Skip to content

Commit 316a3cc

Browse files
committed
[SPARK-57773][SQL] Strip leaked internal metadata in DataSourceV2Relation.create while preserving column IDs
### What changes were proposed in this pull request? `DataSourceV2Relation.create` builds the relation output schema directly from `table.columns().asSchema` (after char/varchar replacement), without removing internal metadata. This PR makes `create` strip internal metadata (the keys in `INTERNAL_METADATA_KEYS`, via `removeInternalMetadata`) from the schema before it becomes the relation's output, **while preserving column IDs**. `removeInternalMetadata` gains a `keepFieldIds: Boolean = false` parameter that skips `FIELD_ID_METADATA_KEY` in the same removal pass (the default preserves behavior for existing callers). `create` calls it with `keepFieldIds = true`: ```scala val schema = removeInternalMetadata( CharVarcharUtils.replaceCharVarcharWithStringInSchema(table.columns.asSchema), keepFieldIds = true) ``` ### Why are the changes needed? `INTERNAL_METADATA_KEYS` exists so that internal-only metadata keys (e.g. `__metadata_col`, `__qualified_access_only`) do not surface to users, and `removeInternalMetadata` is already applied on other schema-producing paths. But `DataSourceV2Relation.create` never applied it, so any internal metadata key that a v2 source attaches to its columns leaks straight onto the relation output and `df.schema`. This hardens `create` so internal-only keys stay internal, consistent with the other paths. Column IDs need special handling: `FIELD_ID_METADATA_KEY` is listed in `INTERNAL_METADATA_KEYS` so that other paths drop it, but it is also deliberately surfaced onto the relation's output (see SPARK-57544). Removing it in `create` would defeat that, so the new `keepFieldIds` flag preserves it in the same pass that strips the rest. ### Does this PR introduce _any_ user-facing change? No. It removes internal-only metadata keys (which were never meant to be user-visible) from the v2 relation output, and preserves the column IDs that are intentionally surfaced. ### How was this patch tested? New unit test in `DataSourceV2RelationSuite` that builds a table whose column carries both a column ID and a leaked internal metadata key, and asserts that `create` strips the internal key from the relation output while preserving the column ID. Verified the test fails without the source change. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 4.8) Closes #56887 from cloud-fan/SPARK-57544-followup. Authored-by: Wenchen Fan <wenchen@databricks.com> Signed-off-by: Wenchen Fan <wenchen@databricks.com>
1 parent 5ebf000 commit 316a3cc

3 files changed

Lines changed: 54 additions & 5 deletions

File tree

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/package.scala

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -251,11 +251,15 @@ package object util extends Logging {
251251
MetadataColumn.PRESERVE_ON_REINSERT
252252
)
253253

254-
def removeInternalMetadata(schema: StructType): StructType = {
254+
def removeInternalMetadata(schema: StructType, keepFieldIds: Boolean = false): StructType = {
255255
StructType(schema.map { field =>
256256
var builder = new MetadataBuilder().withMetadata(field.metadata)
257257
INTERNAL_METADATA_KEYS.foreach { key =>
258-
builder = builder.remove(key)
258+
// Column IDs are listed as internal so most paths drop them, but some paths (e.g. the v2
259+
// relation output) deliberately surface them, so allow callers to keep them.
260+
if (!(keepFieldIds && key == FIELD_ID_METADATA_KEY)) {
261+
builder = builder.remove(key)
262+
}
259263
}
260264
field.copy(metadata = builder.build())
261265
})

sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.plans.QueryPlan
2727
import org.apache.spark.sql.catalyst.plans.logical.{ColumnStat, ExposesMetadataColumns, Histogram, HistogramBin, LeafNode, LogicalPlan, Statistics}
2828
import org.apache.spark.sql.catalyst.streaming.{StreamingSourceIdentifyingName, Unassigned}
2929
import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
30-
import org.apache.spark.sql.catalyst.util.{truncatedString, CharVarcharUtils}
30+
import org.apache.spark.sql.catalyst.util.{removeInternalMetadata, truncatedString, CharVarcharUtils}
3131
import org.apache.spark.sql.connector.catalog.{CatalogPlugin, FunctionCatalog, Identifier, SupportsMetadataColumns, Table, TableCapability, TableCatalog, V2TableUtil}
3232
import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper
3333
import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
@@ -329,7 +329,13 @@ object DataSourceV2Relation {
329329
import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._
330330
// The v2 source may return schema containing char/varchar type. We replace char/varchar
331331
// with "annotated" string type here as the query engine doesn't support char/varchar yet.
332-
val schema = CharVarcharUtils.replaceCharVarcharWithStringInSchema(table.columns.asSchema)
332+
// We also strip internal metadata that may have leaked onto the table columns, so it does
333+
// not surface on the relation's output. Column IDs (FIELD_ID_METADATA_KEY) are an exception:
334+
// although the key is listed in INTERNAL_METADATA_KEYS so that other paths drop it, the
335+
// column-ID feature deliberately surfaces field IDs on the relation's output, so we keep them.
336+
val schema = removeInternalMetadata(
337+
CharVarcharUtils.replaceCharVarcharWithStringInSchema(table.columns.asSchema),
338+
keepFieldIds = true)
333339
DataSourceV2Relation(table, toAttributes(schema), catalog, identifier, options, timeTravelSpec)
334340
}
335341

sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,17 @@
1717

1818
package org.apache.spark.sql.execution.datasources.v2
1919

20+
import java.util
21+
2022
import org.apache.spark.SparkFunSuite
2123
import org.apache.spark.sql.catalyst.catalog.{CatalogColumnStat, CatalogStatistics}
2224
import org.apache.spark.sql.catalyst.plans.logical.{Histogram, HistogramBin}
25+
import org.apache.spark.sql.catalyst.util.FieldMetadataUtils.FIELD_ID_METADATA_KEY
26+
import org.apache.spark.sql.catalyst.util.INTERNAL_METADATA_KEYS
27+
import org.apache.spark.sql.connector.catalog.{Column, Table, TableCapability}
2328
import org.apache.spark.sql.connector.expressions.FieldReference
24-
import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
29+
import org.apache.spark.sql.types.{IntegerType, MetadataBuilder, StringType, StructField, StructType}
30+
import org.apache.spark.sql.util.CaseInsensitiveStringMap
2531

2632
class DataSourceV2RelationSuite extends SparkFunSuite {
2733

@@ -107,4 +113,37 @@ class DataSourceV2RelationSuite extends SparkFunSuite {
107113
val idV2NoHist = colStats.get(FieldReference.column("id"))
108114
assert(!idV2NoHist.histogram().isPresent)
109115
}
116+
117+
test("create strips leaked internal metadata but preserves column IDs") {
118+
// A column carrying both a column ID (surfaced on purpose) and every internal metadata key
119+
// (listed in INTERNAL_METADATA_KEYS), simulating a v2 source that leaks internal metadata.
120+
// The column ID key is excluded here since it is set via `.id(...)` below.
121+
val leakedBuilder = new MetadataBuilder()
122+
INTERNAL_METADATA_KEYS.filter(_ != FIELD_ID_METADATA_KEY).foreach { key =>
123+
leakedBuilder.putString(key, "leaked")
124+
}
125+
val leakedMetadata = leakedBuilder.build()
126+
val table = new Table {
127+
override def name(): String = "t"
128+
override def columns(): Array[Column] = Array(
129+
Column.builderFor("id", IntegerType)
130+
.id("1")
131+
.metadata(leakedMetadata.json)
132+
.build())
133+
override def capabilities(): util.Set[TableCapability] =
134+
util.Set.of[TableCapability]()
135+
}
136+
137+
val relation =
138+
DataSourceV2Relation.create(table, None, None, CaseInsensitiveStringMap.empty())
139+
val field = relation.schema.head
140+
141+
// All leaked internal metadata keys are stripped from the relation output ...
142+
INTERNAL_METADATA_KEYS.filter(_ != FIELD_ID_METADATA_KEY).foreach { key =>
143+
assert(!field.metadata.contains(key), s"internal metadata key $key should be stripped")
144+
}
145+
// ... but the column ID is preserved.
146+
assert(field.id.contains("1"))
147+
assert(field.metadata.contains(FIELD_ID_METADATA_KEY))
148+
}
110149
}

0 commit comments

Comments
 (0)