From 48658c24adace34079805ba17d5a382d46c143a4 Mon Sep 17 00:00:00 2001 From: Peter Toth Date: Sat, 18 Jul 2026 21:47:34 +0200 Subject: [PATCH 1/2] [SQL] Move MergeSubplans and PlanMerger to a new sql/core planmerging package Relocate the MergeSubplans optimizer rule and its PlanMerger helper from sql/catalyst (org.apache.spark.sql.catalyst.optimizer) to a new sql/core package org.apache.spark.sql.execution.planmerging. The ScalarSubqueryReference and NonGroupingAggregateReference placeholder expressions stay in catalyst (MergeSubplansReferences.scala) because BloomFilterMightContain references them. Also move the existing sql/core PlanMergeSuite into the new package as PlanMergingSuite, and move MergeSubplansSuite alongside the rule. --- .../optimizer/MergeSubplansReferences.scala | 65 +++++++++++++++++++ .../spark/sql/execution/SparkOptimizer.scala | 1 + .../planmerging}/MergeSubplans.scala | 45 ++----------- .../execution/planmerging}/PlanMerger.scala | 6 +- .../spark/sql/InjectRuntimeFilterSuite.scala | 2 +- .../planmerging}/MergeSubplansSuite.scala | 2 +- .../planmerging/PlanMergingSuite.scala} | 5 +- 7 files changed, 78 insertions(+), 48 deletions(-) create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplansReferences.scala rename sql/{catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer => core/src/main/scala/org/apache/spark/sql/execution/planmerging}/MergeSubplans.scala (89%) rename sql/{catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer => core/src/main/scala/org/apache/spark/sql/execution/planmerging}/PlanMerger.scala (99%) rename sql/{catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer => core/src/test/scala/org/apache/spark/sql/execution/planmerging}/MergeSubplansSuite.scala (99%) rename sql/core/src/test/scala/org/apache/spark/sql/{PlanMergeSuite.scala => execution/planmerging/PlanMergingSuite.scala} (99%) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplansReferences.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplansReferences.scala new file mode 100644 index 0000000000000..d7f649001d79d --- /dev/null +++ b/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) +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala index 54158d5bb4a94..d15ce3c1f60e3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala +++ b/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( diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplans.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/MergeSubplans.scala similarity index 89% rename from sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplans.scala rename to sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/MergeSubplans.scala index 037abc207298a..45c06f86bd416 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplans.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/MergeSubplans.scala @@ -15,16 +15,16 @@ * limitations under the License. */ -package org.apache.spark.sql.catalyst.optimizer +package org.apache.spark.sql.execution.planmerging import scala.collection.mutable.ArrayBuffer import org.apache.spark.sql.catalyst.expressions._ -import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, CTERelationDef, CTERelationRef, LeafNode, LogicalPlan, OneRowRelation, Project, Subquery, WithCTE} +import org.apache.spark.sql.catalyst.optimizer.{NonGroupingAggregateReference, ScalarSubqueryReference} +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, CTERelationDef, CTERelationRef, LogicalPlan, OneRowRelation, Project, Subquery, WithCTE} import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.catalyst.trees.TreePattern.{AGGREGATE, CTE, NO_GROUPING_AGGREGATE_REFERENCE, SCALAR_SUBQUERY, SCALAR_SUBQUERY_REFERENCE, TreePattern} +import org.apache.spark.sql.catalyst.trees.TreePattern.{AGGREGATE, CTE, NO_GROUPING_AGGREGATE_REFERENCE, SCALAR_SUBQUERY, SCALAR_SUBQUERY_REFERENCE} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.DataType /** * This rule tries to merge multiple subplans that have one row result. This can be either the plan @@ -340,40 +340,3 @@ object MergeSubplans extends Rule[LogicalPlan] { } } -/** - * 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) -} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala similarity index 99% rename from sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala rename to sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala index 1c43f91cee9dc..f216a3349f117 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala @@ -15,7 +15,7 @@ * limitations under the License. */ -package org.apache.spark.sql.catalyst.optimizer +package org.apache.spark.sql.execution.planmerging import scala.collection.mutable @@ -62,8 +62,8 @@ object PlanMerger { // Global counter for generating unique names for propagated filter attributes across all // PlanMerger instances. - private[optimizer] val curId = new java.util.concurrent.atomic.AtomicLong() - private[optimizer] def newId: Long = curId.getAndIncrement() + private[planmerging] val curId = new java.util.concurrent.atomic.AtomicLong() + private[planmerging] def newId: Long = curId.getAndIncrement() } /** diff --git a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala index 53b0132b29cec..85e5461d5a431 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala @@ -19,10 +19,10 @@ package org.apache.spark.sql import org.apache.spark.sql.catalyst.expressions.{Alias, BloomFilterMightContain, Literal} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, BloomFilterAggregate} -import org.apache.spark.sql.catalyst.optimizer.MergeSubplans import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LogicalPlan} import org.apache.spark.sql.execution.{ReusedSubqueryExec, SubqueryExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEPropagateEmptyRelation} +import org.apache.spark.sql.execution.planmerging.MergeSubplans import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{IntegerType, StructType} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplansSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala similarity index 99% rename from sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplansSuite.scala rename to sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala index dfa17f926c5a7..3e63d70e03908 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplansSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala @@ -15,7 +15,7 @@ * limitations under the License. */ -package org.apache.spark.sql.catalyst.optimizer +package org.apache.spark.sql.execution.planmerging import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.dsl.plans._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/PlanMergeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/PlanMergingSuite.scala similarity index 99% rename from sql/core/src/test/scala/org/apache/spark/sql/PlanMergeSuite.scala rename to sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/PlanMergingSuite.scala index e1109f20e6040..7440c5badd679 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/PlanMergeSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/PlanMergingSuite.scala @@ -15,14 +15,15 @@ * limitations under the License. */ -package org.apache.spark.sql +package org.apache.spark.sql.execution.planmerging +import org.apache.spark.sql.Row import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -class PlanMergeSuite extends SharedSparkSession +class PlanMergingSuite extends SharedSparkSession with AdaptiveSparkPlanHelper { import testImplicits._ From dfab083ff4d1794050440c439195a7074368901a Mon Sep 17 00:00:00 2001 From: Peter Toth Date: Sun, 19 Jul 2026 16:25:26 +0200 Subject: [PATCH 2/2] [SPARK-40259][SQL] Support merging equivalent DataSource V2 scans in MergeSubplans Add a generic, Spark-side merge of two DataSource V2 scans of the same table that differ only in projected columns and/or pushed filters, so subplans reading the same V2 table are fused into one scan (as already happens for V1 file sources). A source opts in with the new SupportsScanMerging marker; PlanMerger rebuilds the merged scan by driving the real V2ScanRelationPushDown, so the iterative PartitionPredicate pass runs and the (equal) strict filters come back fully enforced. DataSourceV2ScanRelation.pushedFilters now records the complete set of fully-pushed filters (referencing the relation's pre-pruning output), fixing a pre-existing bug where a filter on a pruned-out column silently vanished from the field. hasBlockingPushdown blocks the merge whenever a pushdown cannot be reproduced by rebuilding the scan: aggregate/join/limit/offset/top-N/sample, and now a fully-pushed non-deterministic filter (dropped from the deterministic-only pushedFilters, so it cannot be re-applied on the rebuilt scan). --- .../connector/read/SupportsScanMerging.java | 49 ++ .../datasources/v2/DataSourceV2Relation.scala | 30 +- ...emoryScanMergingPartitionFilterTable.scala | 87 +++ .../v2/V2ScanRelationPushDown.scala | 104 ++- .../execution/planmerging/PlanMerger.scala | 202 +++++- .../sql/connector/DataSourceV2Suite.scala | 62 +- .../planmerging/DSv2PlanMergingSuite.scala | 121 ++++ .../planmerging/MergeSubplansSuite.scala | 630 +++++++++++++++++- 8 files changed, 1231 insertions(+), 54 deletions(-) create mode 100644 sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsScanMerging.java create mode 100644 sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryScanMergingPartitionFilterTable.scala create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/DSv2PlanMergingSuite.scala 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..971aca12cc354 --- /dev/null +++ b/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. + *

+ * 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. + *

+ * 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. + *

+ * 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 { } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala index 0b7c829939ff7..1562efd8bd95a 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala @@ -158,8 +158,17 @@ 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, @@ -167,10 +176,13 @@ case class DataSourceV2ScanRelation( 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)) ) } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryScanMergingPartitionFilterTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryScanMergingPartitionFilterTable.scala new file mode 100644 index 0000000000000..e7a5079f17869 --- /dev/null +++ b/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() +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala index 14d9f754f95c5..33652889ff14c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala @@ -21,7 +21,7 @@ import java.util.Locale import scala.collection.mutable -import org.apache.spark.{SparkException, SparkIllegalArgumentException} +import org.apache.spark.SparkException import org.apache.spark.internal.LogKeys.{AGGREGATE_FUNCTIONS, COLUMN_NAMES, GROUP_BY_EXPRS, JOIN_CONDITION, JOIN_TYPE, POST_SCAN_FILTERS, PUSHED_FILTERS, RELATION_NAME, RELATION_OUTPUT} import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.expressions.{aggregate, Alias, And, Attribute, AttributeMap, AttributeReference, AttributeSet, Cast, Expression, ExpressionSet, ExprId, IntegerLiteral, Literal, NamedExpression, PredicateHelper, ProjectionOverSchema, SortOrder, SubqueryExpression} @@ -69,6 +69,24 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { } } + /** + * Rebuilds a single scan for a Spark-side scan merge (see + * [[org.apache.spark.sql.connector.read.SupportsScanMerging]]): runs this rule on a synthetic + * `Project(projectList, Filter(conditions, relation))` and returns the resulting + * [[DataSourceV2ScanRelation]]. This lets `PlanMerger` fuse two scans of the same table by + * expressing what the merged scan should project and filter, while the pushdown lifecycle (filter + * translation, column pruning, the iterative PartitionPredicate second pass) stays owned here. + * `conditions` are pushed as a single `Filter` directly above the relation; the caller decides + * which of them must come back fully enforced. + */ + def rebuildScan( + relation: DataSourceV2Relation, + projectList: Seq[NamedExpression], + conditions: Seq[Expression]): Option[DataSourceV2ScanRelation] = { + val child = conditions.reduceOption(And).map(Filter(_, relation)).getOrElse(relation) + apply(Project(projectList, child)).collectFirst { case s: DataSourceV2ScanRelation => s } + } + private def collapseGroupedSumOfCount(plan: LogicalPlan): LogicalPlan = { val excludedRules = SQLConf.get.optimizerExcludedRules.toSeq.flatMap(Utils.stringToSeq) if (excludedRules.contains(CollapseGroupedSumOfCount.ruleName)) { @@ -111,12 +129,21 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { val postScanFilters = postScanFiltersWithoutSubquery ++ normalizedFiltersWithSubquery // Compute the pushed filter expressions: the normalized filters that were fully pushed - // down (i.e., not in postScanFilters). These are stored on the scan relation for - // potential future use in constraint propagation. + // down (i.e., not in postScanFilters). These are stored on the scan relation for potential + // future use in constraint propagation, and are read by a Spark-side scan merge (see + // SupportsScanMerging) to compare and re-enforce a scan's filters. val postScanFilterSet = ExpressionSet(postScanFiltersWithoutSubquery) sHolder.pushedFilterExpressions = normalizedFiltersWithoutSubquery .filterNot(postScanFilterSet.contains) .filter(_.deterministic) + // A non-deterministic filter the source fully enforced is dropped from the + // deterministic-only pushedFilterExpressions above, so a scan merge could neither see nor + // re-apply it. Record it so hasBlockingPushdown blocks the merge. ExpressionSet membership + // can't detect this (a non-deterministic expression never compares equal), so match the + // post-scan filters by identity: a non-deterministic conjunct absent from them was pushed. + sHolder.pushedNonDeterministicFilter = normalizedFiltersWithoutSubquery + .filterNot(_.deterministic) + .exists(f => !postScanFiltersWithoutSubquery.exists(_.fastEquals(f))) logInfo( log""" @@ -826,6 +853,39 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { } } + /** + * Whether a pushdown happened that a Spark-side scan merge (see + * [[org.apache.spark.sql.connector.read.SupportsScanMerging]]) cannot reproduce by rebuilding the + * scan from a fresh `ScanBuilder` and re-applying the pushed filters and pruned columns. When + * true, the resulting [[DataSourceV2ScanRelation]] must not be fused with another scan. + * + * A scan carrying a pushed aggregate, join, or variant extraction is never mergeable; those are + * built by dedicated rules ([[buildScanWithPushedAggregate]] / [[buildScanWithPushedJoin]] / + * [[buildScanWithPushedVariants]]) that mark the scan blocking directly, so this helper does not + * repeat the check. It covers the remaining plain-scan path ([[pruneColumns]]), where the + * non-reproducible pushdowns are: + * - `pushedLimit` -- pushed LIMIT + * - `pushedOffset` -- pushed OFFSET + * - `pushedSample` -- pushed table sample + * - `sortOrders` -- pushed sort order (top-N / ordering) + * - `pushedNonDeterministicFilter` -- a non-deterministic filter the source fully enforced. + * It is dropped from `pushedFilterExpressions` (deterministic-only), so a rebuild could + * neither see nor re-apply it. + * + * Reproducible (re-applied by the merge), so NOT blocking: `output` (column pruning, via + * `SupportsPushDownRequiredColumns`) and `pushedPredicates` / `pushedFilterExpressions` + * (deterministic filters, re-pushed via `SupportsPushDownV2Filters`). + * + * When a new pushdown capability adds state to [[ScanBuilderHolder]], classify it here (or, if it + * has its own build rule, mark that scan blocking there). + */ + private def hasBlockingPushdown(holder: ScanBuilderHolder): Boolean = + holder.pushedLimit.isDefined || + holder.pushedOffset.isDefined || + holder.pushedSample.isDefined || + holder.sortOrders.nonEmpty || + holder.pushedNonDeterministicFilter + def buildScanWithPushedAggregate(plan: LogicalPlan): LogicalPlan = plan.transform { case holder: ScanBuilderHolder if holder.pushedAggregate.isDefined => // No need to do column pruning because only the aggregate columns are used as @@ -838,7 +898,9 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { val wrappedScan = getWrappedScan(scan, holder) // Note: holder.pushedFilterExpressions is not propagated here because the output schema // changes to aggregate columns. When validConstraints is wired up, this needs revisiting. - val scanRelation = DataSourceV2ScanRelation(holder.relation, wrappedScan, realOutput) + val scanRelation = DataSourceV2ScanRelation(holder.relation, wrappedScan, realOutput, + // A pushed aggregate cannot be reproduced by rebuilding the scan, so it is never mergeable. + hasMergeBlockingPushdown = true) val projectList = realOutput.zip(holder.output).map { case (a1, a2) => // The data source may return columns with arbitrary data types and it's safer to cast them // to the expected data type. @@ -857,7 +919,9 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { val wrappedScan = getWrappedScan(scan, holder) // Note: holder.pushedFilterExpressions is not propagated here because the output schema // changes with pushed join. When validConstraints is wired up, this needs revisiting. - val scanRelation = DataSourceV2ScanRelation(holder.relation, wrappedScan, realOutput) + val scanRelation = DataSourceV2ScanRelation(holder.relation, wrappedScan, realOutput, + // A pushed join cannot be reproduced by rebuilding the scan, so it is never mergeable. + hasMergeBlockingPushdown = true) // When join is pushed down, the real output is going to be, for example, // SALARY_01234#0, NAME_ab123#1, DEPT_cd123#2. @@ -885,7 +949,9 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { val wrappedScan = getWrappedScan(scan, holder) // Note: holder.pushedFilterExpressions is not propagated here because the output schema // changes with variant extraction. When validConstraints is wired up, this needs revisiting. - val scanRelation = DataSourceV2ScanRelation(holder.relation, wrappedScan, realOutput) + val scanRelation = DataSourceV2ScanRelation(holder.relation, wrappedScan, realOutput, + // Pushed variant extraction cannot be reproduced by rebuilding the scan; never mergeable. + hasMergeBlockingPushdown = true) // Create projection to map real output to expected output (with transformed types) val outputProjection = realOutput.zip(holder.output).map { case (realAttr, expectedAttr) => @@ -941,19 +1007,16 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { case projectionOverSchema(newExpr) => newExpr } - // Remap pushed filter attributes to the pruned output schema and drop filters - // whose references are no longer in the pruned output. Catch FIELD_NOT_FOUND - // because ProjectionOverSchema throws when a pushed filter references a nested - // struct field that was pruned from the schema. - val remappedPushedFilters = sHolder.pushedFilterExpressions.flatMap { filter => - try Some(projectionFunc(filter)) - catch { - case e: SparkIllegalArgumentException if e.getCondition == "FIELD_NOT_FOUND" => - None - } - }.filter(_.references.subsetOf(AttributeSet(output))) + // Record the fully-pushed filter expressions on the scan relation, keeping their references + // to the relation's (pre-pruning) output. These include filters on columns that were pruned + // out of the scan output -- e.g. an unselected partition column the source still enforces + // internally. PlanMerger needs this complete set to compare two scans' filters and to + // re-enforce them when it rebuilds a merged scan; remapping to the pruned output (as the + // post-scan filters below are) would silently drop pruned-out filter columns and make the + // merge unsound. See DataSourceV2ScanRelation.pushedFilters. val scanRelation = DataSourceV2ScanRelation(sHolder.relation, wrappedScan, output, - pushedFilters = remappedPushedFilters) + pushedFilters = sHolder.pushedFilterExpressions, + hasMergeBlockingPushdown = hasBlockingPushdown(sHolder)) val finalFilters = normalizedFilters.map(projectionFunc) // bottom-most filters are put in the left of the list. @@ -1192,6 +1255,11 @@ case class ScanBuilderHolder( var pushedVariants: Option[VariantInRelation] = None var pushedFilterExpressions: Seq[Expression] = Seq.empty + + // Set when a non-deterministic filter was fully pushed to the source. Such a filter is dropped + // from `pushedFilterExpressions` (deterministic-only), so a Spark-side scan merge cannot see or + // re-apply it; `hasBlockingPushdown` treats this as merge-blocking. + var pushedNonDeterministicFilter: Boolean = false } // A wrapper for v1 scan to carry the translated filters and the handled ones, along with diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala index f216a3349f117..a66071a039be3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala @@ -19,11 +19,13 @@ package org.apache.spark.sql.execution.planmerging 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, AttributeSet, 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.connector.read.SupportsScanMerging +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation, V2ScanRelationPushDown} import org.apache.spark.sql.internal.SQLConf /** @@ -327,9 +329,11 @@ class PlanMerger( // Comparing the canonicalized form is required to ignore different forms of the same // expression. if (mappedNPCondition.canonicalized == cp.condition.canonicalized) { - // Identical conditions: the filter node itself adds no new discrimination between - // the two sides, so we keep it unchanged and pass the child's mappings up. - val mergedPlan = Filter(cp.condition, mergedChild) + // Identical conditions: the filter node adds no new discrimination between the two + // sides, so keep it unchanged. But if it sits above a merged DSv2 scan, recover the + // row-group pruning for its condition on that scan (Phase 2). + val prunedChild = rePushMergedScan(mergedChild, Some(cp.condition)) + val mergedPlan = Filter(cp.condition, prunedChild) Some(TryMergeResult(mergedPlan, npMapping, npFilter, cpFilter)) } else if (filterPropagationSupported && symmetricFilterPropagationEnabled) { if (cp.getTagValue(PlanMerger.MERGED_FILTER_TAG).isDefined) { @@ -353,20 +357,32 @@ class PlanMerger( case a: Alias if a.child.canonicalized == newNPCondition.canonicalized => a.toAttribute } - existingNPFilter match { + val (newProjectList, newCondition, newNPFilterOut) = existingNPFilter match { case Some(reusedFilter) => - val newFilter = cp.withNewChildren(Seq(mergedChild)) - Some(TryMergeResult(newFilter, npMapping, Some((reusedFilter, false)), None)) + // np matches an existing side: no new alias, OR condition unchanged. + (childProject.projectList, cp.condition, (reusedFilter, false)) case None => val newNPFilterAlias = Alias(newNPCondition, s"propagatedFilter_${PlanMerger.newId}")() - val newNPFilter = newNPFilterAlias.toAttribute - val newProject = childProject.copy( - projectList = childProject.projectList ++ Seq(newNPFilterAlias)) - val newFilter = Filter(Or(cp.condition, newNPFilter), newProject) - newFilter.copyTagsFrom(cp) - Some(TryMergeResult(newFilter, npMapping, Some((newNPFilter, true)), None)) + (childProject.projectList :+ newNPFilterAlias, + Or(cp.condition, newNPFilterAlias.toAttribute): Expression, + (newNPFilterAlias.toAttribute, true)) } + // Phase 2: the leaf re-merge rebuilt the scan with strict filters only, dropping + // the OR pruning established in earlier rounds, so re-establish it here from ALL + // propagated conditions, not just the new side's. Only the aliases the OR + // condition references are filter sides; other aliases in the Project are + // computed columns, not filters. + val conditions = newProjectList.collect { + case a: Alias if newCondition.references.contains(a.toAttribute) => a.child + } + val prunedChild = + rePushMergedScan(childProject.child, conditions.reduceOption(Or)) + val newProject = childProject.copy( + projectList = newProjectList, child = prunedChild) + val newFilter = Filter(newCondition, newProject) + newFilter.copyTagsFrom(cp) + Some(TryMergeResult(newFilter, npMapping, Some(newNPFilterOut), None)) } else { // First-time filter propagation: alias both sides' conditions as boolean // attributes in a new Project below the Filter, and set the Filter condition @@ -377,6 +393,13 @@ class PlanMerger( val newNPCondition = npFilter.fold(mappedNPCondition) { case (f, _) => And(f, mappedNPCondition) } val newCPCondition = cpFilter.fold(cp.condition)(And(_, cp.condition)) + // Phase 2: the OR-widen moves both conditions into a boolean Project above the + // merged scan, so the scan itself would read the full table. Recover row-group + // pruning by re-pushing OR(np condition, cp condition) onto the merged DSv2 scan + // (the Filter above still enforces exactness). The pruning is derived from the + // scan-level conditions, not the propagated filter attributes in newNP/newCP. + val prunedChild = rePushMergedScan( + mergedChild, Some(Or(mappedNPCondition, cp.condition))) val newNPFilterAlias = Alias(newNPCondition, s"propagatedFilter_${PlanMerger.newId}")() val newCPFilterAlias = @@ -384,8 +407,8 @@ class PlanMerger( val newNPFilter = newNPFilterAlias.toAttribute val newCPFilter = newCPFilterAlias.toAttribute val project = Project( - mergedChild.output.toList ++ Seq(newNPFilterAlias, newCPFilterAlias), - mergedChild) + prunedChild.output.toList ++ Seq(newNPFilterAlias, newCPFilterAlias), + prunedChild) val newFilter = Filter(Or(newNPFilter, newCPFilter), project) newFilter.copyTagsFrom(cp) newFilter.setTagValue(PlanMerger.MERGED_FILTER_TAG, ()) @@ -483,11 +506,160 @@ class PlanMerger( case _ => None } + case (np: DataSourceV2ScanRelation, cp: DataSourceV2ScanRelation) => + tryMergeScanRelations(np, cp) + // Otherwise merging is not possible. case _ => None }) } + /** + * Whether the two scans pushed semantically equal filters. The two scans reference their columns + * with different `ExprId`s, so the new plan's pushed filters are first remapped onto the cached + * plan's attributes by column name before comparing. The remap is built from the relations' full + * output (not the scans' pruned output) because `pushedFilters` may reference columns pruned out + * of the scan output (e.g. an unselected partition column). The caller has already checked the + * relations are canonically equal, so every column maps by name. Trivially true when neither side + * pushed any filter (the common best-effort case). + */ + private def samePushedFilters( + np: DataSourceV2ScanRelation, + cp: DataSourceV2ScanRelation): Boolean = { + val cpOutputByName = cp.relation.output.map(a => a.name -> a).toMap + val npToCp = AttributeMap[Attribute]( + np.relation.output.flatMap(a => cpOutputByName.get(a.name).map(a -> _))) + ExpressionSet(np.pushedFilters.map(mapAttributes(_, npToCp))) == ExpressionSet(cp.pushedFilters) + } + + /** + * Phase 1 of the DSv2 scan merge: fuse two scans of the same table that differ only in projected + * columns (and carry the same strict pushed filters) into a single scan reading the union of + * their columns. The connector opts in via + * [[org.apache.spark.sql.connector.read.SupportsScanMerging]]; Spark runs the real DSv2 pushdown + * ([[V2ScanRelationPushDown]]) on a synthetic `Filter` over the relation, extracts the merged + * scan, and verifies the (equal) strict filters remain fully enforced. The merged relation's + * `pushedFilters` (computed by that pushdown) let Phase 2 (the [[Filter]] merge) re-push + * them when it rebuilds the scan for row-group pruning. A differing post-scan filter is handled + * by that Phase 2 propagation, not here. Phase 1 has no fallback, so any anomaly (a strict filter + * the rebuilt scan does not fully enforce, an unexpected output schema) results in `None` (no + * merge): it must be correct on its own. + */ + private def tryMergeScanRelations( + np: DataSourceV2ScanRelation, + cp: DataSourceV2ScanRelation): Option[TryMergeResult] = { + val mergeable = + // Same table, options, catalog and identifier: the relation's canonical form covers all of + // these (options compares by content via `CaseInsensitiveStringMap.equals`). + np.relation.canonicalized == cp.relation.canonicalized && + // Neither side carries a pushdown a rebuilt scan cannot reproduce (aggregate, join, + // variant extraction, limit, offset, top-N or sample). + !np.hasMergeBlockingPushdown && !cp.hasMergeBlockingPushdown && + // Reported partitioning/ordering (e.g. storage-partitioned join, reported sort) is not + // reconstructed by the rebuilt scan, so decline the merge rather than silently drop it. + // Merging these can be added as a follow-up. + np.keyGroupedPartitioning.isEmpty && cp.keyGroupedPartitioning.isEmpty && + np.ordering.isEmpty && cp.ordering.isEmpty && + // Both scans opt in to Spark-side merging. + np.scan.isInstanceOf[SupportsScanMerging] && cp.scan.isInstanceOf[SupportsScanMerging] && + // Both pushed the same strict filters, so re-pushing reproduces both sides' row sets. + samePushedFilters(np, cp) + + if (!mergeable) { + return None + } + + val relation = cp.relation + val cpNames = cp.output.map(_.name).toSet + val npOnly = np.output.filterNot(a => cpNames.contains(a.name)) + // cp columns keep cp's exprIds; np-only columns keep np's exprIds. + val unionAttrs = cp.output ++ npOnly + + // Build the merged scan enforcing the (equal) strict filters over the union of columns. No + // extra pruning here -- that is Phase 2's job, once the post-scan Filter conditions are known. + buildMergedScan(relation, unionAttrs, cp.pushedFilters, pruning = None).flatMap { scan => + // The rebuilt scan reuses the relation's exprIds, which match cp's for the shared columns, so + // cp's references stay valid without remapping; np's references are remapped via npMapping. + val scanOutByName = scan.output.map(a => a.name -> a).toMap + val npMapping = + AttributeMap[Attribute](np.output.flatMap(a => scanOutByName.get(a.name).map(a -> _))) + if (npMapping.size == np.output.size) { + Some(TryMergeResult(scan, npMapping)) + } else { + None + } + } + } + + /** + * Rebuilds the merged DSv2 scan via [[V2ScanRelationPushDown.rebuildScan]], projecting `columns` + * and filtering by `strict` (plus best-effort `pruning`). This reuses the production pushdown end + * to end -- the same filter translation, column pruning, determinism/subquery handling and + * iterative PartitionPredicate second pass -- rather than reimplementing a slice of it here. + * + * `strict` filters must come back fully enforced (present in the rebuilt scan's `pushedFilters`); + * otherwise `None`, because nothing above the leaf re-checks it. `pruning` conditions + * are best-effort, but only sound ones are offered to the source: a pruning condition is dropped + * unless it is deterministic (a non-deterministic predicate the source prunes on would drop rows + * the enclosing Filter cannot recover) and references only the relation's own columns (propagated + * boolean filter attributes are not columns of the relation). + */ + private def buildMergedScan( + relation: DataSourceV2Relation, + columns: Seq[AttributeReference], + strict: Seq[Expression], + pruning: Option[Expression]): Option[DataSourceV2ScanRelation] = { + val relationOut = AttributeSet(relation.output) + // Strict filters must be expressible over the relation; if not, we cannot rebuild safely. + if (!strict.forall(_.references.subsetOf(relationOut))) { + return None + } + val relOutByName = relation.output.map(a => a.name -> a).toMap + val projectList = columns.flatMap(a => relOutByName.get(a.name)) + if (projectList.size != columns.size) { + return None + } + // strict is enforced; pruning (a single best-effort OR condition from the caller) is dropped + // unless it is both deterministic and expressible over the relation. A non-deterministic + // pruning predicate is unsound here: the source may prune rows using its own evaluation (e.g. a + // separate rand() draw), and the enclosing Filter re-checks exactness with a different + // evaluation, so pruned rows would be lost. Such pruning is dropped wholesale rather than + // weakened -- best-effort pruning degrades gracefully. Pruning that references non-relation + // attributes (propagated boolean filter aliases) is likewise dropped. + val conds = strict ++ + pruning.filter(p => p.deterministic && p.references.subsetOf(relationOut)) + V2ScanRelationPushDown.rebuildScan(relation, projectList, conds).filter { scan => + // Every intended-strict filter must be fully enforced by the rebuilt scan (nothing above + // re-checks it), and the scan must produce exactly the requested union of columns. + val pushedSet = ExpressionSet(scan.pushedFilters) + strict.forall(pushedSet.contains) && + scan.output.length == columns.length && + columns.forall(a => scan.output.exists(_.name == a.name)) + } + } + + /** + * Phase 2 of the DSv2 scan merge: recover the row-group pruning the OR-widen would lose. + * Per-side post-scan conditions move into a boolean Project above the merged scan, so the scan + * itself would read the full table. This walks through Project nodes to the merged + * [[org.apache.spark.sql.connector.read.SupportsScanMerging]] scan and rebuilds it (via the real + * pushdown) with `pruning` pushed as a best-effort filter on top of the strict filters it already + * enforces; the enclosing Filter still enforces exactness. Phase 2 never aborts the merge: if the + * rebuild cannot keep the strict filters strict, it leaves the Phase 1 scan in place -- correct, + * just without the extra pruning. `pruning` is None when there is nothing to prune on. + */ + private def rePushMergedScan( + mergedChild: LogicalPlan, + pruning: Option[Expression]): LogicalPlan = mergedChild match { + case scanRel: DataSourceV2ScanRelation if scanRel.scan.isInstanceOf[SupportsScanMerging] => + buildMergedScan(scanRel.relation, scanRel.output, scanRel.pushedFilters, pruning) + .map(rebuilt => scanRel.copy(scan = rebuilt.scan)) + .getOrElse(scanRel) + case p: Project => + p.withNewChildren(Seq(rePushMergedScan(p.child, pruning))) + case other => other + } + // 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/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala index c6f3c1c3886c6..55035eb69eeb6 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala @@ -1271,33 +1271,34 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper "pushedFilters should not contain the unsupported filter on column j") } - test("pushedFilters are remapped by ProjectionOverSchema after nested schema pruning") { + test("pushedFilters keep the relation-level struct type after nested schema pruning") { val df = spark.read.format(classOf[NestedSchemaDataSourceV2].getName).load() // NestedSchemaScanBuilder pushes GreaterThan on "s.a". - // Selecting only s.a triggers nested schema pruning: s goes from struct to struct. + // Selecting only s.a triggers nested schema pruning: the scan output narrows s to struct, + // but pushedFilters records the fully-pushed filter against the relation's pre-pruning schema, + // so its struct column keeps the full type struct. val q = df.select($"s.a").filter($"s.a" > 3) checkAnswer(q, (4 until 10).map(i => Row(i))) val scanRelation = getScanRelation(q) assert(scanRelation.pushedFilters.nonEmpty, "pushedFilters should be non-empty") - // Find the struct attribute referenced by the pushed filter. - // Before remapping it would have type struct; after remapping, struct. val structAttrs = scanRelation.pushedFilters .flatMap(_.collect { case a: AttributeReference if a.name == "s" => a }) assert(structAttrs.nonEmpty, "pushed filter should reference struct column s") - val prunedStructType = structAttrs.head.dataType.asInstanceOf[StructType] - assert(prunedStructType.fieldNames.toSeq == Seq("a"), - s"struct column in pushed filter should be pruned to struct but was $prunedStructType") + val structType = structAttrs.head.dataType.asInstanceOf[StructType] + assert(structType.fieldNames.toSeq == Seq("a", "b"), + s"struct column in pushed filter should keep type struct but was $structType") } - test("pushedFilters drops filters referencing pruned nested struct fields") { + test("pushedFilters keep filters referencing pruned nested struct fields") { // Disable constraint propagation so IsNotNull(s.a) is not added as a post-scan // filter (it would keep field a alive in the struct). withSQLConf(SQLConf.CONSTRAINT_PROPAGATION_ENABLED.key -> "false") { val df = spark.read.format(classOf[NestedSchemaDataSourceV2].getName).load() - // Filter on s.a but select only s.b. Column pruning narrows s to struct, - // so the pushed filter on s.a can't be remapped and should be dropped. + // Filter on s.a but select only s.b. Column pruning narrows the scan output's s to struct, + // but pushedFilters keeps the fully-pushed filter on s.a (against the relation schema) so a + // later scan merge can re-enforce it. val q = df.filter($"s.a" > 3).select($"s.b") checkAnswer(q, (4 until 10).map(i => Row(-i))) @@ -1306,8 +1307,8 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper filter.collect { case a: AttributeReference if a.name == "s" => a } .flatMap(_.dataType.asInstanceOf[StructType].fieldNames) } - assert(!referencedStructFields.contains("a"), - "pushedFilters should not reference pruned nested field a") + assert(referencedStructFields.contains("a"), + "pushedFilters should keep the filter referencing nested field a") } } @@ -1336,6 +1337,32 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper "Canonicalized instances with equivalent pushedFilters should be equal") } + test("scan canonicalization distinguishes different pushedFilters") { + // Negative counterpart to the test above: two scans of the same relation whose pushedFilters + // differ must NOT canonicalize equal. This inequality is what lets MergeSubplans tell the two + // scans apart; if they compared equal it could treat them as identical and apply one side's + // filter to both -- a wrong-answer bug. + val table = new SimpleDataSourceV2().getTable(CaseInsensitiveStringMap.empty()) + + val relation1 = DataSourceV2Relation.create( + table, None, None, CaseInsensitiveStringMap.empty()) + val relation2 = DataSourceV2Relation.create( + table, None, None, CaseInsensitiveStringMap.empty()) + val scan1 = relation1.table.asReadable.newScanBuilder(relation1.options).build() + val scan2 = relation2.table.asReadable.newScanBuilder(relation2.options).build() + + val filter1 = CatalystGreaterThan(relation1.output.head, CatalystLiteral(3)) + val filter2 = CatalystGreaterThan(relation2.output.head, CatalystLiteral(5)) + + val scanRelation1 = DataSourceV2ScanRelation(relation1, scan1, relation1.output, + pushedFilters = Seq(filter1)) + val scanRelation2 = DataSourceV2ScanRelation(relation2, scan2, relation2.output, + pushedFilters = Seq(filter2)) + + assert(scanRelation1.canonicalized != scanRelation2.canonicalized, + "Canonicalized instances with different pushedFilters must not be equal") + } + test("pushedFilters excludes non-deterministic filters") { val df = spark.read.format(classOf[AdvancedDataSourceV2].getName).load() // i > 3 is pushable and deterministic; rand() > 0.5 is non-deterministic and not pushable. @@ -1376,20 +1403,23 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper "non-deterministic filter should be retained as a post-scan Filter") } - test("pushedFilters drops filters referencing pruned columns") { + test("pushedFilters keep filters referencing pruned columns") { // Disable constraint propagation so IsNotNull(i) is not added (it would keep // column i in the scan output). This simulates a connector that pushes IsNotNull. withSQLConf(SQLConf.CONSTRAINT_PROPAGATION_ENABLED.key -> "false") { val df = spark.read.format(classOf[AdvancedDataSourceV2].getName).load() - // i > 3 is fully pushed; selecting only j causes column pruning to drop i. + // i > 3 is fully pushed; selecting only j causes column pruning to drop i from the scan + // output, but pushedFilters keeps the fully-pushed filter on i (against the relation schema) + // so a later scan merge can re-enforce it. val q = df.filter($"i" > 3).select($"j") checkAnswer(q, (4 until 10).map(i => Row(-i))) val scanRelation = getScanRelation(q) assert(!scanRelation.output.exists(_.name == "i"), "column i should be pruned from scan output") - assert(scanRelation.pushedFilters.isEmpty, - "pushedFilters should drop filters referencing pruned columns") + val referencedCols = scanRelation.pushedFilters.flatMap(_.references.map(_.name)).toSet + assert(referencedCols.contains("i"), + "pushedFilters should keep the filter referencing the pruned column i") } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/DSv2PlanMergingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/DSv2PlanMergingSuite.scala new file mode 100644 index 0000000000000..1ece21e2a521d --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/DSv2PlanMergingSuite.scala @@ -0,0 +1,121 @@ +/* + * 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.execution.planmerging + +import org.scalatest.BeforeAndAfter + +import org.apache.spark.sql.{DataFrame, QueryTest, Row} +import org.apache.spark.sql.connector.FakeV2ProviderWithCustomSchema +import org.apache.spark.sql.connector.catalog.InMemoryScanMergingPartitionFilterCatalog +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation +import org.apache.spark.sql.test.SharedSparkSession + +/** + * End-to-end test for the DSv2 scan-merge gap this change closes: a source whose (equal) filter is + * strict only via the iterative PartitionPredicate second pass. When [[PlanMerger]] rebuilds the + * merged scan it drives the real + * [[org.apache.spark.sql.execution.datasources.v2.V2ScanRelationPushDown]] (Approach 2), so the + * second pass runs and the filter comes back strict -- the merge proceeds. A hand-rolled single + * push-predicates call would only run the first pass, see the filter come back post-scan, + * mis-classify it as non-strict and decline the merge (leaving two scans). + */ +class DSv2PlanMergingSuite extends QueryTest with SharedSparkSession + with BeforeAndAfter { + + private val v2Source = classOf[FakeV2ProviderWithCustomSchema].getName + private val tbl = "scanmerge.t" + + before { + spark.conf.set("spark.sql.catalog.scanmerge", + classOf[InMemoryScanMergingPartitionFilterCatalog].getName) + } + + after { + spark.sessionState.catalogManager.reset() + spark.conf.unset("spark.sql.catalog.scanmerge") + } + + private def v2Scans(df: DataFrame): Seq[DataSourceV2ScanRelation] = + df.queryExecution.optimizedPlan.collectWithSubqueries { + case s: DataSourceV2ScanRelation => s + } + + test("SPARK-40259: merge two DSv2 scans whose filter is strict only via the second pass") { + withTable(tbl) { + sql(s"CREATE TABLE $tbl (part_col string, c1 int, c2 int) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"INSERT INTO $tbl VALUES ('a', 1, 10), ('a', 2, 20), ('b', 3, 30)") + + // `part_col IN ('a')` is untranslatable/returned in the first pass and only accepted (strict) + // in the iterative PartitionPredicate second pass. The two scalar subqueries differ only in + // their projected data column, so PlanMerger fuses them into a single scan reading {c1, c2}. + val df = sql( + s""" + |SELECT + | (SELECT max(c1) FROM $tbl WHERE part_col IN ('a')) AS m1, + | (SELECT max(c2) FROM $tbl WHERE part_col IN ('a')) AS m2 + |""".stripMargin) + + // Correctness: the merged scan must still enforce the partition filter. Reading all + // partitions would give (3, 30) instead of the filtered (2, 20). + checkAnswer(df, Row(2, 20)) + + // The merged subquery is referenced once per scalar subquery, so the logical plan duplicates + // it (physical planning reuses it). Dedupe by canonical form: a successful merge leaves a + // single distinct scan reading the union of both columns; declining would leave two distinct + // scans, one per column. + val scans = v2Scans(df) + assert(scans.nonEmpty, s"expected a DSv2 scan:\n${df.queryExecution.optimizedPlan}") + assert(scans.map(_.canonicalized).distinct.length == 1, + s"the two scans should be fused into one:\n${df.queryExecution.optimizedPlan}") + val scan = scans.head + assert(scan.output.map(_.name).toSet == Set("c1", "c2"), + s"the merged scan should read the union of both columns; got ${scan.output}") + // The filter must be re-enforced strictly on the merged scan (present in pushedFilters), + // which only happens because the rebuild runs the iterative second pass. + assert(scan.pushedFilters.exists(_.references.exists(_.name == "part_col")), + s"the part_col filter should be re-pushed strict onto the merged scan; " + + s"got pushedFilters=${scan.pushedFilters.mkString("[", ", ", "]")}") + } + } + + test("SPARK-40259: do not merge DSv2 scans with different strict partition filters") { + withTable(tbl) { + sql(s"CREATE TABLE $tbl (part_col string, c1 int, c2 int) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"INSERT INTO $tbl VALUES ('a', 1, 10), ('a', 2, 20), ('b', 3, 30)") + + // Same table and shape, but different (strict, fully-enforced) partition filters. The strict + // filters are unequal, so the two scans must NOT fuse -- otherwise one side's partition would + // be read for both, giving a wrong answer. + val df = sql( + s""" + |SELECT + | (SELECT max(c1) FROM $tbl WHERE part_col IN ('a')) AS m1, + | (SELECT max(c2) FROM $tbl WHERE part_col IN ('b')) AS m2 + |""".stripMargin) + + checkAnswer(df, Row(2, 30)) + + val scans = v2Scans(df) + assert(scans.map(_.canonicalized).distinct.length == 2, + s"scans with different partition filters must not be fused:\n" + + df.queryExecution.optimizedPlan) + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala index 3e63d70e03908..0b34c2e24fbd2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala @@ -19,11 +19,18 @@ package org.apache.spark.sql.execution.planmerging import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.dsl.plans._ -import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, CreateNamedStruct, GetStructField, If, Literal, Or, ScalarSubquery} +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Ascending, Attribute, CreateNamedStruct, Expression, ExprId, GetStructField, If, Literal, Or, ScalarSubquery, SortOrder} import org.apache.spark.sql.catalyst.plans._ import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.rules._ +import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes +import org.apache.spark.sql.connector.catalog.{SupportsRead, Table, TableCapability} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownLimit, SupportsPushDownRequiredColumns, SupportsPushDownV2Filters, SupportsScanMerging} +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation, V2ScanRelationPushDown} import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap class MergeSubplansSuite extends PlanTest { @@ -1966,4 +1973,625 @@ class MergeSubplansSuite extends PlanTest { comparePlans(Optimize.execute(originalQuery.analyze), originalQuery.analyze) } } + + // ---- SPARK-40259: generic DSv2 scan merge ---- + + private val v2Table = new TestV2Table(StructType(Seq( + StructField("a", IntegerType), StructField("b", IntegerType), StructField("c", StringType)))) + + /** A `DataSourceV2ScanRelation` over `table` projecting only the given columns. */ + private def v2ScanReadingOn(table: TestV2Table, cols: Seq[String]): DataSourceV2ScanRelation = { + val fullOutput = toAttributes(table.schema()) + val relation = + DataSourceV2Relation(table, fullOutput, None, None, CaseInsensitiveStringMap.empty()) + val output = cols.map(c => fullOutput.find(_.name == c).get) + val scan = TestV2Scan(StructType(output.map(a => StructField(a.name, a.dataType, a.nullable)))) + DataSourceV2ScanRelation(relation, scan, output) + } + + /** A `DataSourceV2ScanRelation` over [[v2Table]] projecting only the given columns. */ + private def v2ScanReading(cols: String*): DataSourceV2ScanRelation = + v2ScanReadingOn(v2Table, cols) + + /** Like [[v2ScanReading]] but the scan does NOT implement `SupportsScanMerging`. */ + private def v2ScanReadingNoMerge(cols: String*): DataSourceV2ScanRelation = { + val s = v2ScanReading(cols: _*) + s.copy(scan = TestV2ScanNoMerge(s.scan.readSchema())) + } + + private def v2Scans(plan: LogicalPlan): Seq[DataSourceV2ScanRelation] = + plan.collectWithSubqueries { case s: DataSourceV2ScanRelation => s } + + /** + * Normalizes a merged plan so `comparePlans` can match it: (1) drops each DSv2 scan's + * dynamically-built Scan to a schema-only placeholder (the pushed describe() strings are asserted + * separately), and (2) resets the nested DataSourceV2Relation's output exprIds to a positional + * scheme. The reset is needed because `DataSourceV2ScanRelation` is a `LeafNode`, so its + * `relation` is a constructor arg rather than a child -- PlanTest's `normalizeExprIds` never + * recurses into it, and those exprIds otherwise differ between the expected and actual plans. + */ + private def normalizeScans(plan: LogicalPlan): LogicalPlan = plan.transformWithSubqueries { + case s: DataSourceV2ScanRelation => + val fixedRelation = s.relation.copy( + output = s.relation.output.zipWithIndex.map { case (a, i) => a.withExprId(ExprId(i)) }) + s.copy(relation = fixedRelation, scan = TestV2Scan(s.scan.readSchema())) + } + + // Normalize merged DSv2 scans before the standard comparison (a no-op on plans without them), + // so tests can call plain comparePlans. See normalizeScans for why this is needed. + override protected def comparePlans( + plan1: LogicalPlan, + plan2: LogicalPlan, + checkAnalysis: Boolean = true): Unit = + super.comparePlans(normalizeScans(plan1), normalizeScans(plan2), checkAnalysis) + + test("SPARK-40259: merge DSv2 scans that differ only in projected columns") { + val sub1 = ScalarSubquery(v2ScanReading("a").groupBy()(sum($"a").as("sum_a"))) + val sub2 = ScalarSubquery(v2ScanReading("b").groupBy()(sum($"b").as("sum_b"))) + val originalQuery = testRelation.select(sub1, sub2) + + // Expected: the two scans fuse into one reading {a, b}, feeding a single aggregate. No filters, + // so no OR-widen propagation. + val mergedScan = v2ScanReadingOn(v2Table, Seq("a", "b")) + val mergedSubquery = mergedScan + .groupBy()(sum($"a").as("sum_a"), sum($"b").as("sum_b")) + .select(CreateNamedStruct(Seq( + Literal("sum_a"), $"sum_a", + Literal("sum_b"), $"sum_b")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze) + } + + test("SPARK-40259: do not merge DSv2 scans when a pushdown is merge-blocking") { + val sub1 = ScalarSubquery(v2ScanReading("a").groupBy()(sum($"a").as("sum_a"))) + val blockingScan = v2ScanReading("b").copy(hasMergeBlockingPushdown = true) + val sub2 = ScalarSubquery(blockingScan.groupBy()(sum($"b").as("sum_b"))) + val originalQuery = testRelation.select(sub1, sub2) + + // A merge-blocking pushdown declines the merge: the plan is left unchanged. + comparePlans(Optimize.execute(originalQuery.analyze), originalQuery.analyze) + } + + test("SPARK-40259: merge DSv2 scans with different filters via OR-widen propagation") { + val sub1 = ScalarSubquery( + v2ScanReading("a").where($"a" > 1).groupBy()(sum($"a").as("sum_a"))) // cp + val sub2 = ScalarSubquery( + v2ScanReading("b").where($"b" > 2).groupBy()(sum($"b").as("sum_b"))) // np + val originalQuery = testRelation.select(sub1, sub2) + + // Expected: the two scans fuse into one reading {a, b}; the differing filters become + // propagatedFilter aliases OR-widened in a Filter, and each side's aggregate carries its + // filter as a FILTER clause (np = sub2 gets id 0, cp = sub1 gets id 1). + val mergedScan = v2ScanReadingOn(v2Table, Seq("a", "b")) + val npFilterAlias = Alias($"b" > 2, "propagatedFilter_0")() + val cpFilterAlias = Alias($"a" > 1, "propagatedFilter_1")() + val npFilter = npFilterAlias.toAttribute + val cpFilter = cpFilterAlias.toAttribute + val mergedSubquery = mergedScan + .select(mergedScan.output ++ Seq(npFilterAlias, cpFilterAlias): _*) + .where(Or(npFilter, cpFilter)) + .groupBy()( + sum($"a", Some(cpFilter)).as("sum_a"), + sum($"b", Some(npFilter)).as("sum_b")) + .select(CreateNamedStruct(Seq( + Literal("sum_a"), $"sum_a", + Literal("sum_b"), $"sum_b")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + // Additional: the merged scan re-pushes OR(a > 1, b > 2) for row-group pruning, which the + // structural comparison above normalizes out. + val pushed = v2Scans(optimized).head.scan.asInstanceOf[TestV2Scan].pushed + assert(pushed.nonEmpty && pushed.mkString.contains("a") && pushed.mkString.contains("b"), + s"expected an OR pruning predicate over a and b pushed to the merged scan, got: $pushed") + } + } + + test("SPARK-40259: merge proceeds without pruning when the source rejects the pruning") { + val rejecting = new TestV2Table( + StructType(Seq(StructField("a", IntegerType), StructField("b", IntegerType))), + acceptsFilters = false) + val s1 = v2ScanReadingOn(rejecting, Seq("a")) + val s2 = v2ScanReadingOn(rejecting, Seq("b")) + val sub1 = ScalarSubquery(s1.where(s1.output.head > 1).groupBy()(sum($"a").as("sum_a"))) + val sub2 = ScalarSubquery(s2.where(s2.output.head > 2).groupBy()(sum($"b").as("sum_b"))) + val originalQuery = testRelation.select(sub1, sub2) + + // Expected: the same OR-widen merge as the accepting case -- only the pushed pruning differs + // (a rejecting source records none). np = sub2 -> id 0, cp = sub1 -> id 1. + val mergedScan = v2ScanReadingOn(rejecting, Seq("a", "b")) + val npFilterAlias = Alias($"b" > 2, "propagatedFilter_0")() + val cpFilterAlias = Alias($"a" > 1, "propagatedFilter_1")() + val npFilter = npFilterAlias.toAttribute + val cpFilter = cpFilterAlias.toAttribute + val mergedSubquery = mergedScan + .select(mergedScan.output ++ Seq(npFilterAlias, cpFilterAlias): _*) + .where(Or(npFilter, cpFilter)) + .groupBy()( + sum($"a", Some(cpFilter)).as("sum_a"), + sum($"b", Some(npFilter)).as("sum_b")) + .select(CreateNamedStruct(Seq( + Literal("sum_a"), $"sum_a", + Literal("sum_b"), $"sum_b")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + // Phase 2 degrades gracefully: the merge happens; a rejecting source records no pruning. + assert(v2Scans(optimized).head.scan.asInstanceOf[TestV2Scan].pushed.isEmpty, + "a rejecting source must not record any pushed pruning predicate") + } + } + + test("SPARK-40259: do not merge when a strict pushed filter is not re-enforced on rebuild") { + // The scans claim a strict filter on "a" (in pushedFilters), but the source enforces nothing + // strictly (empty strictColumns -> everything is best-effort). Re-pushing the "strict" filter + // would leave it merely best-effort with nothing above to re-check it, so the merge must abort. + val bestEffort = new TestV2Table(StructType(Seq( + StructField("a", IntegerType), StructField("b", IntegerType), StructField("c", StringType)))) + val s1 = v2ScanReadingOn(bestEffort, Seq("a", "b")) + val s2 = v2ScanReadingOn(bestEffort, Seq("a", "c")) + val sub1 = ScalarSubquery( + s1.copy(pushedFilters = Seq(s1.output.find(_.name == "a").get > 0)) + .groupBy()(sum($"b").as("sum_b"))) + val sub2 = ScalarSubquery( + s2.copy(pushedFilters = Seq(s2.output.find(_.name == "a").get > 0)) + .groupBy()(sum($"c").as("max_c"))) + val originalQuery = testRelation.select(sub1, sub2) + + // A strict filter the rebuilt scan cannot re-enforce declines the merge: plan left unchanged. + comparePlans(Optimize.execute(originalQuery.analyze), originalQuery.analyze) + } + + test("SPARK-40259: do not merge DSv2 scans that do not opt into SupportsScanMerging") { + val sub1 = ScalarSubquery(v2ScanReadingNoMerge("a").groupBy()(sum($"a").as("sum_a"))) + val sub2 = ScalarSubquery(v2ScanReadingNoMerge("b").groupBy()(sum($"b").as("sum_b"))) + val originalQuery = testRelation.select(sub1, sub2) + + // Scans that do not implement SupportsScanMerging decline the merge: plan left unchanged. + comparePlans(Optimize.execute(originalQuery.analyze), originalQuery.analyze) + } + + test("SPARK-40259: merge DSv2 scans with identical filters re-pushes the pruning") { + // Identical post-scan filter (a > 1), differing columns: the two scans fuse and the identical + // condition is re-pushed to the merged scan for row-group pruning (Phase 2, single condition). + val sub1 = ScalarSubquery( + v2ScanReading("a", "b").where($"a" > 1).groupBy()(sum($"b").as("sum_b"))) + val sub2 = ScalarSubquery( + v2ScanReading("a", "c").where($"a" > 1).groupBy()(sum($"c").as("sum_c"))) + val originalQuery = testRelation.select(sub1, sub2) + + // Expected: the identical filter a > 1 kept as a single Filter above the merged scan {a, b, c}; + // no OR-widen (the conditions match). The aggregates read b and c. + val mergedScan = v2ScanReadingOn(v2Table, Seq("a", "b", "c")) + val mergedSubquery = mergedScan.where($"a" > 1) + .groupBy()(sum($"b").as("sum_b"), sum($"c").as("sum_c")) + .select(CreateNamedStruct(Seq( + Literal("sum_b"), $"sum_b", + Literal("sum_c"), $"sum_c")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + // Additional: the identical a > 1 is re-pushed to the merged scan for pruning. + val pushed = v2Scans(optimized).head.scan.asInstanceOf[TestV2Scan].pushed + assert(pushed.nonEmpty && pushed.mkString.contains("a"), + s"expected the identical filter a > 1 to be re-pushed to the merged scan, got: $pushed") + } + + test("SPARK-40259: merge three DSv2 scans with differing filters re-pushes the full OR") { + // Three-way merge exercises the tagged (MERGED_FILTER_TAG) branch: the leaf re-merge rebuilds + // the scan strict-only each round, so the tagged branch must re-establish the full 3-way OR. + val t = new TestV2Table(StructType(Seq(StructField("a", IntegerType), + StructField("b", IntegerType), StructField("d", IntegerType)))) + val sub1 = ScalarSubquery( + v2ScanReadingOn(t, Seq("a")).where($"a" > 1).groupBy()(sum($"a").as("s1"))) + val sub2 = ScalarSubquery( + v2ScanReadingOn(t, Seq("b")).where($"b" > 2).groupBy()(sum($"b").as("s2"))) + val sub3 = ScalarSubquery( + v2ScanReadingOn(t, Seq("d")).where($"d" > 3).groupBy()(sum($"d").as("s3"))) + val originalQuery = testRelation.select(sub1, sub2, sub3) + + // Expected: one merged scan reading {a, b, d}; step 1 merges sub1 (cp) + sub2 (np) -> + // propagatedFilter_0 = b > 2 (np), _1 = a > 1 (cp); step 2 merges sub3 (np) -> _2 = d > 3, + // extending the OR. Each aggregate carries its side's FILTER. + val mergedScan = v2ScanReadingOn(t, Seq("a", "b", "d")) + val npFilter0Alias = Alias($"b" > 2, "propagatedFilter_0")() + val cpFilter0Alias = Alias($"a" > 1, "propagatedFilter_1")() + val npFilter1Alias = Alias($"d" > 3, "propagatedFilter_2")() + val npFilter0 = npFilter0Alias.toAttribute + val cpFilter0 = cpFilter0Alias.toAttribute + val npFilter1 = npFilter1Alias.toAttribute + val mergedSubquery = mergedScan + .select(mergedScan.output ++ Seq(npFilter0Alias, cpFilter0Alias, npFilter1Alias): _*) + .where(Or(Or(npFilter0, cpFilter0), npFilter1)) + .groupBy()( + sum($"a", Some(cpFilter0)).as("s1"), + sum($"b", Some(npFilter0)).as("s2"), + sum($"d", Some(npFilter1)).as("s3")) + .select(CreateNamedStruct(Seq( + Literal("s1"), $"s1", + Literal("s2"), $"s2", + Literal("s3"), $"s3")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1), + extractorExpression(0, analyzedMergedSubquery.output, 2)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + // Additional: the merged scan carries the full 3-way OR pruning over a, b and d. + val pushed = v2Scans(optimized).head.scan.asInstanceOf[TestV2Scan].pushed.mkString + assert(pushed.contains("a") && pushed.contains("b") && pushed.contains("d"), + s"expected an OR pruning over a, b and d pushed to the merged scan, got: $pushed") + } + } + + test("SPARK-40259: a non-deterministic filter conjunct is not pushed as a pruning predicate") { + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + // Bare non-grouping aggregates (joined) go through the Aggregate extraction path, which -- + // unlike the ScalarSubquery path -- does not gate on determinism, so the OR-widen pruning + // predicate carries a non-deterministic conjunct into the scan-merge Phase 2. + val agg1 = v2ScanReading("a").where(($"a" > 1) && (rand(0) < Literal(0.5))) + .groupBy()(sum($"a").as("s1")) + val agg2 = v2ScanReading("b").where($"b" > 2).groupBy()(sum($"b").as("s2")) + val originalQuery = agg1.join(agg2) + + val scans = v2Scans(Optimize.execute(originalQuery.analyze)) + assert(scans.length == 1, "the scans should still be fused") + val pushed = scans.head.scan.asInstanceOf[TestV2Scan].pushed.mkString + // Pruning predicate `(a > 1 AND rand() < 0.5) OR (b > 2)` is non-deterministic as a whole. + // A source that prunes on it uses its own rand() draw, while the enclosing Filter re-checks + // exactness with a different draw, so pruned rows would be lost. Best-effort pruning degrades + // gracefully: the predicate is dropped rather than weakened, so nothing is pushed and no + // rand() reaches the source. The merge itself is unaffected. + assert(!pushed.contains("RAND"), + s"the non-deterministic rand() conjunct must not be pushed as pruning, got: $pushed") + assert(pushed.isEmpty, + s"a non-deterministic pruning predicate must be dropped wholesale, got: $pushed") + } + } + + test("SPARK-40259: merge DSv2 scans that pushed the same strict filters (re-push path)") { + // A source that fully enforces (strict) predicates on column "a". + val strictOnA = new TestV2Table( + StructType(Seq(StructField("a", IntegerType), StructField("b", IntegerType), + StructField("c", StringType))), + strictColumns = Set("a")) + val s1 = v2ScanReadingOn(strictOnA, Seq("a", "b")) + val s2 = v2ScanReadingOn(strictOnA, Seq("a", "c")) + val sub1 = ScalarSubquery( + s1.copy(pushedFilters = Seq(s1.output.find(_.name == "a").get > 0)) + .groupBy()(sum($"b").as("sum_b"))) + val sub2 = ScalarSubquery( + s2.copy(pushedFilters = Seq(s2.output.find(_.name == "a").get > 0)) + .groupBy()(sum($"c").as("max_c"))) + val originalQuery = testRelation.select(sub1, sub2) + + // Expected: one merged scan reading {a, b, c} that keeps the strict a > 0 in pushedFilters; the + // two aggregates read b and c. No post-scan filter, so no OR-widen. + val mergedScan0 = v2ScanReadingOn(strictOnA, Seq("a", "b", "c")) + val mergedScan = mergedScan0.copy( + pushedFilters = Seq(mergedScan0.output.find(_.name == "a").get > 0)) + val mergedSubquery = mergedScan + .groupBy()(sum($"b").as("sum_b"), sum($"c").as("max_c")) + .select(CreateNamedStruct(Seq( + Literal("sum_b"), $"sum_b", + Literal("max_c"), $"max_c")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + // Additional: the strict a > 0 is re-enforced on the merged scan (normalized out above). + assert(v2Scans(optimized).head.pushedFilters.nonEmpty, + "the merged relation must carry the strict pushed filters") + } + + test("SPARK-40259: merge scans with equal strict filters and differing post-scan filters") { + // Mixed source: "p" is fully enforced (strict), "a"/"b" are best-effort. Both sides push + // the same strict filter p > 0 and carry a differing post-scan filter (a > 1 / b > 2). The + // scans fuse on the equal strict p; the differing post-scan filters OR-widen above the + // merged scan, so the rebuild re-pushes p > 0 (strict) and (a > 1 OR b > 2) (pruning) at + // once -- the only path that drives buildMergedScan with both non-empty. + val mixed = new TestV2Table( + StructType(Seq(StructField("p", IntegerType), StructField("a", IntegerType), + StructField("b", IntegerType))), + strictColumns = Set("p")) + val s1 = v2ScanReadingOn(mixed, Seq("p", "a")) + val s2 = v2ScanReadingOn(mixed, Seq("p", "b")) + val sub1 = ScalarSubquery( + s1.copy(pushedFilters = Seq(s1.output.find(_.name == "p").get > 0)) + .where($"a" > 1).groupBy()(sum($"a").as("s1"))) + val sub2 = ScalarSubquery( + s2.copy(pushedFilters = Seq(s2.output.find(_.name == "p").get > 0)) + .where($"b" > 2).groupBy()(sum($"b").as("s2"))) + val originalQuery = testRelation.select(sub1, sub2) + + // Expected: one merged scan reading {p, a, b} that keeps p > 0 strict (pushedFilters), with + // the differing post-scan filters OR-widened above it (np = sub2 -> id 0, cp = sub1 -> id 1). + val mergedScan0 = v2ScanReadingOn(mixed, Seq("p", "a", "b")) + val mergedScan = mergedScan0.copy( + pushedFilters = Seq(mergedScan0.output.find(_.name == "p").get > 0)) + val npFilterAlias = Alias($"b" > 2, "propagatedFilter_0")() + val cpFilterAlias = Alias($"a" > 1, "propagatedFilter_1")() + val npFilter = npFilterAlias.toAttribute + val cpFilter = cpFilterAlias.toAttribute + val mergedSubquery = mergedScan + .select(mergedScan.output ++ Seq(npFilterAlias, cpFilterAlias): _*) + .where(Or(npFilter, cpFilter)) + .groupBy()( + sum($"a", Some(cpFilter)).as("s1"), + sum($"b", Some(npFilter)).as("s2")) + .select(CreateNamedStruct(Seq( + Literal("s1"), $"s1", + Literal("s2"), $"s2")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + // Additional: p stays strict (in pushedFilters) and the OR (a, b) is re-pushed as pruning, + // both normalized out of the structural comparison above. + val mergedResult = v2Scans(optimized).head + assert(mergedResult.pushedFilters.exists(_.references.exists(_.name == "p")), + s"the strict filter on p must stay enforced; got ${mergedResult.pushedFilters}") + val pushed = mergedResult.scan.asInstanceOf[TestV2Scan].pushed.mkString + assert(pushed.contains("a") && pushed.contains("b"), + s"the differing post-scan filters must be re-pushed as OR pruning, got: $pushed") + } + } + + test("SPARK-40259: a fully pushed non-deterministic filter blocks merging") { + // A source that fully enforces every pushed predicate. V2ScanRelationPushDown drops a + // non-deterministic filter from pushedFilters (that field is deterministic-only) but records + // it as a merge-blocking pushdown, since a rebuilt merged scan could neither see nor re-apply + // it -- fusing two such scans would silently drop each side's independent rand() draw. A + // deterministic filter, by contrast, lands in pushedFilters and stays mergeable. + val enforcing = new TestV2Table(StructType(Seq( + StructField("a", IntegerType), StructField("b", IntegerType))), enforceAll = true) + def pushDown(mkCondition: DataSourceV2Relation => Expression): DataSourceV2ScanRelation = { + val relation = DataSourceV2Relation( + enforcing, toAttributes(enforcing.schema()), None, None, CaseInsensitiveStringMap.empty()) + V2ScanRelationPushDown(Filter(mkCondition(relation), relation)) + .collectFirst { case s: DataSourceV2ScanRelation => s }.get + } + + val nonDet = pushDown(_ => rand(0) < Literal(0.5)) + assert(nonDet.hasMergeBlockingPushdown, + "a fully pushed non-deterministic filter must block merging") + assert(nonDet.pushedFilters.isEmpty, + "a non-deterministic filter must not appear in the deterministic pushedFilters") + + val det = pushDown(_.output.head > 0) + assert(!det.hasMergeBlockingPushdown, + "a fully pushed deterministic filter must not block merging") + assert(det.pushedFilters.nonEmpty, + "a fully pushed deterministic filter must land in pushedFilters") + } + + test("SPARK-40259: do not merge DSv2 scans that report key-grouped partitioning or ordering") { + // The rebuilt merged scan does not reconstruct reported partitioning/ordering, so a scan + // reporting either declines the merge (checked on both the np and cp side) -- the plan is left + // unchanged -- rather than silently dropping it. Preserving them across a merge is a deferred + // follow-up. (The plain-scan merge is already covered by the projected-columns test above.) + def assertDeclines(withField: DataSourceV2ScanRelation => DataSourceV2ScanRelation): Unit = + Seq( + (withField(v2ScanReading("a")), v2ScanReading("b")), + (v2ScanReading("a"), withField(v2ScanReading("b")))).foreach { case (npScan, cpScan) => + val q = testRelation.select( + ScalarSubquery(npScan.groupBy()(sum($"a").as("sa"))), + ScalarSubquery(cpScan.groupBy()(sum($"b").as("sb")))) + comparePlans(Optimize.execute(q.analyze), q.analyze) + } + + assertDeclines(s => s.copy(keyGroupedPartitioning = Some(Seq(s.output.head)))) + assertDeclines(s => s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending))))) + } + + test("SPARK-40259: merge DSv2 scans reading identical columns but differing filters") { + // Both scans read only "a", so the column union adds nothing (the np-only set is empty); the + // merge is driven purely by the differing filters, and the OR is still re-pushed for pruning. + // (sub1 is cp, sub2 is np.) + val sub1 = ScalarSubquery(v2ScanReading("a").where($"a" > 1).groupBy()(sum($"a").as("s1"))) + val sub2 = ScalarSubquery(v2ScanReading("a").where($"a" > 2).groupBy()(sum($"a").as("s2"))) + val originalQuery = testRelation.select(sub1, sub2) + + // Expected: one merged scan reading just {a} (np-only set empty), the differing filters + // OR-widened above it (np = sub2 -> id 0, cp = sub1 -> id 1). + val mergedScan = v2ScanReadingOn(v2Table, Seq("a")) + val npFilterAlias = Alias($"a" > 2, "propagatedFilter_0")() + val cpFilterAlias = Alias($"a" > 1, "propagatedFilter_1")() + val npFilter = npFilterAlias.toAttribute + val cpFilter = cpFilterAlias.toAttribute + val mergedSubquery = mergedScan + .select(mergedScan.output ++ Seq(npFilterAlias, cpFilterAlias): _*) + .where(Or(npFilter, cpFilter)) + .groupBy()( + sum($"a", Some(cpFilter)).as("s1"), + sum($"a", Some(npFilter)).as("s2")) + .select(CreateNamedStruct(Seq( + Literal("s1"), $"s1", + Literal("s2"), $"s2")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + val pushed = v2Scans(optimized).head.scan.asInstanceOf[TestV2Scan].pushed.mkString + assert(pushed.contains("a"), s"the OR over a should be re-pushed as pruning, got: $pushed") + } + } + + test("SPARK-40259: merge three DSv2 scans that pushed the same strict filter") { + // Round-2 merge feeds the already-merged relation back through samePushedFilters. All three + // scans carry the same strict a > 0, so they fuse into a single scan reading a, b, c, d. + val strictOnA = new TestV2Table( + StructType(Seq(StructField("a", IntegerType), StructField("b", IntegerType), + StructField("c", IntegerType), StructField("d", IntegerType))), + strictColumns = Set("a")) + def strictSub(col: String): ScalarSubquery = { + val s = v2ScanReadingOn(strictOnA, Seq("a", col)) + ScalarSubquery( + s.copy(pushedFilters = Seq(s.output.find(_.name == "a").get > 0)) + .groupBy()(sum(s.output.find(_.name == col).get).as(s"agg_$col"))) + } + val originalQuery = testRelation.select(strictSub("b"), strictSub("c"), strictSub("d")) + + // Expected: one merged scan reading {a, b, c, d} keeping the shared strict a > 0; three + // aggregates read b, c and d. No post-scan filter, so no OR-widen. + val mergedScan0 = v2ScanReadingOn(strictOnA, Seq("a", "b", "c", "d")) + val mergedScan = mergedScan0.copy( + pushedFilters = Seq(mergedScan0.output.find(_.name == "a").get > 0)) + val mergedSubquery = mergedScan + .groupBy()(sum($"b").as("agg_b"), sum($"c").as("agg_c"), sum($"d").as("agg_d")) + .select(CreateNamedStruct(Seq( + Literal("agg_b"), $"agg_b", + Literal("agg_c"), $"agg_c", + Literal("agg_d"), $"agg_d")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1), + extractorExpression(0, analyzedMergedSubquery.output, 2)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + assert(v2Scans(optimized).head.pushedFilters.exists(_.references.exists(_.name == "a")), + "the merged scan must keep the shared strict filter on a") + } + + test("SPARK-40259: a pushed limit blocks merging (hasBlockingPushdown classification)") { + // hasBlockingPushdown flags a pushed limit/offset/sample/sort as merge-blocking; exercise the + // limit term through the real pushdown (the others are analogous). The stub opts into limit + // pushdown, so V2ScanRelationPushDown records pushedLimit on the resulting scan relation. + val relation = DataSourceV2Relation( + v2Table, toAttributes(v2Table.schema()), None, None, CaseInsensitiveStringMap.empty()) + val pushed = V2ScanRelationPushDown(Limit(Literal(1), relation)) + .collectFirst { case s: DataSourceV2ScanRelation => s }.get + assert(pushed.hasMergeBlockingPushdown, + "a pushed limit must be recorded as a merge-blocking pushdown") + } +} + +/** + * Minimal readable DSv2 table whose scan opts into merging, for scan-merge tests. + * + * A pushed predicate that references only columns in `strictColumns` is fully enforced (strict): + * accepted by the source and NOT returned as a post-scan filter. Any other predicate is + * best-effort (stats): accepted for row-group pruning but returned so it is re-checked above the + * scan. When `acceptsFilters` is false the source rejects everything: nothing is accepted and + * every predicate is returned (used to exercise the "no pruning recovery" path). When `enforceAll` + * is true the source fully enforces every accepted predicate (nothing is returned as a post-scan + * filter), including reference-less ones such as a non-deterministic filter. + */ +private class TestV2Table( + tableSchema: StructType, + acceptsFilters: Boolean = true, + strictColumns: Set[String] = Set.empty, + enforceAll: Boolean = false) + extends Table with SupportsRead { + override def name(): String = "test_v2_table" + override def schema(): StructType = tableSchema + override def capabilities(): java.util.Set[TableCapability] = + java.util.Collections.singleton(TableCapability.BATCH_READ) + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = + new TestV2ScanBuilder(tableSchema, acceptsFilters, strictColumns, enforceAll) +} + +private class TestV2ScanBuilder( + tableSchema: StructType, + acceptsFilters: Boolean, + strictColumns: Set[String], + enforceAll: Boolean) + extends ScanBuilder with SupportsPushDownRequiredColumns with SupportsPushDownV2Filters + with SupportsPushDownLimit { + private var prunedSchema: StructType = tableSchema + private var accepted: Array[Predicate] = Array.empty // strict UNION stats (pushedPredicates) + + override def pruneColumns(requiredSchema: StructType): Unit = prunedSchema = requiredSchema + + // Opts into limit pushdown so V2ScanRelationPushDown records a merge-blocking pushedLimit. + override def pushLimit(limit: Int): Boolean = true + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + if (!acceptsFilters) { + predicates // reject everything: nothing accepted, all returned as post-scan + } else if (enforceAll) { + accepted = predicates + Array.empty // fully enforce every accepted predicate, including reference-less ones + } else { + accepted = predicates + predicates.filterNot(isStrict) // stats predicates are re-checked above the scan + } + } + + private def isStrict(p: Predicate): Boolean = + p.references().nonEmpty && + p.references().forall(r => strictColumns.contains(r.fieldNames().mkString("."))) + + override def pushedPredicates(): Array[Predicate] = accepted + override def build(): Scan = TestV2Scan(prunedSchema, accepted.map(_.describe()).toSeq) +} + +/** `pushed` records the `describe()` of the predicates pushed to the scan, for test assertions. */ +private case class TestV2Scan(schema: StructType, pushed: Seq[String] = Seq.empty) + extends Scan with SupportsScanMerging { + override def readSchema(): StructType = schema +} + +private case class TestV2ScanNoMerge(schema: StructType) extends Scan { + override def readSchema(): StructType = schema }