[SPARK-40259][SQL] Support merging equivalent DataSource V2 scans in MergeSubplans

sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsScanMerging.java

@@ -0,0 +1,49 @@
+/*
+ * 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.spark.sql.connector.read;
+
+import org.apache.spark.annotation.Evolving;
+
+/**
+ * A mix-in interface for {@link Scan} that opts the scan in to Spark-side scan merging.
+ * <p>
+ * By implementing this marker a data source declares a determinism contract: holding the scan
+ * options constant, the set of rows and columns the scan reads is fully determined by the filters
+ * pushed via {@link SupportsPushDownV2Filters} and the columns pruned via
+ * {@link SupportsPushDownRequiredColumns}. Equivalently, rebuilding the scan from a fresh
+ * {@link ScanBuilder} with the same options and re-applying the same pushed filters and pruned
+ * columns yields an equivalent scan.
+ * <p>
+ * Given that contract, Spark may fuse two scans of the same table that differ only in their
+ * projected columns and/or pushed filters into a single scan. Spark builds the merged scan itself:
+ * it obtains a fresh {@link ScanBuilder} from the table, prunes it to the union of both read
+ * schemas, re-pushes the (possibly OR-widened) filters, and builds. The merged scan reads the union
+ * of the two scans' columns and a superset of their rows; each original scan's result is recovered
+ * by a projection and filter applied above it. The connector supplies no merge logic of its own --
+ * this interface has no methods.
+ * <p>
+ * Implementers do not need to reason about pushdowns that are not reproducible this way (a pushed
+ * aggregate, join, variant extraction, limit, offset, top-N, or table sample). Spark tracks those
+ * on its own side while building the scan and never merges a scan that carries one, whether or not
+ * the scan implements this interface. The only obligation of an implementer is the determinism
+ * contract above.
+ *
+ * @since 4.3.0
+ */
+@Evolving
+public interface SupportsScanMerging extends Scan { }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplansReferences.scala

@@ -0,0 +1,65 @@
+/*
+ * 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.spark.sql.catalyst.optimizer
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, ExprId, LeafExpression, Unevaluable}
+import org.apache.spark.sql.catalyst.plans.logical.LeafNode
+import org.apache.spark.sql.catalyst.trees.TreePattern.{NO_GROUPING_AGGREGATE_REFERENCE, SCALAR_SUBQUERY_REFERENCE, TreePattern}
+import org.apache.spark.sql.types.DataType
+
+// The temporal reference placeholders below are produced by the `MergeSubplans` rule (now in
+// sql/core) but must remain in catalyst: `ScalarSubqueryReference` is referenced by the catalyst
+// expression `BloomFilterMightContain`, and catalyst cannot depend on sql/core.
+
+/**
+ * Temporal reference to a subquery which is added to a `PlanMerger`.
+ *
+ * @param level The level of the replaced subquery. It defines the `PlanMerger` instance into which
+ *              the subquery is merged.
+ * @param mergedPlanIndex The index of the merged plan in the `PlanMerger`.
+ * @param outputIndex The index of the output attribute of the merged plan.
+ * @param dataType The dataType of original scalar subquery.
+ * @param exprId The expression id of the original scalar subquery.
+ */
+case class ScalarSubqueryReference(
+    level: Int,
+    mergedPlanIndex: Int,
+    outputIndex: Int,
+    override val dataType: DataType,
+    exprId: ExprId) extends LeafExpression with Unevaluable {
+  override def nullable: Boolean = true
+
+  final override val nodePatterns: Seq[TreePattern] = Seq(SCALAR_SUBQUERY_REFERENCE)
+}
+
+/**
+ * Temporal reference to a non-grouping aggregate which is added to a `PlanMerger`.
+ *
+ * @param level The level of the replaced aggregate. It defines the `PlanMerger` instance into which
+ *              the aggregate is merged.
+ * @param mergedPlanIndex The index of the merged plan in the `PlanMerger`.
+ * @param outputIndices The indices of the output attributes of the merged plan.
+ * @param output The output of original aggregate.
+ */
+case class NonGroupingAggregateReference(
+    level: Int,
+    mergedPlanIndex: Int,
+    outputIndices: Seq[Int],
+    override val output: Seq[Attribute]) extends LeafNode {
+  final override val nodePatterns: Seq[TreePattern] = Seq(NO_GROUPING_AGGREGATE_REFERENCE)
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala

@@ -158,19 +158,31 @@ case class DataSourceV2Relation(
  * @param keyGroupedPartitioning if set, the partitioning expressions that are used to split the
  *                               rows in the scan across different partitions
  * @param ordering if set, the ordering provided by the scan
- * @param pushedFilters Catalyst expressions for filters that were fully pushed to the data
- *                      source and do not appear as post-scan filters
+ * @param pushedFilters Catalyst expressions for filters that were fully pushed to the data source
+ *                      and do not appear as post-scan filters. These reference the relation's
+ *                      (pre-pruning) output, so they may reference columns pruned out of `output`
+ *                      (e.g. an unselected partition column the source enforces internally). This
+ *                      complete set is what lets [[org.apache.spark.sql.execution.planmerging.
+ *                      PlanMerger]] soundly compare and re-enforce a scan's filters when fusing two
+ *                      scans via [[org.apache.spark.sql.connector.read.SupportsScanMerging]].
+ * @param hasMergeBlockingPushdown whether a pushdown happened that cannot be reproduced by
+ *                      re-running the scan builder (aggregate, join, variant extraction, limit,
+ *                      offset, top-N, or sample). When true, this scan must not be fused with
+ *                      another via [[org.apache.spark.sql.connector.read.SupportsScanMerging]].
  */
 case class DataSourceV2ScanRelation(
     relation: DataSourceV2Relation,
     scan: Scan,
     output: Seq[AttributeReference],
     keyGroupedPartitioning: Option[Seq[Expression]] = None,
     ordering: Option[Seq[SortOrder]] = None,
-    pushedFilters: Seq[Expression] = Seq.empty) extends LeafNode with NamedRelation {
+    pushedFilters: Seq[Expression] = Seq.empty,
+    hasMergeBlockingPushdown: Boolean = false) extends LeafNode with NamedRelation {
 
   // TODO: Override validConstraints to return ExpressionSet(pushedFilters) so that pushed
   // filters participate in constraint propagation (InferFiltersFromConstraints, PruneFilters).
+  // Note: pushedFilters may reference columns pruned out of `output`, so constraint use must first
+  // intersect with `outputSet` (a constraint has to reference the node's output).
   // This changes which filters InferFiltersFromConstraints adds or removes (e.g., it may
   // skip adding IsNotNull when the scan already implies it, or infer new filters across
   // joins), so plan stability testing is needed first.
@@ -189,6 +201,14 @@ case class DataSourceV2ScanRelation(
 
   override def name: String = relation.name
 
+  // A leaf relation references no upstream attributes. `pushedFilters` (and, for that matter,
+  // partitioning/ordering) are scan metadata, not references to resolve, and `pushedFilters` may
+  // reference columns pruned out of `output` (e.g. an unselected partition column). Without this
+  // override those would surface as `missingInput`, which the optimizer's plan-change validation
+  // flags as dangling references. `mapExpressions`/`transformExpressions` still rewrite the
+  // metadata expressions -- they iterate the product directly, independent of `references`.
+  override def references: AttributeSet = AttributeSet.empty
+
   override def simpleString(maxFields: Int): String = {
     val outputString = truncatedString(output, "[", ", ", "]", maxFields)
     val nameWithTimeTravelSpec = relation.timeTravelSpec match {
@@ -220,7 +240,9 @@ case class DataSourceV2ScanRelation(
       ordering = ordering.map(
         _.map(o => o.copy(child = QueryPlan.normalizeExpressions(o.child, output)))
       ),
-      pushedFilters = pushedFilters.map(QueryPlan.normalizeExpressions(_, output))
+      // pushedFilters may reference columns pruned out of `output` (see the field doc), so they are
+      // normalized against the relation's full output rather than `output`.
+      pushedFilters = pushedFilters.map(QueryPlan.normalizeExpressions(_, relation.output))
     )
   }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryScanMergingPartitionFilterTable.scala

@@ -0,0 +1,87 @@
+/*
+ * 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.spark.sql.connector.catalog
+
+import java.util
+
+import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.connector.read.{Batch, Scan, ScanBuilder, SupportsScanMerging}
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Catalog that hands out [[InMemoryScanMergingPartitionFilterTable]]s, an iterative-pushdown source
+ * that additionally opts its scans in to Spark-side scan merging. Used to exercise the merge of two
+ * DSv2 scans whose (equal) filter is strict only via the iterative PartitionPredicate second pass.
+ */
+class InMemoryScanMergingPartitionFilterCatalog
+  extends InMemoryTableEnhancedPartitionFilterCatalog {
+  import CatalogV2Implicits._
+
+  override def createTable(
+      ident: Identifier,
+      columns: Array[Column],
+      partitions: Array[Transform],
+      properties: util.Map[String, String]): Table = {
+    if (tables.containsKey(ident)) {
+      throw new TableAlreadyExistsException(ident.asMultipartIdentifier)
+    }
+    InMemoryTableCatalog.maybeSimulateFailedTableCreation(properties)
+    val tableName = s"$name.${ident.quoted}"
+    val table =
+      new InMemoryScanMergingPartitionFilterTable(tableName, columns, partitions, properties)
+    tables.put(ident, table)
+    namespaces.putIfAbsent(ident.namespace.toList, Map())
+    table
+  }
+}
+
+/**
+ * An [[InMemoryEnhancedPartitionFilterTable]] whose built scans also implement
+ * [[SupportsScanMerging]], so [[org.apache.spark.sql.execution.planmerging.PlanMerger]] may fuse
+ * two scans of this table. The scan is wrapped in a thin [[ScanMergingWrapperScan]] rather than
+ * subclassing the base batch scan: the wrapper adds the merge marker and deliberately does not
+ * re-expose `SupportsReportPartitioning`, so a partitioned table does not set the scan relation's
+ * `keyGroupedPartitioning` (whose preservation across a merge is a separate follow-up). This keeps
+ * the fixture focused on the iterative-pushdown behavior under test.
+ */
+class InMemoryScanMergingPartitionFilterTable(
+    name: String,
+    columns: Array[Column],
+    partitioning: Array[Transform],
+    properties: util.Map[String, String])
+  extends InMemoryEnhancedPartitionFilterTable(name, columns, partitioning, properties) {
+
+  override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder =
+    new InMemoryEnhancedPartitionFilterScanBuilder(schema()) {
+      override def build(): Scan = ScanMergingWrapperScan(super.build())
+    }
+}
+
+/**
+ * Thin scan decorator that adds the [[SupportsScanMerging]] marker and delegates reading to the
+ * wrapped scan's batch. It exposes only `readSchema` and `toBatch`, so the scan relation carries no
+ * reported partitioning/ordering/statistics -- exactly the shape [[SupportsScanMerging]] promises
+ * is reproducible from a fresh scan builder.
+ */
+case class ScanMergingWrapperScan(inner: Scan) extends Scan with SupportsScanMerging {
+  override def readSchema(): StructType = inner.readSchema()
+  override def toBatch: Batch = inner.toBatch
+  override def description(): String = inner.description()
+}

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala

@@ -26,6 +26,7 @@ import org.apache.spark.sql.connector.catalog.CatalogManager
 import org.apache.spark.sql.execution.datasources.{MarkSingleTaskExecution, PruneFileSourcePartitions, PushVariantIntoScan, SchemaPruning, V1Writes}
 import org.apache.spark.sql.execution.datasources.v2.{GroupBasedRowLevelOperationScanPlanning, OptimizeMetadataOnlyDeleteFromTable, V2ScanPartitioningAndOrdering, V2ScanRelationPushDown, V2Writes}
 import org.apache.spark.sql.execution.dynamicpruning.{CleanupDynamicPruningFilters, PartitionPruning, RowLevelOperationRuntimeGroupFiltering}
+import org.apache.spark.sql.execution.planmerging.MergeSubplans
 import org.apache.spark.sql.execution.python.{ExtractGroupingPythonUDFFromAggregate, ExtractPythonUDFFromAggregate, ExtractPythonUDFs, ExtractPythonUDTFs}
 
 class SparkOptimizer(