readMilvusCollection is not able to read all of the vectors in a collection.

[nisharyan]

I am trying to use the spark-milvus connector to read all the vector embeddings from a collection.
I created a collection with 100,000 vectors and tried to read from it.

Initially the total vectors read is around 50k usually. Later after flushing the segments, the number of vectors read goes up to 95-98k. But the whole collection cannot be read.

• I am using df.count() to check the number of vectors read.
• I am using readMilvusCollection of MilvusUtils to read the collection.

Am I missing something here?

[xiaofan-luan]

did you try to flush after insert all the data?

[nisharyan]

yes, only inserted once to the collections. Verified that all the segments were flushed using attu

[nisharyan]

Also facing this issue when reading a large collection (1 million records).
All the data was flushed and no data was inserted after flushing.

Exception in thread "main" org.apache.spark.sql.AnalysisException: Cannot resolve column name "vector" among (102, segment_rowid)
        at org.apache.spark.sql.errors.QueryCompilationErrors$.cannotResolveColumnNameAmongFieldsError(QueryCompilationErrors.scala:2265)
        at org.apache.spark.sql.Dataset.org$apache$spark$sql$Dataset$$resolveException(Dataset.scala:259)
        at org.apache.spark.sql.Dataset.$anonfun$resolve$1(Dataset.scala:251)
        at scala.Option.getOrElse(Option.scala:189)
        at org.apache.spark.sql.Dataset.resolve(Dataset.scala:251)
        at org.apache.spark.sql.Dataset.col(Dataset.scala:1417)
        at org.apache.spark.sql.Dataset.apply(Dataset.scala:1384)
        at zilliztech.spark.milvus.MilvusUtils$.$anonfun$readMilvusCollection$8(MilvusUtils.scala:94)
        at zilliztech.spark.milvus.MilvusUtils$.$anonfun$readMilvusCollection$8$adapted(MilvusUtils.scala:79)
        at scala.collection.TraversableLike.$anonfun$map$1(TraversableLike.scala:286)
        at scala.collection.IndexedSeqOptimized.foreach(IndexedSeqOptimized.scala:36)
        at scala.collection.IndexedSeqOptimized.foreach$(IndexedSeqOptimized.scala:33)
        at scala.collection.mutable.ArrayOps$ofLong.foreach(ArrayOps.scala:258)
        at scala.collection.TraversableLike.map(TraversableLike.scala:286)
        at scala.collection.TraversableLike.map$(TraversableLike.scala:279)
        at scala.collection.mutable.ArrayOps$ofLong.map(ArrayOps.scala:258)
        at zilliztech.spark.milvus.MilvusUtils$.$anonfun$readMilvusCollection$7(MilvusUtils.scala:79)
        at scala.collection.TraversableLike.$anonfun$map$1(TraversableLike.scala:286)
        at scala.collection.Iterator.foreach(Iterator.scala:943)
        at scala.collection.Iterator.foreach$(Iterator.scala:943)
        at scala.collection.AbstractIterator.foreach(Iterator.scala:1431)
        at scala.collection.IterableLike.foreach(IterableLike.scala:74)
        at scala.collection.IterableLike.foreach$(IterableLike.scala:73)
        at scala.collection.AbstractIterable.foreach(Iterable.scala:56)
        at scala.collection.TraversableLike.map(TraversableLike.scala:286)
        at scala.collection.TraversableLike.map$(TraversableLike.scala:279)
        at scala.collection.AbstractTraversable.map(Traversable.scala:108)
        at zilliztech.spark.milvus.MilvusUtils$.readMilvusCollection(MilvusUtils.scala:78)
        at com.example.milvus.MilvusBridge.readCollection(Main.scala:158)
        at com.example.milvus.MilvusBridgeExample$.main(Main.scala:285)
        at com.example.milvus.MilvusBridgeExample.main(Main.scala)
        at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke0(Native Method)
        at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke(Unknown Source)
        at java.base/jdk.internal.reflect.DelegatingMethodAccessorImpl.invoke(Unknown Source)
        at java.base/java.lang.reflect.Method.invoke(Unknown Source)
        at org.apache.spark.deploy.JavaMainApplication.start(SparkApplication.scala:52)
        at org.apache.spark.deploy.SparkSubmit.org$apache$spark$deploy$SparkSubmit$$runMain(SparkSubmit.scala:958)
        at org.apache.spark.deploy.SparkSubmit.doRunMain$1(SparkSubmit.scala:180)
        at org.apache.spark.deploy.SparkSubmit.submit(SparkSubmit.scala:203)
        at org.apache.spark.deploy.SparkSubmit.doSubmit(SparkSubmit.scala:90)
        at org.apache.spark.deploy.SparkSubmit$$anon$2.doSubmit(SparkSubmit.scala:1046)
        at org.apache.spark.deploy.SparkSubmit$.main(SparkSubmit.scala:1055)
        at org.apache.spark.deploy.SparkSubmit.main(SparkSubmit.scala)

The field name where I store the embeddings is "vector".

[xiaofan-luan]

Cannot resolve column name "vector" among (102, segment_rowid)

it seems that you it's a code issue here, vector is not some field you select

[xiaofan-luan]

you probably want to share you code to run the connector

[nisharyan]

package com.example.milvus

import org.apache.spark.sql.{DataFrame, SaveMode, SparkSession}
import org.apache.spark.sql.util.CaseInsensitiveStringMap
import zilliztech.spark.milvus.{MilvusOptions, MilvusUtils}
import scala.collection.JavaConverters._
import org.apache.spark.ml.linalg.Vector
import org.apache.spark.sql.functions._

class MilvusBridge(spark: SparkSession) {

  // Method to configure Milvus connection options
  def createMilvusOptions(
    host: String,
    port: Int,
    collectionName: String,
    databaseName: String,
    bucketName: String,
    rootPath: String,
    storageEndpoint: String,
    storageUser: String,
    storagePassword: String,
    username: String = "root",
    password: String = ""
  ): MilvusOptions = {
    val properties = Map(
      MilvusOptions.MILVUS_USERNAME -> username,
      MilvusOptions.MILVUS_PASSWORD -> password,
      MilvusOptions.MILVUS_HOST -> host,
      MilvusOptions.MILVUS_PORT -> port.toString,
      MilvusOptions.MILVUS_COLLECTION_NAME -> collectionName,
      MilvusOptions.MILVUS_DATABASE_NAME -> databaseName,
      MilvusOptions.MILVUS_BUCKET -> bucketName,
      MilvusOptions.MILVUS_ROOTPATH -> rootPath,
      MilvusOptions.MILVUS_FS -> "s3a://",
      MilvusOptions.MILVUS_STORAGE_ENDPOINT -> storageEndpoint,
      MilvusOptions.MILVUS_STORAGE_USER -> storageUser,
      MilvusOptions.MILVUS_STORAGE_PASSWORD -> storagePassword,
      MilvusOptions.MILVUS_STORAGE_USESSL -> "false"
    )

    new MilvusOptions(new CaseInsensitiveStringMap(properties.asJava))
  }

  // Method to read from Milvus collection
  def readCollection(milvusOptions: MilvusOptions): DataFrame = {
    MilvusUtils.readMilvusCollection(spark, milvusOptions)
  }

  // Method to write DataFrame to Milvus
  def writeToMilvus(
    df: DataFrame,
    host: String,
    port: Int,
    collectionName: String,
    databaseName: String,
    vectorField: String,
    vectorDim: String,
    primaryKeyField: String
  ): Unit = {
    val writeOptions = Map(
      "milvus.host" -> host,
      "milvus.port" -> port.toString,
      "milvus.collection.name" -> collectionName,
      "milvus.database.name" -> databaseName,
      "milvus.collection.vectorField" -> vectorField,
      "milvus.collection.vectorDim" -> vectorDim,
      "milvus.collection.primaryKeyField" -> primaryKeyField
    )

    df.write
      .format("milvus")
      .options(writeOptions)
      .mode("append")
      .save()
  }
}

// Companion object to create the bridge
object MilvusBridge {
  def apply(spark: SparkSession): MilvusBridge = new MilvusBridge(spark)
}

class MilvusVectorOperations(spark: SparkSession) {
  import spark.implicits._

  // Function to calculate average vector from a DataFrame containing vectors
  def calculateAverageVector(df: DataFrame, vectorColumnName: String): Array[Float] = {
    // First, convert the vector column to array type if it isn't already
    val arrayDF = df.select(
      col(vectorColumnName).cast("array<float>").as("vector_array")
    )

    // Calculate the sum and count of vectors
    val vectorStats = arrayDF.agg(
      collect_list("vector_array").as("all_vectors"),
      count("*").as("total_vectors")
    ).first()

    val allVectors = vectorStats.getAs[Seq[Seq[Float]]](0)
    val totalVectors = vectorStats.getAs[Long](1)

    // Calculate the average vector
    if (totalVectors > 0) {
      val vectorLength = allVectors.head.length
      val sumVector = new Array[Float](vectorLength)

      // Sum all vectors
      allVectors.foreach { vector =>
        for (i <- vector.indices) {
          sumVector(i) += vector(i)
        }
      }

      // Divide by total count to get average
      sumVector.map(_ / totalVectors)
    } else {
      Array.empty[Float]
    }
  }

  // Function to print vector statistics
  def printVectorStats(df: DataFrame, vectorColumnName: String): Unit = {
    val avgVector = calculateAverageVector(df, vectorColumnName)

    println(s"Vector Statistics:")
    println(s"Total number of vectors: ${df.count()}")
    println(s"Vector dimension: ${avgVector.length}")
    println(s"Average vector: ${avgVector.mkString("[", ", ", "]")}")

    // Calculate magnitude of average vector
    val magnitude = Math.sqrt(avgVector.map(x => x * x).sum)
    println(s"Magnitude of average vector: $magnitude")
  }
}

// Companion object
object MilvusVectorOperations {
  def apply(spark: SparkSession): MilvusVectorOperations = new MilvusVectorOperations(spark)
}


// The class where the execution starts.
object MilvusBridgeExample {
  def main(args: Array[String]): Unit = {
    // Initialize SparkSession
    val spark = SparkSession.builder()
      .appName("MilvusBridge Example")
      .master("local[*]")
      .config("spark.jars", "https://github.com/zilliztech/spark-milvus/raw/1.0.0-SNAPSHOT/output/spark-milvus-1.0.0-SNAPSHOT.jar")
      .getOrCreate()

    // Create MilvusBridge instance
    val bridge = MilvusBridge(spark)

    // Configure Milvus options
    val milvusOptions = bridge.createMilvusOptions(
      host = "localhost",
      port = 19530,
      collectionName = "<collection_name>",
      databaseName = "default",
      bucketName = "<aws_s3_bucket_name>",
      rootPath = "",
      storageEndpoint = "s3.us-east-2.amazonaws.com",
      storageUser = "<aws_user_access_key>",
      storagePassword = "<aws_user_secret_key>"
    )

    // Read from Milvus collection
    val df = bridge.readCollection(milvusOptions)

    // Show the results
    df.show()

    val vectorOps = MilvusVectorOperations(spark)

    // Calculate and print vector statistics
    vectorOps.printVectorStats(df, "vector")

    // The average vector for further processing
    val avgVector = vectorOps.calculateAverageVector(df, "vector")

    // Clean up
    spark.stop()
  }
}