sparknlp.annotator.vector_db.vector_db_connector#

Classes#

class VectorDBConnector(classname='com.johnsnowlabs.ml.ai.VectorDBConnector', java_model=None)[source]#

Connector for storing and retrieving embeddings from vector databases.

This annotator takes embeddings from previous annotators (like BertEmbeddings, SentenceEmbeddings, E5VEmbeddings, etc.) and stores them in a vector database for similarity search and retrieval. Currently supports Pinecone with more providers planned.

Two modality modes are supported via setModalityMode:

• text (default) – expects DOCUMENT, SENTENCE_EMBEDDINGS input columns. Upserted metadata is augmented with modality=text.

• image – expects MAGE, SENTENCE_EMBEDDINGS input columns (e.g. from ImageAssembler + E5VEmbeddings). Upserted metadata is augmented with modality=image, image_origin, image_width, image_height, and image_nChannels. Vector IDs are deterministic UUID-v3 values derived from the image file-path (origin), ensuring stable re-indexing.

Parameters:

provider

Vector database provider. Currently supported: ‘pinecone’

indexName

Name of the index/collection in the vector database

namespace

Namespace/partition within the index (optional)

idColumn

Column name to use as vector ID (if not set, generates UUID; for image mode a stable UUID-v3 derived from the image origin path is used)

metadataColumns

Column names to include as metadata with vectors

batchSize

Number of vectors to upsert in a single batch

modalityMode

Modality mode: ‘text’ (default) or ‘image’

Examples

Text mode example:

>>> import sparknlp
>>> from sparknlp.base import *
>>> from sparknlp.annotator import *
>>> from pyspark.ml import Pipeline

>>> documentAssembler = DocumentAssembler() \
...     .setInputCol("text") \
...     .setOutputCol("document")

>>> embeddings = BertSentenceEmbeddings.pretrained() \
...     .setInputCols(["document"]) \
...     .setOutputCol("sentence_embeddings")

>>> vectorDB = VectorDBConnector() \
...     .setInputCols(["document", "sentence_embeddings"]) \
...     .setOutputCol("vectordb_result") \
...     .setProvider("pinecone") \
...     .setIndexName("my-index") \
...     .setNamespace("production") \
...     .setIdColumn("id") \
...     .setMetadataColumns(["text", "category"]) \
...     .setBatchSize(100)

>>> pipeline = Pipeline().setStages([
...     documentAssembler,
...     embeddings,
...     vectorDB
... ])

>>> data = spark.createDataFrame([
...     ("1", "Spark NLP is great", "tech"),
...     ("2", "Vector databases enable semantic search", "tech")
... ]).toDF("id", "text", "category")

>>> result = pipeline.fit(data).transform(data)

Image mode example:

>>> imageAssembler = ImageAssembler() \
...     .setInputCol("image") \
...     .setOutputCol("image_assembler")

>>> e5vEmbeddings = E5VEmbeddings.pretrained() \
...     .setInputCols(["image_assembler"]) \
...     .setOutputCol("image_embeddings")

>>> vectorDB = VectorDBConnector() \
...     .setInputCols(["image_assembler", "image_embeddings"]) \
...     .setOutputCol("vectordb_result") \
...     .setProvider("pinecone") \
...     .setIndexName("my-multimodal-index") \
...     .setModalityMode("image") \
...     .setBatchSize(50)

>>> pipeline = Pipeline().setStages([
...     imageAssembler,
...     e5vEmbeddings,
...     vectorDB
... ])

name = 'VectorDBConnector'[source]#

inputAnnotatorTypes[source]#

outputAnnotatorType = 'document'[source]#

provider[source]#

indexName[source]#

namespace[source]#

idColumn[source]#

metadataColumns[source]#

batchSize[source]#

modalityMode[source]#

setProvider(value)[source]#

Sets the vector database provider.

Parameters:

valuestr

Vector database provider. Currently supported: ‘pinecone’

setIndexName(value)[source]#

Sets the name of the index/collection in the vector database.

Parameters:

valuestr

Name of the index/collection

setNamespace(value)[source]#

Sets the namespace/partition within the index.

Parameters:

valuestr

Namespace/partition name (optional)

setIdColumn(value)[source]#

Sets the column name to use as vector ID.

Parameters:

valuestr

Column name for vector ID. If not set, UUIDs will be generated.

setMetadataColumns(value)[source]#

Sets the column names to include as metadata with vectors.

Parameters:

valuelist[str]

List of column names to include as metadata

setBatchSize(value)[source]#

Sets the number of vectors to upsert in a single batch.

Parameters:

valueint

Batch size for upsert operations (max 1000)

setModalityMode(value)[source]#

Sets the modality mode for indexing.

Use ‘text’ (default) for DOCUMENT + SENTENCE_EMBEDDINGS pipelines and ‘image’ for IMAGE + SENTENCE_EMBEDDINGS pipelines (e.g. ImageAssembler + E5VEmbeddings). In image mode vector IDs are stable UUID-v3 values derived from the image origin path, and upserted metadata automatically includes modality, image_origin, image_width, image_height, and image_nChannels fields.

Parameters:

valuestr

‘text’ or ‘image’

getModalityMode()[source]#

Gets the current modality mode.

Returns:

str

‘text’ or ‘image’