qlong commented on code in PR #56556:
URL: https://github.com/apache/spark/pull/56556#discussion_r3449705846
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -418,14 +437,61 @@ 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
+ // Variant extraction pushdown is for extractions, not for whole-variant
reads. A bare variant
+ // reference (e.g. `SELECT v`, or a column lifted into the local Project
to feed a `variant_get`
+ // that sits above a Join/Sort/Aggregate barrier this local rewrite cannot
see) makes
+ // `collectRequestedFields` record `RequestedVariantField.fullVariant` --
path `$`, target
+ // VariantType -- meaning "produce the entire value." That is not a real
extraction:
+ // - No I/O benefit: the whole variant is read regardless, so shredding
it to a single
+ // full-variant slot never reads fewer bytes than leaving the column
raw.
+ // - It is the case readers mishandle: a lone full-variant slot
collapses to a boolean
+ // placeholder (see ParquetScan.rewriteVariantPushdownSchema), and
above a barrier the
+ // re-exposed `GetStructField(v_new, i) AS v#orig` alias can be
dropped by
+ // RemoveRedundantAliases, yielding wrong results or an invalid plan.
+ // So when a variant is read ONLY as a whole -- its sole requested field
is fullVariant -- leave
+ // it raw and let normal column pruning read it. Scoped to the sole-field
case on purpose:
+ // - It is exactly the broken shape (a lone full-variant slot is what
collapses / re-exposes).
+ // - When fullVariant coexists with real extractions on the same variant
(e.g.
+ // `SELECT v, variant_get(v, '$.a')`), the shredded struct has >= 2
slots, so the reader
+ // does not collapse it; keeping it shredded preserves the typed-path
pushdown.
+ // This also keeps the rule correct for free as barrier-aware pushdown
lands later: once a
+ // `variant_get` above a Join/Sort/Aggregate becomes visible to the
rewrite it is collected as a
+ // real typed path (not fullVariant), so it is no longer sole-fullVariant
and the column shreds
+ // automatically -- this guard simply steps aside. Runs before the
IsNull/IsNotNull injection
+ // below so presence-only references, which legitimately shred to a tiny
placeholder, are
+ // unaffected.
+ variants.mapping.values.foreach { pathToFields =>
+ pathToFields.filterInPlace { case (_, fields) =>
+ !(fields.size == 1 &&
fields.contains(RequestedVariantField.fullVariant))
Review Comment:
Fixed, targetType.isInstanceOf[VariantType]) is more direct, added tests for
2-arg.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -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
Review Comment:
Thanks for review. Changed to set `sHolder.pushedVariants = Some(new
VariantInRelation())` before retturn.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala:
##########
@@ -794,6 +805,832 @@ 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. v is lifted in only as a bare
reference -> a
+ // whole-variant request. Before the fix this shredded to a lone
full-variant slot, which
Review Comment:
fixed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
