This is an automated email from the ASF dual-hosted git repository.
cloud-fan 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 bd65ed7ce121 [SPARK-57499][SQL] Fix column pruning and invalid plans
in variant extraction pushdown on DSv2 scans
bd65ed7ce121 is described below
commit bd65ed7ce121ef3fa38b7b6dd7e3581c1555d4f1
Author: Qiegang Long <[email protected]>
AuthorDate: Mon Jun 22 09:58:22 2026 -0700
[SPARK-57499][SQL] Fix column pruning and invalid plans in variant
extraction pushdown on DSv2 scans
### What changes were proposed in this pull request?
Three fixes in `pushVariantExtractions` (called
by`V2ScaqqnRelationPushDown.pushDownVariants`):
1. **Guard against double-visit**: Add a `pushedVariants.isEmpty`
sentinel check so the inner `ScanBuilderHolder` leaf visit (caused by
`transformDown` recursing into the child after returning the plan
unchanged) returns immediately. This ensures
`builder.pushVariantExtractions` is called exactly once per holder.
2. **Eager column pruning**: While `projectList` and `filters` are in
scope, call `builder.pruneColumns(requiredSchema)` for builders
implementing `SupportsPushDownRequiredColumns` and trim
`sHolder.output` to the required columns. By the time
`buildScanWithPushedVariants` calls `build()`, the builder already
has the correct pruned schema. This is similiar to how
buildScanWithPushedAggregate works.
3. **Keep whole-variant reads raw**: pushdown is for extractions, not
whole-variant reads. A bare variant reference -- `SELECT v`, or a
column lifted to feed a `variant_get` above a Join/Sort/Aggregate
barrier the local rewrite cannot see -- is recorded as `fullVariant`
(path `$`), meaning "the entire value." Shredding that to a lone
full-variant slot saves no I/O and is mishandled: the Parquet reader
collapses it to a boolean placeholder, and above a barrier the
re-exposed `GetStructField AS v#orig` alias is dropped by
`RemoveRedundantAliases`, giving wrong results or an invalid plan. So
when fullVariant is a variant's only requested field, leave the
column raw; when it coexists with real extractions (e.g. `SELECT v,
variant_get(v, '$.a')`) the >=2-slot struct is not collapsed and
keeps its pushdown. This subsumes the join-key case and shreds
automatically once barrier-aware pushdown makes the `variant_get`
visible as a typed path.
Jira: https://issues.apache.org/jira/browse/SPARK-57499
### Why are the changes needed?
Three bugs on the accepted variant pushdown path:
**Issue 1 (Performance) — column pruning is skipped.**
`buildScanWithPushedVariants` calls
`builder.build()` and replaces the `ScanBuilderHolder` with a
`DataSourceV2ScanRelation`.
The subsequent `pruneColumns` rule matches only `ScanBuilderHolder` nodes,
so it is a
no-op and `builder.pruneColumns()` is never called. The scan reads the full
table schema
including unreferenced columns. For unreferenced `VARIANT` columns this is
especially
costly — each is fully reconstructed from its shredded Parquet tree on
every row.
**Issue 2 (Correctness) — invalid plan / crash on tables with >=2 VARIANT
columns**
pushDownVariants uses transformDown, which recurses into the child
ScanBuilderHolder after returning the plan unchanged. The bare
ScanBuilderHolder matches PhysicalOperation a second time, collecting an
unreferenced sibling VARIANT column as a full-variant request and pushing it to
the builder again. ParquetScanBuilder overwrites its state on every call, so
the second push clobbers the correct extraction from the first. The rewritten
scan then emits a fresh ExprId for the variant whil [...]
This affects any extraction shape — projection, ORDER BY, aggregate, join —
on a table with two or more VARIANT columns. Single-VARIANT tables are
unaffected.
Reproduce on stock spark-4.1.x (path-based views force DSv2):
```
SET spark.sql.sources.useV1SourceList = "";
CREATE TABLE t (a INT, v1 VARIANT, v2 VARIANT) USING PARQUET LOCATION
'/tmp/vt';
INSERT INTO t VALUES
(1, parse_json('{"x":1,"price":3,"name":"x"}'), parse_json('{"y":2}')),
(2, parse_json('{"x":9,"price":1,"name":"z"}'), parse_json('{"y":8}'));
CREATE OR REPLACE TEMPORARY VIEW tv USING parquet OPTIONS (path
'/tmp/vt');
CREATE OR REPLACE TEMPORARY VIEW codes USING parquet OPTIONS (path
'/tmp/vt');
-- All four crash with: [INTERNAL_ERROR_ATTRIBUTE_NOT_FOUND] Could not find
v1#NN in [...]
SELECT variant_get(v1, '$.x', 'int') FROM tv;
SELECT variant_get(v1,'$.name','string') AS nm
FROM tv ORDER BY variant_get(v1,'$.price','int');
SELECT max(variant_get(v1, '$.price', 'int')) FROM tv;
SELECT l.a FROM tv l JOIN codes r
ON variant_get(l.v1,'$.x','int') = variant_get(r.v1,'$.x','int');
[INTERNAL_ERROR_ATTRIBUTE_NOT_FOUND] Could not find v1#21 in
[a#33,v1#34,v2#35]. SQLSTATE: XX000
at
org.apache.spark.sql.catalyst.expressions.BindReferences$.attributeNotFoundException(BoundAttribute.scala:109)
...
```
**Issue 3 (Correctness) -- wrong results / crash when a whole-variant read
is shredded**. A bare
variant reference (a plain `SELECT v`, or a column lifted to feed a
variant_get
above a Join/Sort/Aggregate barrier the local rewrite cannot see) is
recorded as a
full-variant request (path "$"). Shredding it to a lone full-variant slot
is both
useless (the whole value is read regardless) and mishandled:
- The Parquet reader collapses a lone VariantType slot to a boolean
placeholder,
so ORDER BY variant_get(v, '$.price') sorts on the placeholder and
silently
returns the wrong order, and max(variant_get(v, '$.price')) fails to
codegen.
- A join key is re-exposed above the join as GetStructField(v_new, i) AS
v#orig;
RemoveRedundantAliases collapses the alias and the condition references
a
dropped ExprId, failing plan validation.
This issue is masked on stock 4.1 by the Issue 2 binding crash, which fails
first; it surfaces only after that issue is fixed, which is why it has no repro
for stock 4.1. See new unit tests for join, order by, aggregrate.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
- Added new unit tests
- manual testing with spark-sql
### Was this patch authored or co-authored using generative AI tooling?
Co-authored with Claude code (Sonnet 4.6)
Closes #56556 from qlong/SPARK-57499-variant-pushdown-column-pruning.
Authored-by: Qiegang Long <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
.../datasources/v2/V2ScanRelationPushDown.scala | 123 ++-
.../datasources/PushVariantIntoScanSuite.scala | 898 +++++++++++++++++++++
2 files changed, 1006 insertions(+), 15 deletions(-)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala
index 2b291bf3a4db..42afb513f406 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala
@@ -34,12 +34,12 @@ import
org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
import org.apache.spark.sql.connector.expressions.{SortOrder => V2SortOrder}
import org.apache.spark.sql.connector.expressions.aggregate.{Aggregation, Avg,
Count, CountStar, Max, Min, Sum}
import org.apache.spark.sql.connector.expressions.filter.Predicate
-import org.apache.spark.sql.connector.read.{Scan, ScanBuilder,
SupportsPushDownAggregates, SupportsPushDownFilters, SupportsPushDownJoin,
SupportsPushDownVariantExtractions, V1Scan, VariantExtraction}
+import org.apache.spark.sql.connector.read.{Scan, ScanBuilder,
SupportsPushDownAggregates, SupportsPushDownFilters, SupportsPushDownJoin,
SupportsPushDownRequiredColumns, SupportsPushDownVariantExtractions, V1Scan,
VariantExtraction}
import org.apache.spark.sql.execution.datasources.{DataSourceStrategy,
VariantInRelation, VariantMetadata}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.internal.connector.VariantExtractionImpl
import org.apache.spark.sql.sources
-import org.apache.spark.sql.types.{DataType, DecimalType, IntegerType,
StringType, StructField, StructType}
+import org.apache.spark.sql.types.{DataType, DecimalType, IntegerType,
StringType, StructField, StructType, VariantType}
import org.apache.spark.sql.util.SchemaUtils._
import org.apache.spark.util.ArrayImplicits._
import org.apache.spark.util.Utils
@@ -358,11 +358,30 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
case agg: Aggregate => rewriteAggregate(agg)
}
- def pushDownVariants(plan: LogicalPlan): LogicalPlan = plan.transformDown {
- case p@PhysicalOperation(projectList, filters, sHolder @
ScanBuilderHolder(_, _,
- builder: SupportsPushDownVariantExtractions))
- if
conf.getConf(org.apache.spark.sql.internal.SQLConf.PUSH_VARIANT_INTO_SCAN) =>
- pushVariantExtractions(p, projectList, filters, sHolder, builder)
+ // Two-visit protocol explanation:
+ // `transformDown` is pre-order and always recurses into children, because
+ // `pushVariantExtractions` returns `originalPlan` unchanged (same object
reference).
+ // This means a tree like Project(pl, ScanBuilderHolder) gets visited twice:
+ // (1) Outer visit: PhysicalOperation collapses Project->ScanBuilderHolder
and
+ // yields the real (projectList, filters, sHolder). This is the
authoritative
+ // visit where column pruning decisions are made.
+ // (2) Inner leaf visit: PhysicalOperation on a bare ScanBuilderHolder
LeafNode
+ // yields (sHolder.output, Nil, sHolder) -- projectList = full schema
output.
+ // Without a guard, this would add fullVariant for *every* variant
column.
+ //
+ // The `pushedVariants.isEmpty` guard prevents the inner visit from
re-running once
+ // the outer visit has completed (successfully or not). An empty-mapping
sentinel
+ // `Some(new VariantInRelation())` is written by the outer visit when there
is
+ // nothing to push; `buildScanWithPushedVariants` requires
`mapping.nonEmpty` to
+ // fire, so the sentinel correctly bypasses it while still suppressing the
inner visit.
+ def pushDownVariants(plan: LogicalPlan): LogicalPlan = {
+ plan.transformDown {
+ case p@PhysicalOperation(projectList, filters, sHolder @
ScanBuilderHolder(_, _,
+ builder: SupportsPushDownVariantExtractions))
+ if
conf.getConf(org.apache.spark.sql.internal.SQLConf.PUSH_VARIANT_INTO_SCAN) &&
+ sHolder.pushedVariants.isEmpty =>
+ pushVariantExtractions(p, projectList, filters, sHolder, builder)
+ }
}
/**
@@ -418,14 +437,48 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
for ((a, defaultValue) <- schemaAttributes.zip(defaultValues)) {
variants.addVariantFields(a.exprId, a.dataType, defaultValue, Nil)
}
- if (variants.mapping.isEmpty) return originalPlan
+ if (variants.mapping.isEmpty) {
+ // No variant columns in the schema. Mark as attempted so the guard in
+ // pushDownVariants prevents a spurious second visit on the leaf node.
+ sHolder.pushedVariants = Some(new VariantInRelation())
+ return originalPlan
+ }
// Collect requested fields from project list and filters
projectList.foreach(variants.collectRequestedFields)
filters.foreach(variants.collectRequestedFields)
- // If no variant columns remain after collection, return original plan
- if (variants.mapping.forall(_._2.isEmpty)) return originalPlan
+ // Drop a variant whose sole requested field targets VariantType (a
whole-variant read: bare
+ // `SELECT v`, 2-arg `variant_get(v, '$.a')`, or a column lifted to feed a
variant_get above a
+ // barrier). Shredding it to a lone variant-typed slot saves no I/O and is
mishandled: the
+ // reader collapses such a slot to a boolean placeholder while catalyst
keeps it VariantType, so
+ // holder.output and readSchema disagree. Leaving it raw avoids that. The
predicate matches the
+ // reader's collapse condition by target type, not the exact `fullVariant`
value -- the latter
+ // would miss 2-arg variant_get and non-UTC sessions (different path/tz,
same broken slot).
+ // When the field coexists with real extractions the struct has >= 2 slots
and is not collapsed,
+ // so this is scoped to the sole-field case. Runs before the
IsNull/IsNotNull injection so
+ // presence-only references, which legitimately shred to a placeholder,
are unaffected.
+ variants.mapping.values.foreach { pathToFields =>
+ pathToFields.filterInPlace { case (_, fields) =>
+ !(fields.size == 1 &&
fields.head._1.targetType.isInstanceOf[VariantType])
+ }
+ }
+ variants.mapping.filterInPlace { case (_, pathToFields) =>
pathToFields.nonEmpty }
+
+ // If a variant column is referenced only via IsNull/IsNotNull (e.g. WHERE
isnotnull(v)),
+ // collectRequestedFields adds nothing to its field map (the
IsNull/IsNotNull branch is a
+ // no-op). Inject a fullVariant entry for each such column so rewriteType
generates the
+ // placeholder struct and the extraction is pushed on this (outer) visit.
Without this,
+ // the inner leaf visit would fire with sHolder.output as its projectList,
causing all
+ // variant columns -- including unreferenced siblings -- to receive
fullVariant treatment.
+ val referencedAttrs = AttributeSet((projectList ++
filters).flatMap(_.references))
+ for (a <- sHolder.relation.output) {
+ if (variants.mapping.contains(a.exprId) &&
+ variants.mapping(a.exprId).values.forall(_.isEmpty) &&
+ referencedAttrs.contains(a)) {
+ variants.collectRequestedFields(a)
+ }
+ }
// Build individual VariantExtraction for each field access
// Track which extraction corresponds to which (attr, field, ordinal)
@@ -475,20 +528,31 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
}
}
- // Call the API to push down variant extractions
- if (extractionInfo.isEmpty) return originalPlan
+ // No extraction was requested for any variant column. Set the sentinel to
suppress
+ // the inner leaf visit and avoid a spurious empty push.
+ if (extractionInfo.isEmpty) {
+ sHolder.pushedVariants = Some(new VariantInRelation())
+ return originalPlan
+ }
// Companion extractions can only be honored by readers that support
cast-error deferral. If
// none were generated, the pushdown carries only non-strict accesses
(`try_variant_get`, plain
// variant reads, casts to variant/string) that are safe regardless of
deferral support.
- if (hasCompanionExtraction && !builder.supportsDeferCastError()) return
originalPlan
+ if (hasCompanionExtraction && !builder.supportsDeferCastError()) {
+ // Set the sentinel like the other early returns, else the leaf is
re-visited (double-visit).
+ sHolder.pushedVariants = Some(new VariantInRelation())
+ return originalPlan
+ }
val extractions: Array[VariantExtraction] =
extractionInfo.map(_._1).toArray
val pushedResults = builder.pushVariantExtractions(extractions)
// Filter to only the accepted extractions
val acceptedExtractions =
extractionInfo.zip(pushedResults).filter(_._2).map(_._1)
- if (acceptedExtractions.isEmpty) return originalPlan
+ if (acceptedExtractions.isEmpty) {
+ sHolder.pushedVariants = Some(new VariantInRelation())
+ return originalPlan
+ }
// Group accepted extractions by attribute to rebuild the struct schemas
val extractionsByAttr = acceptedExtractions.groupBy(_._2)
@@ -513,6 +577,31 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
sHolder.pushedVariantAttributeMap = attributeMap
sHolder.output = newOutput
+ // Commit the required top-level columns to the builder now, while
projectList and
+ // filters are in scope. This mirrors how rewriteAggregate sets
holder.output to only
+ // the pushed-down aggregate columns before build() -- both patterns
ensure that
+ // buildScanWithPushed* can call build() and zip holder.output with
readSchema() directly.
+ //
+ // projectList/filters reference original ExprIds (this function returns
originalPlan
+ // unchanged). holder.relation.output carries those original ExprIds, so
AttributeSet
+ // matching works directly against it.
+ //
+ // attributeMap maps old ExprId -> new AttributeReference (with a fresh
ExprId for rewritten
+ // variant columns). Invert it to map new ExprId -> old ExprId so we can
filter holder.output
+ // (which carries new ExprIds) against requiredColumns (which carries old
ExprIds).
+ sHolder.builder match {
+ case r: SupportsPushDownRequiredColumns =>
+ val requiredColumns = AttributeSet((projectList ++
filters).flatMap(_.references))
+ val neededRelOutput =
sHolder.relation.output.filter(requiredColumns.contains)
+ val newToOldExprId = attributeMap.map { case (oldId, newAttr) =>
newAttr.exprId -> oldId }
+ val oldExprIdToRelAttr = sHolder.relation.output.map(a => a.exprId ->
a).toMap
+ sHolder.output = sHolder.output.filter { a =>
+
oldExprIdToRelAttr.get(newToOldExprId(a.exprId)).exists(requiredColumns.contains)
+ }
+ r.pruneColumns(neededRelOutput.toStructType)
+ case _ => // builder does not support column pruning; holder.output
stays full-schema
+ }
+
// Return the original plan unchanged - transformation happens in
buildScanWithPushedVariants
originalPlan
}
@@ -794,13 +883,17 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
def buildScanWithPushedVariants(plan: LogicalPlan): LogicalPlan =
plan.transform {
case p@PhysicalOperation(projectList, filters, holder: ScanBuilderHolder)
- if holder.pushedVariants.isDefined =>
+ if holder.pushedVariants.exists(_.mapping.nonEmpty) =>
val variants = holder.pushedVariants.get
val attributeMap = holder.pushedVariantAttributeMap
// Build the scan
val scan = holder.builder.build()
val realOutput = toAttributes(scan.readSchema())
+ assert(realOutput.length == holder.output.length,
+ s"The data source returns ${realOutput.length} columns but the plan
expected " +
+ s"${holder.output.length}:
scan=[${realOutput.map(_.name).mkString(",")}], " +
+ s"plan=[${holder.output.map(_.name).mkString(",")}]")
val wrappedScan = getWrappedScan(scan, holder)
// Note: holder.pushedFilterExpressions is not propagated here because
the output schema
// changes with variant extraction. When validConstraints is wired up,
this needs revisiting.
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala
index d6a9cfc94e8c..67bee55c2275 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala
@@ -472,6 +472,17 @@ abstract class PushVariantIntoScanV2SuiteBase extends
QueryTest with PushVariant
super.sparkConf.set(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key,
vectorizedReaderEnabled.toString)
+ // Locate the single DataSourceV2ScanRelation in an optimized plan,
regardless of how many
+ // Project/Filter nodes wrap it. This keeps scan-content assertions robust
against optimizer
+ // rules (e.g. CollapseProject) that may or may not collapse the identity
outputProjection
+ // that buildScanWithPushedVariants inserts.
+ protected def findScanRelation(plan: LogicalPlan): DataSourceV2ScanRelation
= {
+ val scans = plan.collect { case s: DataSourceV2ScanRelation => s }
+ assert(scans.length == 1,
+ s"Expected exactly one DataSourceV2ScanRelation but found
${scans.length}:\n$plan")
+ scans.head
+ }
+
test(s"V2 test - basic variant field extraction ($readerName)") {
withTempPath { dir =>
val path = dir.getCanonicalPath
@@ -794,6 +805,893 @@ abstract class PushVariantIntoScanV2SuiteBase extends
QueryTest with PushVariant
}
}
+ test(s"V2 test - column pruning: only referenced variant column in scan
($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v1 variant, v2 variant, v3 variant) using
PARQUET " +
+ s"location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}'),
parse_json('{\"c\": 3}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ val df = spark.read.parquet(path)
+ df.createOrReplaceTempView("T_V2")
+ val query = "select variant_get(v1, '$.a', 'int') from T_V2"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ sql(query).queryExecution.optimizedPlan match {
+ case Project(_, scanRelation: DataSourceV2ScanRelation) =>
+ // Only v1 should be in the scan; v2 and v3 must be pruned.
+ assert(scanRelation.output.map(_.name) == Seq("v1"),
+ s"Expected scan output [v1] but got
${scanRelation.output.map(_.name)}")
+ assert(scanRelation.output(0).dataType.isInstanceOf[StructType],
+ "Expected v1 to be rewritten to struct type after extraction
pushdown")
+ case other =>
+ fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other")
+ }
+ }
+ }
+ }
+
+ test(s"V2 test - column pruning: mixed variant/scalar, only variant
referenced ($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v variant, s string, i int) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values (parse_json('{\"a\": 42}'), 'hello',
7)")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ val df = spark.read.parquet(path)
+ df.createOrReplaceTempView("T_V2")
+ val query = "select variant_get(v, '$.a', 'int') from T_V2"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ sql(query).queryExecution.optimizedPlan match {
+ case Project(_, scanRelation: DataSourceV2ScanRelation) =>
+ // s and i are unreferenced; only v should appear in the scan.
+ assert(scanRelation.output.map(_.name) == Seq("v"),
+ s"Expected scan output [v] but got
${scanRelation.output.map(_.name)}")
+ // v is referenced via variant_get, so it is rewritten to an
extraction struct
+ // (slot per path) rather than read as a whole variant.
+ assert(scanRelation.output(0).dataType.isInstanceOf[StructType],
+ s"Expected v rewritten to struct, but got
${scanRelation.output(0).dataType}")
+ case other =>
+ fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other")
+ }
+ }
+ }
+ }
+
+ test(s"V2 test - column pruning: referenced scalar survives alongside
variant ($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v variant, s string) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values (parse_json('{\"a\": 5}'), 'keep_me')")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ val df = spark.read.parquet(path)
+ df.createOrReplaceTempView("T_V2")
+ val query = "select variant_get(v, '$.a', 'int'), s from T_V2"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ sql(query).queryExecution.optimizedPlan match {
+ case Project(_, scanRelation: DataSourceV2ScanRelation) =>
+ // Both v and s are referenced; neither should be pruned.
+ val names = scanRelation.output.map(_.name).toSet
+ assert(names == Set("v", "s"),
+ s"Expected scan output {v, s} but got $names")
+ // v is rewritten to an extraction struct; the scalar s keeps its
original type.
+ val byName = scanRelation.output.map(a => a.name ->
a.dataType).toMap
+ assert(byName("v").isInstanceOf[StructType],
+ s"Expected v rewritten to struct, but got ${byName("v")}")
+ assert(byName("s") == StringType,
+ s"Expected s to remain StringType, but got ${byName("s")}")
+ case other =>
+ fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other")
+ }
+ }
+ }
+ }
+
+ test(s"V2 test - column pruning: filter on second variant keeps it in scan
($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v1 variant, v2 variant) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"a\": 10}'), parse_json('{\"b\": 1}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ val df = spark.read.parquet(path)
+ df.createOrReplaceTempView("T_V2")
+ val query =
+ "select variant_get(v1, '$.a', 'int') from T_V2 " +
+ "where variant_get(v2, '$.b', 'int') > 0"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ // v1 is in SELECT, v2 is in WHERE -- both must survive column pruning.
+ val scanRelation =
findScanRelation(sql(query).queryExecution.optimizedPlan)
+ val names = scanRelation.output.map(_.name).toSet
+ assert(names == Set("v1", "v2"),
+ s"Expected scan output {v1, v2} but got $names")
+ // Both are referenced via variant_get, so both are rewritten to
extraction structs.
+ scanRelation.output.foreach { a =>
+ assert(a.dataType.isInstanceOf[StructType],
+ s"Expected ${a.name} rewritten to struct, but got ${a.dataType}")
+ }
+ }
+ }
+ }
+
+ test(s"V2 test - isnotnull on variant with sibling variant column
($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v1 variant, v2 variant) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values (parse_json('{\"a\": 1}'),
parse_json('{\"b\": 2}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ val df = spark.read.parquet(path)
+ df.createOrReplaceTempView("T_V2")
+ val query = "select 1 from T_V2 where isnotnull(v1)"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ // Only v1 should be in the scan; v2 is unreferenced and must be
pruned.
+ val scanRelation =
findScanRelation(sql(query).queryExecution.optimizedPlan)
+ assert(scanRelation.output.map(_.name) == Seq("v1"),
+ s"Expected scan output [v1] but got
${scanRelation.output.map(_.name)}")
+ assert(scanRelation.output(0).dataType.isInstanceOf[StructType],
+ "Expected v1 to be rewritten to a placeholder struct after pushdown")
+ }
+ }
+ }
+
+ // Aggregate barrier: an aggregate over one variant column (v1) with a
filter on the same
+ // column. The Aggregate node terminates the PhysicalOperation match, so
variant extraction
+ // pushdown only fires on the Filter -> Project -> ScanBuilderHolder subtree
below it. The
+ // $.price path inside max() lives above the barrier and is NOT pushed (v1
stays a full
+ // variant), but the unreferenced sibling v2 is never in the subtree's
projectList/filters,
+ // so top-level column pruning still drops it from the scan.
+ test(s"V2 test - column pruning: aggregate barrier prunes unreferenced
variant ($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v1 variant, v2 variant) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"price\": 10, \"country\": \"china\"}'),
parse_json('{\"b\": 2}')), " +
+ "(parse_json('{\"price\": 20, \"country\": \"japan\"}'),
parse_json('{\"b\": 3}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ val df = spark.read.parquet(path)
+ df.createOrReplaceTempView("T_V2")
+ val query =
+ "select max(variant_get(v1, '$.price', 'long')) from T_V2 " +
+ "where variant_get(v1, '$.country', 'string') = 'china'"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ // v2 is unreferenced anywhere in the query -- it must be pruned out
of the scan even
+ // though the aggregate barrier prevents the $.price extraction from
being pushed.
+ val scanRelation =
findScanRelation(sql(query).queryExecution.optimizedPlan)
+ assert(scanRelation.output.map(_.name) == Seq("v1"),
+ s"Expected scan output [v1] but got
${scanRelation.output.map(_.name)}")
+ }
+ }
+ }
+
+ test(s"V2 test - column pruning: group by + variant_get aggregate, sibling
pruned " +
+ s"($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v variant, v2 variant, name string) " +
+ s"using PARQUET location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"price\": 10, \"country\": \"china\"}'), " +
+ " parse_json('{\"junk\": 1}'), 'widget'), " +
+ "(parse_json('{\"price\": 20, \"country\": \"china\"}'), " +
+ " parse_json('{\"junk\": 2}'), 'widget'), " +
+ "(parse_json('{\"price\": 30, \"country\": \"japan\"}'), " +
+ " parse_json('{\"junk\": 3}'), 'gadget')")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ spark.read.parquet(path).createOrReplaceTempView("T_V2")
+ val query =
+ "select name, max(variant_get(v, '$.price', 'long')) from T_V2 " +
+ "where variant_get(v, '$.country', 'string') = 'china' group by name"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ // v2 is unreferenced and must be pruned. v is referenced via
variant_get in the WHERE
+ // (local, $.country) and via max(variant_get(v, '$.price')) above the
aggregate barrier
+ // (lifted in as a bare reference -> fullVariant). Because the local
$.country extraction is
+ // also present, fullVariant is not the column's only field, so v is
still shredded (a
+ // multi-slot struct does not collapse to a placeholder); it is not
kept raw.
+ val scanRelation =
findScanRelation(sql(query).queryExecution.optimizedPlan)
+ assert(scanRelation.output.map(_.name).toSet == Set("v", "name"),
+ s"Expected scan output {v, name} but got
${scanRelation.output.map(_.name)}")
+ val vAttr = scanRelation.output.find(_.name == "v").get
+ assert(vAttr.dataType.isInstanceOf[StructType],
+ s"Expected v rewritten to struct, but got ${vAttr.dataType}")
+ }
+ }
+ }
+
+ test(s"V2 test - order by variant_get: correct ordering, variant read raw,
sibling pruned " +
+ s"($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v variant, v2 variant, name string) " +
+ s"using PARQUET location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"price\": 3}'), parse_json('{\"junk\": 1}'), 'x'), "
+
+ "(parse_json('{\"price\": 1}'), parse_json('{\"junk\": 2}'), 'z'), "
+
+ "(parse_json('{\"price\": 2}'), parse_json('{\"junk\": 3}'), 'y')")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ spark.read.parquet(path).createOrReplaceTempView("T_V2")
+ // The sort key variant_get(v, '$.price') lives in Sort.order, above
the scan window, so v
+ // is lifted in only as a bare reference -> a whole-variant request ->
v stays raw and the
+ // sort evaluates on the real variant. Compare ORDER-SENSITIVELY: a
shredded full-variant
+ // slot would collapse to a boolean placeholder and silently
mis-order, which the
+ // order-insensitive checkAnswer would not catch.
+ val query = "select name from T_V2 order by variant_get(v, '$.price',
'int')"
+ val expectedOrder = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect().map(_.getString(0)).toList
+ }
+ assert(expectedOrder == List("z", "y", "x"),
+ s"baseline sanity: expected z,y,x but got $expectedOrder")
+ assert(sql(query).collect().map(_.getString(0)).toList ==
expectedOrder,
+ "ORDER BY variant_get produced the wrong row order")
+ val scanRelation =
findScanRelation(sql(query).queryExecution.optimizedPlan)
+ assert(scanRelation.output.map(_.name).toSet == Set("v", "name"),
+ s"Expected scan output {v, name} but got
${scanRelation.output.map(_.name)}")
+ val vAttr = scanRelation.output.find(_.name == "v").get
+ assert(vAttr.dataType.isInstanceOf[VariantType],
+ s"Expected v left as raw VariantType, but got ${vAttr.dataType}")
+ }
+ }
+ }
+
+ test(s"V2 test - aggregate max(variant_get) with no local filter: no codegen
crash, variant " +
+ s"read raw ($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v variant, v2 variant, name string) " +
+ s"using PARQUET location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"price\": 10}'), parse_json('{\"junk\": 1}'), 'a'),
" +
+ "(parse_json('{\"price\": 30}'), parse_json('{\"junk\": 2}'), 'a'),
" +
+ "(parse_json('{\"price\": 20}'), parse_json('{\"junk\": 3}'), 'b')")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ spark.read.parquet(path).createOrReplaceTempView("T_V2")
+ // variant_get(v, '$.price') is inside an aggregate function, above
the aggregate barrier,
+ // with no local filter/projection on v, so v is lifted in only as a
bare reference -> a
+ // whole-variant request. A whole-variant read is kept raw: shredding
it to a lone
+ // full-variant slot would collapse to a boolean placeholder, and
max(variant_get(<boolean>,
+ // ...)) would fail to codegen. Keeping v raw yields correct
aggregates with a valid plan.
+ val query =
+ "select name, max(variant_get(v, '$.price', 'int')) as mx from T_V2
group by name"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows) // a -> 30, b -> 20
+ val scanRelation =
findScanRelation(sql(query).queryExecution.optimizedPlan)
+ val vAttr = scanRelation.output.find(_.name == "v").get
+ assert(vAttr.dataType.isInstanceOf[VariantType],
+ s"Expected v left as raw VariantType, but got ${vAttr.dataType}")
+ assert(!scanRelation.output.exists(_.name == "v2"),
+ s"Expected v2 pruned but got ${scanRelation.output.map(_.name)}")
+ }
+ }
+ }
+
+ test(s"V2 test - join on variant_get key: no crash, variant read raw,
sibling pruned " +
+ s"($readerName)") {
+ withTempPath { dir =>
+ val itemsPath = dir.getCanonicalPath + "/items"
+ val countriesPath = dir.getCanonicalPath + "/countries"
+ withTable("temp_items", "temp_countries") {
+ sql(s"create table temp_items (v variant, v2 variant, name string) " +
+ s"using PARQUET location '$itemsPath'")
+ sql("insert into temp_items values " +
+ "(parse_json('{\"country_code\": \"CN\"}'), parse_json('{\"junk\":
1}'), 'widget'), " +
+ "(parse_json('{\"country_code\": \"JP\"}'), parse_json('{\"junk\":
2}'), 'gadget')")
+ sql(s"create table temp_countries (code string, country_name string) "
+
+ s"using PARQUET location '$countriesPath'")
+ sql("insert into temp_countries values ('CN', 'China'), ('JP',
'Japan')")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ spark.read.parquet(itemsPath).createOrReplaceTempView("ITEMS_V2")
+
spark.read.parquet(countriesPath).createOrReplaceTempView("COUNTRIES_V2")
+ val query =
+ "select i.name, c.country_name from ITEMS_V2 i join COUNTRIES_V2 c "
+
+ "on variant_get(i.v, '$.country_code', 'string') = c.code"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ // A variant_get join key produces a valid plan (no attribute-binding
failure) and
+ // correct results; the key is read raw and the extraction is
evaluated above the scan.
+ checkAnswer(sql(query), expectedRows)
+ val scans = sql(query).queryExecution.optimizedPlan.collect {
+ case s: DataSourceV2ScanRelation => s
+ }
+ val itemsScan = scans.find(_.output.exists(_.name == "v")).getOrElse(
+ fail(s"Could not find the items scan
in:\n${sql(query).queryExecution.optimizedPlan}"))
+ // v is the join key (referenced only via variant_get in the join
condition): it must be
+ // left as a raw variant, not shredded. v2 is unreferenced and must be
pruned.
+ assert(itemsScan.output.map(_.name).toSet == Set("v", "name"),
+ s"Expected items scan output {v, name} but got
${itemsScan.output.map(_.name)}")
+ val vAttr = itemsScan.output.find(_.name == "v").get
+ assert(vAttr.dataType.isInstanceOf[VariantType],
+ s"Expected v left as raw VariantType, but got ${vAttr.dataType}")
+ }
+ }
+ }
+
+ test(s"V2 test - join on variant_get key with local variant filter: no crash
($readerName)") {
+ withTempPath { dir =>
+ val itemsPath = dir.getCanonicalPath + "/items"
+ val countriesPath = dir.getCanonicalPath + "/countries"
+ withTable("temp_items", "temp_countries") {
+ sql(s"create table temp_items (v variant, v2 variant) " +
+ s"using PARQUET location '$itemsPath'")
+ sql("insert into temp_items values " +
+ "(parse_json('{\"country_code\": \"CN\", \"qty\": 5}'),
parse_json('{\"junk\": 1}')), " +
+ "(parse_json('{\"country_code\": \"JP\", \"qty\": 0}'),
parse_json('{\"junk\": 2}'))")
+ sql(s"create table temp_countries (code string, country_name string) "
+
+ s"using PARQUET location '$countriesPath'")
+ sql("insert into temp_countries values ('CN', 'China'), ('JP',
'Japan')")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ spark.read.parquet(itemsPath).createOrReplaceTempView("ITEMS_V2")
+
spark.read.parquet(countriesPath).createOrReplaceTempView("COUNTRIES_V2")
+ // v appears in both a join condition and a local WHERE extraction.
The join reference
+ // forces v to stay raw; the query must still optimize without
crashing and be correct.
+ val query =
+ "select c.country_name from ITEMS_V2 i join COUNTRIES_V2 c " +
+ "on variant_get(i.v, '$.country_code', 'string') = c.code " +
+ "where variant_get(i.v, '$.qty', 'int') > 0"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ }
+ }
+ }
+
+ test(s"V2 test - join on variant_get key through an aliasing projection: no
crash, no " +
+ s"mis-shred ($readerName)") {
+ withTempPath { dir =>
+ val itemsPath = dir.getCanonicalPath + "/items"
+ val countriesPath = dir.getCanonicalPath + "/countries"
+ withTable("temp_items", "temp_countries") {
+ sql(s"create table temp_items (v variant, v2 variant, name string) " +
+ s"using PARQUET location '$itemsPath'")
+ sql("insert into temp_items values " +
+ "(parse_json('{\"country_code\": \"CN\"}'), parse_json('{\"junk\":
1}'), 'widget'), " +
+ "(parse_json('{\"country_code\": \"JP\"}'), parse_json('{\"junk\":
2}'), 'gadget')")
+ sql(s"create table temp_countries (code string, country_name string) "
+
+ s"using PARQUET location '$countriesPath'")
+ sql("insert into temp_countries values ('CN', 'China'), ('JP',
'Japan')")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ spark.read.parquet(itemsPath).createOrReplaceTempView("ITEMS_V2")
+
spark.read.parquet(countriesPath).createOrReplaceTempView("COUNTRIES_V2")
+ // The variant column is aliased (v AS vw) in a subquery before the
join condition reads
+ // it via variant_get. Whether the optimizer inlines the alias or
keeps it, the query must
+ // optimize without crashing, return correct results, and not
over-shred (v2 pruned, the
+ // variant join key not turned into a struct that yields wrong data).
+ val query =
+ "select c.country_name from " +
+ "(select v as vw, name as nm from ITEMS_V2) i join COUNTRIES_V2 c " +
+ "on variant_get(i.vw, '$.country_code', 'string') = c.code"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ // Must not crash during optimization and must be correct.
+ checkAnswer(sql(query), expectedRows)
+ val scans = sql(query).queryExecution.optimizedPlan.collect {
+ case s: DataSourceV2ScanRelation => s
+ }
+ val itemsScan = scans.find(_.output.exists(a => a.name == "v" ||
a.name == "vw"))
+ .getOrElse(fail(
+ s"Could not find the items scan
in:\n${sql(query).queryExecution.optimizedPlan}"))
+ // v2 is unreferenced and must be pruned.
+ assert(!itemsScan.output.exists(_.name == "v2"),
+ s"Expected v2 pruned but items scan was
${itemsScan.output.map(_.name)}")
+ // The variant join key must not be shredded: it is read as a raw
variant.
+ val vAttr = itemsScan.output.find(a => a.name == "v" || a.name ==
"vw").get
+ assert(vAttr.dataType.isInstanceOf[VariantType],
+ s"Expected the variant join key left raw, but got ${vAttr.dataType}")
+ }
+ }
+ }
+
+ // Non-shredded variant: parse_json writes the variant blob without shredded
typed_value
+ // columns, so the Parquet file contains only metadata+value columns. The
logical plan
+ // rewrite (variant_get -> GetStructField) and column pruning both still
apply; the Parquet
+ // reader handles the non-shredded data at physical read time.
+ // All of the tests above already use parse_json (non-shredded) data, so
this comment
+ // documents that the non-shredded case is fully covered by the existing
test suite.
+
+ test(s"V2 test - sole full variant: select v directly stays raw
($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}')), " +
+ "(parse_json('{\"a\": 9}'), parse_json('{\"b\": 8}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ spark.read.parquet(path).createOrReplaceTempView("T_V2")
+ // Selecting the whole variant is the canonical sole-fullVariant case:
its only requested
+ // field is the entire value, which is not an extraction. v must stay
raw (shredding it to a
+ // lone full-variant slot would collapse to a boolean placeholder and
corrupt the output).
+ // v2 is unreferenced and must be pruned.
+ val query = "select v from T_V2"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ val scanRelation =
findScanRelation(sql(query).queryExecution.optimizedPlan)
+ assert(scanRelation.output.map(_.name) == Seq("v"),
+ s"Expected scan output [v] but got
${scanRelation.output.map(_.name)}")
+ assert(scanRelation.output(0).dataType.isInstanceOf[VariantType],
+ s"Expected v left raw, but got ${scanRelation.output(0).dataType}")
+ }
+ }
+ }
+
+ test(s"V2 test - sole full variant on one column, sibling extracted
($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}')), " +
+ "(parse_json('{\"a\": 9}'), parse_json('{\"b\": 8}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ spark.read.parquet(path).createOrReplaceTempView("T_V2")
+ // The guard is per-column: v is read whole (sole fullVariant -> raw)
while v2 has a real
+ // extraction (-> shredded). Both must coexist in the same scan.
+ val query = "select v, variant_get(v2, '$.b', 'int') as b from T_V2"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ val scanRelation =
findScanRelation(sql(query).queryExecution.optimizedPlan)
+ val vAttr = scanRelation.output.find(_.name == "v").get
+ val v2Attr = scanRelation.output.find(_.name == "v2").get
+ assert(vAttr.dataType.isInstanceOf[VariantType],
+ s"Expected v left raw, but got ${vAttr.dataType}")
+ assert(v2Attr.dataType.isInstanceOf[StructType],
+ s"Expected v2 shredded to struct, but got ${v2Attr.dataType}")
+ }
+ }
+ }
+
+ test(s"V2 test - sole 2-arg variant_get (VariantType target) stays raw
($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"a\": {\"x\": 1}}'), parse_json('{\"b\": 2}')), " +
+ "(parse_json('{\"a\": {\"x\": 9}}'), parse_json('{\"b\": 8}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ spark.read.parquet(path).createOrReplaceTempView("T_V2")
+ // 2-arg variant_get has no type argument, so its target is
VariantType. Its requested
+ // field is NOT the exact `fullVariant` sentinel (the path is "$.a",
not "$"), but it still
+ // shreds to a lone variant-typed slot which the reader collapses to a
boolean placeholder.
+ // The strip must catch it by target type, not by exact-value
equality, and keep v raw.
+ val query = "select variant_get(v, '$.a') as a from T_V2"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ val scanRelation =
findScanRelation(sql(query).queryExecution.optimizedPlan)
+ // v is read whole (kept raw); the unreferenced sibling v2 is pruned.
+ assert(scanRelation.output.map(_.name) == Seq("v"),
+ s"Expected scan output [v] but got
${scanRelation.output.map(_.name)}")
+ assert(scanRelation.output(0).dataType.isInstanceOf[VariantType],
+ s"Expected v left raw, but got ${scanRelation.output(0).dataType}")
+ }
+ }
+ }
+
+ test(s"V2 test - sole variant_get('$$') under non-UTC session stays raw
($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}')), " +
+ "(parse_json('{\"a\": 9}'), parse_json('{\"b\": 8}'))")
+ }
+ // A non-UTC session makes variant_get carry a non-UTC timezone, so even
the "$" path is not
+ // equal to the UTC `fullVariant` sentinel. The strip must still keep v
raw -- matching by
+ // target type rather than by the exact sentinel value covers this case.
+ withSQLConf(
+ SQLConf.USE_V1_SOURCE_LIST.key -> "",
+ SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Los_Angeles") {
+ spark.read.parquet(path).createOrReplaceTempView("T_V2")
+ val query = "select variant_get(v, '$') as whole from T_V2"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ val scanRelation =
findScanRelation(sql(query).queryExecution.optimizedPlan)
+ assert(scanRelation.output.map(_.name) == Seq("v"),
+ s"Expected scan output [v] but got
${scanRelation.output.map(_.name)}")
+ assert(scanRelation.output(0).dataType.isInstanceOf[VariantType],
+ s"Expected v left raw, but got ${scanRelation.output(0).dataType}")
+ }
+ }
+ }
+
+ test(s"V2 test - mixed extractions with order by a pushed path: shredded,
correct order " +
+ s"($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"name\": \"x\", \"price\": 3}'), parse_json('{\"j\":
1}')), " +
+ "(parse_json('{\"name\": \"z\", \"price\": 1}'), parse_json('{\"j\":
2}')), " +
+ "(parse_json('{\"name\": \"y\", \"price\": 2}'), parse_json('{\"j\":
3}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ spark.read.parquet(path).createOrReplaceTempView("T_V2")
+ // Two real extractions ($.name, $.price) are in the local Project,
and the ORDER BY lifts a
+ // bare v (-> fullVariant) for the sort. fullVariant is not the only
field, so v stays
+ // shredded (a multi-slot struct with a real full-variant slot, no
placeholder collapse) and
+ // the sort reads that real slot. Compare ORDER-SENSITIVELY.
+ val query = "select variant_get(v, '$.name', 'string') as nm, " +
+ "variant_get(v, '$.price', 'long') as pr " +
+ "from T_V2 order by variant_get(v, '$.price', 'long')"
+ val expectedOrder = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect().map(r => (r.getString(0), r.getLong(1))).toList
+ }
+ assert(expectedOrder == List(("z", 1L), ("y", 2L), ("x", 3L)),
+ s"baseline sanity: $expectedOrder")
+ val actualOrder = sql(query).collect().map(r => (r.getString(0),
r.getLong(1))).toList
+ assert(actualOrder == expectedOrder,
+ "ORDER BY a pushed path produced the wrong row order")
+ val scanRelation =
findScanRelation(sql(query).queryExecution.optimizedPlan)
+ val vAttr = scanRelation.output.find(_.name == "v").get
+ assert(vAttr.dataType.isInstanceOf[StructType],
+ s"Expected v shredded, but got ${vAttr.dataType}")
+ }
+ }
+ }
+
+ test(s"V2 test - order by a pushed path not in the select list: correct
order ($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"name\": \"x\", \"price\": 3}'), parse_json('{\"j\":
1}')), " +
+ "(parse_json('{\"name\": \"z\", \"price\": 1}'), parse_json('{\"j\":
2}')), " +
+ "(parse_json('{\"name\": \"y\", \"price\": 2}'), parse_json('{\"j\":
3}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ spark.read.parquet(path).createOrReplaceTempView("T_V2")
+ // Only $.name is selected (local); the sort key $.price is lifted as
a bare v (full
+ // variant). v shreds to {$.name, full-variant}; the sort reads the
real full-variant slot.
+ // The ORDER-SENSITIVE check guards against a placeholder mis-order.
+ val query = "select variant_get(v, '$.name', 'string') as nm " +
+ "from T_V2 order by variant_get(v, '$.price', 'long')"
+ val expectedOrder = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect().map(_.getString(0)).toList
+ }
+ assert(expectedOrder == List("z", "y", "x"), s"baseline sanity:
$expectedOrder")
+ assert(sql(query).collect().map(_.getString(0)).toList ==
expectedOrder,
+ "ORDER BY a non-selected pushed path produced the wrong row order")
+ }
+ }
+ }
+
+ test(s"V2 test - group by variant_get key: optimal shred, no full-variant
slot ($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"k\": \"a\"}'), parse_json('{\"j\": 1}')), " +
+ "(parse_json('{\"k\": \"a\"}'), parse_json('{\"j\": 2}')), " +
+ "(parse_json('{\"k\": \"b\"}'), parse_json('{\"j\": 3}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ spark.read.parquet(path).createOrReplaceTempView("T_V2")
+ // The grouping expression variant_get(v, '$.k') is pulled into a
Project below the
+ // Aggregate (PullOutGroupingExpressions), so it is visible to the
rewrite and shreds to
+ // the $.k slot -- a real extraction, no full-variant slot and no
bare-v lift. v2 is pruned.
+ val query =
+ "select variant_get(v, '$.k', 'string') as k, count(1) as n " +
+ "from T_V2 group by variant_get(v, '$.k', 'string')"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ val scanRelation =
findScanRelation(sql(query).queryExecution.optimizedPlan)
+ assert(scanRelation.output.map(_.name) == Seq("v"),
+ s"Expected scan output [v] but got
${scanRelation.output.map(_.name)}")
+ val vStruct = scanRelation.output(0).dataType match {
+ case s: StructType => s
+ case other => fail(s"Expected v shredded to struct, but got $other")
+ }
+ assert(!vStruct.fields.exists(_.dataType.isInstanceOf[VariantType]),
+ s"Expected optimal extraction with no full-variant slot, but got
$vStruct")
+ }
+ }
+ }
+
+ test(s"V2 test - distinct variant_get: optimal shred ($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"k\": \"a\"}'), parse_json('{\"j\": 1}')), " +
+ "(parse_json('{\"k\": \"a\"}'), parse_json('{\"j\": 2}')), " +
+ "(parse_json('{\"k\": \"b\"}'), parse_json('{\"j\": 3}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ spark.read.parquet(path).createOrReplaceTempView("T_V2")
+ // DISTINCT becomes an Aggregate whose grouping expression is the
distinct column; like
+ // GROUP BY, the variant_get is pulled below the Aggregate and shreds
to the $.k slot.
+ val query = "select distinct variant_get(v, '$.k', 'string') as k from
T_V2"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ val scanRelation =
findScanRelation(sql(query).queryExecution.optimizedPlan)
+ assert(scanRelation.output(0).dataType.isInstanceOf[StructType],
+ s"Expected v shredded, but got ${scanRelation.output(0).dataType}")
+ }
+ }
+ }
+
+ test(s"V2 test - window order by variant_get: no crash, shredded
($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"k\": \"a\", \"price\": 3}'), parse_json('{\"j\":
1}')), " +
+ "(parse_json('{\"k\": \"a\", \"price\": 1}'), parse_json('{\"j\":
2}')), " +
+ "(parse_json('{\"k\": \"b\", \"price\": 2}'), parse_json('{\"j\":
3}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ spark.read.parquet(path).createOrReplaceTempView("T_V2")
+ // Window is a fourth barrier type. The OVER (ORDER BY variant_get(v,
'$.price')) order key
+ // and the selected variant_get(v, '$.k') are both pushed; the query
must optimize without
+ // crashing and return correct results.
+ val query = "select variant_get(v, '$.k', 'string') as k, " +
+ "row_number() over (order by variant_get(v, '$.price', 'long')) as
rn from T_V2"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ val scanRelation =
findScanRelation(sql(query).queryExecution.optimizedPlan)
+ assert(scanRelation.output.find(_.name ==
"v").get.dataType.isInstanceOf[StructType],
+ "Expected v shredded for the pushed window extractions")
+ }
+ }
+ }
+
+ test(s"V2 test - two variant columns each extracted: both shredded
($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v variant, v2 variant) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}')), " +
+ "(parse_json('{\"a\": 9}'), parse_json('{\"b\": 8}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ spark.read.parquet(path).createOrReplaceTempView("T_V2")
+ // Both variant columns carry real extractions on different paths;
each shreds independently
+ // and both remain in the scan.
+ val query =
+ "select variant_get(v, '$.a', 'int') as a, variant_get(v2, '$.b',
'int') as b from T_V2"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ val scanRelation =
findScanRelation(sql(query).queryExecution.optimizedPlan)
+ assert(scanRelation.output.map(_.name).toSet == Set("v", "v2"),
+ s"Expected scan output {v, v2} but got
${scanRelation.output.map(_.name)}")
+ assert(scanRelation.output.forall(_.dataType.isInstanceOf[StructType]),
+ s"Expected both v and v2 shredded, but got " +
+ s"${scanRelation.output.map(a =>
s"${a.name}:${a.dataType.simpleString}")}")
+ }
+ }
+ }
+
+ test(s"V2 test - column pruning: no variant_get, normal pruning applies
($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v1 variant, v2 variant, s string) using
PARQUET " +
+ s"location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}'), 'hello')")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ val df = spark.read.parquet(path)
+ df.createOrReplaceTempView("T_V2")
+ // No variant_get: pushdown does not fire; normal pruneColumns applies.
+ val query = "select s from T_V2"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ sql(query).queryExecution.optimizedPlan match {
+ case scanRelation: DataSourceV2ScanRelation =>
+ assert(scanRelation.output.map(_.name) == Seq("s"),
+ s"Expected scan output [s] but got
${scanRelation.output.map(_.name)}")
+ case Project(_, scanRelation: DataSourceV2ScanRelation) =>
+ assert(scanRelation.output.map(_.name) == Seq("s"),
+ s"Expected scan output [s] but got
${scanRelation.output.map(_.name)}")
+ case other =>
+ fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other")
+ }
+ }
+ }
+ }
+
+ test(s"V2 test - column pruning: 2 variant cols, only second referenced
($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v1 variant, v2 variant) using PARQUET
location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 99}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ val df = spark.read.parquet(path)
+ df.createOrReplaceTempView("T_V2")
+ val query = "select variant_get(v2, '$.b', 'int') from T_V2"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ sql(query).queryExecution.optimizedPlan match {
+ case Project(_, scanRelation: DataSourceV2ScanRelation) =>
+ assert(scanRelation.output.map(_.name) == Seq("v2"),
+ s"Expected scan output [v2] but got
${scanRelation.output.map(_.name)}")
+ assert(scanRelation.output(0).dataType.isInstanceOf[StructType],
+ s"Expected v2 rewritten to struct, but got
${scanRelation.output(0).dataType}")
+ case other =>
+ fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other")
+ }
+ }
+ }
+ }
+
+ test(s"V2 test - column pruning: 3 variant cols, 2 referenced
($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v1 variant, v2 variant, v3 variant) using
PARQUET " +
+ s"location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}'),
parse_json('{\"c\": 3}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ val df = spark.read.parquet(path)
+ df.createOrReplaceTempView("T_V2")
+ val query = "select variant_get(v1, '$.a', 'int'), variant_get(v3,
'$.c', 'int') from T_V2"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ sql(query).queryExecution.optimizedPlan match {
+ case Project(_, scanRelation: DataSourceV2ScanRelation) =>
+ val names = scanRelation.output.map(_.name).toSet
+ assert(names == Set("v1", "v3"),
+ s"Expected scan output {v1, v3} but got $names")
+ scanRelation.output.foreach { a =>
+ assert(a.dataType.isInstanceOf[StructType],
+ s"Expected ${a.name} rewritten to struct, but got
${a.dataType}")
+ }
+ case other =>
+ fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other")
+ }
+ }
+ }
+ }
+
+ test(s"V2 test - column pruning: 3 variant cols, all 3 referenced
($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v1 variant, v2 variant, v3 variant) using
PARQUET " +
+ s"location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}'),
parse_json('{\"c\": 3}'))")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ val df = spark.read.parquet(path)
+ df.createOrReplaceTempView("T_V2")
+ val query = "select variant_get(v1, '$.a', 'int'), variant_get(v2,
'$.b', 'int'), " +
+ "variant_get(v3, '$.c', 'int') from T_V2"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ sql(query).queryExecution.optimizedPlan match {
+ case Project(_, scanRelation: DataSourceV2ScanRelation) =>
+ val names = scanRelation.output.map(_.name).toSet
+ assert(names == Set("v1", "v2", "v3"),
+ s"Expected all 3 columns in scan but got $names")
+ scanRelation.output.foreach { a =>
+ assert(a.dataType.isInstanceOf[StructType],
+ s"Expected ${a.name} rewritten to struct, but got
${a.dataType}")
+ }
+ case other =>
+ fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other")
+ }
+ }
+ }
+ }
+
+ test(s"V2 test - column pruning: 3 variant cols, none referenced via
variant_get ($readerName)") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ withTable("temp_v1") {
+ sql(s"create table temp_v1 (v1 variant, v2 variant, v3 variant, s
string) " +
+ s"using PARQUET location '$path'")
+ sql("insert into temp_v1 values " +
+ "(parse_json('{\"a\": 1}'), parse_json('{\"b\": 2}'),
parse_json('{\"c\": 3}'), 'hi')")
+ }
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ val df = spark.read.parquet(path)
+ df.createOrReplaceTempView("T_V2")
+ // Only s is selected; no variant_get, so extraction pushdown does not
fire.
+ // Normal pruneColumns removes v1/v2/v3.
+ val query = "select s from T_V2"
+ val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key ->
"false") {
+ sql(query).collect()
+ }
+ checkAnswer(sql(query), expectedRows)
+ sql(query).queryExecution.optimizedPlan match {
+ case scanRelation: DataSourceV2ScanRelation =>
+ assert(scanRelation.output.map(_.name) == Seq("s"),
+ s"Expected scan output [s] but got
${scanRelation.output.map(_.name)}")
+ case Project(_, scanRelation: DataSourceV2ScanRelation) =>
+ assert(scanRelation.output.map(_.name) == Seq("s"),
+ s"Expected scan output [s] but got
${scanRelation.output.map(_.name)}")
+ case other =>
+ fail(s"Unexpected plan shape: ${other.getClass.getName}\n$other")
+ }
+ }
+ }
+ }
+
test(s"V2 No push down for JSON ($readerName)") {
withTempPath { dir =>
val path = dir.getCanonicalPath
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
