Working with Vector stores

Vector store types

Managed Knowledge Banks

A managed Knowledge Bank uses a vector store index/table that DSS creates and updates through the embedding flow. These KBs are typically created and written by an Embed Documents or Embed Dataset recipe.

Out of the box, Knowledge Banks are created with a Vector Store called Chroma. This does not require any setup, and provides good performance even for quite a large corpus.

As an alternative, other no-setup Vector Stores are available: Milvus (local), Qdrant, and FAISS.

For more advanced use cases, you may wish to use a dedicated Vector Store. Dataiku supports several third-party vector stores that require you to set up a dedicated connection beforehand:

  • Azure AI Search

  • Elasticsearch

  • OpenSearch, including AWS OpenSearch services (both managed cluster & serverless)

  • Milvus (remote)

  • Pinecone

  • pgvector (based on a PostgreSQL Connection)

  • Vertex Vector Search (based on a Google Cloud Storage Connection)

When creating the Embedding recipe, select the desired vector store type, then select your connection. You can also change the vector store type later, by editing the settings of the Knowledge Bank.

DSS provides a default name for the index (may be known as “table”, “collection” or “service” depending on the type of vector store) that you can update if needed.

Unmanaged Knowledge Banks

An unmanaged Knowledge Bank uses an existing remote vector store index or service. DSS connects to that remote store for retrieval, but does not create or manage the underlying index/table.

Create a new Knowledge Bank, choose the connection and the index/table to read from and the embedding model to use (when applicable, in which case it must be compatible with the embedding model that was used to index the underlying data). You can now use this Knowledge Bank for retrieval in Retrieval-Augmented LLMs, or in Agent KB Search tools.

Dataiku supports unmanaged Knowledge Banks on the following vector store types:

  • Azure AI Search

  • Elasticsearch

  • OpenSearch, including AWS OpenSearch services (both managed cluster & serverless)

  • Milvus (remote)

  • Pinecone

  • pgvector (based on a PostgreSQL Connection)

  • Vertex Vector Search (based on a Google Cloud Storage Connection)

  • Snowflake Cortex Search (based on a Snowflake connection)

  • Databricks AI Search (based on a Databricks connection)

Snowflake Cortex Search from a dataset

Snowflake Cortex offers a Search Service feature, that automatically maintains an up-to-date searchable index of some source table, re-indexing it when it changes.

You can use an unmanaged Snowflake Knowledge Bank by creating a KB, selecting a Cortex connection, and specifying an existing Cortex Search service.

You can alternatively have DSS create a the Cortex Search service and unmanaged Knowledge Bank from a selected Snowflake dataset. In this case the link between the dataset and the Knowledge Bank is shown in the Flow.

To use this:

  • the source must be a Snowflake dataset,

  • its Snowflake connection must have Allow knowledge banks enabled,

  • its Snowflake connection must have the required privileges to create a Cortex Search service in the dataset’s database & schema.

The Cortex Search service is created in the Snowflake dataset’s underlying database and schema, and uses the warehouse defined on the dataset’s Snowflake connection.

By default, DSS uses Full refresh mode. Incremental refresh is also available. For incremental refresh, CHANGE_TRACKING must be enabled on the source dataset’s underlying table.

When the source dataset is rebuilt, DSS synchronizes the settings of the linked Cortex Search service, possibly recreating the service to keep it aligned with the dataset schema, or resuming suspended indexing for incremental refresh when CHANGE_TRACKING is enabled.

If you later change the connection or service name in the Knowledge Bank settings, the Knowledge Bank is detached from its source dataset and then behaves like a standard unmanaged Knowledge Bank.

Connections

When setting up an Elasticsearch, OpenSearch, Google Cloud Storage, PostgreSQL, Snowflake, or Databricks connection for Knowledge Bank usage, you must enable the Allow knowledge banks setting in the connection panel.

Limitations

  • You need an Elasticsearch version >=7.14.0 to store a Knowledge Bank.

  • Elasticsearch >=8.0.0 and <8.8.0 supports only embeddings of size smaller than 1024. Embedding models generating larger embedding vectors will not work.

  • Milvus (local) does not support empty values in metadata. Dataiku fills empty values with defaults depending on the type (NaN for numbers, False for booleans, and empty string for other types).

  • Milvus (local and remote) does not support adding new metadata columns to already-built Knowledge Banks created from an Embed Documents recipe. New metadata columns are ignored until the Knowledge Bank is cleared and rebuilt.

  • Milvus (local and remote) supports switching update method from Smart sync to Append only after clearing the Knowledge Bank.

  • Only Private key authentication is supported for Google Cloud Storage connections used for Knowledge bank usage.

  • Smart update methods (Smart sync and Upsert) are not supported on the following vector store types: Pinecone, AWS OpenSearch serverless.

  • Note that, after running the embedding recipe, remote vector stores might take some time to update their indexing data in their respective user interfaces.

  • pgvector supports filtering on metadata columns with names containing only alphanumeric characters and underscore _.

  • In pgvector, metadata columns with names containing only lowercase letters and underscore _ are stored in individual columns, while other columns (with names containing uppercase letters or other symbols) are nested into a json column called metadata.