test: multivec py (#31)

Author: Anush008

.gitignore

@@ -22,4 +22,5 @@ poetry.lock
 .pytest_cache/
 *_pycache__
 
-senv
+senv
+.venv

pom.xml

@@ -6,7 +6,7 @@
 	<modelVersion>4.0.0</modelVersion>
 	<groupId>io.qdrant</groupId>
 	<artifactId>spark</artifactId>
-	<version>2.3.1</version>
+	<version>2.3.2</version>
 	<name>qdrant-spark</name>
 	<url>https://github.com/qdrant/qdrant-spark</url>
 	<description>An Apache Spark connector for the Qdrant vector database</description>
@@ -18,7 +18,7 @@
 	</licenses>
 	<developers>
 		<developer>
-			<name>Anush Shetty</name>
+			<name>Anush008</name>
 			<email>anush.shetty@qdrant.com</email>
 			<organization>Qdrant</organization>
 			<organizationUrl>http://qdrant.tech</organizationUrl>

src/test/python/conftest.py

@@ -89,6 +89,11 @@ def qdrant():
                 size=QDRANT_EMBEDDING_DIM,
                 distance=QDRANT_DISTANCE,
             ),
+            "multi": models.VectorParams(
+                size=QDRANT_EMBEDDING_DIM,
+                distance=QDRANT_DISTANCE,
+                multivector_config=models.MultiVectorConfig(comparator=models.MultiVectorComparator.MAX_SIM)
+            )
         },
         sparse_vectors_config={
             "sparse": models.SparseVectorParams(),

src/test/python/requirements.txt

@@ -1,4 +1,4 @@
 pyspark==3.5.1
 pytest==8.2.0
-qdrant-client==1.9.0
+qdrant-client==1.10.1
 testcontainers==4.4.0

src/test/python/schema.py

@@ -102,5 +102,6 @@
         StructField("dense_vector", ArrayType(FloatType()), nullable=False),
         StructField("sparse_indices", ArrayType(IntegerType()), nullable=False),
         StructField("sparse_values", ArrayType(FloatType()), nullable=False),
+        StructField("multi", ArrayType(ArrayType(FloatType())), nullable=False),
     ]
 )

src/test/python/test_qdrant_ingest.py

@@ -197,6 +197,29 @@ def test_upsert_multiple_sparse_dense_vectors(
         qdrant.client.count(qdrant.collection_name).count == df.count()
     ), "Uploaded points count is not equal to the dataframe count"
 
+def test_upsert_multi_vector(
+    qdrant: Qdrant, spark_session: SparkSession
+):
+    df = (
+        spark_session.read.schema(schema)
+        .option("multiline", "true")
+        .json(str(input_file_path))
+    )
+    opts = {
+        "qdrant_url": qdrant.url,
+        "collection_name": qdrant.collection_name,
+        "multi_vector_fields": "multi",
+        "multi_vector_names": "multi",
+        "schema": df.schema.json(),
+        "api_key": qdrant.api_key,
+    }
+
+    df.write.format("io.qdrant.spark.Qdrant").options(**opts).mode("append").save()
+
+    assert (
+        qdrant.client.count(qdrant.collection_name).count == df.count()
+    ), "Uploaded points count is not equal to the dataframe count"
+
 
 # Test an upsert without vectors. All the dataframe fields will be treated as payload
 def test_upsert_without_vectors(qdrant: Qdrant, spark_session: SparkSession):