switch to Client6 RC2 version

Author: antas-marcin

build.sbt

@@ -16,7 +16,7 @@ ThisBuild / scalafixDependencies += "org.scalalint" %% "rules" % "0.2.1" % "runt
 
 lazy val sparkVersion = "4.0.1"
 lazy val grpcNettyShadedVersion = "1.76.0"
-lazy val weaviateClient6Version = "6.0.0-M2"
+lazy val weaviateClient6Version = "6.0.0-RC2"
 lazy val scalaCollectionCompatVersion = "2.14.0"
 lazy val scalatestVersion = "3.2.19"
 lazy val gsonVersion = "2.13.2"

src/main/scala/WeaviateDataWriter.scala

@@ -2,12 +2,12 @@ package io.weaviate.spark
 
 import com.google.gson.reflect.TypeToken
 import com.google.gson.{Gson, JsonSyntaxException}
-import io.weaviate.client6.v1.api.collections.data.Reference
+import io.weaviate.client6.v1.api.collections.WeaviateObject
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage}
 import org.apache.spark.sql.types._
-import io.weaviate.client6.v1.api.collections.{CollectionConfig, ObjectMetadata, Vectors, WeaviateObject}
+import io.weaviate.client6.v1.api.collections.{CollectionConfig, Vectors}
 import org.apache.spark.sql.catalyst.util.{ArrayData, GenericArrayData}
 
 import java.util.{Map => JavaMap}
@@ -20,23 +20,25 @@ case class WeaviateCommitMessage(msg: String) extends WriterCommitMessage
 
 case class WeaviateDataWriter(weaviateOptions: WeaviateOptions, schema: StructType)
   extends DataWriter[InternalRow] with Serializable with Logging {
-  var batch = mutable.Map[String, WeaviateObject[JavaMap[String, Object], Reference, ObjectMetadata]]()
+  var batch = mutable.Map[String, WeaviateObject[JavaMap[String, Object]]]()
   private lazy val weaviateClass = weaviateOptions.getCollectionConfig()
 
   override def write(record: InternalRow): Unit = {
     val weaviateObject = buildWeaviateObject(record, weaviateClass)
-    batch += (weaviateObject.uuid() -> weaviateObject)
+    batch += (weaviateObject.uuid -> weaviateObject)
 
     if (batch.size >= weaviateOptions.batchSize) writeBatch()
   }
 
   def writeBatch(retries: Int = weaviateOptions.retries): Unit = {
     if (batch.isEmpty) return
 
-    val consistencyLevel = weaviateOptions.consistencyLevel
     val client = weaviateOptions.getClient()
 
-    val collection = client.collections.use(weaviateOptions.className).withTenant(weaviateOptions.tenant).withConsistencyLevel(consistencyLevel)
+    val collection = client.collections
+      .use(weaviateOptions.className)
+      .withTenant(weaviateOptions.tenant)
+      .withConsistencyLevel(weaviateOptions.consistencyLevel)
 
     val results = collection.data.insertMany(batch.values.toList.asJava)
 
@@ -76,11 +78,9 @@ case class WeaviateDataWriter(weaviateOptions: WeaviateOptions, schema: StructTy
     }
   }
 
-  private[spark] def buildWeaviateObject(record: InternalRow, collectionConfig: CollectionConfig = null): WeaviateObject[java.util.Map[String, Object], Reference, ObjectMetadata] = {
-    var builder: WeaviateObject.Builder[java.util.Map[String, Object], Reference, ObjectMetadata] = new WeaviateObject.Builder()
-    builder = builder.collection(weaviateOptions.className)
+  private[spark] def buildWeaviateObject(record: InternalRow, collectionConfig: CollectionConfig = null): WeaviateObject[java.util.Map[String, Object]] = {
+    val builder: WeaviateObject.Builder[java.util.Map[String, Object]] = new WeaviateObject.Builder()
 
-    val metadata = new ObjectMetadata.Builder()
     val properties = mutable.Map[String, AnyRef]()
     var vector: Array[Float] = null
     val vectors = mutable.Map[String, Array[Float]]()
@@ -103,13 +103,13 @@ case class WeaviateDataWriter(weaviateOptions: WeaviateOptions, schema: StructTy
 
           multiVectors += (weaviateOptions.multiVectors(key) -> multiVector)
         }
-        case weaviateOptions.id => metadata.uuid(record.getString(field._2))
+        case weaviateOptions.id => builder.uuid(record.getString(field._2))
         case _ => properties(field._1.name) = getPropertyValue(field._2, record, field._1.dataType, false, field._1.name, collectionConfig)
       }
     )
 
     if (weaviateOptions.id == null) {
-      metadata.uuid(java.util.UUID.randomUUID.toString)
+      builder.uuid(java.util.UUID.randomUUID.toString)
     }
 
     val allVectors = ListBuffer.empty[Vectors]
@@ -124,9 +124,8 @@ case class WeaviateDataWriter(weaviateOptions: WeaviateOptions, schema: StructTy
       val arr = multiVectors.map { case (key, multiVector) => Vectors.of(key, multiVector) }.toArray
       allVectors ++= arr
     }
-    metadata.vectors(allVectors.toSeq : _*)
 
-    builder.properties(properties.asJava).metadata(metadata.build()).build()
+    builder.tenant(weaviateOptions.tenant).properties(properties.asJava).vectors(allVectors.toSeq : _*).build()
   }
 
   def getPropertyValue(index: Int, record: InternalRow, dataType: DataType, parseObjectArrayItem: Boolean, propertyName: String, collectionConfig: CollectionConfig): AnyRef = {
@@ -216,7 +215,7 @@ case class WeaviateDataWriter(weaviateOptions: WeaviateOptions, schema: StructTy
           })
         }
         objList.asJava
-      case default => throw new SparkDataTypeNotSupported(s"DataType ${default} is not supported by Weaviate")
+      case default => throw SparkDataTypeNotSupported(s"DataType ${default} is not supported by Weaviate")
     }
   }
 

src/main/scala/WeaviateOptions.scala

@@ -90,7 +90,7 @@ class WeaviateOptions(config: CaseInsensitiveStringMap) extends Serializable {
     if (oidcUsername.trim().nonEmpty && oidcPassword.trim().nonEmpty) {
       config.authentication(Authentication.resourceOwnerPassword(oidcUsername, oidcPassword, oidcScopes.asJava))
     } else if (oidcClientSecret.trim().nonEmpty) {
-      config.authentication(Authentication.clientCredentials(oidcClientId, oidcClientSecret, oidcScopes.asJava))
+      config.authentication(Authentication.clientCredentials(oidcClientSecret, oidcScopes.asJava))
     } else if (oidcAccessToken.trim().nonEmpty) {
       config.authentication(Authentication.bearerToken(oidcAccessToken, oidcRefreshToken, oidcAccessTokenLifetime))
     } else if (apiKey.trim().nonEmpty) {

src/test/python/test_spark.py

@@ -33,7 +33,7 @@ def get_connector_version():
 
 connector_version = os.environ.get("CONNECTOR_VERSION", get_connector_version())
 scala_version = os.environ.get("SCALA_VERSION", "2.13")
-weaviate_version = os.environ.get("WEAVIATE_VERSION", "1.30.3")
+weaviate_version = os.environ.get("WEAVIATE_VERSION", "1.32.17")
 spark_connector_jar_path = os.environ.get(
     "CONNECTOR_JAR_PATH", f"target/scala-{scala_version}/spark-connector-assembly-{connector_version}.jar"
 )
@@ -69,7 +69,7 @@ def weaviate_client():
                      "CLUSTER_HOSTNAME": "node1",
                      "PERSISTENCE_DATA_PATH": "./data"},
     )
-    time.sleep(0.5)
+    time.sleep(2)
     wclient = weaviate.Client('http://localhost:8080')
     test_class_name = "TestWillBeRemoved"
     retries = 3

src/test/scala/SparkIntegrationTest.scala

@@ -5,7 +5,7 @@ import org.apache.spark.sql.execution.streaming.MemoryStream
 import org.apache.spark.sql.{AnalysisException, DataFrame, Encoder, Encoders}
 import org.scalatest.BeforeAndAfter
 import org.scalatest.funsuite.AnyFunSuite
-import io.weaviate.client6.v1.api.collections.{Property}
+import io.weaviate.client6.v1.api.collections.Property
 import io.weaviate.client6.v1.internal.ObjectBuilder
 import org.apache.spark.SparkException
 

src/test/scala/TestWeaviateDataWriter.scala

@@ -29,8 +29,7 @@ class TestWeaviateDataWriter extends AnyFunSuite {
     assert(weaviateObject.properties().get("title").equals("Sam"))
     assert(weaviateObject.properties().get("content") == "Sam")
     assert(weaviateObject.properties().get("wordCount") == 5)
-    // how to get tenant?
-//    assert(weaviateObject. == "TenantA")
+    assert(weaviateObject.tenant() == "TenantA")
   }
 
   test("Test Build Weaviate Object without supplied ID") {
@@ -52,7 +51,7 @@ class TestWeaviateDataWriter extends AnyFunSuite {
     assert(weaviateObject.properties().get("content") == "Sam")
     assert(weaviateObject.properties().get("wordCount") == 5)
     assert(weaviateObject.uuid() != null)
-//    assert(weaviateObject.getTenant == null)
+    assert(weaviateObject.tenant() == null)
   }
 
   test("Test Build Weaviate Object with supplied ID") {
@@ -78,7 +77,7 @@ class TestWeaviateDataWriter extends AnyFunSuite {
     assert(weaviateObject.properties().get("content") == "Sam")
     assert(weaviateObject.properties().get("wordCount") == 5)
     assert(weaviateObject.uuid() == uuid)
-//    assert(weaviateObject.getTenant == null)
+    assert(weaviateObject.tenant() == null)
   }
 
   test("Test Build Weaviate Object with DateString") {
@@ -103,7 +102,7 @@ class TestWeaviateDataWriter extends AnyFunSuite {
     assert(weaviateObject.properties().get("content") == "Sam")
     assert(weaviateObject.properties().get("wordCount") == 5)
     assert(weaviateObject.properties().get("date") == "2022-11-18T00:00:00Z")
-//    assert(weaviateObject.getTenant == null)
+    assert(weaviateObject.tenant() == null)
   }
 
   test("Test Build Weaviate Object with Unsupported Data types") {
@@ -180,9 +179,9 @@ class TestWeaviateDataWriter extends AnyFunSuite {
     assert(weaviateObject.properties().get("content") == "Sam")
     assert(weaviateObject.properties().get("wordCount") == 5)
     assert(weaviateObject.uuid() == uuid)
-    assert(weaviateObject.metadata().vectors() != null)
-    assert(weaviateObject.metadata().vectors().getDefaultSingle().sameElements(embedding))
-//    assert(weaviateObject.getTenant == null)
+    assert(weaviateObject.vectors() != null)
+    assert(weaviateObject.vectors().getDefaultSingle().sameElements(embedding))
+    assert(weaviateObject.tenant() == null)
   }
 
   test("Test Build Weaviate Object with vectors") {
@@ -214,10 +213,10 @@ class TestWeaviateDataWriter extends AnyFunSuite {
     assert(weaviateObject.properties().get("content") == "Sam")
     assert(weaviateObject.properties().get("wordCount") == 5)
     assert(weaviateObject.uuid() == uuid)
-    assert(weaviateObject.metadata().vectors() != null)
-    assert(weaviateObject.metadata().vectors().getSingle("v1").sameElements(embedding1))
-    assert(weaviateObject.metadata().vectors().getSingle("v2").sameElements(embedding2))
-//    assert(weaviateObject.getTenant == null)
+    assert(weaviateObject.vectors() != null)
+    assert(weaviateObject.vectors().getSingle("v1").sameElements(embedding1))
+    assert(weaviateObject.vectors().getSingle("v2").sameElements(embedding2))
+    assert(weaviateObject.tenant() == null)
   }
 
   test("Test Build Weaviate Object with vectors and multi vectors") {
@@ -250,12 +249,12 @@ class TestWeaviateDataWriter extends AnyFunSuite {
     assert(weaviateObject.properties().get("content") == "Sam")
     assert(weaviateObject.properties().get("wordCount") == 5)
     assert(weaviateObject.uuid() == uuid)
-    assert(weaviateObject.metadata().vectors() != null)
-    assert(weaviateObject.metadata().vectors().getSingle("v1").sameElements(embedding1))
-    assert(weaviateObject.metadata().vectors().getMulti("colbert").length == 2)
-    assert(weaviateObject.metadata().vectors().getMulti("colbert")(0).sameElements(colbert(0)))
-    assert(weaviateObject.metadata().vectors().getMulti("colbert")(1).sameElements(colbert(1)))
-//    assert(weaviateObject.getTenant == null)
+    assert(weaviateObject.vectors() != null)
+    assert(weaviateObject.vectors().getSingle("v1").sameElements(embedding1))
+    assert(weaviateObject.vectors().getMulti("colbert").length == 2)
+    assert(weaviateObject.vectors().getMulti("colbert")(0).sameElements(colbert(0)))
+    assert(weaviateObject.vectors().getMulti("colbert")(1).sameElements(colbert(1)))
+    assert(weaviateObject.tenant() == null)
   }
 
   test("Test Build Weaviate Object with geo coordinates") {
@@ -287,6 +286,6 @@ class TestWeaviateDataWriter extends AnyFunSuite {
     assert(weaviateObject.uuid() == uuid)
     assert(weaviateObject.properties().get("title") == "title")
     assert(weaviateObject.properties().get("geo") != null)
-//    assert(weaviateObject.getTenant == null)
+    assert(weaviateObject.tenant() == null)
   }
 }

src/test/scala/WeaviateDocker.scala

@@ -25,7 +25,7 @@ object WeaviateDocker {
   var retries = 10
 
   def start(vectorizerModule: String = "none", enableModules: String = "text2vec-openai"): Int = {
-    val weaviateVersion = "1.30.6"
+    val weaviateVersion = "1.34.0"
     val docker_run =
       s"""docker run -d --name=weaviate-test-container-will-be-deleted
 -p 8080:8080
@@ -64,6 +64,7 @@ semitechnologies/weaviate:$weaviateVersion"""
       }.getOrElse(false)
     }
 
+    Thread.sleep(2000L)
     val maxAttempts = 10
     for (_ <- 1 to maxAttempts) {
       if (checkReadinessProbe) {