[SPARK-58082][SQL] Add DSv2 relation and row-level command tree patterns

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 =>

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 &&

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 =>
 

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]

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 {

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)

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 = {

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))
+  }
+}

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))
+  }
 }

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) =>

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) =>

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)