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 b09282e7e1 [spark] support get row id in vector search (#8475)
b09282e7e1 is described below
commit b09282e7e11379f8199630073577bbabb921c96c
Author: Stefanietry <[email protected]>
AuthorDate: Mon Jul 6 18:14:34 2026 +0800
[spark] support get row id in vector search (#8475)
Purpose: Support get _row_id when process vector search
Linked issue: https://github.com/apache/paimon/issues/8474
---
.../apache/paimon/reader/ScoreRecordIterator.java | 12 ++++
.../globalindex/IndexedSplitRecordReader.java | 14 ++++-
.../paimon/spark/PaimonRecordReaderIterator.scala | 17 ++++--
.../plans/logical/PaimonTableValuedFunctions.scala | 6 +-
.../paimon/spark/execution/PaimonStrategy.scala | 12 ++--
.../apache/paimon/spark/SparkMultimodalITCase.java | 69 +++++++++++++++++++++-
6 files changed, 109 insertions(+), 21 deletions(-)
diff --git
a/paimon-common/src/main/java/org/apache/paimon/reader/ScoreRecordIterator.java
b/paimon-common/src/main/java/org/apache/paimon/reader/ScoreRecordIterator.java
index 908f35404e..10e4548141 100644
---
a/paimon-common/src/main/java/org/apache/paimon/reader/ScoreRecordIterator.java
+++
b/paimon-common/src/main/java/org/apache/paimon/reader/ScoreRecordIterator.java
@@ -30,6 +30,8 @@ public interface ScoreRecordIterator<T> extends
RecordReader.RecordIterator<T> {
float returnedScore();
+ long returnedRowId();
+
@Override
default <R> ScoreRecordIterator<R> transform(Function<T, R> function) {
ScoreRecordIterator<T> thisIterator = this;
@@ -39,6 +41,11 @@ public interface ScoreRecordIterator<T> extends
RecordReader.RecordIterator<T> {
return thisIterator.returnedScore();
}
+ @Override
+ public long returnedRowId() {
+ return thisIterator.returnedRowId();
+ }
+
@Nullable
@Override
public R next() throws IOException {
@@ -65,6 +72,11 @@ public interface ScoreRecordIterator<T> extends
RecordReader.RecordIterator<T> {
return thisIterator.returnedScore();
}
+ @Override
+ public long returnedRowId() {
+ return thisIterator.returnedRowId();
+ }
+
@Nullable
@Override
public T next() throws IOException {
diff --git
a/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplitRecordReader.java
b/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplitRecordReader.java
index b57e2c5403..3778f37cb9 100644
---
a/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplitRecordReader.java
+++
b/paimon-core/src/main/java/org/apache/paimon/globalindex/IndexedSplitRecordReader.java
@@ -59,18 +59,26 @@ public class IndexedSplitRecordReader implements
RecordReader<InternalRow> {
return new ScoreRecordIterator<InternalRow>() {
private float score = Float.NaN;
+ private long rowId;
@Override
public float returnedScore() {
return score;
}
+ @Override
+ public long returnedRowId() {
+ return rowId;
+ }
+
@Override
public InternalRow next() throws IOException {
InternalRow row = iterator.next();
- if (row != null && rowIdToScore != null) {
- Long rowId = row.getLong(rowIdIndex);
- this.score = rowIdToScore.get(rowId);
+ if (row != null && rowIdIndex >= 0) {
+ this.rowId = row.getLong(rowIdIndex);
+ if (rowIdToScore != null) {
+ this.score = rowIdToScore.get(rowId);
+ }
if (projectedRow != null) {
projectedRow.replaceRow(row);
return projectedRow;
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 ac06698fe9..706715f7ed 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
@@ -20,9 +20,10 @@ package org.apache.paimon.spark
import org.apache.paimon.data.{BinaryString, GenericRow, InternalRow =>
PaimonInternalRow, JoinedRow}
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, SEARCH_SCORE_COLUMN}
+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.table.source.{DataSplit, Split}
import org.apache.paimon.utils.{CloseableIterator, Preconditions}
@@ -48,10 +49,12 @@ case class PaimonRecordReaderIterator(
private val needMetadata = metadataColumns.nonEmpty
private val needPathAndIndexMetadata =
metadataColumns.exists(c => PATH_AND_INDEX_META_COLUMNS.contains(c.name))
- private val needScoreMetadata = {
- metadataColumns.exists(_.name == SEARCH_SCORE_COLUMN)
- }
- Preconditions.checkArgument(!needScoreMetadata || metadataColumns.size == 1)
+ private val needVectorSearchMetadata =
reader.isInstanceOf[IndexedSplitRecordReader] &&
+ metadataColumns.exists(c => c.name == ROW_ID_COLUMN || c.name ==
SEARCH_SCORE_COLUMN)
+
+ Preconditions.checkArgument(
+ !needVectorSearchMetadata ||
+ metadataColumns.forall(c => c.name == ROW_ID_COLUMN || c.name ==
SEARCH_SCORE_COLUMN))
private val metadataRow: GenericRow =
GenericRow.of(Array.fill(metadataColumns.size)(null.asInstanceOf[AnyRef]):
_*)
@@ -126,7 +129,7 @@ case class PaimonRecordReaderIterator(
while (!stop) {
val dataRow = currentIterator.next()
if (dataRow != null) {
- if (needScoreMetadata) {
+ if (needVectorSearchMetadata) {
updateScoreMetadata(
currentIterator.asInstanceOf[ScoreRecordIterator[PaimonInternalRow]])
currentResult = joinedRow.replace(dataRow, metadataRow)
@@ -179,6 +182,8 @@ case class PaimonRecordReaderIterator(
metadataColumns.zipWithIndex.foreach {
case (metadataColumn, index) =>
metadataColumn.name match {
+ case PaimonMetadataColumn.ROW_ID_COLUMN =>
+ metadataRow.setField(index, fileRecordIterator.returnedRowId())
case PaimonMetadataColumn.SEARCH_SCORE_COLUMN =>
metadataRow.setField(index, fileRecordIterator.returnedScore())
}
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 7c24bdd5f0..d3d9194150 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
@@ -844,11 +844,11 @@ case class DynamicVectorSearchRelation(
relationOutput: Seq[Attribute])
extends LeafNode {
- private lazy val outputWithScore: Seq[Attribute] =
+ private lazy val outputWithMetaFields: Seq[Attribute] =
relationOutput ++
- Seq(PaimonMetadataColumn.SEARCH_SCORE.toAttribute)
+ Seq(PaimonMetadataColumn.ROW_ID.toAttribute,
PaimonMetadataColumn.SEARCH_SCORE.toAttribute)
- override def output: Seq[Attribute] = outputWithScore
+ override def output: Seq[Attribute] = outputWithMetaFields
}
case class LateralVectorSearch(
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 a559fe4ada..43e2e11264 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
@@ -315,14 +315,12 @@ case class LateralVectorSearchExec(
readerTracker: LateralVectorSearchReaderTracker):
LateralVectorSearchContext = {
val rowType = innerTable.rowType()
val readFieldNames = vectorSearchOutput
- .filterNot(_.name == PaimonMetadataColumn.SEARCH_SCORE_COLUMN)
+ .filterNot(
+ attr =>
+ attr.name == PaimonMetadataColumn.SEARCH_SCORE_COLUMN ||
+ attr.name == SpecialFields.ROW_ID.name())
.map(_.name)
- val readFieldNamesWithRowId =
- if (readFieldNames.contains(SpecialFields.ROW_ID.name())) {
- readFieldNames
- } else {
- readFieldNames :+ SpecialFields.ROW_ID.name()
- }
+ val readFieldNamesWithRowId = readFieldNames :+ SpecialFields.ROW_ID.name()
val rowTypeWithRowId = SpecialFields.rowTypeWithRowId(rowType)
val readRowType = rowType.project(readFieldNames.asJava)
val readRowTypeWithRowId = SpecialFields.rowTypeWithRowId(readRowType)
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 3a4dd41b16..340122a12a 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
@@ -32,6 +32,7 @@ import org.junit.jupiter.api.io.TempDir;
import java.io.IOException;
import java.util.List;
+import java.util.Map;
import java.util.stream.Collectors;
import static org.assertj.core.api.Assertions.assertThat;
@@ -143,19 +144,68 @@ public class SparkMultimodalITCase {
"select gid, sid, embs from my_db1.vector_test
where date = '20260420' and embs is not null;")
.collectAsList();
assertThat(rows).hasSize(8);
+ Map<Long, Long> baseRowIds =
+ spark.sql("select gid, _row_id from my_db1.vector_test where
date = '20260420'")
+ .collectAsList().stream()
+ .collect(Collectors.toMap(row -> row.getLong(0), row
-> row.getLong(1)));
+ assertThat(baseRowIds).hasSize(8);
rows =
spark.sql(
"select gid, sid, embs from
vector_search('my_db1.vector_test', 'embs', array(1.0f, 2.0f, 3.0f, 4.0f), 5)
where date = '20260420'")
.collectAsList();
assertThat(rows).hasSize(5);
+
+ // **vector search with row id */
+ String vectorSearchWithRowIdSql =
+ "select gid, sid, embs, _row_id AS _row_id "
+ + "from vector_search('my_db1.vector_test', 'embs',
array(1.0f, 2.0f, 3.0f, 4.0f), 5) "
+ + "where date = '20260420'";
+ Dataset<Row> df = spark.sql(vectorSearchWithRowIdSql);
+ assertThat(df.columns()).hasSize(4);
+ assertThat(df.columns()).contains("_row_id");
+ rows = df.collectAsList();
+ assertThat(rows).hasSize(5);
+ assertThat(rows.stream().noneMatch(row -> row.isNullAt(3))).isTrue();
+ assertThat(
+ rows.stream()
+ .allMatch(
+ row ->
+ baseRowIds
+ .get(row.getLong(0))
+
.equals(row.getLong(3))))
+ .isTrue();
+
+ // **vector search with row id and score */
+ String vectorSearchWithRowIdAndScoreSql =
+ "select gid, sid, embs, _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(vectorSearchWithRowIdAndScoreSql);
+ assertThat(df.columns()).hasSize(5);
+ assertThat(df.columns()).contains("_row_id", "__paimon_search_score");
+ rows = df.collectAsList();
+ assertThat(rows).hasSize(5);
+ assertThat(rows.stream().allMatch(row -> !row.isNullAt(3) &&
!row.isNullAt(4))).isTrue();
+ assertThat(
+ rows.stream()
+ .allMatch(
+ row ->
+ baseRowIds
+ .get(row.getLong(0))
+
.equals(row.getLong(3))))
+ .isTrue();
+
+ // **vector search with score */
String vectorSearchSql =
"select gid, sid, embs, __paimon_search_score "
+ "from vector_search('my_db1.vector_test', 'embs',
array(1.0f, 2.0f, 3.0f, 4.0f), 5) "
+ "where date = '20260420'";
- Dataset<Row> df = spark.sql(vectorSearchSql);
+ df = spark.sql(vectorSearchSql);
assertThat(df.columns()).hasSize(4);
rows = df.collectAsList();
assertThat(rows).hasSize(5);
+
+ // ** distribute vector search */
spark.sql("SET
`spark.paimon.vector-search.distribute.enabled`=`true`");
spark.sql("SET `spark.paimon.global-index.thread-num`=`1`");
List<Row> compareRows = spark.sql(vectorSearchSql).collectAsList();
@@ -170,13 +220,28 @@ public class SparkMultimodalITCase {
.collect(Collectors.toList()));
spark.close();
+ // ** lateral vector search */
spark = builder.getOrCreate();
spark.sql("SET
`spark.paimon.vector-search.distribute.enabled`=`false`");
rows =
spark.sql(
- "SELECT q.gid AS query_gid, q.embs AS
query_embs, r.gid AS result_gid FROM my_db1.vector_test AS q, LATERAL (SELECT
gid FROM vector_search('my_db1.vector_test', 'embs', q.embs, 5)) AS r WHERE
q.`date` = '20260420';")
+ "SELECT q.gid AS query_gid, q.embs AS
query_embs, "
+ + "r.gid AS result_gid, r._row_id AS
result_row_id "
+ + "FROM my_db1.vector_test AS q, "
+ + "LATERAL (SELECT gid, _row_id "
+ + "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().noneMatch(row -> row.isNullAt(3))).isTrue();
+ assertThat(
+ rows.stream()
+ .allMatch(
+ row ->
+ baseRowIds
+ .get(row.getLong(2))
+
.equals(row.getLong(3))))
+ .isTrue();
assertThat(
rows.stream()
.collect(
