This is an automated email from the ASF dual-hosted git repository.
gengliangwang pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 138da638d12f [SPARK-58082][SQL] Add DSv2 relation and row-level
command tree patterns
138da638d12f is described below
commit 138da638d12fdf8c81e08876f95dfd9f22b01f0b
Author: yyanyy <[email protected]>
AuthorDate: Tue Aug 18 13:51:54 2026 -0700
[SPARK-58082][SQL] Add DSv2 relation and row-level command tree patterns
### What changes were proposed in this pull request?
Add Catalyst tree-pattern identity bits for the DSv2 relation leaves and the
row-level command nodes, and use them to prune the optimizer/analyzer rules
that match those nodes.
- Add `DATA_SOURCE_V2_RELATION` and `DATA_SOURCE_V2_SCAN_RELATION` to
`TreePattern`, and attach them to `DataSourceV2Relation` and
`DataSourceV2ScanRelation` respectively. `DATA_SOURCE_V2_SCAN_RELATION`
is used
as a pruning guard (see below); `DATA_SOURCE_V2_RELATION` is attached to
the
relation leaf but has no consumer rule in this PR (reserved for future
pruning).
- Add `DELETE_FROM_TABLE`, `UPDATE_TABLE`, `MERGE_INTO_TABLE`,
`REPLACE_DATA`
and `WRITE_DELTA`. Since every command already carries the shared
`COMMAND`
bit (declared `final` on the `Command` trait), `Command` now exposes a
`nodePatternsInternal()` hook, mirroring how `PlanExpression` combines
`PLAN_EXPRESSION` with a per-subquery bit, so a command node contributes
its
own identity bit while `COMMAND` stays guaranteed-present.
`DeleteFromTable`,
`UpdateTable`, `MergeIntoTable`, `ReplaceData` and `WriteDelta` override
it.
- Guard the previously-unguarded rules that key off these nodes with
`containsPattern` / `containsAnyPattern`:
`RewriteDeleteFromTable`, `RewriteUpdateTable`, `RewriteMergeIntoTable`,
`V2ScanPartitioningAndOrdering`,
`GroupBasedRowLevelOperationScanPlanning`,
`OptimizeMetadataOnlyDeleteFromTable` and
`RowLevelOperationRuntimeGroupFiltering`.
- Add tree-pattern contract tests.
### Why are the changes needed?
A tree-pattern bit is an O(1) bitset lookup on a plan's cached
`treePatternBits`.
Rules that today run an unconditional `transform`/`resolveOperators` walk
over
every node, on every invocation, even for plans that contain none of the
target
nodes, can instead cheaply skip whole subtrees that provably don't contain
the
node they rewrite. The DSv2 scan/relation leaves and the row-level command
nodes
previously had no identity pattern (the command nodes carried only the
coarse
`COMMAND` bit shared by all commands), so rules matching them either ran
unguarded or could only prune on the broad `COMMAND` bit. This adds the
missing
per-node bits and applies them as pruning guards, matching the existing
idiom
used throughout the optimizer (for example `ReplaceExceptWithFilter` on
`EXCEPT`,
`ResolveRowLevelCommandAssignments` on `COMMAND`).
### Does this PR introduce any user-facing change?
No. This only adds internal Catalyst tree-pattern metadata and rule pruning
guards. Query results and plans are unchanged; the guarded rules fire on
exactly
the same nodes as before, just skipping trees that cannot contain them.
### How was this patch tested?
- New unit tests pinning the pattern contract:
`DataSourceV2RelationSuite` (asserts the relation/scan-relation identity
bits)
and a new `V2CommandTreePatternSuite` (asserts each of the five row-level
command nodes -- `DeleteFromTable`, `UpdateTable`, `MergeIntoTable`,
`ReplaceData` and `WriteDelta` -- carries both `COMMAND` and its own bit).
- Regression suites covering the guarded rules, confirming the pruning
guards do
not drop any rule firing:
- Row-level operations (Group/Delta-based Delete/Update/Merge): 347
tests, all
passing.
- Scan pushdown / partitioning (`DataSourceV2Suite`,
`KeyGroupedPartitioningSuite`, `ProjectedOrderingAndPartitioningSuite`,
`DataSourceV2TableSampleSuite`): all passing, including `TABLESAMPLE
SYSTEM`
validation over DSv1 relations.
- V2 write analysis (`V2AppendData*`, `V2Overwrite*` ANSI/Strict suites):
388 tests, all passing, confirming the `Seq(COMMAND) ++
nodePatternsInternal()`
composition does not disturb existing `COMMAND`-pruned analysis rules.
- `./build/sbt catalyst/compile sql/compile` and `git diff --check`.
### Was this patch authored or co-authored using generative AI tooling?
yes, opus 4.8
Closes #57178 from yyanyy/spark-dsv2-tree-patterns-20260708.
Authored-by: yyanyy <[email protected]>
Signed-off-by: Gengliang Wang <[email protected]>
(cherry picked from commit deeadfea019fde91af4a82fc2de3651edb508c6e)
Signed-off-by: Gengliang Wang <[email protected]>
---
.../catalyst/analysis/RewriteDeleteFromTable.scala | 4 +-
.../catalyst/analysis/RewriteMergeIntoTable.scala | 4 +-
.../sql/catalyst/analysis/RewriteUpdateTable.scala | 4 +-
.../spark/sql/catalyst/plans/logical/Command.scala | 7 +-
.../sql/catalyst/plans/logical/v2Commands.scala | 11 +++
.../spark/sql/catalyst/trees/TreePatterns.scala | 7 ++
.../datasources/v2/DataSourceV2Relation.scala | 5 ++
.../plans/logical/V2CommandTreePatternSuite.scala | 97 ++++++++++++++++++++++
.../datasources/v2/DataSourceV2RelationSuite.scala | 25 ++++++
.../GroupBasedRowLevelOperationScanPlanning.scala | 4 +-
.../v2/OptimizeMetadataOnlyDeleteFromTable.scala | 4 +-
.../v2/V2ScanPartitioningAndOrdering.scala | 7 +-
.../RowLevelOperationRuntimeGroupFiltering.scala | 4 +-
13 files changed, 174 insertions(+), 9 deletions(-)
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 c8795b8ad9b1..8bb55b1822b1 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}
@@ -36,7 +37,8 @@ import
org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, Extr
*/
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 714f815161ae..1f8e9df4a110 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}
@@ -42,7 +43,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 207b3e721704..352ddc0b0acb 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 bd277e92d11d..e054c955cc77 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 85f3c726c7ec..b816016a3ec8 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}
@@ -448,6 +449,8 @@ case class ReplaceData(
override protected def withNewChildInternal(newChild: LogicalPlan):
ReplaceData = {
copy(query = newChild)
}
+
+ override protected def nodePatternsInternal(): Seq[TreePattern] =
Seq(REPLACE_DATA)
}
/**
@@ -558,6 +561,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
@@ -1098,6 +1103,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)
}
/**
@@ -1138,6 +1145,8 @@ case class UpdateTable(
case r: NamedRelation => r.skipSchemaResolution
case _ => false
}
+
+ override protected def nodePatternsInternal(): Seq[TreePattern] =
Seq(UPDATE_TABLE)
}
/**
@@ -1229,6 +1238,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 0679d0c332df..55b9f0621fd6 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 1c7a7abe4a8a..da2b3a54d74c 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
@@ -27,6 +27,7 @@ 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.plans.logical.statsEstimation.EstimationUtils
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}
@@ -145,6 +146,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)
}
/**
@@ -225,6 +228,8 @@ case class DataSourceV2ScanRelation(
filterAttrs.toImmutableArraySeq, this))
}
+ override val nodePatterns: Seq[TreePattern] =
Seq(DATA_SOURCE_V2_SCAN_RELATION)
+
override def name: String = relation.name
// A leaf relation references no upstream attributes. `pushedFilters` (and,
for that matter,
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 000000000000..3b20eae1a75e
--- /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 388af4c4c5b4..79fd057f037e 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
@@ -26,6 +26,7 @@ import
org.apache.spark.sql.catalyst.expressions.AttributeReference
import org.apache.spark.sql.catalyst.plans.SQLHelper
import org.apache.spark.sql.catalyst.plans.logical.{Histogram, HistogramBin}
import
org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils
+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}
@@ -507,6 +508,30 @@ class DataSourceV2RelationSuite extends SparkFunSuite with
SQLHelper {
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))
+ }
}
private class FakeTableWithSchema(
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 0cb7967cb736..a41aad05d435 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 c052edea53e0..0a0afb73dbb2 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 479154e7a285..165425aea19d 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 87139a3a20e1..ae7b432d72f5 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.expressions.NamedReference
import org.apache.spark.sql.connector.read.{Scan, SupportsRuntimeV2Filtering}
import org.apache.spark.sql.connector.write.RowLevelOperation.Command.{DELETE,
MERGE, UPDATE}
@@ -51,7 +52,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.filterAttributes) =>
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
