This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 4137611284 [spark] skip table lookup for meta-columns-only vector
search since row ids and scores are already available from the index result
(#8485)
4137611284 is described below
commit 4137611284a339ed70f625d132f2f62738da3d80
Author: Stefanietry <[email protected]>
AuthorDate: Tue Jul 7 15:20:49 2026 +0800
[spark] skip table lookup for meta-columns-only vector search since row ids
and scores are already available from the index result (#8485)
---
.../apache/paimon/spark/PaimonScanBuilder.scala | 16 +++++
.../org/apache/paimon/spark/PaimonBaseScan.scala | 54 ++++++++++-------
.../paimon/spark/PaimonRecordReaderIterator.scala | 6 +-
.../apache/paimon/spark/PaimonScanBuilder.scala | 18 +++++-
.../plans/logical/PaimonTableValuedFunctions.scala | 2 +-
.../paimon/spark/execution/PaimonStrategy.scala | 31 ++++++++++
.../spark/read/VectorSearchResultUtils.scala | 70 ++++++++++++++++++++++
.../paimon/spark/schema/PaimonMetadataColumn.scala | 3 +
.../apache/paimon/spark/SparkMultimodalITCase.java | 41 +++++++++++++
9 files changed, 216 insertions(+), 25 deletions(-)
diff --git
a/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
b/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
index 5253cc9894..e451a433cd 100644
---
a/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
+++
b/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
@@ -18,6 +18,7 @@
package org.apache.paimon.spark
+import org.apache.paimon.spark.read.{PaimonLocalScan, VectorSearchResultUtils}
import org.apache.paimon.table.InnerTable
import org.apache.spark.sql.connector.read.Scan
@@ -34,6 +35,21 @@ class PaimonScanBuilder(val table: InnerTable) extends
PaimonBaseScanBuilder {
(ftst.origin(), None, None, Option(ftst.fullTextSearch()))
case _ => (table, pushedVectorSearch, None, pushedFullTextSearch)
}
+ if (
+ vectorSearch.isDefined &&
+
VectorSearchResultUtils.isVectorSearchMetaOnly(requiredSchema.fieldNames.toSeq)
+ ) {
+ val result = PaimonBaseScan.evalVectorSearch(
+ actualTable,
+ vectorSearch.get,
+ pushedPartitionFilters,
+ pushedDataFilters)
+ return PaimonLocalScan(
+ VectorSearchResultUtils.toRows(result, requiredSchema),
+ requiredSchema,
+ actualTable,
+ pushedPartitionFilters)
+ }
PaimonScan(
actualTable,
requiredSchema,
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala
index c3dd135b9d..9249507467 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala
@@ -21,7 +21,7 @@ package org.apache.paimon.spark
import org.apache.paimon.CoreOptions
import org.apache.paimon.globalindex.GlobalIndexResult
import org.apache.paimon.partition.PartitionPredicate
-import org.apache.paimon.predicate.PredicateBuilder
+import org.apache.paimon.predicate.{Predicate, PredicateBuilder, VectorSearch}
import org.apache.paimon.spark.metric.SparkMetricRegistry
import org.apache.paimon.spark.read.{BaseScan, BatchReadTagCleanupListener,
PaimonSupportsRuntimeFiltering, SparkHybridSearchBuilderImpl,
SparkVectorSearchBuilderImpl}
import org.apache.paimon.spark.sources.PaimonMicroBatchStream
@@ -83,25 +83,11 @@ abstract class PaimonBaseScan(table: InnerTable)
}
private def evalVectorSearch(): GlobalIndexResult = {
- val vectorSearch = pushedVectorSearch.get
- val vectorSearchBuilder =
- if (CoreOptions.fromMap(table.options).vectorSearchDistributeEnabled()) {
- new SparkVectorSearchBuilderImpl(table)
- } else {
- table.newVectorSearchBuilder()
- }
- val vectorBuilder = vectorSearchBuilder
- .withVector(vectorSearch.vector())
- .withVectorColumn(vectorSearch.fieldName())
- .withLimit(vectorSearch.limit())
- .withOptions(vectorSearch.options())
- if (pushedPartitionFilters.nonEmpty) {
-
vectorBuilder.withPartitionFilter(PartitionPredicate.and(pushedPartitionFilters.asJava))
- }
- if (pushedDataFilters.nonEmpty) {
- vectorBuilder.withFilter(PredicateBuilder.and(pushedDataFilters.asJava))
- }
- vectorBuilder.newVectorRead().read(vectorBuilder.newVectorScan().scan())
+ PaimonBaseScan.evalVectorSearch(
+ table,
+ pushedVectorSearch.get,
+ pushedPartitionFilters,
+ pushedDataFilters)
}
private def evalHybridSearch(): GlobalIndexResult = {
@@ -182,3 +168,31 @@ abstract class PaimonBaseScan(table: InnerTable)
}
}
}
+
+object PaimonBaseScan {
+
+ private[spark] def evalVectorSearch(
+ table: InnerTable,
+ vectorSearch: VectorSearch,
+ pushedPartitionFilters: Seq[PartitionPredicate],
+ pushedDataFilters: Seq[Predicate]): GlobalIndexResult = {
+ val vectorSearchBuilder =
+ if (CoreOptions.fromMap(table.options).vectorSearchDistributeEnabled()) {
+ new SparkVectorSearchBuilderImpl(table)
+ } else {
+ table.newVectorSearchBuilder()
+ }
+ val vectorBuilder = vectorSearchBuilder
+ .withVector(vectorSearch.vector())
+ .withVectorColumn(vectorSearch.fieldName())
+ .withLimit(vectorSearch.limit())
+ .withOptions(vectorSearch.options())
+ if (pushedPartitionFilters.nonEmpty) {
+
vectorBuilder.withPartitionFilter(PartitionPredicate.and(pushedPartitionFilters.asJava))
+ }
+ if (pushedDataFilters.nonEmpty) {
+ vectorBuilder.withFilter(PredicateBuilder.and(pushedDataFilters.asJava))
+ }
+ vectorBuilder.newVectorRead().read(vectorBuilder.newVectorScan().scan())
+ }
+}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala
index 706715f7ed..64b0b5166c 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonRecordReaderIterator.scala
@@ -23,7 +23,7 @@ import org.apache.paimon.fs.Path
import org.apache.paimon.globalindex.IndexedSplitRecordReader
import org.apache.paimon.reader.{FileRecordIterator, RecordReader,
ScoreRecordIterator}
import org.apache.paimon.spark.schema.PaimonMetadataColumn
-import
org.apache.paimon.spark.schema.PaimonMetadataColumn.{PARTITION_AND_BUCKET_META_COLUMNS,
PATH_AND_INDEX_META_COLUMNS, ROW_ID_COLUMN, SEARCH_SCORE_COLUMN}
+import
org.apache.paimon.spark.schema.PaimonMetadataColumn.{PARTITION_AND_BUCKET_META_COLUMNS,
PATH_AND_INDEX_META_COLUMNS, VECTOR_SEARCH_META_COLUMN_NAMES}
import org.apache.paimon.table.source.{DataSplit, Split}
import org.apache.paimon.utils.{CloseableIterator, Preconditions}
@@ -50,11 +50,11 @@ case class PaimonRecordReaderIterator(
private val needPathAndIndexMetadata =
metadataColumns.exists(c => PATH_AND_INDEX_META_COLUMNS.contains(c.name))
private val needVectorSearchMetadata =
reader.isInstanceOf[IndexedSplitRecordReader] &&
- metadataColumns.exists(c => c.name == ROW_ID_COLUMN || c.name ==
SEARCH_SCORE_COLUMN)
+ metadataColumns.exists(c =>
VECTOR_SEARCH_META_COLUMN_NAMES.contains(c.name))
Preconditions.checkArgument(
!needVectorSearchMetadata ||
- metadataColumns.forall(c => c.name == ROW_ID_COLUMN || c.name ==
SEARCH_SCORE_COLUMN))
+ metadataColumns.forall(c =>
VECTOR_SEARCH_META_COLUMN_NAMES.contains(c.name)))
private val metadataRow: GenericRow =
GenericRow.of(Array.fill(metadataColumns.size)(null.asInstanceOf[AnyRef]):
_*)
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
index 23f0924f1f..6d7f894bea 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
@@ -22,7 +22,7 @@ import org.apache.paimon.partition.PartitionPredicate
import org.apache.paimon.predicate._
import org.apache.paimon.predicate.SortValue.{NullOrdering, SortDirection}
import
org.apache.paimon.spark.aggregate.AggregatePushDownUtils.tryPushdownAggregation
-import org.apache.paimon.spark.read.{PaimonLocalScan,
PaimonSupportsPushDownVariantExtractions}
+import org.apache.paimon.spark.read.{PaimonLocalScan,
PaimonSupportsPushDownVariantExtractions, VectorSearchResultUtils}
import org.apache.paimon.table.{FileStoreTable, InnerTable}
import org.apache.spark.sql.connector.expressions
@@ -139,6 +139,22 @@ class PaimonScanBuilder(val table: InnerTable)
case _ => (table, pushedVectorSearch, None, pushedFullTextSearch)
}
+ if (
+ vectorSearch.isDefined &&
+
VectorSearchResultUtils.isVectorSearchMetaOnly(requiredSchema.fieldNames.toSeq)
+ ) {
+ val result = PaimonBaseScan.evalVectorSearch(
+ actualTable,
+ vectorSearch.get,
+ pushedPartitionFilters,
+ pushedDataFilters)
+ return PaimonLocalScan(
+ VectorSearchResultUtils.toRows(result, requiredSchema),
+ requiredSchema,
+ actualTable,
+ pushedPartitionFilters)
+ }
+
PaimonScan(
actualTable,
requiredSchema,
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/plans/logical/PaimonTableValuedFunctions.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/plans/logical/PaimonTableValuedFunctions.scala
index d3d9194150..851acf933e 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/plans/logical/PaimonTableValuedFunctions.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/plans/logical/PaimonTableValuedFunctions.scala
@@ -846,7 +846,7 @@ case class DynamicVectorSearchRelation(
private lazy val outputWithMetaFields: Seq[Attribute] =
relationOutput ++
- Seq(PaimonMetadataColumn.ROW_ID.toAttribute,
PaimonMetadataColumn.SEARCH_SCORE.toAttribute)
+ PaimonMetadataColumn.VECTOR_SEARCH_META_COLUMNS.map(_.toAttribute)
override def output: Seq[Attribute] = outputWithMetaFields
}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala
index 43e2e11264..b3c23e03d6 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala
@@ -28,6 +28,7 @@ import org.apache.paimon.spark.catalog.{SparkBaseCatalog,
SupportView}
import org.apache.paimon.spark.catalyst.analysis.ResolvedPaimonView
import
org.apache.paimon.spark.catalyst.plans.logical.{CopyIntoLocationCommand,
CopyIntoLocationSource, CopyIntoTableCommand, CreateOrReplaceTagCommand,
CreatePaimonView, DeleteTagCommand, DropPaimonView, LateralVectorSearch,
PaimonCallCommand, PaimonDropPartitions, PaimonTableValuedFunctions,
RenameTagCommand, ResolvedIdentifier, ShowPaimonViews, ShowTagsCommand,
TruncatePaimonTableWithFilter}
import org.apache.paimon.spark.data.SparkInternalRow
+import org.apache.paimon.spark.read.VectorSearchResultUtils
import org.apache.paimon.spark.schema.PaimonMetadataColumn
import org.apache.paimon.table.{InnerTable, SpecialFields, Table}
import org.apache.paimon.table.source.{BatchVectorSearchBuilder,
InnerTableScan, ReadBuilder, VectorScan}
@@ -360,6 +361,8 @@ case class LateralVectorSearchExec(
scoreMetadataColumns,
sparkRow,
rowIdOrdinal = resultRowType.getFieldIndex(SpecialFields.ROW_ID.name()),
+ metaColumnsOnly =
+
VectorSearchResultUtils.isVectorSearchMetaOnly(vectorSearchOutput.map(_.name)),
projectionInputOrdinals = vectorSearchOutput.map {
attr =>
if (attr.name == PaimonMetadataColumn.SEARCH_SCORE_COLUMN) {
@@ -431,6 +434,9 @@ case class LateralVectorSearchExec(
s"Batch vector search returned ${globalIndexResults.size} results for
${queries.size} " +
"query vectors. The result count must match the query count."
)
+ if (context.metaColumnsOnly) {
+ return searchMetaColumns(queries, globalIndexResults, context)
+ }
val rowIdToMatches = createRowIdToMatches(queries, globalIndexResults)
val batchGlobalIndexResult =
createBatchGlobalIndexResult(globalIndexResults)
val scan = context.readBuilder
@@ -466,6 +472,30 @@ case class LateralVectorSearchExec(
}
}
+ private def searchMetaColumns(
+ queries: Seq[LateralVectorSearchQuery],
+ globalIndexResults: Seq[GlobalIndexResult],
+ context: LateralVectorSearchContext): Iterator[(InternalRow,
InternalRow)] = {
+ queries.zip(globalIndexResults).iterator.flatMap {
+ case (query, result) =>
+ val scoreGetter = result match {
+ case scored: ScoredGlobalIndexResult => Some(scored.scoreGetter())
+ case _ => None
+ }
+ result.results().iterator().asScala.map {
+ rowId =>
+ val values = vectorSearchOutput
+ .map(attr => VectorSearchResultUtils.valueOf(attr.name, rowId,
scoreGetter))
+ .toArray
+ val projectedRow = context.rightProjection(
+ new JoinedRow(
+ query.outerRow,
+ new GenericInternalRow(values.asInstanceOf[Array[Any]])))
+ (query.outerRow, projectedRow)
+ }
+ }
+ }
+
private def projectRightRow(
rightRow: InternalRow,
searchMatch: LateralVectorSearchMatch,
@@ -574,6 +604,7 @@ case class LateralVectorSearchExec(
scoreMetadataColumns: Seq[PaimonMetadataColumn],
sparkRow: SparkInternalRow,
rowIdOrdinal: Int,
+ metaColumnsOnly: Boolean,
projectionInputOrdinals: Seq[Int],
rightProjection: UnsafeProjection,
batchSize: Int,
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/VectorSearchResultUtils.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/VectorSearchResultUtils.scala
new file mode 100644
index 0000000000..e224d0b154
--- /dev/null
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/VectorSearchResultUtils.scala
@@ -0,0 +1,70 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.paimon.spark.read
+
+import org.apache.paimon.globalindex.{GlobalIndexResult,
ScoredGlobalIndexResult}
+import
org.apache.paimon.spark.catalyst.plans.logical.DynamicVectorSearchRelation
+import org.apache.paimon.spark.schema.PaimonMetadataColumn
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
+import org.apache.spark.sql.types.StructType
+
+import scala.collection.JavaConverters._
+
+/** Utilities for returning vector search results without reading table rows.
*/
+object VectorSearchResultUtils {
+
+ def isVectorSearchMetaOnly(fieldNames: Seq[String]): Boolean = {
+ fieldNames.nonEmpty && fieldNames.forall(
+ PaimonMetadataColumn.VECTOR_SEARCH_META_COLUMN_NAMES.contains)
+ }
+
+ def toRows(result: GlobalIndexResult, requiredSchema: StructType):
Array[InternalRow] = {
+ val fieldNames = requiredSchema.fieldNames
+ val scoreGetter = result match {
+ case scored: ScoredGlobalIndexResult => Some(scored.scoreGetter())
+ case _ => None
+ }
+ result
+ .results()
+ .iterator()
+ .asScala
+ .map {
+ rowId =>
+ new GenericInternalRow(
+ fieldNames.map(valueOf(_, rowId,
scoreGetter)).asInstanceOf[Array[Any]])
+ }
+ .toArray
+ }
+
+ def valueOf(
+ fieldName: String,
+ rowId: Long,
+ scoreGetter: Option[org.apache.paimon.globalindex.ScoreGetter]): Any = {
+ fieldName match {
+ case PaimonMetadataColumn.ROW_ID_COLUMN => rowId
+ case PaimonMetadataColumn.SEARCH_SCORE_COLUMN =>
+ scoreGetter.map(_.score(rowId)).getOrElse(Float.NaN)
+ case _ =>
+ throw new IllegalArgumentException(
+ s"Field $fieldName cannot be returned directly from vector search
result.")
+ }
+ }
+}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/schema/PaimonMetadataColumn.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/schema/PaimonMetadataColumn.scala
index d61dfb8c59..41c6cdaaaf 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/schema/PaimonMetadataColumn.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/schema/PaimonMetadataColumn.scala
@@ -61,6 +61,7 @@ object PaimonMetadataColumn {
val PATH_AND_INDEX_META_COLUMNS: Seq[String] = Seq(FILE_PATH_COLUMN,
ROW_INDEX_COLUMN)
val PARTITION_AND_BUCKET_META_COLUMNS: Seq[String] = Seq(PARTITION_COLUMN,
BUCKET_COLUMN)
val ROW_TRACKING_META_COLUMNS: Seq[String] = Seq(ROW_ID_COLUMN,
SEQUENCE_NUMBER_COLUMN)
+ val VECTOR_SEARCH_META_COLUMN_NAMES: Seq[String] = Seq(ROW_ID_COLUMN,
SEARCH_SCORE_COLUMN)
val SUPPORTED_METADATA_COLUMNS: Seq[String] = Seq(
ROW_INDEX_COLUMN,
@@ -92,6 +93,8 @@ object PaimonMetadataColumn {
val SEARCH_SCORE: PaimonMetadataColumn =
PaimonMetadataColumn(Integer.MAX_VALUE - 106, SEARCH_SCORE_COLUMN,
FloatType)
+ val VECTOR_SEARCH_META_COLUMNS: Seq[PaimonMetadataColumn] = Seq(ROW_ID,
SEARCH_SCORE)
+
def dvMetaCols: Seq[PaimonMetadataColumn] = Seq(FILE_PATH, ROW_INDEX)
def get(metadataColumn: String, partitionType: StructType):
PaimonMetadataColumn = {
diff --git
a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkMultimodalITCase.java
b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkMultimodalITCase.java
index 340122a12a..8fc7d603b1 100644
---
a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkMultimodalITCase.java
+++
b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkMultimodalITCase.java
@@ -195,6 +195,21 @@ public class SparkMultimodalITCase {
.equals(row.getLong(3))))
.isTrue();
+ // **vector search with metadata columns only */
+ String vectorSearchWithMetadataColumnsOnlySql =
+ "select _row_id AS _row_id, __paimon_search_score "
+ + "from vector_search('my_db1.vector_test', 'embs',
array(1.0f, 2.0f, 3.0f, 4.0f), 5) "
+ + "where date = '20260420'";
+ df = spark.sql(vectorSearchWithMetadataColumnsOnlySql);
+ assertThat(df.columns()).hasSize(2);
+ assertThat(df.columns()).contains("_row_id", "__paimon_search_score");
+ rows = df.collectAsList();
+ assertThat(rows).hasSize(5);
+ assertThat(rows.stream().allMatch(row -> !row.isNullAt(0) &&
!row.isNullAt(1))).isTrue();
+ assertThat(rows.stream().allMatch(row ->
baseRowIds.containsValue(row.getLong(0))))
+ .isTrue();
+ assertThat(rows.stream().map(row ->
row.getLong(0)).collect(Collectors.toSet())).hasSize(5);
+
// **vector search with score */
String vectorSearchSql =
"select gid, sid, embs, __paimon_search_score "
@@ -256,6 +271,32 @@ public class SparkMultimodalITCase {
.containsEntry(6L, 5L)
.containsEntry(7L, 5L)
.containsEntry(8L, 5L);
+
+ // ** lateral vector search with metadata columns only in subquery */
+ rows =
+ spark.sql(
+ "SELECT q.gid AS query_gid, "
+ + "r._row_id AS result_row_id, "
+ + "r.__paimon_search_score AS
result_score "
+ + "FROM my_db1.vector_test AS q, "
+ + "LATERAL (SELECT _row_id,
__paimon_search_score "
+ + "FROM
vector_search('my_db1.vector_test', 'embs', q.embs, 5)) AS r "
+ + "WHERE q.`date` = '20260420';")
+ .collectAsList();
+ assertThat(rows).hasSize(40);
+ assertThat(rows.stream().allMatch(row -> !row.isNullAt(0) &&
!row.isNullAt(1))).isTrue();
+ assertThat(rows.stream().allMatch(row ->
baseRowIds.containsValue(row.getLong(1))))
+ .isTrue();
+ Map<Long, java.util.Set<Long>> rowIdsPerQueryGid =
+ rows.stream()
+ .collect(
+ Collectors.groupingBy(
+ row -> row.getLong(0),
+ Collectors.mapping(
+ row -> row.getLong(1),
Collectors.toSet())));
+ assertThat(rowIdsPerQueryGid).hasSize(8);
+ assertThat(rowIdsPerQueryGid.values().stream().allMatch(rowIds ->
rowIds.size() == 5))
+ .isTrue();
spark.close();
spark = builder.getOrCreate();
