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 a21653a011b34..bc1624bc9aebf 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 @@ -1633,6 +1633,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 " + @@ -7896,6 +7908,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