From 511e5d1a332e3c4129b476a4bbd23eab1fcb5a98 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 2 Jun 2026 15:45:10 +0800 Subject: [PATCH] [SPARK-57205][SQL] Add SupportsScanMerging to merge equivalent V2 file scans ### What changes were proposed in this pull request? This adds a `SupportsScanMerging` mix-in for DSv2 `Scan` and teaches `MergeSubplans` (through `PlanMerger`) to use it. When two scans of the same table differ only in their projected columns and/or pushed filters, the optimizer asks the source to fuse them into one scan that covers both, and then deduplicates the subplans built on top. The interface is implemented once on the `FileScan` base, so every built-in file format (Parquet, ORC, CSV, JSON, Text, Avro) participates. Two scans merge when they read the same data (same file index, schema, options and partition filters); the merged scan reads the union of the two read schemas. When the pushed data filters differ, the merged scan widens them to `OR(f1, f2)` and reads a superset -- the exact per-side predicate is still enforced by the post-scan `Filter`, which `MergeSubplans` turns into per-aggregate `FILTER (WHERE ...)` clauses. That widening is off by default, gated by `spark.sql.files.scanMerge.ignorePushedDataFilters`. Merging is declined when an aggregate is pushed into the scan, since there is then no post-scan `Filter` to separate the two sides. ### Why are the changes needed? V1 file sources already get this. Their filter and column pushdown happens during physical planning, so when `MergeSubplans` runs the logical plan still has plain `Filter`/`Project` nodes over identical relation leaves, which the existing rules merge. DSv2 bakes pushdown into `DataSourceV2ScanRelation` during logical optimization, so two subqueries that differ only in `WHERE` or `SELECT` become structurally different leaves and cannot be merged. For example TPC-DS q9 (fifteen scalar subqueries over `store_sales`) collapses to a single scan under V1 but reads the table fifteen times under V2. This closes that gap. ### Does this PR introduce _any_ user-facing change? No. The connector interface and the config are internal, and the config defaults to off, so the default plan is unchanged. ### How was this patch tested? New tests in `PlanMergeSuite` run each query through both the V1 and V2 file-source paths, with AQE on and off, and assert identical results and identical merge structure. The merge-semantics cases cover differing columns, differing data filters, differing partition filters, filter propagation through a join, and a multi-subquery composition, plus negative cases; the format-coverage cases run a representative query over Parquet, ORC and JSON. `MergeSubplansSuite`, `DataSourceV2Suite` and the Parquet/ORC V2 aggregate-pushdown suites still pass. I also checked that V1 and V2 produce identical results and merge structure on the affected TPC-DS queries (q2, q9, q28, q59, q77a, q88, q90) at scale factor 1. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Opus 4.8 --- .../connector/read/SupportsScanMerging.java | 72 ++++ .../sql/catalyst/optimizer/PlanMerger.scala | 94 ++++- .../apache/spark/sql/internal/SQLConf.scala | 15 + .../execution/datasources/v2/FileScan.scala | 87 +++- .../datasources/v2/orc/OrcScan.scala | 3 + .../datasources/v2/parquet/ParquetScan.scala | 12 + .../org/apache/spark/sql/PlanMergeSuite.scala | 385 ++++++++++++++++++ 7 files changed, 664 insertions(+), 4 deletions(-) create mode 100644 sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsScanMerging.java diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsScanMerging.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsScanMerging.java new file mode 100644 index 0000000000000..d0dba9d1715db --- /dev/null +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsScanMerging.java @@ -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.spark.sql.connector.read; + +import java.util.Optional; + +import org.apache.spark.annotation.Evolving; +import org.apache.spark.sql.connector.catalog.SupportsRead; + +/** + * A mix-in interface for {@link Scan}. Data sources can implement this interface to allow + * Spark's optimizer to fuse two {@link Scan}s of the same {@link SupportsRead} table into a + * single {@link Scan} covering both: it reads the union of their projected columns and a + * superset of their rows, from which each original scan's result can be recovered by a + * projection and filter applied above the merged scan. + *

+ * Spark calls {@link #mergeWith(SupportsScanMerging, SupportsRead)} when it detects that two + * scans target the same table and read structurally compatible data (e.g. same file + * index, schema, partitioning) but differ in pushed predicates, projections, or other + * scan-level state. Implementations decide whether the two scans can be safely combined + * and, if so, return the merged scan via {@link Optional#of}. Returning {@link Optional#empty} + * declines the merge. + *

+ * Contract for implementations: + *

+ * + * @since 4.3.0 + */ +@Evolving +public interface SupportsScanMerging extends Scan { + + /** + * Attempts to merge this scan with {@code other} into a single equivalent scan. + * + * @param other the other scan to merge with; guaranteed to also implement + * {@link SupportsScanMerging} and to be a scan of the same {@code table} + * @param table the {@link SupportsRead} table that owns both scans, used to obtain + * a {@link ScanBuilder} for constructing the merged scan + * @return {@link Optional#of} the merged scan if merging is supported, or + * {@link Optional#empty} to decline + */ + Optional mergeWith(SupportsScanMerging other, SupportsRead table); +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala index 1c43f91cee9dc..25707eab37ea3 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala @@ -19,11 +19,15 @@ package org.apache.spark.sql.catalyst.optimizer import scala.collection.mutable -import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeMap, Expression, If, Literal, NamedExpression, Or} +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeMap, AttributeReference, Expression, ExpressionSet, If, Literal, NamedExpression, Or} import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression import org.apache.spark.sql.catalyst.plans.{Cross, Inner, JoinType, LeftAnti, LeftOuter, LeftSemi, RightOuter} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project} import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes +import org.apache.spark.sql.connector.catalog.SupportsRead +import org.apache.spark.sql.connector.read.SupportsScanMerging +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation import org.apache.spark.sql.internal.SQLConf /** @@ -483,11 +487,99 @@ class PlanMerger( case _ => None } + case (np: DataSourceV2ScanRelation, cp: DataSourceV2ScanRelation) => + tryMergeV2ScanRelations(np, cp) + // Otherwise merging is not possible. case _ => None }) } + /** Attempt to merge two V2 scan relations via the [[SupportsScanMerging]] connector + * interface. Both scans must reference the same underlying relation and have + * canonicalized-identical relation-level `pushedFilters` (see the gate below). The + * merged scan's output is the union of both sides' read schemas; cp's attribute IDs + * are preserved for shared columns and np's attribute IDs are mapped to the + * corresponding merged attribute. + * + * How this reaches V1 parity for filter differences: for sources whose pushdown is + * best-effort (e.g. Parquet row-group filters), the exact predicate remains as a + * post-scan [[Filter]] node ABOVE the scan, and the relation-level `pushedFilters` + * field is empty. So two branches that differ only in their WHERE clause look like + * `Filter(p1) over scan` vs `Filter(p2) over scan` -- the existing `(Filter, Filter)` + * symmetric-propagation cases OR-widen the exact filter and emit per-side aggregate + * `FILTER (WHERE ...)` clauses, exactly as for V1. This leaf case only merges the + * scans underneath, OR-widening their best-effort pushed filter (via the connector's + * `mergeWith`) so the merged scan reads a superset of either side's rows. The + * per-side correctness is then restored by the Filter nodes above. + */ + private def tryMergeV2ScanRelations( + np: DataSourceV2ScanRelation, + cp: DataSourceV2ScanRelation): Option[TryMergeResult] = { + // Both scans must reference the same underlying relation. + if (np.relation.canonicalized != cp.relation.canonicalized) return None + + // Relation-level pushedFilters must match. For best-effort sources these are empty + // (the exact predicate stays as a post-scan Filter, handled by the (Filter, Filter) + // cases above). For sources that report EXACT pushdown (no residual Filter node), + // differing pushedFilters here would have no Filter above to re-apply per-side + // predicates, so OR-widening the scan would change results -- this gate + // conservatively declines that case. + if (ExpressionSet(np.pushedFilters) != ExpressionSet(cp.pushedFilters)) return None + + // Both scans must implement SupportsScanMerging. + val (npm, cpm) = (np.scan, cp.scan) match { + case (a: SupportsScanMerging, b: SupportsScanMerging) => (a, b) + case _ => return None + } + + // The underlying table must be SupportsRead so we can call newScanBuilder. + val table = np.relation.table match { + case sr: SupportsRead => sr + case _ => return None + } + + val mergedOpt = npm.mergeWith(cpm, table) + if (!mergedOpt.isPresent) return None + val mergedScan = mergedOpt.get() + + // Build merged output: preserve cp's attribute IDs for shared columns, allocate + // fresh attribute IDs for np-exclusive columns. + val realOutput = toAttributes(mergedScan.readSchema()) + val cpByName = cp.output.map(a => a.name -> a).toMap + val mergedOutput = realOutput.map { ra => + cpByName.get(ra.name) match { + case Some(cpa) => + AttributeReference(ra.name, ra.dataType, ra.nullable, ra.metadata)( + cpa.exprId, cpa.qualifier) + case None => ra + } + } + + // np.output -> mergedOutput by name. If any np-output column is missing from + // the merged scan, the merge is not usable. AttributeMap is invariant in its + // value type, so widen the pairs to (Attribute, Attribute) explicitly. + val mergedByName = mergedOutput.map(a => a.name -> a).toMap + val mappingOpts: Seq[Option[(Attribute, Attribute)]] = + np.output.map(npa => mergedByName.get(npa.name).map(npa -> _)) + if (mappingOpts.exists(_.isEmpty)) return None + + val npMapping = AttributeMap(mappingOpts.flatten) + + // Preserve the original pushedFilters (np.pushedFilters == cp.pushedFilters by + // the check above). Partitioning and ordering hints from V2ScanPartitioningAndOrdering + // are intentionally dropped: a merged scan with a wider read schema may have a + // different output partitioning than either input, and re-running pushdown via + // the connector's mergeWith does not preserve those hints. Consumers that need + // them can be re-fixed by a later pass of V2ScanPartitioningAndOrdering. + val mergedRelation = DataSourceV2ScanRelation( + np.relation, + mergedScan, + mergedOutput, + pushedFilters = cp.pushedFilters) + Some(TryMergeResult(mergedRelation, npMapping)) + } + // Returns true when a filter attribute originating from `fromLeft` child of a join with // `joinType` can be safely propagated through that join to a parent Aggregate. // diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 0aed28e92558f..e4cf5111fa477 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -1620,6 +1620,18 @@ object SQLConf { .booleanConf .createWithDefault(false) + val FILE_SCAN_MERGE_IGNORE_PUSHED_DATA_FILTERS = + buildConf("spark.sql.files.scanMerge.ignorePushedDataFilters") + .internal() + .doc("When true, allow file-source `Scan`s with different pushed data filters to be merged " + + "via the SupportsScanMerging connector API. The merged scan widens the filter to " + + "OR(f1, f2) and per-scan filters are reapplied at consumers (e.g. via aggregate FILTER " + + "clauses). When false, scans must have identical pushed data filters to merge.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .booleanConf + .createWithDefault(false) + val PARQUET_WRITE_LEGACY_FORMAT = buildConf("spark.sql.parquet.writeLegacyFormat") .doc("If true, data will be written in a way of Spark 1.4 and earlier. For example, decimal " + "values will be written in Apache Parquet's fixed-length byte array format, which other " + @@ -7870,6 +7882,9 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def parquetAggregatePushDown: Boolean = getConf(PARQUET_AGGREGATE_PUSHDOWN_ENABLED) + def fileScanMergeIgnorePushedDataFilters: Boolean = + getConf(FILE_SCAN_MERGE_IGNORE_PUSHED_DATA_FILTERS) + def orcFilterPushDown: Boolean = getConf(ORC_FILTER_PUSHDOWN_ENABLED) def orcAggregatePushDown: Boolean = getConf(ORC_AGGREGATE_PUSHDOWN_ENABLED) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala index 5348f9ab6df62..16962d2f7b3c3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala @@ -16,7 +16,7 @@ */ package org.apache.spark.sql.execution.datasources.v2 -import java.util.{Locale, OptionalLong} +import java.util.{Locale, Optional, OptionalLong} import org.apache.hadoop.fs.Path @@ -25,10 +25,11 @@ import org.apache.spark.internal.LogKeys.{PATH, REASON} import org.apache.spark.internal.config.IO_WARNING_LARGEFILETHRESHOLD import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.SQLConfHelper -import org.apache.spark.sql.catalyst.expressions.{AttributeSet, Expression, ExpressionSet} +import org.apache.spark.sql.catalyst.expressions.{And, AttributeSet, Expression, ExpressionSet, Or} import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes +import org.apache.spark.sql.connector.catalog.SupportsRead import org.apache.spark.sql.connector.read._ import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.execution.PartitionedFileUtil @@ -37,12 +38,14 @@ import org.apache.spark.sql.internal.{SessionStateHelper, SQLConf} import org.apache.spark.sql.internal.connector.SupportsMetadata import org.apache.spark.sql.sources.Filter import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.util.Utils trait FileScan extends Scan with Batch with SupportsReportStatistics with SupportsMetadata + with SupportsScanMerging with SQLConfHelper with Logging { /** @@ -56,6 +59,9 @@ trait FileScan extends Scan def fileIndex: PartitioningAwareFileIndex + /** The scan options, used to obtain a fresh [[ScanBuilder]] when merging. */ + def options: CaseInsensitiveStringMap + def dataSchema: StructType /** @@ -89,7 +95,7 @@ trait FileScan extends Scan protected def seqToString(seq: Seq[Any]): String = seq.mkString("[", ", ", "]") - private lazy val (normalizedPartitionFilters, normalizedDataFilters) = { + protected lazy val (normalizedPartitionFilters, normalizedDataFilters) = { val partitionFilterAttributes = AttributeSet(partitionFilters).map(a => a.name -> a).toMap val normalizedPartitionFilters = ExpressionSet(partitionFilters.map( QueryPlan.normalizeExpressions(_, toAttributes(fileIndex.partitionSchema) @@ -211,6 +217,81 @@ trait FileScan extends Scan a.sortBy(_.hashCode()).sameElements(b.sortBy(_.hashCode())) } + // ===== SupportsScanMerging ================================================= + // + // Two file scans of the same table can be fused into one when they read the same data and + // differ only in projected columns and/or pushed data filters. This brings V2 file sources to + // parity with V1, where MergeSubplans already merges such subplans (V1 keeps the Filter/Project + // nodes in the logical plan, so the relation leaves are identical). The logic here is + // format-agnostic; format-specific state is handled via the hooks below. + + /** + * Format-specific scan state (beyond fileIndex / schema / options / partition filters / data + * filters / projected columns) that must match for two scans to be mergeable. The default + * requires no extra state. Overridden e.g. by ParquetScan to require equal variant extractions. + * `other` is guaranteed to be the same concrete class as `this`. + */ + protected def canMergeScanStateWith(other: FileScan): Boolean = true + + /** + * Whether this scan has pushed-down aggregation. When true, the scan emits aggregated rows with + * no post-scan Filter to reconcile per-side predicates, so merging is declined. Overridden by + * ParquetScan / OrcScan which support aggregate pushdown. + */ + protected def hasAggregatePushedDown: Boolean = false + + private def mergeEligible(o: FileScan): Boolean = + getClass == o.getClass && + fileIndex == o.fileIndex && + dataSchema == o.dataSchema && + options == o.options && + normalizedPartitionFilters == o.normalizedPartitionFilters && + !hasAggregatePushedDown && !o.hasAggregatePushedDown && + canMergeScanStateWith(o) + + override def mergeWith( + other: SupportsScanMerging, + table: SupportsRead): Optional[SupportsScanMerging] = other match { + case o: FileScan if mergeEligible(o) => + // Strict when the two scans already agree on pushed data filters: rebuild with the union of + // read schemas, rows unchanged. Otherwise (relaxed) widen the best-effort pushed filter to + // OR(f1, f2) so the merged scan reads a superset; the exact per-side predicate is reapplied + // by the post-scan Filter node (split into aggregate FILTER clauses by PlanMerger). The + // relaxed path is gated by `spark.sql.files.scanMerge.ignorePushedDataFilters`. + val strict = normalizedDataFilters == o.normalizedDataFilters + if (!strict && !conf.fileScanMergeIgnorePushedDataFilters) { + Optional.empty() + } else { + table.newScanBuilder(options) match { + case builder: FileScanBuilder => + // Partition filters are guaranteed equal (see `mergeEligible`), so push them as-is -- + // never widen, which would needlessly produce OR(pf, pf) and obscure partition + // pruning. Only data filters use the strict/relaxed distinction. + builder.pushFilters(mergedDataFilters(o, strict) ++ partitionFilters) + builder.pruneColumns(readSchema().merge(o.readSchema())) + builder.build() match { + case merged: SupportsScanMerging => Optional.of(merged) + case _ => Optional.empty() + } + case _ => Optional.empty() + } + } + case _ => Optional.empty() + } + + // Strict: the two data-filter lists are equivalent, so push this scan's. Relaxed: widen to + // OR(AND(df1), AND(df2)) so the merged scan reads a superset; an empty list on either side + // means that side reads everything, so the union also reads everything (no data filter). + private def mergedDataFilters(o: FileScan, strict: Boolean): Seq[Expression] = { + if (strict) { + dataFilters + } else if (dataFilters.nonEmpty && o.dataFilters.nonEmpty) { + Seq(Or(dataFilters.reduce(And), o.dataFilters.reduce(And))) + } else { + Seq.empty + } + } + private val isCaseSensitive = conf.caseSensitiveAnalysis private def normalizeName(name: String): String = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/orc/OrcScan.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/orc/OrcScan.scala index 6242cd3ca2c62..dde91ed7597e0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/orc/OrcScan.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/orc/OrcScan.scala @@ -104,4 +104,7 @@ case class OrcScan( Map("PushedAggregation" -> pushedAggregationsStr) ++ Map("PushedGroupBy" -> pushedGroupByStr) } + + // ORC supports aggregate pushdown, which disables scan merging (see FileScan). + override protected def hasAggregatePushedDown: Boolean = pushedAggregate.nonEmpty } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/parquet/ParquetScan.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/parquet/ParquetScan.scala index d0c7859964e09..556eec4fca369 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/parquet/ParquetScan.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/parquet/ParquetScan.scala @@ -217,4 +217,16 @@ case class ParquetScan( Map("PushedGroupBy" -> pushedGroupByStr) ++ Map("PushedVariantExtractions" -> variantExtractionStr) } + + // SupportsScanMerging hooks (generic logic in FileScan): Parquet disables merging when an + // aggregate is pushed, and requires equal variant extractions to merge. + override protected def hasAggregatePushedDown: Boolean = pushedAggregate.nonEmpty + + override protected def canMergeScanStateWith(other: FileScan): Boolean = other match { + case o: ParquetScan => + java.util.Arrays.equals( + pushedVariantExtractions.asInstanceOf[Array[Object]], + o.pushedVariantExtractions.asInstanceOf[Array[Object]]) + case _ => false + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/PlanMergeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/PlanMergeSuite.scala index e1109f20e6040..92abc9cafb8c5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/PlanMergeSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/PlanMergeSuite.scala @@ -17,6 +17,7 @@ package org.apache.spark.sql +import org.apache.spark.sql.catalyst.optimizer.MergeSubplans import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper} import org.apache.spark.sql.internal.SQLConf @@ -396,6 +397,390 @@ class PlanMergeSuite extends SharedSparkSession } } + // =========================================================================== + // SupportsScanMerging: V2 file-source scan merging, at parity with V1. + // + // The merge logic lives in the FileScan base, so it is format-agnostic. Tests are organized in + // two groups: + // 1. "merge semantics" -- exhaustive coverage of the merge taxonomy (differing columns / + // data filters / partition filters / through-join, plus negatives). These use Parquet as + // the representative format; repeating every shape on every format would only re-test the + // shared FileScan logic. + // 2. "format coverage" -- one representative query run across every built-in file format, to + // confirm each concrete Scan inherits the merge (its SupportsScanMerging hooks + `options` + // wiring work). + // Each test runs through the V1 and V2 file-source paths (USE_V1_SOURCE_LIST) and asserts they + // produce identical results and identical merge structure. + // =========================================================================== + + private val allFileSources = Seq("avro", "csv", "json", "kafka", "orc", "text", "parquet") + + // USE_V1_SOURCE_LIST routes a source through V1 when it is listed (its V2 path is disabled) and + // through V2 when omitted. `allViaV1` routes every source through V1; `viaV2(fmt)` routes `fmt` + // through V2 and everything else through V1. + private val allViaV1: String = allFileSources.mkString(",") + private def viaV2(format: String): String = allFileSources.filterNot(_ == format).mkString(",") + + // Counts (SubqueryExec, ReusedSubqueryExec) in the executed plan -- a merge + // collapses N subqueries into 1 SubqueryExec + (N-1) ReusedSubqueryExec. + private def subqueryCounts(df: DataFrame): (Int, Int) = { + val plan = df.queryExecution.executedPlan + val subqueryIds = collectWithSubqueries(plan) { case s: SubqueryExec => s.id } + val reusedSubqueryIds = collectWithSubqueries(plan) { + case rs: ReusedSubqueryExec => rs.child.id + } + (subqueryIds.size, reusedSubqueryIds.size) + } + + /** + * Runs `queryFn(table)` through the V1 file-source path and the V2 file-source path and + * asserts that both produce identical results AND identical merge behavior (same + * SubqueryExec / ReusedSubqueryExec counts). This is the core V1/V2 capability-parity + * guard: a query that merges under V1 must merge identically under V2, and vice versa. + * + * Uses Parquet as the representative format -- the merge logic is format-agnostic (lives in + * FileScan), so the merge-semantics tests built on this helper need only one format. Per-format + * coverage is provided separately by the "format coverage" tests. + */ + private def assertV1V2Parity( + data: DataFrame, + query: String => String, + expected: Seq[Row], + expectedSubqueries: Int, + expectedReused: Int, + extraConfs: Seq[(String, String)] = Seq.empty): Unit = { + Seq(false, true).foreach { enableAQE => + Seq("V1" -> allViaV1, "V2" -> viaV2("parquet")).foreach { case (label, v1List) => + withTempPath { path => + val pathStr = path.getAbsolutePath + data.write.mode("overwrite").parquet(pathStr) + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> enableAQE.toString, + SQLConf.USE_V1_SOURCE_LIST.key -> v1List) ++ extraConfs): _*) { + val df = spark.sql(query(s"parquet.`$pathStr`")) + withClue(s"[$label, AQE=$enableAQE] results: ") { + checkAnswer(df, expected) + } + val (nSub, nReused) = subqueryCounts(df) + withClue(s"[$label, AQE=$enableAQE] SubqueryExec count: ") { + assert(nSub == expectedSubqueries) + } + withClue(s"[$label, AQE=$enableAQE] ReusedSubqueryExec count: ") { + assert(nReused == expectedReused) + } + } + } + } + } + } + + + /** + * Like [[assertV1V2Parity]] but compares the V1 path directly against the V2 path instead of + * against hard-coded expectations. Asserts: + * 1. V1 and V2 return identical rows. + * 2. V1 and V2 have identical merge structure (same SubqueryExec / ReusedSubqueryExec counts). + * 3. The merge actually fired -- the merged plan has strictly fewer distinct SubqueryExec + * than the same query with MergeSubplans excluded. Without this the parity check would + * pass vacuously if the optimization silently stopped firing in BOTH paths. + * Used for composition-style queries where the exact merged counts are not worth predicting -- + * the point is that V2 behaves exactly like V1, and that a merge genuinely happens. + * + * Note: the "merge fired" guard checks for a reduction in distinct SubqueryExec, so this helper + * is only suitable for scalar-subquery / joined-subquery merges (not main-plan merges that + * leave no SubqueryExec, e.g. UNION-of-aggregates). + */ + private def assertV1V2Consistent( + data: DataFrame, + query: String => String, + extraConfs: Seq[(String, String)] = Seq.empty, + format: String = "parquet"): Unit = { + Seq(false, true).foreach { enableAQE => + withTempPath { path => + val pathStr = path.getAbsolutePath + data.write.mode("overwrite").format(format).save(pathStr) + def run(v1List: String, confs: Seq[(String, String)]): (Seq[Row], (Int, Int)) = { + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> enableAQE.toString, + SQLConf.USE_V1_SOURCE_LIST.key -> v1List) ++ confs): _*) { + val df = spark.sql(query(s"$format.`$pathStr`")) + (df.collect().toSeq, subqueryCounts(df)) + } + } + val (v1Rows, v1Counts) = run(allViaV1, extraConfs) + val (v2Rows, v2Counts) = run(viaV2(format), extraConfs) + // Baseline with the merge rule excluded, to prove the merge actually fired. + val (_, offCounts) = run( + allViaV1, extraConfs :+ (SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> MergeSubplans.ruleName)) + withClue(s"[AQE=$enableAQE] V1 vs V2 rows: ") { + assert(v1Rows.map(_.toString).sorted == v2Rows.map(_.toString).sorted) + } + withClue(s"[AQE=$enableAQE] V1 vs V2 (SubqueryExec, ReusedSubqueryExec) counts: ") { + assert(v1Counts == v2Counts) + } + withClue(s"[AQE=$enableAQE] merge did not fire (merged $v1Counts vs off $offCounts): ") { + assert(v1Counts._1 < offCounts._1) + } + } + } + } + + test("V1/V2 parity: many scalar subqueries over one relation, differing in filter and " + + "aggregate") { + // Integration shape: scalar subqueries over a single relation where some share a filter but + // compute different aggregates (strict / column-union leaf merge) and others differ in their + // filter (relaxed / OR-widen leaf merge). 3 buckets x {count, avg(disc), avg(paid)}. + // (This is the structure of TPC-DS q9.) + val data = spark.range(60).selectExpr( + "id + 1 as qty", "cast(id as double) as disc", "cast(id * 2 as double) as paid") + assertV1V2Consistent( + data, + tbl => + s""" + |SELECT + | CASE WHEN (SELECT count(*) FROM $tbl WHERE qty BETWEEN 1 AND 20) > 10 + | THEN (SELECT avg(disc) FROM $tbl WHERE qty BETWEEN 1 AND 20) + | ELSE (SELECT avg(paid) FROM $tbl WHERE qty BETWEEN 1 AND 20) END b1, + | CASE WHEN (SELECT count(*) FROM $tbl WHERE qty BETWEEN 21 AND 40) > 10 + | THEN (SELECT avg(disc) FROM $tbl WHERE qty BETWEEN 21 AND 40) + | ELSE (SELECT avg(paid) FROM $tbl WHERE qty BETWEEN 21 AND 40) END b2, + | CASE WHEN (SELECT count(*) FROM $tbl WHERE qty BETWEEN 41 AND 60) > 10 + | THEN (SELECT avg(disc) FROM $tbl WHERE qty BETWEEN 41 AND 60) + | ELSE (SELECT avg(paid) FROM $tbl WHERE qty BETWEEN 41 AND 60) END b3 + """.stripMargin, + extraConfs = Seq( + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.FILE_SCAN_MERGE_IGNORE_PUSHED_DATA_FILTERS.key -> "true")) + } + + // ---- Format coverage ------------------------------------------------------- + // One representative query (differing columns + differing filter, exercising both the strict + // column-union and relaxed OR-widen leaf paths) run across every built-in file format, to + // confirm each concrete FileScan inherits the merge. Parquet is included so the same shape is + // checked uniformly. CSV/Text are omitted because their all-string round-trip is not + // type-faithful for these aggregates; their scans share the identical FileScan merge path. + Seq("parquet", "orc", "json").foreach { fmt => + test(s"V1/V2 parity, format coverage ($fmt): differing columns and filter") { + val data = spark.range(100).selectExpr( + "id as a", "id * 2 as b", "cast(id % 5 as int) as c") + assertV1V2Consistent( + data, + tbl => + s""" + |SELECT + | (SELECT sum(a) FROM $tbl WHERE c = 1), + | (SELECT sum(b) FROM $tbl WHERE c = 2) + """.stripMargin, + extraConfs = Seq( + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.FILE_SCAN_MERGE_IGNORE_PUSHED_DATA_FILTERS.key -> "true"), + format = fmt) + } + } + + test("V1/V2 parity: merge subqueries differing only in projected columns") { + val data = spark.range(100).selectExpr("id as a", "id * 2 as b", "id % 5 as c") + assertV1V2Parity( + data, + tbl => + s""" + |SELECT + | (SELECT sum(a) FROM $tbl WHERE c = 1), + | (SELECT sum(b) FROM $tbl WHERE c = 1) + """.stripMargin, + expected = Row( + (0 until 100).filter(_ % 5 == 1).sum, + (0 until 100).filter(_ % 5 == 1).map(_ * 2L).sum) :: Nil, + expectedSubqueries = 1, + expectedReused = 1) + } + + test("V1/V2 parity: merge subqueries differing in filter (symmetric propagation)") { + val data = spark.range(100).selectExpr("id as a", "id % 5 as c") + assertV1V2Parity( + data, + tbl => + s""" + |SELECT + | (SELECT sum(a) FROM $tbl WHERE c = 1), + | (SELECT sum(a) FROM $tbl WHERE c = 2) + """.stripMargin, + expected = Row( + (0 until 100).filter(_ % 5 == 1).sum, + (0 until 100).filter(_ % 5 == 2).sum) :: Nil, + expectedSubqueries = 1, + expectedReused = 1, + extraConfs = Seq( + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.FILE_SCAN_MERGE_IGNORE_PUSHED_DATA_FILTERS.key -> "true")) + } + + test("V1/V2 parity: merge subqueries differing in BOTH column and filter") { + val data = spark.range(100).selectExpr("id as a", "id * 2 as b", "id % 5 as c") + assertV1V2Parity( + data, + tbl => + s""" + |SELECT + | (SELECT sum(a) FROM $tbl WHERE c = 1), + | (SELECT sum(b) FROM $tbl WHERE c = 2) + """.stripMargin, + expected = Row( + (0 until 100).filter(_ % 5 == 1).sum, + (0 until 100).filter(_ % 5 == 2).map(_ * 2L).sum) :: Nil, + expectedSubqueries = 1, + expectedReused = 1, + extraConfs = Seq( + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.FILE_SCAN_MERGE_IGNORE_PUSHED_DATA_FILTERS.key -> "true")) + } + + test("V1/V2 parity: no merge for filter difference when symmetric propagation off") { + val data = spark.range(100).selectExpr("id as a", "id % 5 as c") + assertV1V2Parity( + data, + tbl => + s""" + |SELECT + | (SELECT sum(a) FROM $tbl WHERE c = 1), + | (SELECT sum(a) FROM $tbl WHERE c = 2) + """.stripMargin, + expected = Row( + (0 until 100).filter(_ % 5 == 1).sum, + (0 until 100).filter(_ % 5 == 2).sum) :: Nil, + expectedSubqueries = 2, + expectedReused = 0, + extraConfs = Seq( + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "false")) + } + + test("V1/V2 parity: merge subqueries differing in partition filter") { + val expP1 = (0 until 100).filter(_ % 5 == 1).sum + val expP2 = (0 until 100).filter(_ % 5 == 2).sum + Seq(false, true).foreach { enableAQE => + Seq("V1" -> allViaV1, "V2" -> viaV2("parquet")).foreach { case (label, v1List) => + withTempPath { path => + val pathStr = path.getAbsolutePath + spark.range(100).selectExpr("id as a", "id % 5 as p") + .write.mode("overwrite").partitionBy("p").parquet(pathStr) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> enableAQE.toString, + SQLConf.USE_V1_SOURCE_LIST.key -> v1List, + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.FILE_SCAN_MERGE_IGNORE_PUSHED_DATA_FILTERS.key -> "true") { + val df = spark.sql( + s""" + |SELECT + | (SELECT sum(a) FROM parquet.`$pathStr` WHERE p = 1), + | (SELECT sum(a) FROM parquet.`$pathStr` WHERE p = 2) + """.stripMargin) + withClue(s"[$label, AQE=$enableAQE] results: ") { + checkAnswer(df, Row(expP1, expP2) :: Nil) + } + } + } + } + } + } + + test("V1/V2 parity: same partition filter, differing data filter (relaxed path, partitioned)") { + // Both subqueries share the partition filter (p = 1) but differ in a data-column filter + // (c = 1 vs c = 2) over PARTITIONED data. This is the relaxed leaf-merge path on partitioned + // input: only the data filter is OR-widened; the (equal) partition filter is pushed as-is. + // Expected: p=1 (odd ids) AND c=1 -> ids 1,11,...,91; p=1 AND c=2 -> ids 7,17,...,97. + val expC1 = (0 until 100).filter(id => id % 2 == 1 && id % 5 == 1).sum + val expC2 = (0 until 100).filter(id => id % 2 == 1 && id % 5 == 2).sum + Seq(false, true).foreach { enableAQE => + withTempPath { path => + val pathStr = path.getAbsolutePath + spark.range(100) + .selectExpr("id as a", "cast(id % 5 as int) as c", "cast(id % 2 as int) as p") + .write.mode("overwrite").partitionBy("p").parquet(pathStr) + def run(v1List: String, extra: Seq[(String, String)]): (Seq[Row], (Int, Int)) = + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> enableAQE.toString, + SQLConf.USE_V1_SOURCE_LIST.key -> v1List, + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.FILE_SCAN_MERGE_IGNORE_PUSHED_DATA_FILTERS.key -> "true") ++ extra): _*) { + val df = spark.sql( + s""" + |SELECT + | (SELECT sum(a) FROM parquet.`$pathStr` WHERE p = 1 AND c = 1), + | (SELECT sum(a) FROM parquet.`$pathStr` WHERE p = 1 AND c = 2) + """.stripMargin) + withClue(s"[AQE=$enableAQE, $v1List] results: ") { + checkAnswer(df, Row(expC1, expC2) :: Nil) + } + (df.collect().toSeq, subqueryCounts(df)) + } + val (_, v1Counts) = run(allViaV1, Seq.empty) + val (_, v2Counts) = run(viaV2("parquet"), Seq.empty) + val (_, offCounts) = + run(allViaV1, Seq(SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> MergeSubplans.ruleName)) + withClue(s"[AQE=$enableAQE] V1 vs V2 merge counts: ") { assert(v1Counts == v2Counts) } + withClue(s"[AQE=$enableAQE] merge did not fire ($v1Counts vs off $offCounts): ") { + assert(v1Counts._1 < offCounts._1) + } + } + } + } + + test("V1/V2 parity: joined aggregate subqueries differing in a join-side filter (through-join)") { + // Integration shape: several count(*) subqueries over the same fact joined to dimensions, + // each differing only in a filter on a DIMENSION (here `td.t_hour`). The differing filter is + // on a join input, so merging requires through-join filter propagation. Under V2 the dimension + // scans differ in their pushed filter (OR-widened by the leaf merge) while the fact/other-dim + // scans merge strictly -- exercising the V2 leaf merge together with join-crossing propagation. + // (This is the structure of TPC-DS q88.) + Seq(false, true).foreach { enableAQE => + withTempDir { dir => + val ssPath = new java.io.File(dir, "ss").getAbsolutePath + val hdPath = new java.io.File(dir, "hd").getAbsolutePath + val tdPath = new java.io.File(dir, "td").getAbsolutePath + // td: t_time_sk 0..29, hours cycling 8/9/10. + spark.range(30).selectExpr("id as t_time_sk", "8 + cast(id % 3 as int) as t_hour") + .write.mode("overwrite").parquet(tdPath) + // hd: demo 0..4, dep_count = demo. + spark.range(5).selectExpr("id as hd_demo_sk", "cast(id as int) as hd_dep_count") + .write.mode("overwrite").parquet(hdPath) + // ss: fact linking to td and hd. + spark.range(300).selectExpr( + "cast(id % 30 as int) as ss_sold_time_sk", "cast(id % 5 as int) as ss_hdemo_sk") + .write.mode("overwrite").parquet(ssPath) + + def run(v1List: String, extra: Seq[(String, String)]): (Seq[String], (Int, Int)) = + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> enableAQE.toString, + SQLConf.USE_V1_SOURCE_LIST.key -> v1List, + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_THROUGH_JOIN_ENABLED.key -> "true", + SQLConf.FILE_SCAN_MERGE_IGNORE_PUSHED_DATA_FILTERS.key -> "true") ++ extra): _*) { + def bucket(hour: Int): String = + s"""(SELECT count(*) FROM parquet.`$ssPath` ss, parquet.`$hdPath` hd, + | parquet.`$tdPath` td + | WHERE ss.ss_hdemo_sk = hd.hd_demo_sk AND ss.ss_sold_time_sk = td.t_time_sk + | AND hd.hd_dep_count = 3 AND td.t_hour = $hour)""".stripMargin + val df = spark.sql( + s"SELECT ${bucket(8)} h8, ${bucket(9)} h9, ${bucket(10)} h10") + (df.collect().map(_.toString).sorted.toSeq, subqueryCounts(df)) + } + val (v1Rows, v1Counts) = run(allViaV1, Seq.empty) + val (v2Rows, v2Counts) = run(viaV2("parquet"), Seq.empty) + val (_, offCounts) = + run(allViaV1, Seq(SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> MergeSubplans.ruleName)) + withClue(s"[AQE=$enableAQE] V1 vs V2 rows: ") { assert(v1Rows == v2Rows) } + withClue(s"[AQE=$enableAQE] V1 vs V2 merge counts: ") { assert(v1Counts == v2Counts) } + withClue(s"[AQE=$enableAQE] merge did not fire (merged $v1Counts vs off $offCounts): ") { + assert(v1Counts._1 < offCounts._1) + } + } + } + } + test("SPARK-56677: Merge scalar subqueries with filter propagation through Join") { // subquery1 has no filter; subquery2 filters on b > 1 (a column from the right side of the join // that is not part of the join condition). Predicate pushdown can only push this filter to