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..3b20eae1a75ed --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/V2CommandTreePatternSuite.scala @@ -0,0 +1,97 @@ +/* + * 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 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 + * 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) + + // 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)) + } + + 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)) + } + + 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/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/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)