[spark] Repartition limited lateral vector search input

docs/generated/spark_connector_configuration.html

@@ -86,6 +86,12 @@
             <td>Boolean</td>
             <td>Whether to adjust the target split size based on pruned (projected) columns. If enabled, split size estimation uses only the columns actually being read.</td>
         </tr>
+        <tr>
+            <td><h5>vector-search.lateral-join.parallelism</h5></td>
+            <td style="word-wrap: break-word;">16</td>
+            <td>Integer</td>
+            <td>Parallelism used to repartition a single-partition LIMIT input before executing a lateral vector search.</td>
+        </tr>
         <tr>
             <td><h5>write.data-evolution.update-conflict-retry.max-attempts</h5></td>
             <td style="word-wrap: break-word;">20</td>

paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/spark/sql/paimon/shims/MinorVersionShim.scala

@@ -20,6 +20,7 @@ package org.apache.spark.sql.paimon.shims
 
 import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression}
 import org.apache.spark.sql.catalyst.plans.logical.{CTERelationRef, LogicalPlan, MergeAction, MergeIntoTable}
+import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, Distribution}
 
 object MinorVersionShim {
 
@@ -29,6 +30,9 @@ object MinorVersionShim {
       output: Seq[Attribute],
       isStreaming: Boolean): CTERelationRef = CTERelationRef(cteId, resolved, output)
 
+  def createClusteredDistribution(expressions: Seq[Expression], numPartitions: Int): Distribution =
+    ClusteredDistribution(expressions, Some(numPartitions))
+
   def createMergeIntoTable(
       targetTable: LogicalPlan,
       sourceTable: LogicalPlan,

paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/spark/sql/paimon/shims/MinorVersionShim.scala

@@ -20,6 +20,7 @@ package org.apache.spark.sql.paimon.shims
 
 import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression}
 import org.apache.spark.sql.catalyst.plans.logical.{CTERelationRef, LogicalPlan, MergeAction, MergeIntoTable}
+import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, Distribution}
 
 object MinorVersionShim {
 
@@ -29,6 +30,12 @@ object MinorVersionShim {
       output: Seq[Attribute],
       isStreaming: Boolean): CTERelationRef = CTERelationRef(cteId, resolved, output)
 
+  def createClusteredDistribution(expressions: Seq[Expression], numPartitions: Int): Distribution =
+    ClusteredDistribution(
+      expressions,
+      requireAllClusterKeys = false,
+      requiredNumPartitions = Some(numPartitions))
+
   def createMergeIntoTable(
       targetTable: LogicalPlan,
       sourceTable: LogicalPlan,

paimon-spark/paimon-spark-4.0/src/main/scala/org/apache/spark/sql/paimon/shims/Spark4Shim.scala

@@ -39,6 +39,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
 import org.apache.spark.sql.catalyst.parser.ParserInterface
 import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Assignment, ColumnDefinition, CTERelationRef, InsertAction, LogicalPlan, MergeAction, MergeIntoTable, MergeRows, SubqueryAlias, TableSpec, UnresolvedWith, UpdateAction}
 import org.apache.spark.sql.catalyst.plans.logical.MergeRows.Keep
+import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, Distribution}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.{ArrayData, GeneratedColumn, IdentityColumn, ResolveDefaultColumns}
 import org.apache.spark.sql.connector.catalog.{CatalogV2Util, Column, Identifier, StagingTableCatalog, Table, TableCatalog}
@@ -243,6 +244,14 @@ class Spark4Shim extends SparkShim {
     CTERelationRef(cteId, resolved, output.toSeq, isStreaming)
   }
 
+  override def createClusteredDistribution(
+      expressions: Seq[Expression],
+      numPartitions: Int): Distribution =
+    ClusteredDistribution(
+      expressions,
+      requireAllClusterKeys = false,
+      requiredNumPartitions = Some(numPartitions))
+
   override def supportsHashAggregate(
       aggregateBufferAttributes: Seq[Attribute],
       groupingExpression: Seq[Expression]): Boolean = {

paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java

@@ -39,6 +39,14 @@ public class SparkConnectorOptions {
                     .withDescription(
                             "If true, map Paimon TIMESTAMP to Spark TIMESTAMP instead of TIMESTAMP_NTZ.");
 
+    public static final ConfigOption<Integer> VECTOR_SEARCH_LATERAL_JOIN_PARALLELISM =
+            key("vector-search.lateral-join.parallelism")
+                    .intType()
+                    .defaultValue(16)
+                    .withDescription(
+                            "Parallelism used to repartition a single-partition LIMIT input before "
+                                    + "executing a lateral vector search.");
+
     public static final ConfigOption<Boolean> MERGE_SCHEMA =
             key("write.merge-schema")
                     .booleanType()

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/RepartitionLateralVectorSearchInput.scala

@@ -0,0 +1,91 @@
+/*
+ * 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.SparkConnectorOptions
+import org.apache.paimon.spark.catalyst.plans.logical.LateralVectorSearch
+import org.apache.paimon.spark.util.OptionUtils
+
+import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, CTERelationRef, GlobalLimit, HintInfo, Join, LogicalPlan, Repartition, RepartitionOperation, ResolvedHint, UnaryNode, WithCTE}
+import org.apache.spark.sql.catalyst.rules.Rule
+
+/** Restores parallelism lost by a global limit before executing a lateral vector search. */
+object RepartitionLateralVectorSearchInput extends Rule[LogicalPlan] {
+
+  override def apply(plan: LogicalPlan): LogicalPlan = {
+    val cteDefinitions = plan
+      .collect { case withCTE: WithCTE => withCTE.cteDefs }
+      .flatten
+      .map(definition => definition.id -> definition.child)
+      .toMap
+
+    plan.transformUp {
+      case lateralVectorSearch: LateralVectorSearch
+          if hasUnrepartitionedGlobalLimit(lateralVectorSearch.left, cteDefinitions, Set.empty) =>
+        lateralVectorSearch.copy(
+          left = Repartition(parallelism, shuffle = true, lateralVectorSearch.left))
+    }
+  }
+
+  private[spark] def parallelism: Int = {
+    val value =
+      OptionUtils
+        .getOptionString(SparkConnectorOptions.VECTOR_SEARCH_LATERAL_JOIN_PARALLELISM)
+        .toInt
+    require(
+      value > 0,
+      s"spark.paimon.${SparkConnectorOptions.VECTOR_SEARCH_LATERAL_JOIN_PARALLELISM.key()} " +
+        s"must be positive, but got $value")
+    value
+  }
+
+  private def hasUnrepartitionedGlobalLimit(
+      plan: LogicalPlan,
+      cteDefinitions: Map[Long, LogicalPlan],
+      visitedCTEs: Set[Long]): Boolean = plan match {
+    case repartition: RepartitionOperation if repartition.shuffle => false
+    case repartition: RepartitionOperation =>
+      hasUnrepartitionedGlobalLimit(repartition.child, cteDefinitions, visitedCTEs)
+    case _: GlobalLimit => true
+    case reference: CTERelationRef if !visitedCTEs.contains(reference.cteId) =>
+      cteDefinitions
+        .get(reference.cteId)
+        .exists(hasUnrepartitionedGlobalLimit(_, cteDefinitions, visitedCTEs + reference.cteId))
+    case join: Join
+        if hasBroadcastHint(join.hint.rightHint) || hasResolvedBroadcastHint(join.right) =>
+      hasUnrepartitionedGlobalLimit(join.left, cteDefinitions, visitedCTEs)
+    case join: Join
+        if hasBroadcastHint(join.hint.leftHint) || hasResolvedBroadcastHint(join.left) =>
+      hasUnrepartitionedGlobalLimit(join.right, cteDefinitions, visitedCTEs)
+    case unary: UnaryNode =>
+      hasUnrepartitionedGlobalLimit(unary.child, cteDefinitions, visitedCTEs)
+    case _ => false
+  }
+
+  private def hasBroadcastHint(hint: Option[HintInfo]): Boolean = {
+    hint.flatMap(_.strategy).contains(BROADCAST)
+  }
+
+  private def hasResolvedBroadcastHint(plan: LogicalPlan): Boolean = {
+    plan.exists {
+      case hint: ResolvedHint => hasBroadcastHint(Some(hint.hints))
+      case _ => false
+    }
+  }
+}

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala

@@ -27,6 +27,7 @@ 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.optimizer.RepartitionLateralVectorSearchInput
 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.format.PaimonFormatTable
@@ -43,11 +44,15 @@ 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, PredicateHelper, UnsafeProjection}
+import org.apache.spark.sql.catalyst.optimizer.BuildRight
 import org.apache.spark.sql.catalyst.plans.logical.{AddPartitions, CreateTableAsSelect, DescribeRelation, DropPartitions, LogicalPlan, RepairTable, ReplaceTable, ReplaceTableAsSelect, ShowCreateTable}
+import org.apache.spark.sql.catalyst.plans.physical.{Distribution, UnspecifiedDistribution}
 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.{GlobalLimitExec, PaimonDescribeTableExec, SparkPlan, SparkStrategy, UnaryExecNode}
 import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Implicits, DataSourceV2Relation}
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec}
 import org.apache.spark.sql.execution.shim.{PaimonCreateTableAsSelectStrategy, PaimonReplaceTableAsSelectStrategy, PaimonReplaceTableStrategy}
 import org.apache.spark.sql.paimon.shims.SparkShimLoader
 
@@ -309,6 +314,30 @@ case class LateralVectorSearchExec(
 
   override def output: Seq[Attribute] = child.output ++ projectOutput
 
+  // Statistics-based broadcast selection is only known after physical planning. Request a
+  // distribution here so EnsureRequirements can restore the streamed LIMIT side's parallelism.
+  override def requiredChildDistribution: Seq[Distribution] = {
+    if (hasUnrepartitionedGlobalLimit(child)) {
+      Seq(
+        SparkShimLoader.shim.createClusteredDistribution(
+          child.output,
+          RepartitionLateralVectorSearchInput.parallelism))
+    } else {
+      Seq(UnspecifiedDistribution)
+    }
+  }
+
+  private def hasUnrepartitionedGlobalLimit(plan: SparkPlan): Boolean = plan match {
+    case _: ShuffleExchangeLike => false
+    case _: GlobalLimitExec => true
+    case join: BroadcastHashJoinExec =>
+      hasUnrepartitionedGlobalLimit(if (join.buildSide == BuildRight) join.left else join.right)
+    case join: BroadcastNestedLoopJoinExec =>
+      hasUnrepartitionedGlobalLimit(if (join.buildSide == BuildRight) join.left else join.right)
+    case unary: UnaryExecNode => hasUnrepartitionedGlobalLimit(unary.child)
+    case _ => false
+  }
+
   @transient override lazy val producedAttributes: AttributeSet = {
     AttributeSet(vectorSearchOutput ++ output.filterNot(attr => inputSet.contains(attr)))
   }

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, PushDownLateralVectorSearchFilter}
+import org.apache.paimon.spark.catalyst.optimizer.{MergePaimonScalarSubqueries, OptimizeMetadataOnlyDeleteFromPaimonTable, PushDownLateralVectorSearchFilter, RepartitionLateralVectorSearchInput}
 import org.apache.paimon.spark.catalyst.plans.logical.PaimonTableValuedFunctions
 import org.apache.paimon.spark.commands.BucketExpression
 import org.apache.paimon.spark.execution.{OldCompatibleStrategy, PaimonStrategy}
@@ -104,6 +104,7 @@ class PaimonSparkSessionExtensions extends (SparkSessionExtensions => Unit) {
     // TODO: Enable MAP selected-key pushdown after core reader supports
     // __PAIMON_MAP_SELECTED_KEYS read type.
     extensions.injectOptimizerRule(_ => MergePaimonScalarSubqueries)
+    extensions.injectOptimizerRule(_ => RepartitionLateralVectorSearchInput)
     extensions.injectOptimizerRule(_ => PushDownLateralVectorSearchFilter)
 
     // planner extensions

paimon-spark/paimon-spark-common/src/main/scala/org/apache/spark/sql/paimon/shims/SparkShim.scala

@@ -33,6 +33,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression}
 import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
 import org.apache.spark.sql.catalyst.parser.ParserInterface
 import org.apache.spark.sql.catalyst.plans.logical.{Assignment, CTERelationRef, InsertAction, LogicalPlan, MergeAction, MergeIntoTable, SubqueryAlias, TableSpec, UnresolvedWith, UpdateAction}
+import org.apache.spark.sql.catalyst.plans.physical.Distribution
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.ArrayData
 import org.apache.spark.sql.connector.catalog.{Column, Identifier, StagingTableCatalog, Table, TableCatalog}
@@ -159,6 +160,8 @@ trait SparkShim {
       output: Seq[Attribute],
       isStreaming: Boolean): CTERelationRef
 
+  def createClusteredDistribution(expressions: Seq[Expression], numPartitions: Int): Distribution
+
   def supportsHashAggregate(
       aggregateBufferAttributes: Seq[Attribute],
       groupingExpression: Seq[Expression]): Boolean