diff --git a/data_chain/stores/database/database.py b/data_chain/stores/database/database.py index d0cd3961d26281fa4a6b9cc586369df078b779fe..f0b67f0737de3ff8d2481c977d2c9ec808f749a6 100644 --- a/data_chain/stores/database/database.py +++ b/data_chain/stores/database/database.py @@ -366,7 +366,8 @@ class DocumentEntity(Base): server_default=func.current_timestamp(), onupdate=func.current_timestamp() ) - __table_args__ = ( + if config['DATABASE_TYPE'].lower() == 'postgres': + __table_args__ = ( Index('abstract_ts_vector_index', abstract_ts_vector, postgresql_using='gin'), Index( @@ -375,9 +376,21 @@ class DocumentEntity(Base): postgresql_using='hnsw', postgresql_with={'m': 32, 'ef_construction': 200}, postgresql_ops={'abstract_vector': 'vector_cosine_ops'} - ), - Index('abstract_bm25_index', abstract, postgresql_using='bm25') - ) + ) + ) + else: + __table_args__ = ( + Index('abstract_ts_vector_index', + abstract_ts_vector, postgresql_using='gin'), + Index( + 'abstract_vector_index', + abstract_vector, + postgresql_using='hnsw', + postgresql_with={'m': 32, 'ef_construction': 200}, + postgresql_ops={'abstract_vector': 'vector_cosine_ops'} + ), + Index('abstract_bm25_index', abstract, postgresql_using='bm25') + ) class ChunkEntity(Base): @@ -411,7 +424,8 @@ class ChunkEntity(Base): TIMESTAMP(timezone=True), server_default=func.current_timestamp(), onupdate=func.current_timestamp()) - __table_args__ = ( + if config['DATABASE_TYPE'].lower() == 'postgres': + __table_args__ = ( Index('text_ts_vector_index', text_ts_vector, postgresql_using='gin'), Index( 'text_vector_index', @@ -419,9 +433,20 @@ class ChunkEntity(Base): postgresql_using='hnsw', postgresql_with={'m': 32, 'ef_construction': 200}, postgresql_ops={'text_vector': 'vector_cosine_ops'} - ), - Index('text_bm25_index', text, postgresql_using='bm25') + ) ) + else: + __table_args__ = ( + Index('text_ts_vector_index', text_ts_vector, postgresql_using='gin'), + Index( + 'text_vector_index', + text_vector, + postgresql_using='hnsw', + postgresql_with={'m': 32, 'ef_construction': 200}, + postgresql_ops={'text_vector': 'vector_cosine_ops'} + ), + Index('text_bm25_index', text, postgresql_using='bm25') + ) class ImageEntity(Base):