(spark) branch master updated: [SPARK-60111][SQL] Preserve tags when rebuilding DataSource V2 relations

[ge...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

gengliangwang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new 35686db92365 [SPARK-60111][SQL] Preserve tags when rebuilding DataSource V2 relations
35686db92365 is described below

commit 35686db92365eb5aa74686f65176ff1a24dcd33f
Author: yyanyy <yy...@gmail.com>
AuthorDate: Mon Oct 19 21:39:33 2026 -0700

    [SPARK-60111][SQL] Preserve tags when rebuilding DataSource V2 relations
    
    ### What changes were proposed in this pull request?
    
    This pull request preserves `TreeNode` tags when a `DataSourceV2Relation` is rebuilt in two places:
    
    * `DataSourceV2Relation.withMetadataColumns`
    * `Analyzer.ResolveOutputRelation` when it cleans and marks a DataSource V2 write target
    
    It also adds regression tests for both rebuild paths.
    
    ### Why are the changes needed?
    
    These paths use the Scala case-class `copy` method directly. Unlike `TreeNode.makeCopy` and the tree transformation helpers, the generated `copy` method does not call `copyTagsFrom`. As a result, tags attached to the original relation are silently absent from the rebuilt relation.
    
    Preserving the tags is consistent with the existing `LogicalRelation.withMetadataColumns` behavior and with the `TreeNode` contract for semantically equivalent replacements.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No. This change preserves internal logical-plan metadata.
    
    ### How was this patch tested?
    
    Added the following regression tests:
    
    * `DataSourceV2RelationSuite: withMetadataColumns preserves tags`
    * `V2ResolveOutputRelationSuite: ResolveOutputRelation preserves tags on a rebuilt DSv2 write target`
    
    Before the production changes, each new test failed because the rebuilt relation had no tag. After the changes, the two relevant suites completed with 30 tests succeeded and 0 failed:
    
    `build/sbt 'catalyst/testOnly org.apache.spark.sql.execution.datasources.v2.DataSourceV2RelationSuite org.apache.spark.sql.catalyst.analysis.V2ResolveOutputRelationSuite'`
    
    The Catalyst compile and test scalastyle checks also passed with 0 errors:
    
    `build/sbt 'catalyst/scalastyle' 'catalyst/Test/scalastyle'`
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: OpenAI Codex
    
    Closes #59330 from yyanyy/spark-preserve-dsv2-tags-20261008.
    
    Authored-by: yyanyy <yy...@gmail.com>
    Signed-off-by: Gengliang Wang <ge...@apache.org>
---
 .../spark/sql/catalyst/analysis/Analyzer.scala     |  4 ++-
 .../datasources/v2/DataSourceV2Relation.scala      |  4 ++-
 .../catalyst/analysis/V2WriteAnalysisSuite.scala   | 30 +++++++++++++++++++++-
 .../datasources/v2/DataSourceV2RelationSuite.scala | 30 +++++++++++++++++++---
 4 files changed, 62 insertions(+), 6 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
index 1c24d24c6653..b1f7eded0b76 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -4095,8 +4095,10 @@ class Analyzer(
           val cleanedTable = v2Write.table match {
             case r: DataSourceV2Relation =>
               val cleaned = r.output.map(CharVarcharUtils.cleanAttrMetadata)
-              r.copy(output =
+              val cleanedRelation = r.copy(output =
                 GeneratedColumn.markAutoFilledGeneratedColumns(cleaned, autoFilledGenCols))
+              cleanedRelation.copyTagsFrom(r)
+              cleanedRelation
             case other => other
           }
           v2Write.withNewQuery(projection).withNewTable(cleanedTable)
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 da080055bbd0..7ca3dea43ffc 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
@@ -137,7 +137,9 @@ case class DataSourceV2Relation(
   def withMetadataColumns(): DataSourceV2Relation = {
     val newMetadata = metadataOutput.filterNot(outputSet.contains)
     if (newMetadata.nonEmpty) {
-      copy(output = output ++ newMetadata)
+      val newRelation = copy(output = output ++ newMetadata)
+      newRelation.copyTagsFrom(this)
+      newRelation
     } else {
       this
     }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/V2WriteAnalysisSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/V2WriteAnalysisSuite.scala
index 3b4da14c7b15..24cdbd0cf0df 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/V2WriteAnalysisSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/V2WriteAnalysisSuite.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
-import java.util.Locale
+import java.util.{Locale, Set => JSet}
 
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
@@ -25,10 +25,38 @@ import org.apache.spark.sql.catalyst.expressions.{Alias, ArrayTransform, Attribu
 import org.apache.spark.sql.catalyst.expressions.objects.AssertNotNull
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
 import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.connector.catalog.{Column, Table, TableCapability}
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.internal.SQLConf.StoreAssignmentPolicy
 import org.apache.spark.sql.types._
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+class V2ResolveOutputRelationSuite extends AnalysisTest {
+
+  test("ResolveOutputRelation preserves tags on a rebuilt DSv2 write target") {
+    val table = new Table {
+      override def name(): String = "table-name"
+      override def columns(): Array[Column] = Array(Column.create("id", LongType))
+      override def capabilities(): JSet[TableCapability] =
+        JSet.of(TableCapability.BATCH_WRITE)
+    }
+    val target =
+      DataSourceV2Relation.create(table, None, None, CaseInsensitiveStringMap.empty())
+    val tag = TreeNodeTag[Unit]("resolve-output-relation-test")
+    target.setTagValue(tag, ())
+    val query = LocalRelation(AttributeReference("id", IntegerType)())
+    val write = AppendData.byName(target, query)
+
+    val resolved = getAnalyzer.ResolveOutputRelation(write).asInstanceOf[AppendData]
+    val resolvedTarget = resolved.table.asInstanceOf[DataSourceV2Relation]
+
+    assert(resolvedTarget ne target)
+    assert(resolvedTarget.getTagValue(tag).contains(()))
+  }
+}
 
 class V2AppendDataANSIAnalysisSuite extends V2ANSIWriteAnalysisSuiteBase {
   override def byName(table: NamedRelation, query: LogicalPlan): LogicalPlan = {
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 2d0272ff14ef..8f4d300be61d 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,15 +26,15 @@ import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.SQLHelper
 import org.apache.spark.sql.catalyst.plans.logical.{Filter, Histogram, HistogramBin, Project}
 import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils
-import org.apache.spark.sql.catalyst.trees.TreePattern
+import org.apache.spark.sql.catalyst.trees.{TreeNodeTag, 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.catalog.{Column, MetadataColumn, SupportsMetadataColumns, Table, TableCapability}
 import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
 import org.apache.spark.sql.connector.read.{Scan, Statistics => V2Statistics, SupportsReportStatistics}
 import org.apache.spark.sql.connector.read.colstats.ColumnStatistics
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.{IntegerType, MetadataBuilder, StringType, StructField, StructType}
+import org.apache.spark.sql.types.{DataType, IntegerType, MetadataBuilder, StringType, StructField, StructType}
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
 
 class DataSourceV2RelationSuite extends SparkFunSuite with SQLHelper {
@@ -724,6 +724,30 @@ class DataSourceV2RelationSuite extends SparkFunSuite with SQLHelper {
     assert(field.metadata.contains(FIELD_ID_METADATA_KEY))
   }
 
+  test("withMetadataColumns preserves tags") {
+    val table = new Table with SupportsMetadataColumns {
+      override def name(): String = "t"
+      override def columns(): Array[Column] = Array(Column.create("id", IntegerType))
+      override def metadataColumns(): Array[MetadataColumn] =
+        Array(new MetadataColumn {
+          override def name(): String = "index"
+          override def dataType(): DataType = IntegerType
+          override def isNullable(): Boolean = false
+        })
+      override def capabilities(): util.Set[TableCapability] =
+        util.Set.of[TableCapability]()
+    }
+    val relation =
+      DataSourceV2Relation.create(table, None, None, CaseInsensitiveStringMap.empty())
+    val tag = TreeNodeTag[Unit]("metadata-columns-test")
+    relation.setTagValue(tag, ())
+
+    val relationWithMetadata = relation.withMetadataColumns()
+
+    assert(relationWithMetadata ne relation)
+    assert(relationWithMetadata.getTagValue(tag).contains(()))
+  }
+
   test("nodePatterns declare the DSv2 relation identity tree patterns") {
     val table = new Table {
       override def name(): String = "t"


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org