From 7e652bfbf162d4a1b3bba5de06dd6d9e17749a04 Mon Sep 17 00:00:00 2001 From: Yan Yan Date: Fri, 10 Jul 2026 01:31:27 +0000 Subject: [PATCH 1/2] [SPARK][SQL] Add DSv2 relation and row-level command tree patterns Add Catalyst tree-pattern identity bits for the post-pushdown DSv2 scan leaf, the pre-pushdown DSv2 relation, and the row-level command nodes, and use them to prune the optimizer/analyzer rules that match them. - Add DATA_SOURCE_V2_RELATION and DATA_SOURCE_V2_SCAN_RELATION and attach them to DataSourceV2Relation / DataSourceV2ScanRelation. - Add DELETE_FROM_TABLE, UPDATE_TABLE, MERGE_INTO_TABLE, REPLACE_DATA and WRITE_DELTA. Command now exposes a nodePatternsInternal() hook (mirroring PlanExpression) so command nodes combine the shared COMMAND bit with their own identity bit; DeleteFromTable/UpdateTable/MergeIntoTable/ReplaceData/ WriteDelta override it. - Guard the previously-unguarded rules with containsPattern/containsAnyPattern: RewriteDeleteFromTable, RewriteUpdateTable, RewriteMergeIntoTable, V2ScanRelationPushDown, V2ScanPartitioningAndOrdering, GroupBasedRowLevelOperationScanPlanning, OptimizeMetadataOnlyDeleteFromTable, RowLevelOperationRuntimeGroupFiltering. - Add tree-pattern contract tests (DataSourceV2RelationSuite, V2CommandTreePatternSuite). --- .../analysis/RewriteDeleteFromTable.scala | 4 +- .../analysis/RewriteMergeIntoTable.scala | 4 +- .../analysis/RewriteUpdateTable.scala | 4 +- .../sql/catalyst/plans/logical/Command.scala | 7 ++- .../catalyst/plans/logical/v2Commands.scala | 11 ++++ .../sql/catalyst/trees/TreePatterns.scala | 7 +++ .../datasources/v2/DataSourceV2Relation.scala | 5 ++ .../logical/V2CommandTreePatternSuite.scala | 56 +++++++++++++++++++ .../v2/DataSourceV2RelationSuite.scala | 26 +++++++++ ...upBasedRowLevelOperationScanPlanning.scala | 4 +- .../OptimizeMetadataOnlyDeleteFromTable.scala | 4 +- .../v2/V2ScanPartitioningAndOrdering.scala | 7 ++- .../v2/V2ScanRelationPushDown.scala | 7 +++ ...wLevelOperationRuntimeGroupFiltering.scala | 4 +- 14 files changed, 141 insertions(+), 9 deletions(-) create mode 100644 sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/V2CommandTreePatternSuite.scala diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteDeleteFromTable.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteDeleteFromTable.scala index f8881e2077103..b7c4499586513 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteDeleteFromTable.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteDeleteFromTable.scala @@ -20,6 +20,7 @@ package org.apache.spark.sql.catalyst.analysis import org.apache.spark.sql.catalyst.expressions.{Alias, EqualNullSafe, Expression, Literal, Not} import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral import org.apache.spark.sql.catalyst.plans.logical.{DeleteFromTable, Filter, LogicalPlan, Project, ReplaceData, WriteDelta} +import org.apache.spark.sql.catalyst.trees.TreePattern.DELETE_FROM_TABLE import org.apache.spark.sql.catalyst.util.RowDeltaUtils._ import org.apache.spark.sql.connector.catalog.{SupportsDeleteV2, SupportsRowLevelOperations, TruncatableTable} import org.apache.spark.sql.connector.write.{RowLevelOperationTable, SupportsDelta} @@ -37,7 +38,8 @@ import org.apache.spark.sql.util.CaseInsensitiveStringMap */ object RewriteDeleteFromTable extends RewriteRowLevelCommand { - override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators { + override def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperatorsWithPruning( + _.containsPattern(DELETE_FROM_TABLE)) { case d @ DeleteFromTable(aliasedTable, cond) if d.resolved => EliminateSubqueryAliases(aliasedTable) match { case ExtractV2Table(_: TruncatableTable) if cond == TrueLiteral => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala index e11a985d85ddb..682aa4eb3825f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala @@ -24,6 +24,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression import org.apache.spark.sql.catalyst.plans.{FullOuter, Inner, JoinType, LeftAnti, LeftOuter, RightOuter} import org.apache.spark.sql.catalyst.plans.logical.{DeleteAction, Filter, HintInfo, InsertAction, InsertOnlyMerge, Join, JoinHint, LogicalPlan, MergeAction, MergeIntoTable, MergeRows, NO_BROADCAST_AND_REPLICATION, Project, ReplaceData, UpdateAction, WriteDelta} import org.apache.spark.sql.catalyst.plans.logical.MergeRows.{Copy, Delete, Discard, Insert, Instruction, Keep, ROW_ID, Split, Update} +import org.apache.spark.sql.catalyst.trees.TreePattern.MERGE_INTO_TABLE import org.apache.spark.sql.catalyst.util.RowDeltaUtils.{COPY_OPERATION, INSERT_OPERATION, OPERATION_COLUMN, UPDATE_OPERATION} import org.apache.spark.sql.connector.catalog.SupportsRowLevelOperations import org.apache.spark.sql.connector.write.{RowLevelOperationTable, SupportsDelta} @@ -43,7 +44,8 @@ object RewriteMergeIntoTable extends RewriteRowLevelCommand with PredicateHelper private final val ROW_FROM_SOURCE = "__row_from_source" private final val ROW_FROM_TARGET = "__row_from_target" - override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators { + override def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperatorsWithPruning( + _.containsPattern(MERGE_INTO_TABLE)) { // aligned is false when schema evolution is pending (see ResolveRowLevelCommandAssignments) case m @ MergeIntoTable(aliasedTable, source, cond, matchedActions, notMatchedActions, notMatchedBySourceActions, _) if m.resolved && m.rewritable && m.aligned && diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala index 207b3e7217040..352ddc0b0acb1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala @@ -20,6 +20,7 @@ package org.apache.spark.sql.catalyst.analysis import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, EqualNullSafe, Expression, If, Literal, MetadataAttribute, Not, SubqueryExpression} import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral import org.apache.spark.sql.catalyst.plans.logical.{Assignment, Expand, Filter, LogicalPlan, Project, ReplaceData, Union, UpdateTable, WriteDelta} +import org.apache.spark.sql.catalyst.trees.TreePattern.UPDATE_TABLE import org.apache.spark.sql.catalyst.util.RowDeltaUtils._ import org.apache.spark.sql.connector.catalog.SupportsRowLevelOperations import org.apache.spark.sql.connector.write.{RowLevelOperationTable, SupportsDelta} @@ -34,7 +35,8 @@ import org.apache.spark.sql.types.IntegerType */ object RewriteUpdateTable extends RewriteRowLevelCommand { - override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators { + override def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperatorsWithPruning( + _.containsPattern(UPDATE_TABLE)) { case u @ UpdateTable(aliasedTable, assignments, cond) if u.resolved && u.rewritable && u.aligned => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/Command.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/Command.scala index bd277e92d11d2..e054c955cc774 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/Command.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/Command.scala @@ -35,7 +35,12 @@ trait Command extends LogicalPlan { // is created. That said, the statistics of a command is useless. Here we just return a dummy // statistics to avoid unnecessary statistics calculation of command's children. override def stats: Statistics = Statistics.DUMMY - final override val nodePatterns: Seq[TreePattern] = Seq(COMMAND) + // Every command carries the shared COMMAND pattern. Keeping this `final` guarantees no subclass + // can drop COMMAND; subclasses add their own identity pattern(s) via `nodePatternsInternal()`. + final override val nodePatterns: Seq[TreePattern] = Seq(COMMAND) ++ nodePatternsInternal() + + // Subclasses can override this to contribute additional identity tree patterns. + protected def nodePatternsInternal(): Seq[TreePattern] = Seq() } trait LeafCommand extends Command with LeafLike[LogicalPlan] diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala index eed6fd7ca6caf..d031a18201529 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala @@ -26,6 +26,7 @@ import org.apache.spark.sql.catalyst.catalog.CatalogTypes.TablePartitionSpec import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.plans.DescribeCommandSchema import org.apache.spark.sql.catalyst.trees.BinaryLike +import org.apache.spark.sql.catalyst.trees.TreePattern.{DELETE_FROM_TABLE, MERGE_INTO_TABLE, REPLACE_DATA, TreePattern, UPDATE_TABLE, WRITE_DELTA} import org.apache.spark.sql.catalyst.types.DataTypeUtils import org.apache.spark.sql.catalyst.util._ import org.apache.spark.sql.catalyst.util.TypeUtils.{ordinalNumber, toSQLExpr} @@ -447,6 +448,8 @@ case class ReplaceData( override protected def withNewChildInternal(newChild: LogicalPlan): ReplaceData = { copy(query = newChild) } + + override protected def nodePatternsInternal(): Seq[TreePattern] = Seq(REPLACE_DATA) } /** @@ -557,6 +560,8 @@ case class WriteDelta( override protected def withNewChildInternal(newChild: LogicalPlan): WriteDelta = { copy(query = newChild) } + + override protected def nodePatternsInternal(): Seq[TreePattern] = Seq(WRITE_DELTA) } trait V2CreateTableAsSelectPlan @@ -1084,6 +1089,8 @@ case class DeleteFromTable( override def child: LogicalPlan = table override protected def withNewChildInternal(newChild: LogicalPlan): DeleteFromTable = copy(table = newChild) + + override protected def nodePatternsInternal(): Seq[TreePattern] = Seq(DELETE_FROM_TABLE) } /** @@ -1123,6 +1130,8 @@ case class UpdateTable( case r: NamedRelation => r.skipSchemaResolution case _ => false } + + override protected def nodePatternsInternal(): Seq[TreePattern] = Seq(UPDATE_TABLE) } /** @@ -1214,6 +1223,8 @@ case class MergeIntoTable( newRight: LogicalPlan): MergeIntoTable = { copy(targetTable = newLeft, sourceTable = newRight) } + + override protected def nodePatternsInternal(): Seq[TreePattern] = Seq(MERGE_INTO_TABLE) } object MergeIntoTable { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala index 94b4666a88a8d..bd4c36f3802b3 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala @@ -144,6 +144,9 @@ object TreePattern extends Enumeration { val COLLECT_METRICS: Value = Value val COMMAND: Value = Value val CTE: Value = Value + val DATA_SOURCE_V2_RELATION: Value = Value + val DATA_SOURCE_V2_SCAN_RELATION: Value = Value + val DELETE_FROM_TABLE: Value = Value val DESERIALIZE_TO_OBJECT: Value = Value val DF_DROP_COLUMNS: Value = Value val DISTINCT_LIKE: Value = Value @@ -163,6 +166,7 @@ object TreePattern extends Enumeration { val LIMIT: Value = Value val LOCAL_RELATION: Value = Value val LOGICAL_QUERY_STAGE: Value = Value + val MERGE_INTO_TABLE: Value = Value val METRIC_VIEW_PLACEHOLDER: Value = Value val NATURAL_LIKE_JOIN: Value = Value val NEAREST_BY_JOIN: Value = Value @@ -179,6 +183,7 @@ object TreePattern extends Enumeration { val RELATION_TIME_TRAVEL: Value = Value val REPARTITION_OPERATION: Value = Value val REBALANCE_PARTITIONS: Value = Value + val REPLACE_DATA: Value = Value val RESOLVED_METRIC_VIEW: Value = Value val SEQUENTIAL_STREAMING_UNION: Value = Value val SERIALIZE_FROM_OBJECT: Value = Value @@ -189,10 +194,12 @@ object TreePattern extends Enumeration { val UNION: Value = Value val UNPIVOT: Value = Value val UPDATE_EVENT_TIME_WATERMARK_COLUMN: Value = Value + val UPDATE_TABLE: Value = Value val TYPED_FILTER: Value = Value val WINDOW: Value = Value val WINDOW_GROUP_LIMIT: Value = Value val WITH_WINDOW_DEFINITION: Value = Value + val WRITE_DELTA: Value = Value val ZIP: Value = Value // Unresolved Plan patterns (Alphabetically ordered) 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..65e4c05335325 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 @@ -26,6 +26,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, Attri import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.plans.logical.{ColumnStat, ExposesMetadataColumns, Histogram, HistogramBin, LeafNode, LogicalPlan, Statistics} import org.apache.spark.sql.catalyst.streaming.{StreamingSourceIdentifyingName, Unassigned} +import org.apache.spark.sql.catalyst.trees.TreePattern.{DATA_SOURCE_V2_RELATION, DATA_SOURCE_V2_SCAN_RELATION, TreePattern} import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.catalyst.util.{removeInternalMetadata, truncatedString, CharVarcharUtils} import org.apache.spark.sql.connector.catalog.{CatalogPlugin, FunctionCatalog, Identifier, SupportsMetadataColumns, Table, TableCapability, TableCatalog, V2TableUtil} @@ -143,6 +144,8 @@ case class DataSourceV2Relation( table.capabilities.contains(TableCapability.AUTOMATIC_SCHEMA_EVOLUTION) def isVersioned: Boolean = table.version != null + + override val nodePatterns: Seq[TreePattern] = Seq(DATA_SOURCE_V2_RELATION) } /** @@ -187,6 +190,8 @@ case class DataSourceV2ScanRelation( case _ => AttributeSet.empty } + override val nodePatterns: Seq[TreePattern] = Seq(DATA_SOURCE_V2_SCAN_RELATION) + override def name: String = relation.name override def simpleString(maxFields: Int): String = { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/V2CommandTreePatternSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/V2CommandTreePatternSuite.scala new file mode 100644 index 0000000000000..6f0ba676aa6ef --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/V2CommandTreePatternSuite.scala @@ -0,0 +1,56 @@ +/* + * 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.plans.logical + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.catalyst.dsl.expressions._ +import org.apache.spark.sql.catalyst.expressions.Literal +import org.apache.spark.sql.catalyst.trees.TreePattern + +/** + * Pins the tree-pattern identity contract for the DSv2 row-level command nodes: each node carries + * both the shared `COMMAND` bit and its own identity bit, so rules can prune on either. + */ +class V2CommandTreePatternSuite extends SparkFunSuite { + + private val target = LocalRelation($"a".int, $"b".int) + private val source = LocalRelation($"c".int, $"d".int) + + test("DeleteFromTable declares COMMAND and DELETE_FROM_TABLE") { + val plan = DeleteFromTable(target, Literal.TrueLiteral) + assert(plan.containsAllPatterns(TreePattern.COMMAND, TreePattern.DELETE_FROM_TABLE)) + } + + test("UpdateTable declares COMMAND and UPDATE_TABLE") { + val plan = UpdateTable(target, Seq.empty, Some(Literal.TrueLiteral)) + assert(plan.containsAllPatterns(TreePattern.COMMAND, TreePattern.UPDATE_TABLE)) + } + + test("MergeIntoTable declares COMMAND and MERGE_INTO_TABLE") { + val plan = MergeIntoTable( + target, + source, + mergeCondition = $"a" === $"c", + matchedActions = Seq(DeleteAction(None)), + notMatchedActions = Seq(InsertAction(None, + Seq(Assignment($"a", $"c"), Assignment($"b", $"d")))), + notMatchedBySourceActions = Seq(DeleteAction(None)), + withSchemaEvolution = false) + assert(plan.containsAllPatterns(TreePattern.COMMAND, TreePattern.MERGE_INTO_TABLE)) + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala index 46ed870cb3221..0acefec943da1 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala @@ -22,10 +22,12 @@ import java.util import org.apache.spark.SparkFunSuite import org.apache.spark.sql.catalyst.catalog.{CatalogColumnStat, CatalogStatistics} import org.apache.spark.sql.catalyst.plans.logical.{Histogram, HistogramBin} +import org.apache.spark.sql.catalyst.trees.TreePattern import org.apache.spark.sql.catalyst.util.FieldMetadataUtils.FIELD_ID_METADATA_KEY import org.apache.spark.sql.catalyst.util.INTERNAL_METADATA_KEYS import org.apache.spark.sql.connector.catalog.{Column, Table, TableCapability} import org.apache.spark.sql.connector.expressions.FieldReference +import org.apache.spark.sql.connector.read.Scan import org.apache.spark.sql.types.{IntegerType, MetadataBuilder, StringType, StructField, StructType} import org.apache.spark.sql.util.CaseInsensitiveStringMap @@ -146,4 +148,28 @@ class DataSourceV2RelationSuite extends SparkFunSuite { assert(field.id.contains("1")) assert(field.metadata.contains(FIELD_ID_METADATA_KEY)) } + + test("nodePatterns declare the DSv2 relation identity tree patterns") { + val table = new Table { + override def name(): String = "t" + override def columns(): Array[Column] = + Array(Column.create("id", IntegerType)) + override def capabilities(): util.Set[TableCapability] = + util.Set.of[TableCapability]() + } + + val relation = + DataSourceV2Relation.create(table, None, None, CaseInsensitiveStringMap.empty()) + assert(relation.containsPattern(TreePattern.DATA_SOURCE_V2_RELATION)) + + val scan = new Scan { + override def readSchema(): StructType = relation.schema + } + val scanRelation = DataSourceV2ScanRelation(relation, scan, relation.output) + assert(scanRelation.containsPattern(TreePattern.DATA_SOURCE_V2_SCAN_RELATION)) + // The scan leaf must not inherit the pre-pushdown relation's identity pattern from its + // `relation` field (which is a case-class arg, not a child), so pruning on the two bits + // stays distinct. + assert(!scanRelation.containsPattern(TreePattern.DATA_SOURCE_V2_RELATION)) + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala index 0cb7967cb7369..a41aad05d4351 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala @@ -23,6 +23,7 @@ import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral import org.apache.spark.sql.catalyst.planning.{GroupBasedRowLevelOperation, PhysicalOperation} import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan, ReplaceData} import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.REPLACE_DATA import org.apache.spark.sql.connector.expressions.filter.{Predicate => V2Filter} import org.apache.spark.sql.connector.read.ScanBuilder import org.apache.spark.sql.connector.write.RowLevelOperation.Command.MERGE @@ -40,7 +41,8 @@ object GroupBasedRowLevelOperationScanPlanning extends Rule[LogicalPlan] with Pr import DataSourceV2Implicits._ - override def apply(plan: LogicalPlan): LogicalPlan = plan transformDown { + override def apply(plan: LogicalPlan): LogicalPlan = plan.transformDownWithPruning( + _.containsPattern(REPLACE_DATA)) { // push down the filter from the command condition instead of the filter in the rewrite plan, // which is negated for data sources that only support replacing groups of data (e.g. files) case GroupBasedRowLevelOperation(rd: ReplaceData, cond, _, relation: DataSourceV2Relation) => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeMetadataOnlyDeleteFromTable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeMetadataOnlyDeleteFromTable.scala index c02fdec1fab60..5632fd21764f3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeMetadataOnlyDeleteFromTable.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeMetadataOnlyDeleteFromTable.scala @@ -21,6 +21,7 @@ import org.apache.spark.sql.catalyst.expressions.{Expression, PredicateHelper, S import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral import org.apache.spark.sql.catalyst.plans.logical.{DeleteFromTable, DeleteFromTableWithFilters, LogicalPlan, ReplaceData, RowLevelWrite, WriteDelta} import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.{REPLACE_DATA, WRITE_DELTA} import org.apache.spark.sql.connector.catalog.{SupportsDeleteV2, TruncatableTable} import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.connector.write.RowLevelOperation @@ -37,7 +38,8 @@ import org.apache.spark.util.ArrayImplicits._ */ object OptimizeMetadataOnlyDeleteFromTable extends Rule[LogicalPlan] with PredicateHelper { - override def apply(plan: LogicalPlan): LogicalPlan = plan transform { + override def apply(plan: LogicalPlan): LogicalPlan = plan.transformWithPruning( + _.containsAnyPattern(REPLACE_DATA, WRITE_DELTA)) { case RewrittenRowLevelCommand(rowLevelPlan, DELETE, cond, relation: DataSourceV2Relation) => relation.table match { case table: SupportsDeleteV2 if !SubqueryExpression.hasSubquery(cond) => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitioningAndOrdering.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitioningAndOrdering.scala index 479154e7a2852..165425aea19dc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitioningAndOrdering.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitioningAndOrdering.scala @@ -21,6 +21,7 @@ import org.apache.spark.internal.LogKeys.CLASS_NAME import org.apache.spark.sql.catalyst.expressions.V2ExpressionUtils import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.DATA_SOURCE_V2_SCAN_RELATION import org.apache.spark.sql.connector.read.{SupportsReportOrdering, SupportsReportPartitioning} import org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning, UnknownPartitioning} import org.apache.spark.util.ArrayImplicits._ @@ -40,7 +41,8 @@ object V2ScanPartitioningAndOrdering extends Rule[LogicalPlan] with Logging { } } - private def partitioning(plan: LogicalPlan) = plan.transformDown { + private def partitioning(plan: LogicalPlan) = plan.transformDownWithPruning( + _.containsPattern(DATA_SOURCE_V2_SCAN_RELATION)) { case d @ ExtractV2ScanInfo(relation, scan: SupportsReportPartitioning, _) if d.keyGroupedPartitioning.isEmpty => val catalystPartitioning = scan.outputPartitioning() match { @@ -68,7 +70,8 @@ object V2ScanPartitioningAndOrdering extends Rule[LogicalPlan] with Logging { d.copy(keyGroupedPartitioning = catalystPartitioning) } - private def ordering(plan: LogicalPlan) = plan.transformDown { + private def ordering(plan: LogicalPlan) = plan.transformDownWithPruning( + _.containsPattern(DATA_SOURCE_V2_SCAN_RELATION)) { case d @ ExtractV2ScanInfo(relation, scan: SupportsReportOrdering, _) => val ordering = V2ExpressionUtils.toCatalystOrdering(scan.outputOrdering(), relation, relation.funCatalog) 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..7e4a3c36c8193 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 @@ -30,6 +30,7 @@ import org.apache.spark.sql.catalyst.optimizer.{CollapseGroupedSumOfCount, Colla import org.apache.spark.sql.catalyst.planning.{PhysicalOperation, ScanOperation} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LeafNode, Limit, LimitAndOffset, LocalLimit, LogicalPlan, Offset, OffsetAndLimit, Project, Sample, SampleMethod, Sort} import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.DATA_SOURCE_V2_RELATION import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.connector.expressions.{SortOrder => V2SortOrder} import org.apache.spark.sql.connector.expressions.aggregate.{Aggregation, Avg, Count, CountStar, Max, Min, Sum} @@ -48,6 +49,12 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { import DataSourceV2Implicits._ def apply(plan: LogicalPlan): LogicalPlan = { + // All pushdown sub-rules enter through a `DataSourceV2Relation` (createScanBuilder), so skip + // the whole fold when the plan has no such relation. + if (!plan.containsPattern(DATA_SOURCE_V2_RELATION)) { + return plan + } + val pushdownRules = Seq[LogicalPlan => LogicalPlan] ( createScanBuilder, pushDownSample, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala index 9f8409efa360e..9b7035fd5b92e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala @@ -24,6 +24,7 @@ import org.apache.spark.sql.catalyst.optimizer.RewritePredicateSubquery import org.apache.spark.sql.catalyst.planning.{DeltaBasedRowLevelOperation, GroupBasedRowLevelOperation} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LogicalPlan, RowLevelWrite} import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.{REPLACE_DATA, WRITE_DELTA} import org.apache.spark.sql.connector.read.SupportsRuntimeV2Filtering import org.apache.spark.sql.connector.write.RowLevelOperation.Command.{DELETE, MERGE, UPDATE} import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Implicits, DataSourceV2Relation, DataSourceV2ScanRelation, ExtractV2Scan} @@ -49,7 +50,8 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla import DataSourceV2Implicits._ - override def apply(plan: LogicalPlan): LogicalPlan = plan transformDown { + override def apply(plan: LogicalPlan): LogicalPlan = plan.transformDownWithPruning( + _.containsAnyPattern(REPLACE_DATA, WRITE_DELTA)) { case GroupBasedRowLevelOperation(replaceData, _, Some(cond), ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) if canInjectGroupFilters(cond, scan) => injectGroupFilters(replaceData, cond, scan) From 818095512cbabbaa7b42a102df2f48339090c7ce Mon Sep 17 00:00:00 2001 From: Yan Yan Date: Fri, 10 Jul 2026 22:44:01 +0000 Subject: [PATCH 2/2] [SPARK-58082][SQL] Address review: drop V2ScanRelationPushDown guard, add ReplaceData/WriteDelta tests - Remove the DATA_SOURCE_V2_RELATION pruning guard from V2ScanRelationPushDown. The guard was too aggressive: pushDownSample also validates SYSTEM sampling that cannot be pushed down (e.g. over a DSv1 relation, which has no DataSourceV2Relation) and throws UNSUPPORTED_FEATURE.TABLESAMPLE_SYSTEM_NO_SCAN. With the guard, that validation was skipped and DSv1 SYSTEM sampling surfaced a SparkException at execution instead of the expected AnalysisException, breaking DataSourceV2TableSampleSuite "TABLESAMPLE SYSTEM on DSv1 table errors". Restore the rule to its original unconditional fold. The DATA_SOURCE_V2_RELATION bit stays available for other rules (e.g. V2ScanPartitioningAndOrdering). - Add ReplaceData and WriteDelta cases to V2CommandTreePatternSuite, pinning the COMMAND + REPLACE_DATA / WRITE_DELTA identity bits that GroupBasedRowLevelOperationScanPlanning, OptimizeMetadataOnlyDeleteFromTable and RowLevelOperationRuntimeGroupFiltering prune on. --- .../logical/V2CommandTreePatternSuite.scala | 41 +++++++++++++++++++ .../v2/V2ScanRelationPushDown.scala | 7 ---- 2 files changed, 41 insertions(+), 7 deletions(-) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/V2CommandTreePatternSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/V2CommandTreePatternSuite.scala index 6f0ba676aa6ef..3b20eae1a75ed 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/V2CommandTreePatternSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/V2CommandTreePatternSuite.scala @@ -17,10 +17,18 @@ package org.apache.spark.sql.catalyst.plans.logical +import java.util + import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.catalyst.ProjectingInternalRow import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.expressions.Literal import org.apache.spark.sql.catalyst.trees.TreePattern +import org.apache.spark.sql.catalyst.util.{ReplaceDataProjections, WriteDeltaProjections} +import org.apache.spark.sql.connector.catalog.{Column, Table, TableCapability} +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation +import org.apache.spark.sql.types.{IntegerType, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap /** * Pins the tree-pattern identity contract for the DSv2 row-level command nodes: each node carries @@ -31,6 +39,19 @@ class V2CommandTreePatternSuite extends SparkFunSuite { private val target = LocalRelation($"a".int, $"b".int) private val source = LocalRelation($"c".int, $"d".int) + // A minimal `NamedRelation` for the row-level write nodes, whose `nodePatternsInternal()` + // returns a compile-time constant, so no resolution is needed to pin the identity bit. + private val v2Relation: DataSourceV2Relation = { + val table = new Table { + override def name(): String = "t" + override def columns(): Array[Column] = Array(Column.create("a", IntegerType)) + override def capabilities(): util.Set[TableCapability] = util.Set.of[TableCapability]() + } + DataSourceV2Relation.create(table, None, None, CaseInsensitiveStringMap.empty()) + } + + private val emptyRow = ProjectingInternalRow(new StructType(), IndexedSeq.empty) + test("DeleteFromTable declares COMMAND and DELETE_FROM_TABLE") { val plan = DeleteFromTable(target, Literal.TrueLiteral) assert(plan.containsAllPatterns(TreePattern.COMMAND, TreePattern.DELETE_FROM_TABLE)) @@ -53,4 +74,24 @@ class V2CommandTreePatternSuite extends SparkFunSuite { withSchemaEvolution = false) assert(plan.containsAllPatterns(TreePattern.COMMAND, TreePattern.MERGE_INTO_TABLE)) } + + test("ReplaceData declares COMMAND and REPLACE_DATA") { + val plan = ReplaceData( + v2Relation, + Literal.TrueLiteral, + source, + v2Relation, + ReplaceDataProjections(emptyRow, None)) + assert(plan.containsAllPatterns(TreePattern.COMMAND, TreePattern.REPLACE_DATA)) + } + + test("WriteDelta declares COMMAND and WRITE_DELTA") { + val plan = WriteDelta( + v2Relation, + Literal.TrueLiteral, + source, + v2Relation, + WriteDeltaProjections(None, emptyRow, None)) + assert(plan.containsAllPatterns(TreePattern.COMMAND, TreePattern.WRITE_DELTA)) + } } 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 7e4a3c36c8193..14d9f754f95c5 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 @@ -30,7 +30,6 @@ import org.apache.spark.sql.catalyst.optimizer.{CollapseGroupedSumOfCount, Colla import org.apache.spark.sql.catalyst.planning.{PhysicalOperation, ScanOperation} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LeafNode, Limit, LimitAndOffset, LocalLimit, LogicalPlan, Offset, OffsetAndLimit, Project, Sample, SampleMethod, Sort} import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.catalyst.trees.TreePattern.DATA_SOURCE_V2_RELATION import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.connector.expressions.{SortOrder => V2SortOrder} import org.apache.spark.sql.connector.expressions.aggregate.{Aggregation, Avg, Count, CountStar, Max, Min, Sum} @@ -49,12 +48,6 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { import DataSourceV2Implicits._ def apply(plan: LogicalPlan): LogicalPlan = { - // All pushdown sub-rules enter through a `DataSourceV2Relation` (createScanBuilder), so skip - // the whole fold when the plan has no such relation. - if (!plan.containsPattern(DATA_SOURCE_V2_RELATION)) { - return plan - } - val pushdownRules = Seq[LogicalPlan => LogicalPlan] ( createScanBuilder, pushDownSample,