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 ffaebaec51 [spark] support lateral inner join for vector search (#8252)
ffaebaec51 is described below
commit ffaebaec51f12b4f5e2b2064512d834c2742cc43
Author: Stefanietry <[email protected]>
AuthorDate: Wed Jun 24 08:50:06 2026 +0800
[spark] support lateral inner join for vector search (#8252)
---
docs/generated/core_configuration.html | 18 +-
.../main/java/org/apache/paimon/CoreOptions.java | 12 +
.../PaimonIncompatibleResolutionRules.scala | 9 +
.../PushDownLateralVectorSearchFilter.scala | 72 ++++
.../plans/logical/PaimonTableValuedFunctions.scala | 262 +++++++++++++-
.../paimon/spark/execution/PaimonStrategy.scala | 387 ++++++++++++++++++++-
.../extensions/PaimonSparkSessionExtensions.scala | 3 +-
.../apache/paimon/spark/SparkMultimodalITCase.java | 23 ++
.../spark/sql/TableValuedFunctionsTest.scala | 142 +++++++-
9 files changed, 911 insertions(+), 17 deletions(-)
diff --git a/docs/generated/core_configuration.html
b/docs/generated/core_configuration.html
index f525a00b64..1e82e01041 100644
--- a/docs/generated/core_configuration.html
+++ b/docs/generated/core_configuration.html
@@ -1668,6 +1668,12 @@ If the data size allocated for the sorting task is
uneven,which may lead to perf
<td>Boolean</td>
<td>Whether to process distributed vector search.</td>
</tr>
+ <tr>
+ <td><h5>vector-search.lateral-join.batch-size</h5></td>
+ <td style="word-wrap: break-word;">256</td>
+ <td>Integer</td>
+ <td>The batch size for lateral vector search. Each batch executes
vector topK search and table lookup for multiple query vectors.</td>
+ </tr>
<tr>
<td><h5>vector.file.format</h5></td>
<td style="word-wrap: break-word;">(none)</td>
@@ -1734,12 +1740,6 @@ If the data size allocated for the sorting task is
uneven,which may lead to perf
<td>Boolean</td>
<td>If set to true, compactions and snapshot expiration will be
skipped. This option is used along with dedicated compact jobs.</td>
</tr>
- <tr>
- <td><h5>write.sequence-number-init-mode</h5></td>
- <td style="word-wrap: break-word;">scan</td>
- <td><p>Enum</p></td>
- <td>Specify how to initialize the next sequence number for primary
key table writers.<br /><br />Possible values:<ul><li>"scan": initialize by
scanning existing file metadata.</li><li>"snapshot": initialize from the
maximum sequence number recorded in snapshot properties, which can avoid
scanning existing file metadata in write-only mode.</li></ul></td>
- </tr>
<tr>
<td><h5>write.batch-memory</h5></td>
<td style="word-wrap: break-word;">128 mb</td>
@@ -1752,6 +1752,12 @@ If the data size allocated for the sorting task is
uneven,which may lead to perf
<td>Integer</td>
<td>Write batch size for any file format if it supports.</td>
</tr>
+ <tr>
+ <td><h5>write.sequence-number-init-mode</h5></td>
+ <td style="word-wrap: break-word;">scan</td>
+ <td><p>Enum</p></td>
+ <td>Specify how to initialize the next sequence number for primary
key table writers.<br /><br />Possible values:<ul><li>"scan": initialize by
scanning existing file metadata.</li><li>"snapshot": initialize from the
maximum sequence number recorded in snapshot properties, which can avoid
scanning existing file metadata in write-only mode.</li></ul></td>
+ </tr>
<tr>
<td><h5>zorder.var-length-contribution</h5></td>
<td style="word-wrap: break-word;">8</td>
diff --git a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
index 5eb86a56bc..8ca39caf66 100644
--- a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
+++ b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
@@ -2658,6 +2658,14 @@ public class CoreOptions implements Serializable {
.defaultValue(false)
.withDescription("Whether to process distributed vector
search.");
+ public static final ConfigOption<Integer>
VECTOR_SEARCH_LATERAL_JOIN_BATCH_SIZE =
+ key("vector-search.lateral-join.batch-size")
+ .intType()
+ .defaultValue(256)
+ .withDescription(
+ "The batch size for lateral vector search. Each
batch executes vector "
+ + "topK search and table lookup for
multiple query vectors.");
+
@Immutable
public static final ConfigOption<Boolean> PK_CLUSTERING_OVERRIDE =
key("pk-clustering-override")
@@ -4155,6 +4163,10 @@ public class CoreOptions implements Serializable {
return options.get(VECTOR_SEARCH_DISTRIBUTE_ENABLED);
}
+ public int vectorSearchLateralJoinBatchSize() {
+ return options.get(VECTOR_SEARCH_LATERAL_JOIN_BATCH_SIZE);
+ }
+
/** Specifies the merge engine for table with primary key. */
public enum MergeEngine implements DescribedEnum {
DEDUPLICATE("deduplicate", "De-duplicate and keep the last row."),
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/analysis/PaimonIncompatibleResolutionRules.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/analysis/PaimonIncompatibleResolutionRules.scala
index 9824597c1e..a56116545a 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/analysis/PaimonIncompatibleResolutionRules.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/analysis/PaimonIncompatibleResolutionRules.scala
@@ -21,6 +21,8 @@ package org.apache.paimon.spark.catalyst.analysis
import
org.apache.paimon.spark.catalyst.plans.logical.{PaimonTableValuedFunctions,
PaimonTableValueFunction}
import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.expressions.LateralSubquery
+import org.apache.spark.sql.catalyst.plans.logical.LateralJoin
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.rules.Rule
@@ -32,6 +34,13 @@ case class PaimonIncompatibleResolutionRules(session:
SparkSession) extends Rule
case func: PaimonTableValueFunction if func.args.forall(_.resolved) =>
PaimonTableValuedFunctions.resolvePaimonTableValuedFunction(session,
func)
+ case lateralJoin @ LateralJoin(left, lateralSubquery: LateralSubquery,
joinType, condition)
+ if left.resolved && lateralSubquery.plan.resolved &&
+
PaimonTableValuedFunctions.containsDynamicVectorSearch(lateralSubquery.plan) =>
+ PaimonTableValuedFunctions
+ .resolveLateralVectorSearch(left, lateralSubquery.plan, joinType,
condition)
+ .getOrElse(lateralJoin)
+
}
}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownLateralVectorSearchFilter.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownLateralVectorSearchFilter.scala
new file mode 100644
index 0000000000..de5baf0b35
--- /dev/null
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownLateralVectorSearchFilter.scala
@@ -0,0 +1,72 @@
+/*
+ * 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.catalyst.optimizer
+
+import org.apache.paimon.spark.catalyst.plans.logical.{LateralVectorSearch,
PaimonTableValuedFunctions}
+
+import org.apache.spark.sql.catalyst.expressions.{And, PredicateHelper}
+import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan}
+import org.apache.spark.sql.catalyst.rules.Rule
+
+/** Pushes filters on the query side below lateral vector search. */
+object PushDownLateralVectorSearchFilter extends Rule[LogicalPlan] with
PredicateHelper {
+
+ override def apply(plan: LogicalPlan): LogicalPlan = plan.transform {
+ case filter @ Filter(condition, lvs: LateralVectorSearch) =>
+ val predicates = splitConjunctivePredicates(condition)
+ val (pushDownToLeft, otherPredicates) = predicates.partition {
+ predicate => predicate.deterministic &&
predicate.references.subsetOf(lvs.child.outputSet)
+ }
+ val (pushDownToSearch, stayUp) = otherPredicates.partition {
+ predicate =>
+ predicate.deterministic &&
+ predicate.references.nonEmpty &&
+ predicate.references.subsetOf(lvs.searchFilterOutputSet) &&
+ PaimonTableValuedFunctions
+ .convertLateralVectorSearchFilters(
+ lvs.innerTable,
+ lvs.vectorSearchOutput,
+ lvs.projectList,
+ lvs.projectOutput,
+ Seq(predicate))
+ .isDefined
+ }
+
+ if (pushDownToLeft.isEmpty && pushDownToSearch.isEmpty) {
+ filter
+ } else {
+ val lvsWithPushedLeft = if (pushDownToLeft.isEmpty) {
+ lvs
+ } else {
+ lvs.copy(left = Filter(buildBalancedPredicate(pushDownToLeft, And),
lvs.child))
+ }
+ val newLateralVectorSearch = if (pushDownToSearch.isEmpty) {
+ lvsWithPushedLeft
+ } else {
+ lvsWithPushedLeft.copy(
+ searchFilters = lvsWithPushedLeft.searchFilters ++
pushDownToSearch)
+ }
+ if (stayUp.isEmpty) {
+ newLateralVectorSearch
+ } else {
+ Filter(buildBalancedPredicate(stayUp, And), newLateralVectorSearch)
+ }
+ }
+ }
+}
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 7d41d7f1b5..ad3675f369 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
@@ -20,19 +20,21 @@ package org.apache.paimon.spark.catalyst.plans.logical
import org.apache.paimon.CoreOptions
import org.apache.paimon.globalindex.HybridSearchRanker
-import org.apache.paimon.predicate.{FullTextQuery, FullTextSearch,
HybridSearch, HybridSearchRoute, VectorSearch}
-import org.apache.paimon.spark.SparkTable
+import org.apache.paimon.predicate.{FullTextQuery, FullTextSearch,
HybridSearch, HybridSearchRoute, Predicate, VectorSearch}
+import org.apache.paimon.spark.{SparkTable, SparkTypeUtils,
SparkV2FilterConverter}
import
org.apache.paimon.spark.catalyst.plans.logical.PaimonTableValuedFunctions._
+import org.apache.paimon.spark.schema.PaimonMetadataColumn
import org.apache.paimon.table.{DataTable, FullTextSearchTable,
HybridSearchTable, InnerTable, VectorSearchTable}
import
org.apache.paimon.table.source.snapshot.TimeTravelUtil.InconsistentTagBucketException
-import org.apache.spark.sql.PaimonUtils.createDataset
+import org.apache.spark.sql.PaimonUtils.{createDataset, normalizeExprs,
toAttributes, translateFilterV2}
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.FunctionIdentifier
import org.apache.spark.sql.catalyst.analysis.FunctionRegistryBase
import
org.apache.spark.sql.catalyst.analysis.TableFunctionRegistry.TableFunctionBuilder
-import org.apache.spark.sql.catalyst.expressions.{Attribute, CreateArray,
CreateMap, CreateNamedStruct, Expression, ExpressionInfo, Literal}
-import org.apache.spark.sql.catalyst.plans.logical.{LeafNode, LogicalPlan}
+import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute,
AttributeSet, CreateArray, CreateMap, CreateNamedStruct, Expression,
ExpressionInfo, ExprId, Literal, OuterReference}
+import org.apache.spark.sql.catalyst.plans.{Inner, JoinType}
+import org.apache.spark.sql.catalyst.plans.logical.{Filter, LeafNode,
LogicalPlan, Project, SubqueryAlias, UnaryNode}
import org.apache.spark.sql.catalyst.util.MapData
import org.apache.spark.sql.connector.catalog.{Identifier, Table, TableCatalog}
import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
@@ -40,9 +42,17 @@ import org.apache.spark.sql.util.CaseInsensitiveStringMap
import org.apache.spark.unsafe.types.UTF8String
import scala.collection.JavaConverters._
+import scala.collection.mutable.ArrayBuffer
+import scala.util.control.NonFatal
object PaimonTableValuedFunctions {
+ private case class DynamicVectorSearchExtraction(
+ relation: DynamicVectorSearchRelation,
+ projectList: Seq[Expression],
+ projectOutput: Seq[Attribute],
+ searchFilters: Seq[Expression])
+
val INCREMENTAL_QUERY = "paimon_incremental_query"
val INCREMENTAL_BETWEEN_TIMESTAMP = "paimon_incremental_between_timestamp"
val INCREMENTAL_TO_AUTO_TAG = "paimon_incremental_to_auto_tag"
@@ -153,6 +163,9 @@ object PaimonTableValuedFunctions {
argsWithoutTable: Seq[Expression]): LogicalPlan = {
sparkTable match {
case st @ SparkTable(innerTable: InnerTable) =>
+ if (vsq.hasOuterReference(argsWithoutTable)) {
+ return vsq.createDynamicVectorSearch(innerTable, argsWithoutTable)
+ }
val vectorSearch = vsq.createVectorSearch(innerTable, argsWithoutTable)
val vectorSearchTable = VectorSearchTable.create(innerTable,
vectorSearch)
DataSourceV2Relation.create(
@@ -189,6 +202,138 @@ object PaimonTableValuedFunctions {
}
}
+ def resolveLateralVectorSearch(
+ left: LogicalPlan,
+ right: LogicalPlan,
+ joinType: JoinType,
+ condition: Option[Expression]): Option[LogicalPlan] = {
+ extractDynamicVectorSearch(right) match {
+ case None if containsDynamicVectorSearch(right) =>
+ throw new UnsupportedOperationException(
+ "LATERAL vector_search only supports SELECT <columns> FROM
vector_search(...) " +
+ "or SELECT <columns> FROM vector_search(...) WHERE <searched-table
predicate>.")
+ case None =>
+ None
+ case Some(_) if joinType != Inner =>
+ throw new RuntimeException(
+ s"LATERAL vector_search only supports INNER join, but got:
$joinType.")
+ case Some(
+ DynamicVectorSearchExtraction(relation, projectList,
projectOutput, searchFilters)) =>
+ val vectorSearchOutput = vectorSearchOutputForProject(relation,
projectList, searchFilters)
+ if (
+ searchFilters.nonEmpty && convertLateralVectorSearchFilters(
+ relation.innerTable,
+ vectorSearchOutput,
+ projectList,
+ projectOutput,
+ searchFilters).isEmpty
+ ) {
+ throw new UnsupportedOperationException(
+ "LATERAL vector_search only supports deterministic subquery
predicates " +
+ "convertible to Paimon predicates on searched-table columns.")
+ }
+ val lateralVectorSearch =
+ LateralVectorSearch(
+ left,
+ relation.innerTable,
+ relation.columnName,
+ relation.queryVectorExpr,
+ relation.limit,
+ relation.options,
+ vectorSearchOutput,
+ projectList,
+ projectOutput,
+ searchFilters
+ )
+ Some(condition.map(Filter(_,
lateralVectorSearch)).getOrElse(lateralVectorSearch))
+ }
+ }
+
+ def convertLateralVectorSearchFilters(
+ innerTable: InnerTable,
+ vectorSearchOutput: Seq[Attribute],
+ projectList: Seq[Expression],
+ projectOutput: Seq[Attribute],
+ filters: Seq[Expression]): Option[Seq[Predicate]] = {
+ val converter = SparkV2FilterConverter(innerTable.rowType())
+ val projectionByExprId = projectList
+ .zip(projectOutput)
+ .map { case (project, outputAttr) => outputAttr.exprId ->
stripAlias(project) }
+ .toMap
+ try {
+ val predicates = ArrayBuffer[Predicate]()
+ normalizeExprs(filters.map(rewriteSearchFilter(_, projectionByExprId)),
vectorSearchOutput)
+ .flatMap(splitConjunctivePredicatesForFilter)
+ .foreach {
+ filter =>
+ translateFilterV2(filter).flatMap(converter.convert(_)) match {
+ case Some(predicate) => predicates += predicate
+ case None => return None
+ }
+ }
+ Some(predicates.toSeq)
+ } catch {
+ case NonFatal(_) => None
+ }
+ }
+
+ private def rewriteSearchFilter(
+ filter: Expression,
+ projectionByExprId: Map[ExprId, Expression]): Expression = {
+ filter.transform { case attr: Attribute =>
projectionByExprId.getOrElse(attr.exprId, attr) }
+ }
+
+ private def stripAlias(expression: Expression): Expression = {
+ expression match {
+ case Alias(child, _) => child
+ case other => other
+ }
+ }
+
+ private def splitConjunctivePredicatesForFilter(condition: Expression):
Seq[Expression] = {
+ condition match {
+ case And(left, right) =>
+ splitConjunctivePredicatesForFilter(left) ++
splitConjunctivePredicatesForFilter(right)
+ case other => other :: Nil
+ }
+ }
+
+ private def vectorSearchOutputForProject(
+ relation: DynamicVectorSearchRelation,
+ projectList: Seq[Expression],
+ searchFilters: Seq[Expression]): Seq[Attribute] = {
+ val projectReferences =
+ AttributeSet.fromAttributeSets((projectList ++
searchFilters).map(_.references))
+ relation.output.filter(projectReferences.contains)
+ }
+
+ private def extractDynamicVectorSearch(
+ plan: LogicalPlan): Option[DynamicVectorSearchExtraction] = {
+ plan match {
+ case SubqueryAlias(_, child) =>
+ extractDynamicVectorSearch(child).map {
+ extraction => extraction.copy(projectOutput = plan.output)
+ }
+ case Project(projectList, Filter(condition, relation:
DynamicVectorSearchRelation))
+ if projectList.forall(_.resolved) && condition.resolved =>
+ Some(DynamicVectorSearchExtraction(relation, projectList, plan.output,
Seq(condition)))
+ case Project(projectList, relation: DynamicVectorSearchRelation)
+ if projectList.forall(_.resolved) =>
+ Some(DynamicVectorSearchExtraction(relation, projectList, plan.output,
Nil))
+ case Filter(condition, relation: DynamicVectorSearchRelation) if
condition.resolved =>
+ Some(
+ DynamicVectorSearchExtraction(relation, relation.output,
relation.output, Seq(condition)))
+ case relation: DynamicVectorSearchRelation =>
+ Some(DynamicVectorSearchExtraction(relation, relation.output,
relation.output, Nil))
+ case _ => None
+ }
+ }
+
+ def containsDynamicVectorSearch(plan: LogicalPlan): Boolean = {
+ plan.isInstanceOf[DynamicVectorSearchRelation] || plan.children.exists(
+ containsDynamicVectorSearch)
+ }
+
private def resolveFullTextSearchQuery(
sparkTable: Table,
sparkCatalog: TableCatalog,
@@ -447,6 +592,48 @@ case class VectorSearchQuery(override val args:
Seq[Expression])
}
value.toString
}
+
+ def hasOuterReference(argsWithoutTable: Seq[Expression]): Boolean = {
+ val queryVector = argsWithoutTable(1)
+ (argsWithoutTable.size == 3 || argsWithoutTable.size == 4) &&
+ (queryVector.references.nonEmpty || containsOuterReference(queryVector))
+ }
+
+ private def containsOuterReference(expr: Expression): Boolean = {
+ expr.isInstanceOf[OuterReference] ||
expr.children.exists(containsOuterReference)
+ }
+
+ def createDynamicVectorSearch(
+ innerTable: InnerTable,
+ argsWithoutTable: Seq[Expression]): DynamicVectorSearchRelation = {
+ if (argsWithoutTable.size != 3 && argsWithoutTable.size != 4) {
+ throw new RuntimeException(
+ s"$VECTOR_SEARCH needs three or four parameters after table_name: " +
+ s"column_name, query_vector, limit[, options]. " +
+ 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 limit = parsePositiveLimit(argsWithoutTable(2).eval())
+ val options: Map[String, String] =
+ if (argsWithoutTable.size == 4) {
+ extractOptions(argsWithoutTable(3))
+ } else {
+ Map.empty[String, String]
+ }
+ DynamicVectorSearchRelation(
+ innerTable,
+ columnName,
+ argsWithoutTable(1),
+ limit,
+ options,
+ toAttributes(SparkTypeUtils.fromPaimonRowType(innerTable.rowType())))
+ }
}
/**
@@ -637,6 +824,71 @@ case class HybridSearchQuery(override val args:
Seq[Expression])
}
+case class DynamicVectorSearchRelation(
+ innerTable: InnerTable,
+ columnName: String,
+ queryVectorExpr: Expression,
+ limit: Int,
+ options: Map[String, String],
+ relationOutput: Seq[Attribute])
+ extends LeafNode {
+
+ private lazy val outputWithScore: Seq[Attribute] =
+ relationOutput ++
+ Seq(PaimonMetadataColumn.SEARCH_SCORE.toAttribute)
+
+ override def output: Seq[Attribute] = outputWithScore
+}
+
+case class LateralVectorSearch(
+ left: LogicalPlan,
+ innerTable: InnerTable,
+ columnName: String,
+ queryVectorExpr: Expression,
+ limit: Int,
+ options: Map[String, String],
+ vectorSearchOutput: Seq[Attribute],
+ projectList: Seq[Expression],
+ projectOutput: Seq[Attribute],
+ searchFilters: Seq[Expression] = Nil)
+ extends UnaryNode {
+
+ override def child: LogicalPlan = left
+
+ override def output: Seq[Attribute] = left.output ++ projectOutput
+
+ lazy val searchFilterOutputSet: AttributeSet = {
+ val tableOutputSet = AttributeSet(
+ vectorSearchOutput.filterNot(_.name ==
PaimonMetadataColumn.SEARCH_SCORE_COLUMN))
+ AttributeSet(projectList.zip(projectOutput).collect {
+ case (expr, attr) if isSearchFilterAttribute(expr, tableOutputSet) =>
+ attr
+ })
+ }
+
+ private def isSearchFilterAttribute(
+ expression: Expression,
+ tableOutputSet: AttributeSet): Boolean = {
+ expression match {
+ case Alias(attr: Attribute, _) => tableOutputSet.contains(attr)
+ case attr: Attribute => tableOutputSet.contains(attr)
+ case _ => false
+ }
+ }
+
+ override lazy val producedAttributes: AttributeSet = {
+ AttributeSet(vectorSearchOutput ++ output.filterNot(attr =>
inputSet.contains(attr)))
+ }
+
+ override lazy val references: AttributeSet = {
+ AttributeSet.fromAttributeSets(expressions.map(_.references)) --
producedAttributes
+ }
+
+ override protected def withNewChildInternal(newChild: LogicalPlan):
LogicalPlan = {
+ copy(left = newChild)
+ }
+}
+
/**
* Plan for the [[FULL_TEXT_SEARCH]] table-valued function.
*
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 321e61f2cb..d6eeee11db 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
@@ -18,18 +18,30 @@
package org.apache.paimon.spark.execution
+import org.apache.paimon.CoreOptions
+import org.apache.paimon.globalindex.{GlobalIndexResult,
ScoredGlobalIndexResult}
import org.apache.paimon.partition.PartitionPredicate
-import org.apache.paimon.spark.{SparkCatalog, SparkGenericCatalog, SparkTable,
SparkUtils}
+import
org.apache.paimon.partition.PartitionPredicate.splitPartitionPredicatesAndDataPredicates
+import org.apache.paimon.predicate.{Predicate, PredicateBuilder}
+import org.apache.paimon.spark.{PaimonRecordReaderIterator, SparkCatalog,
SparkGenericCatalog, SparkTable, SparkUtils}
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, PaimonCallCommand,
PaimonDropPartitions, RenameTagCommand, ResolvedIdentifier, ShowPaimonViews,
ShowTagsCommand, TruncatePaimonTableWithFilter}
-import org.apache.paimon.table.Table
+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.schema.PaimonMetadataColumn
+import org.apache.paimon.table.{InnerTable, SpecialFields, Table}
+import org.apache.paimon.table.source.{BatchVectorSearchBuilder,
InnerTableScan, ReadBuilder, VectorScan}
+import org.apache.paimon.types.RowType
+import org.apache.paimon.utils.RoaringNavigableMap64
+import org.apache.spark.TaskContext
+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.{Expression,
GenericInternalRow, PredicateHelper}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet,
Expression, GenericInternalRow, JoinedRow, OuterReference, 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}
import org.apache.spark.sql.execution.{PaimonDescribeTableExec, SparkPlan,
SparkStrategy}
import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Implicits,
DataSourceV2Relation}
@@ -37,6 +49,7 @@ import
org.apache.spark.sql.execution.shim.{PaimonCreateTableAsSelectStrategy, P
import org.apache.spark.sql.paimon.shims.SparkShimLoader
import scala.collection.JavaConverters._
+import scala.collection.mutable.ArrayBuffer
case class PaimonStrategy(spark: SparkSession)
extends SparkStrategy
@@ -61,6 +74,20 @@ case class PaimonStrategy(spark: SparkSession)
val input = buildInternalRow(args)
PaimonCallExec(c.output, procedure, input) :: Nil
+ case lvs: LateralVectorSearch =>
+ LateralVectorSearchExec(
+ lvs.innerTable,
+ lvs.columnName,
+ lvs.queryVectorExpr,
+ lvs.limit,
+ lvs.options,
+ lvs.vectorSearchOutput,
+ lvs.projectList,
+ lvs.projectOutput,
+ lvs.searchFilters,
+ planLater(lvs.left)
+ ) :: Nil
+
case t @ ShowTagsCommand(PaimonCatalogAndIdentifier(catalog, ident)) =>
ShowTagsExec(catalog, ident, t.output) :: Nil
@@ -215,3 +242,355 @@ case class PaimonStrategy(spark: SparkSession)
SparkShimLoader.shim.classicApi.recacheByPlan(spark, v2Relation)
}
}
+
+case class LateralVectorSearchExec(
+ innerTable: InnerTable,
+ columnName: String,
+ queryVectorExpr: Expression,
+ limit: Int,
+ options: Map[String, String],
+ vectorSearchOutput: Seq[Attribute],
+ projectList: Seq[Expression],
+ projectOutput: Seq[Attribute],
+ searchFilters: Seq[Expression],
+ child: SparkPlan)
+ extends SparkPlan
+ with PredicateHelper {
+
+ override def children: Seq[SparkPlan] = Seq(child)
+
+ override def output: Seq[Attribute] = child.output ++ projectOutput
+
+ @transient override lazy val producedAttributes: AttributeSet = {
+ AttributeSet(vectorSearchOutput ++ output.filterNot(attr =>
inputSet.contains(attr)))
+ }
+
+ @transient
+ override lazy val references: AttributeSet = {
+ AttributeSet.fromAttributeSets(expressions.map(_.references)) --
producedAttributes
+ }
+
+ override protected def withNewChildrenInternal(newChildren:
IndexedSeq[SparkPlan]): SparkPlan = {
+ copy(child = newChildren.head)
+ }
+
+ 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 rightProjection =
+ UnsafeProjection.create(strippedProjectList, child.output ++
vectorSearchOutput)
+ val joinedRow = new JoinedRow
+ val readerTracker = new LateralVectorSearchReaderTracker
+ Option(TaskContext.get())
+ .foreach(_.addTaskCompletionListener[Unit](_ =>
readerTracker.closeCurrent()))
+ val searchContext = createSearchContext(rightProjection, readerTracker)
+ val batchSize = searchContext.batchSize
+
+ outerRows.map(_.copy()).grouped(batchSize).flatMap {
+ outerRowBatch =>
+ val searchBatch = ArrayBuffer[LateralVectorSearchQuery]()
+ outerRowBatch.foreach {
+ outerRow =>
+ toFloatArray(queryVectorProjection(outerRow).get(0,
strippedQueryExpr.dataType))
+ .foreach(
+ queryVector => searchBatch +=
LateralVectorSearchQuery(outerRow, queryVector))
+ }
+
+ if (searchBatch.isEmpty) {
+ Iterator.empty
+ } else {
+ search(searchBatch.toVector, searchContext).map {
+ case (outerRow, rightRow) =>
+ joinedRow(outerRow, rightRow)
+ joinedRow.copy()
+ }
+ }
+ }
+ }
+ }
+
+ private def createSearchContext(
+ rightProjection: UnsafeProjection,
+ readerTracker: LateralVectorSearchReaderTracker):
LateralVectorSearchContext = {
+ val rowType = innerTable.rowType()
+ val readFieldNames = vectorSearchOutput
+ .filterNot(_.name == PaimonMetadataColumn.SEARCH_SCORE_COLUMN)
+ .map(_.name)
+ val readFieldNamesWithRowId =
+ if (readFieldNames.contains(SpecialFields.ROW_ID.name())) {
+ readFieldNames
+ } else {
+ readFieldNames :+ SpecialFields.ROW_ID.name()
+ }
+ val rowTypeWithRowId = SpecialFields.rowTypeWithRowId(rowType)
+ val readRowType = rowType.project(readFieldNames.asJava)
+ val readRowTypeWithRowId = SpecialFields.rowTypeWithRowId(readRowType)
+ val readBuilder = innerTable
+ .newReadBuilder()
+ .withReadType(rowTypeWithRowId.project(readFieldNamesWithRowId.asJava))
+ val scoreMetadataColumns =
+ if (vectorSearchOutput.exists(_.name ==
PaimonMetadataColumn.SEARCH_SCORE_COLUMN)) {
+ Seq(PaimonMetadataColumn.SEARCH_SCORE)
+ } else {
+ Seq.empty
+ }
+ val resultRowType =
+ if (scoreMetadataColumns.isEmpty) {
+ readRowTypeWithRowId
+ } else {
+ new RowType(
+ (readRowTypeWithRowId.getFields.asScala ++ scoreMetadataColumns.map(
+ _.toPaimonDataField)).asJava)
+ }
+ val sparkRow = SparkInternalRow.create(resultRowType)
+ val vectorSearchBuilder = innerTable
+ .newBatchVectorSearchBuilder()
+ .withVectorColumn(columnName)
+ .withLimit(limit)
+ .withOptions(options.asJava)
+ pushSearchFilters(readBuilder, vectorSearchBuilder)
+
+ val vectorPlan = vectorSearchBuilder.newVectorScan().scan()
+ val batchSize =
+ Math.max(1, new
CoreOptions(innerTable.options()).vectorSearchLateralJoinBatchSize())
+
+ LateralVectorSearchContext(
+ readBuilder,
+ vectorSearchBuilder,
+ vectorPlan,
+ scoreMetadataColumns,
+ sparkRow,
+ rowIdOrdinal = resultRowType.getFieldIndex(SpecialFields.ROW_ID.name()),
+ projectionInputOrdinals = vectorSearchOutput.map {
+ attr =>
+ if (attr.name == PaimonMetadataColumn.SEARCH_SCORE_COLUMN) {
+ -1
+ } else {
+ resultRowType.getFieldIndex(attr.name)
+ }
+ },
+ rightProjection,
+ batchSize,
+ readerTracker
+ )
+ }
+
+ private def pushSearchFilters(
+ readBuilder: ReadBuilder,
+ vectorSearchBuilder: BatchVectorSearchBuilder): Unit = {
+ val predicates = convertSearchFilters()
+ if (predicates.nonEmpty) {
+ val split = splitPartitionPredicatesAndDataPredicates(
+ predicates.asJava,
+ innerTable.rowType(),
+ innerTable.partitionKeys())
+ if (split.getLeft.isPresent) {
+ val partitionFilter = split.getLeft.get()
+ readBuilder.withPartitionFilter(partitionFilter)
+ vectorSearchBuilder.withPartitionFilter(partitionFilter)
+ }
+ if (!split.getRight.isEmpty) {
+ val dataFilter = PredicateBuilder.and(split.getRight)
+ readBuilder.withFilter(dataFilter)
+ vectorSearchBuilder.withFilter(dataFilter)
+ }
+ }
+ }
+
+ private def convertSearchFilters(): Seq[Predicate] = {
+ if (searchFilters.isEmpty) {
+ Seq.empty
+ } else {
+ PaimonTableValuedFunctions
+ .convertLateralVectorSearchFilters(
+ innerTable,
+ vectorSearchOutput,
+ projectList,
+ projectOutput,
+ searchFilters)
+ .getOrElse {
+ throw new UnsupportedOperationException(
+ s"Cannot convert searched-table predicates for LATERAL
vector_search: $searchFilters")
+ }
+ }
+ }
+
+ private def search(
+ queries: Seq[LateralVectorSearchQuery],
+ context: LateralVectorSearchContext): Iterator[(InternalRow,
InternalRow)] = {
+ val vectors = queries.map(_.queryVector).toArray
+ val globalIndexResults = context.vectorSearchBuilder
+ .withVectors(vectors)
+ .newBatchVectorRead()
+ .readBatch(context.vectorPlan)
+ .asScala
+ .toVector
+ // Batch vector search must return one result per input query vector and
preserve the input
+ // order, because createRowIdToMatches pairs each result with its original
outer row by index.
+ require(
+ globalIndexResults.size == queries.size,
+ s"Batch vector search returned ${globalIndexResults.size} results for
${queries.size} " +
+ "query vectors. The result count must match the query count."
+ )
+ val rowIdToMatches = createRowIdToMatches(queries, globalIndexResults)
+ val batchGlobalIndexResult =
createBatchGlobalIndexResult(globalIndexResults)
+ val scan = context.readBuilder
+ .newScan()
+ .withGlobalIndexResult(batchGlobalIndexResult)
+ .asInstanceOf[InnerTableScan]
+ val read = context.readBuilder.newRead()
+
+ scan.plan().splits().asScala.iterator.flatMap {
+ split =>
+ val reader =
+ PaimonRecordReaderIterator(read.createReader(split),
context.scoreMetadataColumns, split)
+ val readerState = context.readerTracker.track(reader)
+ new Iterator[Iterator[(InternalRow, InternalRow)]] {
+ override def hasNext: Boolean = {
+ val hasNext = reader.hasNext
+ if (!hasNext) {
+ readerState.closeOnce()
+ }
+ hasNext
+ }
+
+ override def next(): Iterator[(InternalRow, InternalRow)] = {
+ val rightRow = context.sparkRow.replace(reader.next())
+ val rowId = rightRow.getLong(context.rowIdOrdinal)
+ rowIdToMatches.getOrElse(rowId, Seq.empty).iterator.map {
+ searchMatch =>
+ val projectedRow = projectRightRow(rightRow, searchMatch,
context)
+ (searchMatch.outerRow, projectedRow)
+ }
+ }
+ }.flatMap(identity)
+ }
+ }
+
+ private def projectRightRow(
+ rightRow: InternalRow,
+ searchMatch: LateralVectorSearchMatch,
+ context: LateralVectorSearchContext): InternalRow = {
+ val values = new Array[Any](vectorSearchOutput.size)
+ vectorSearchOutput.zipWithIndex.foreach {
+ case (attr, index) =>
+ val ordinal = context.projectionInputOrdinals(index)
+ values(index) = if (ordinal < 0) {
+ searchMatch.score
+ } else {
+ rightRow.get(ordinal, attr.dataType)
+ }
+ }
+ context.rightProjection(new JoinedRow(searchMatch.outerRow, new
GenericInternalRow(values)))
+ }
+
+ private def createRowIdToMatches(
+ queries: Seq[LateralVectorSearchQuery],
+ globalIndexResults: Seq[GlobalIndexResult]): Map[Long,
Seq[LateralVectorSearchMatch]] = {
+ val rowIdToMatches =
+ scala.collection.mutable.LinkedHashMap[Long,
ArrayBuffer[LateralVectorSearchMatch]]()
+ queries.zip(globalIndexResults).foreach {
+ case (query, result) =>
+ val scoreGetter = result match {
+ case scored: ScoredGlobalIndexResult => Some(scored.scoreGetter())
+ case _ => None
+ }
+ result.results().iterator().asScala.foreach {
+ rowId =>
+ rowIdToMatches.getOrElseUpdate(rowId, ArrayBuffer()) +=
+ LateralVectorSearchMatch(
+ query.outerRow,
+ scoreGetter.map(_.score(rowId)).getOrElse(Float.NaN))
+ }
+ }
+ rowIdToMatches.iterator.map { case (rowId, matches) => rowId ->
matches.toSeq }.toMap
+ }
+
+ private def createBatchGlobalIndexResult(
+ globalIndexResults: Seq[GlobalIndexResult]): GlobalIndexResult = {
+ val rowIds = new RoaringNavigableMap64()
+ globalIndexResults.foreach(result => rowIds.or(result.results()))
+ GlobalIndexResult.create(rowIds)
+ }
+
+ private def toFloatArray(value: Any): Option[Array[Float]] = {
+ value match {
+ case null => None
+ case arrayData: ArrayData => Some(arrayData.toFloatArray())
+ case _ =>
+ throw new RuntimeException(s"Cannot extract query vector from
expression value: $value")
+ }
+ }
+
+ private class LateralVectorSearchReaderTracker {
+ @volatile private var currentReader: LateralVectorSearchReaderState = _
+
+ def track(reader: PaimonRecordReaderIterator):
LateralVectorSearchReaderState = {
+ val state = new LateralVectorSearchReaderState(reader, this)
+ this.synchronized {
+ currentReader = state
+ }
+ state
+ }
+
+ def clear(state: LateralVectorSearchReaderState): Unit = {
+ this.synchronized {
+ if (currentReader eq state) {
+ currentReader = null
+ }
+ }
+ }
+
+ def closeCurrent(): Unit = {
+ val reader = currentReader
+ if (reader != null) {
+ reader.closeOnce()
+ }
+ }
+ }
+
+ private class LateralVectorSearchReaderState(
+ reader: PaimonRecordReaderIterator,
+ tracker: LateralVectorSearchReaderTracker) {
+ private var closed = false
+
+ def closeOnce(): Unit = {
+ this.synchronized {
+ if (!closed) {
+ closed = true
+ try {
+ reader.close()
+ } finally {
+ tracker.clear(this)
+ }
+ }
+ }
+ }
+ }
+
+ private case class LateralVectorSearchContext(
+ readBuilder: ReadBuilder,
+ vectorSearchBuilder: BatchVectorSearchBuilder,
+ vectorPlan: VectorScan.Plan,
+ scoreMetadataColumns: Seq[PaimonMetadataColumn],
+ sparkRow: SparkInternalRow,
+ rowIdOrdinal: Int,
+ projectionInputOrdinals: Seq[Int],
+ rightProjection: UnsafeProjection,
+ batchSize: Int,
+ readerTracker: LateralVectorSearchReaderTracker)
+
+ private case class LateralVectorSearchQuery(outerRow: InternalRow,
queryVector: Array[Float])
+
+ private case class LateralVectorSearchMatch(outerRow: InternalRow, score:
Float)
+}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala
index e433a5f7d4..61481e201c 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala
@@ -19,7 +19,7 @@
package org.apache.paimon.spark.extensions
import org.apache.paimon.spark.catalyst.analysis.{PaimonAnalysis,
PaimonDeleteTable, PaimonFunctionResolver, PaimonIncompatibleResolutionRules,
PaimonMergeInto, PaimonPostHocResolutionRules, PaimonProcedureResolver,
PaimonUpdateTable, PaimonViewResolver, ReplacePaimonFunctions,
RewriteUpsertTable}
-import
org.apache.paimon.spark.catalyst.optimizer.{MergePaimonScalarSubqueries,
OptimizeMetadataOnlyDeleteFromPaimonTable}
+import
org.apache.paimon.spark.catalyst.optimizer.{MergePaimonScalarSubqueries,
OptimizeMetadataOnlyDeleteFromPaimonTable, PushDownLateralVectorSearchFilter}
import
org.apache.paimon.spark.catalyst.plans.logical.PaimonTableValuedFunctions
import org.apache.paimon.spark.commands.BucketExpression
import org.apache.paimon.spark.execution.{OldCompatibleStrategy,
PaimonStrategy}
@@ -102,6 +102,7 @@ class PaimonSparkSessionExtensions extends
(SparkSessionExtensions => Unit) {
extensions.injectOptimizerRule(spark => ReplacePaimonFunctions(spark))
extensions.injectOptimizerRule(_ =>
OptimizeMetadataOnlyDeleteFromPaimonTable)
extensions.injectOptimizerRule(_ => MergePaimonScalarSubqueries)
+ extensions.injectOptimizerRule(_ => PushDownLateralVectorSearchFilter)
// planner extensions
extensions.injectPlannerStrategy(spark => PaimonStrategy(spark))
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 af3401ea2a..3a4dd41b16 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
@@ -170,6 +170,29 @@ public class SparkMultimodalITCase {
.collect(Collectors.toList()));
spark.close();
+ 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';")
+ .collectAsList();
+ assertThat(rows).hasSize(40);
+ assertThat(
+ rows.stream()
+ .collect(
+ Collectors.groupingBy(
+ row -> row.getLong(0),
Collectors.counting())))
+ .hasSize(8)
+ .containsEntry(1L, 5L)
+ .containsEntry(2L, 5L)
+ .containsEntry(3L, 5L)
+ .containsEntry(4L, 5L)
+ .containsEntry(5L, 5L)
+ .containsEntry(6L, 5L)
+ .containsEntry(7L, 5L)
+ .containsEntry(8L, 5L);
+ spark.close();
+
spark = builder.getOrCreate();
spark.sql("DROP TABLE IF EXISTS `my_db1`.`vector_test`;");
spark.close();
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 8114766a99..94565a9ad2 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
@@ -21,10 +21,11 @@ package org.apache.paimon.spark.sql
import org.apache.paimon.data.{BinaryString, GenericRow, Timestamp}
import org.apache.paimon.manifest.ManifestCommittable
import org.apache.paimon.spark.PaimonHiveTestBase
-import
org.apache.paimon.spark.catalyst.plans.logical.PaimonTableValuedFunctions
+import org.apache.paimon.spark.catalyst.plans.logical.{LateralVectorSearch,
PaimonTableValuedFunctions}
import org.apache.paimon.utils.DateTimeUtils
import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.plans.logical.Filter
import java.time.LocalDateTime
import java.util.Collections
@@ -41,6 +42,145 @@ class TableValuedFunctionsTest extends PaimonHiveTestBase {
assert(error.getMessage.contains("Limit must be no greater than"))
}
+ test("lateral vector search preserves subquery alias qualifiers") {
+ withTable("vector_search_source", "vector_search_result") {
+ spark.sql("""
+ |CREATE TABLE vector_search_source (gid BIGINT, embs
ARRAY<FLOAT>, dt STRING)
+ |USING paimon
+ |TBLPROPERTIES (
+ | 'vector.file.format' = 'lance',
+ | 'vector-field' = 'embs',
+ | 'field.embs.vector-dim' = '3',
+ | 'row-tracking.enabled' = 'true',
+ | 'data-evolution.enabled' = 'true')
+ |PARTITIONED BY (dt)
+ |""".stripMargin)
+ spark.sql("""
+ |CREATE TABLE vector_search_result (
+ | query_gid BIGINT,
+ | query_embs ARRAY<FLOAT>,
+ | result_gid BIGINT,
+ | result_embs ARRAY<FLOAT>,
+ | score FLOAT,
+ | dt STRING)
+ |USING paimon
+ |PARTITIONED BY (dt)
+ |""".stripMargin)
+
+ val insertOptimizedPlan = spark
+ .sql("""
+ |SELECT q.gid AS query_gid, q.embs AS query_embs,
+ | r.gid AS result_gid, r.embs AS result_embs,
+ | r.__paimon_search_score AS score
+ |FROM vector_search_source AS q,
+ |LATERAL (
+ | SELECT gid, embs, __paimon_search_score
+ | FROM vector_search('vector_search_source', 'embs', q.embs, 5)
+ |) AS r
+ |WHERE q.dt = '20260608'
+ |""".stripMargin)
+ .queryExecution
+ .optimizedPlan
+ val lateralVectorSearches = insertOptimizedPlan.collect {
+ case lvs: LateralVectorSearch => lvs
+ }
+ assert(lateralVectorSearches.size == 1, insertOptimizedPlan.toString)
+
+ val optimizedPlanWithoutScore = spark
+ .sql("""
+ |SELECT q.gid AS query_gid, r.embs AS result_embs
+ |FROM vector_search_source AS q,
+ |LATERAL (
+ | SELECT embs
+ | FROM vector_search('vector_search_source', 'embs', q.embs, 5)
+ |) AS r
+ |""".stripMargin)
+ .queryExecution
+ .optimizedPlan
+ assert(
+ optimizedPlanWithoutScore.exists(_.isInstanceOf[LateralVectorSearch]),
+ optimizedPlanWithoutScore.toString)
+
+ val analyzedPlanWithJoinCondition = spark
+ .sql("""
+ |SELECT q.gid AS query_gid, r.result_gid, r.score
+ |FROM vector_search_source AS q,
+ |LATERAL (
+ | SELECT gid AS result_gid, __paimon_search_score AS score
+ | FROM vector_search('vector_search_source', 'embs', q.embs, 5)
+ |) AS r
+ |WHERE q.gid = r.result_gid AND r.score >= 0.0
+ |""".stripMargin)
+ .queryExecution
+ .analyzed
+ val lateralVectorSearchFilters = analyzedPlanWithJoinCondition.collect {
+ case filter @ Filter(_, _: LateralVectorSearch) => filter
+ }
+ assert(lateralVectorSearchFilters.size == 1,
analyzedPlanWithJoinCondition.toString)
+ assert(
+ lateralVectorSearchFilters.head.condition.references
+ .subsetOf(lateralVectorSearchFilters.head.child.outputSet),
+ analyzedPlanWithJoinCondition.toString
+ )
+
+ val optimizedPlanWithSearchFilter = spark
+ .sql("""
+ |SELECT q.gid AS query_gid, r.result_gid, r.dt
+ |FROM vector_search_source AS q,
+ |LATERAL (
+ | SELECT gid AS result_gid, dt
+ | FROM vector_search('vector_search_source', 'embs', q.embs, 5)
+ |) AS r
+ |WHERE r.dt = '20260608'
+ |""".stripMargin)
+ .queryExecution
+ .optimizedPlan
+ val lateralVectorSearchesWithSearchFilter =
optimizedPlanWithSearchFilter.collect {
+ case lvs: LateralVectorSearch => lvs
+ }
+ assert(
+ lateralVectorSearchesWithSearchFilter.size == 1,
+ optimizedPlanWithSearchFilter.toString)
+ assert(
+ lateralVectorSearchesWithSearchFilter.head.searchFilters.nonEmpty,
+ optimizedPlanWithSearchFilter.toString)
+
+ val optimizedPlanWithSubqueryFilter = spark
+ .sql("""
+ |SELECT q.gid AS query_gid, r.result_gid
+ |FROM vector_search_source AS q,
+ |LATERAL (
+ | SELECT gid AS result_gid
+ | FROM vector_search('vector_search_source', 'embs', q.embs, 5)
+ | WHERE dt = '20260608'
+ |) AS r
+ |""".stripMargin)
+ .queryExecution
+ .optimizedPlan
+ val lateralVectorSearchesWithSubqueryFilter =
optimizedPlanWithSubqueryFilter.collect {
+ case lvs: LateralVectorSearch => lvs
+ }
+ assert(
+ lateralVectorSearchesWithSubqueryFilter.size == 1,
+ optimizedPlanWithSubqueryFilter.toString)
+ assert(
+ lateralVectorSearchesWithSubqueryFilter.head.searchFilters.nonEmpty,
+ optimizedPlanWithSubqueryFilter.toString)
+
+ val constantVectorPlan = spark
+ .sql("""
+ |SELECT gid
+ |FROM vector_search(
+ | 'vector_search_source', 'embs', array(1.0f, 2.0f, 3.0f), 5)
+ |""".stripMargin)
+ .queryExecution
+ .optimizedPlan
+ assert(
+ !constantVectorPlan.exists(_.isInstanceOf[LateralVectorSearch]),
+ constantVectorPlan.toString)
+ }
+ }
+
withPk.foreach {
hasPk =>
bucketModes.foreach {
