add test

Author: jerry-024

paimon-core/src/main/java/org/apache/paimon/table/VectorSearchTable.java

@@ -24,21 +24,20 @@
 import org.apache.paimon.table.source.InnerTableScan;
 import org.apache.paimon.types.RowType;
 
-import javax.annotation.Nullable;
-
 import java.util.List;
 import java.util.Map;
 
 /**
- * A table wrapper to hold vector search information. This is used by Spark engine to pass vector
- * search pushdown information from logical plan optimization to physical plan execution.
+ * A table wrapper to hold vector search information. This is used to pass vector search pushdown
+ * information from logical plan optimization to physical plan execution. For now, it is only used
+ * by internal for Spark engine.
  */
 public class VectorSearchTable implements ReadonlyTable {
 
     private final InnerTable origin;
     private final VectorSearch vectorSearch;
 
-    VectorSearchTable(InnerTable origin, VectorSearch vectorSearch) {
+    private VectorSearchTable(InnerTable origin, VectorSearch vectorSearch) {
         this.origin = origin;
         this.vectorSearch = vectorSearch;
     }
@@ -47,7 +46,6 @@ public static VectorSearchTable create(InnerTable origin, VectorSearch vectorSea
         return new VectorSearchTable(origin, vectorSearch);
     }
 
-    @Nullable
     public VectorSearch vectorSearch() {
         return vectorSearch;
     }
@@ -93,7 +91,7 @@ public InnerTableRead newRead() {
 
     @Override
     public InnerTableScan newScan() {
-        return origin.newScan();
+        throw new UnsupportedOperationException();
     }
 
     @Override

paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/PaimonScan.scala

@@ -39,7 +39,7 @@ case class PaimonScan(
     pushedDataFilters: Seq[Predicate],
     override val pushedLimit: Option[Int],
     override val pushedTopN: Option[TopN],
-    override val pushedVectorSearch: Option[VectorSearch] = None,
+    override val pushedVectorSearch: Option[VectorSearch],
     bucketedScanDisabled: Boolean = false)
   extends PaimonBaseScan(table)
   with SupportsRuntimeFiltering

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScan.scala

@@ -41,7 +41,7 @@ case class PaimonScan(
     pushedDataFilters: Seq[Predicate],
     override val pushedLimit: Option[Int],
     override val pushedTopN: Option[TopN],
-    override val pushedVectorSearch: Option[VectorSearch] = None,
+    override val pushedVectorSearch: Option[VectorSearch],
     bucketedScanDisabled: Boolean = false)
   extends PaimonBaseScan(table)
   with SupportsReportPartitioning

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala

@@ -130,7 +130,13 @@ class PaimonScanBuilder(val table: InnerTable)
       case None =>
         val (actualTable, vectorSearch) = table match {
           case vst: org.apache.paimon.table.VectorSearchTable =>
-            (vst.origin(), Some(vst.vectorSearch()))
+            val tableVectorSearch = Option(vst.vectorSearch())
+            val vs = (tableVectorSearch, pushedVectorSearch) match {
+              case (Some(_), _) => tableVectorSearch
+              case (None, Some(_)) => pushedVectorSearch
+              case (None, None) => None
+            }
+            (vst.origin(), vs)
           case _ => (table, pushedVectorSearch)
         }
 

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/plans/logical/PaimonTableValuedFunctions.scala

@@ -117,7 +117,7 @@ object PaimonTableValuedFunctions {
       argsWithoutTable: Seq[Expression]): LogicalPlan = {
     sparkTable match {
       case st @ SparkTable(innerTable: InnerTable) =>
-        val vectorSearch = vsq.createVectorSearch(argsWithoutTable)
+        val vectorSearch = vsq.createVectorSearch(innerTable, argsWithoutTable)
         val vectorSearchTable = VectorSearchTable.create(innerTable, vectorSearch)
         DataSourceV2Relation.create(
           st.copy(table = vectorSearchTable),
@@ -126,7 +126,8 @@ object PaimonTableValuedFunctions {
           CaseInsensitiveStringMap.empty())
       case _ =>
         throw new RuntimeException(
-          s"vector_search only supports Paimon tables, got ${sparkTable.getClass.getName}")
+          "vector_search only supports Paimon SparkTable backed by InnerTable, " +
+            s"but got table implementation: ${sparkTable.getClass.getName}")
     }
   }
 
@@ -258,21 +259,32 @@ case class VectorSearchQuery(override val args: Seq[Expression])
     Map.empty
   }
 
-  def createVectorSearch(argsWithoutTable: Seq[Expression]): VectorSearch = {
-    assert(
-      argsWithoutTable.size == 3,
-      s"$VECTOR_SEARCH needs four parameters: table_name, column_name, query_vector, limit. " +
-        s"Got ${argsWithoutTable.size + 1} parameters."
-    )
-
+  def createVectorSearch(
+      innerTable: InnerTable,
+      argsWithoutTable: Seq[Expression]): VectorSearch = {
+    if (argsWithoutTable.size != 3) {
+      throw new RuntimeException(
+        s"$VECTOR_SEARCH needs three parameters after table_name: column_name, query_vector, limit. " +
+          s"Got ${argsWithoutTable.size} parameters after table_name."
+      )
+    }
     val columnName = argsWithoutTable.head.eval().toString
+    if (!innerTable.rowType().containsField(columnName)) {
+      throw new RuntimeException(
+        s"Column $columnName does not exist in table ${innerTable.name()}"
+      )
+    }
     val queryVector = extractQueryVector(argsWithoutTable(1))
     val limit = argsWithoutTable(2).eval() match {
       case i: Int => i
       case l: Long => l.toInt
       case other => throw new RuntimeException(s"Invalid limit type: ${other.getClass.getName}")
     }
-
+    if (limit <= 0) {
+      throw new IllegalArgumentException(
+        s"Limit must be a positive integer, but got: $limit"
+      )
+    }
     new VectorSearch(queryVector, limit, columnName)
   }
 
@@ -281,7 +293,7 @@ case class VectorSearchQuery(override val args: Seq[Expression])
       case Literal(arrayData, _) if arrayData != null =>
         val arr = arrayData.asInstanceOf[org.apache.spark.sql.catalyst.util.ArrayData]
         arr.toFloatArray()
-      case CreateArray(elements, _) =>
+      case CreateArray(elements, _) if elements != null =>
         elements.map {
           case Literal(v: Float, _) => v
           case Literal(v: Double, _) => v.toFloat

paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/BaseVectorSearchPushDownTest.scala

@@ -55,6 +55,46 @@ class BaseVectorSearchPushDownTest extends PaimonSparkTestBase with StreamTest {
 
       // Should return results (actual filtering depends on vector index)
       assert(result.nonEmpty)
+
+      // Test invalid limit (negative)
+      val ex1 = intercept[Exception] {
+        spark
+          .sql("""
+                 |SELECT * FROM vector_search('T', 'v', array(1.0f, 0.0f, 0.0f), -3)
+                 |""".stripMargin)
+          .collect()
+      }
+      assert(ex1.getMessage.contains("Limit must be a positive integer"))
+
+      // Test invalid limit (zero)
+      val ex2 = intercept[Exception] {
+        spark
+          .sql("""
+                 |SELECT * FROM vector_search('T', 'v', array(1.0f, 0.0f, 0.0f), 0)
+                 |""".stripMargin)
+          .collect()
+      }
+      assert(ex2.getMessage.contains("Limit must be a positive integer"))
+
+      // Test missing parameters
+      val ex3 = intercept[Exception] {
+        spark
+          .sql("""
+                 |SELECT * FROM vector_search('T', 'v', array(1.0f, 0.0f, 0.0f))
+                 |""".stripMargin)
+          .collect()
+      }
+      assert(ex3.getMessage.contains("vector_search needs three parameters after table_name"))
+
+      // Test non-existent column
+      val ex4 = intercept[Exception] {
+        spark
+          .sql("""
+                 |SELECT * FROM vector_search('T', 'non_existent_col', array(1.0f, 0.0f, 0.0f), 3)
+                 |""".stripMargin)
+          .collect()
+      }
+      assert(ex4.getMessage.nonEmpty)
     }
   }
 }