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 e5335fe343 [spark] Fix column pruning for lateral vector search query
vector (#8425)
e5335fe343 is described below
commit e5335fe34395ccc34672d315f9e67f26d28e385e
Author: Zouxxyy <[email protected]>
AuthorDate: Thu Jul 2 13:34:37 2026 +0800
[spark] Fix column pruning for lateral vector search query vector (#8425)
- Strip `OuterReference` from `queryVectorExpr` and `projectList` when
creating `LateralVectorSearch`, so that the default `references`
computation sees the query vector column and
column pruning does not remove it from the left plan
- This also removes redundant `OuterReference` stripping in the physical
execution
---
.../plans/logical/PaimonTableValuedFunctions.scala | 6 ++++--
.../paimon/spark/execution/PaimonStrategy.scala | 17 ++++-------------
.../paimon/spark/sql/TableValuedFunctionsTest.scala | 19 +++++++++++++++++++
3 files changed, 27 insertions(+), 15 deletions(-)
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 ee3a930386..c3f3116b5a 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
@@ -232,16 +232,18 @@ object PaimonTableValuedFunctions {
"LATERAL vector_search only supports deterministic subquery
predicates " +
"convertible to Paimon predicates on searched-table columns.")
}
+ val stripOuterRef: Expression => Expression =
+ _.transform { case OuterReference(a) => a.toAttribute }
val lateralVectorSearch =
LateralVectorSearch(
left,
relation.innerTable,
relation.columnName,
- relation.queryVectorExpr,
+ stripOuterRef(relation.queryVectorExpr),
relation.limit,
relation.options,
vectorSearchOutput,
- projectList,
+ projectList.map(stripOuterRef),
projectOutput,
searchFilters
)
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 d6eeee11db..a559fe4ada 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
@@ -39,7 +39,7 @@ import org.apache.spark.rdd.RDD
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.analysis.{ResolvedNamespace,
ResolvedTable}
-import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet,
Expression, GenericInternalRow, JoinedRow, OuterReference, PredicateHelper,
UnsafeProjection}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet,
Expression, GenericInternalRow, JoinedRow, PredicateHelper, UnsafeProjection}
import org.apache.spark.sql.catalyst.plans.logical.{CreateTableAsSelect,
DescribeRelation, LogicalPlan, ReplaceTable, ReplaceTableAsSelect,
ShowCreateTable}
import org.apache.spark.sql.catalyst.util.ArrayData
import org.apache.spark.sql.connector.catalog.{Identifier,
PaimonLookupCatalog, TableCatalog}
@@ -277,18 +277,9 @@ case class LateralVectorSearchExec(
override protected def doExecute(): RDD[InternalRow] = {
child.execute().mapPartitions {
outerRows =>
- val strippedQueryExpr = queryVectorExpr.transform {
- case OuterReference(namedExpression) => namedExpression.toAttribute
- }
- val queryVectorProjection =
UnsafeProjection.create(Seq(strippedQueryExpr), child.output)
- val strippedProjectList = projectList.map {
- project =>
- project.transform {
- case OuterReference(namedExpression) =>
namedExpression.toAttribute
- }
- }
+ val queryVectorProjection =
UnsafeProjection.create(Seq(queryVectorExpr), child.output)
val rightProjection =
- UnsafeProjection.create(strippedProjectList, child.output ++
vectorSearchOutput)
+ UnsafeProjection.create(projectList, child.output ++
vectorSearchOutput)
val joinedRow = new JoinedRow
val readerTracker = new LateralVectorSearchReaderTracker
Option(TaskContext.get())
@@ -301,7 +292,7 @@ case class LateralVectorSearchExec(
val searchBatch = ArrayBuffer[LateralVectorSearchQuery]()
outerRowBatch.foreach {
outerRow =>
- toFloatArray(queryVectorProjection(outerRow).get(0,
strippedQueryExpr.dataType))
+ toFloatArray(queryVectorProjection(outerRow).get(0,
queryVectorExpr.dataType))
.foreach(
queryVector => searchBatch +=
LateralVectorSearchQuery(outerRow, queryVector))
}
diff --git
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/TableValuedFunctionsTest.scala
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/TableValuedFunctionsTest.scala
index 94565a9ad2..a44711d5b9 100644
---
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/TableValuedFunctionsTest.scala
+++
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/TableValuedFunctionsTest.scala
@@ -167,6 +167,25 @@ class TableValuedFunctionsTest extends PaimonHiveTestBase {
lateralVectorSearchesWithSubqueryFilter.head.searchFilters.nonEmpty,
optimizedPlanWithSubqueryFilter.toString)
+ val optimizedPlanWithoutQueryVector = spark
+ .sql("""
+ |SELECT q.gid AS query_gid, r.gid AS result_gid
+ |FROM vector_search_source AS q,
+ |LATERAL (
+ | SELECT gid
+ | FROM vector_search('vector_search_source', 'embs', q.embs, 5)
+ |) AS r
+ |""".stripMargin)
+ .queryExecution
+ .optimizedPlan
+ val lateralVectorSearchesWithoutQueryVector =
+ optimizedPlanWithoutQueryVector.collect { case lvs:
LateralVectorSearch => lvs }
+ assert(
+ lateralVectorSearchesWithoutQueryVector.size == 1,
+ "Query vector column not in outer SELECT should not break lateral
vector search: " +
+ optimizedPlanWithoutQueryVector.toString
+ )
+
val constantVectorPlan = spark
.sql("""
|SELECT gid
