From ef2a92e89780191afabd2298e40b11c14934be87 Mon Sep 17 00:00:00 2001 From: Qiegang Long Date: Thu, 9 Jul 2026 20:04:23 -0400 Subject: [PATCH] [SPARK-58089][SQL] Push variant extractions through Aggregate/Sort/Join `PushVariantIntoScan` (v1) and `V2ScanRelationPushDown` (v2) both rely on `PhysicalOperation`, which collapses only a contiguous `Project`/`Filter` chain and stops at `Aggregate`, `Sort`, and `Join`. As a result, `variant_get` expressions embedded in aggregate function arguments, sort-order expressions, or join conditions are invisible to the pushdown rules, and the full variant column is read raw instead of being shredded to the requested typed fields. When an `Aggregate` or `Sort` sits above a `Join`, the barrier compounds: even a `Project` hoisted above the join tree is still unreachable from the scan side. Example queries that fail to shred without this fix: -- Aggregate: extraction in agg arg / GROUP BY SELECT AVG(variant_get(data, '$.qty', 'double')) FROM store_sales GROUP BY variant_get(data, '$.id', 'string') -- Join: extraction in ON condition SELECT ss.k FROM store_sales ss JOIN date_dim d ON ss.date_sk = variant_get(d.data, '$.sk', 'int') -- Aggregate over Join (TPC-DS Q26 shape) SELECT AVG(variant_get(ss.data, '$.qty', 'double')) FROM store_sales ss JOIN date_dim d ON ss.date_sk = d.date_sk JOIN store st ON ss.store_sk = st.store_sk GROUP BY variant_get(ss.data, '$.id', 'string') In all cases the scan emits `data:struct<0:variant>` (full blob) instead of `data:struct<0:double, 1:string>`. Introduce a new optimizer rule `PullOutVariantExtractions`, registered as the first rule in `earlyScanPushDownRules` (before `SchemaPruning`), which hoists `variant_get` calls out of `Aggregate`, `Sort`, and `Join` into a `Project` directly below the operator so the downstream pushdown rules can see them. - **Aggregate**: hoist extractions from aggregate function arguments into a `Project` below the `Aggregate`. The `Aggregate` defines its own output so the raw variant column is not passed through. - **Sort / Join**: match the `Project` sitting directly above the `Sort`/`Join` (its `references` give the live-above set) and drop any variant column no longer referenced once its extraction is hoisted. - **Push-through-Join**: hoisting above a `Join` is not enough because `PhysicalOperation` stops there. The new `pushSideAliases` helper recursively routes each hoisted `_ve` alias down through the join tree to the side whose output owns the referenced attribute, landing it in a `Project` directly above the scan. Handles any depth of chained joins in one pass.for review --- .../apache/spark/sql/internal/SQLConf.scala | 13 + .../spark/sql/execution/SparkOptimizer.scala | 6 +- .../PullOutVariantExtractions.scala | 468 +++++++++ .../PushVariantIntoScanSuite.scala | 951 +++++++++++++++++- 4 files changed, 1404 insertions(+), 34 deletions(-) create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PullOutVariantExtractions.scala diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 22f5c8416e050..2232f9bbe087c 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -6490,6 +6490,19 @@ object SQLConf { .booleanConf .createWithDefault(true) + val PUSH_VARIANT_INTO_SCAN_PULL_OUT_EXTRACTIONS = + buildConf("spark.sql.variant.pushVariantIntoScan.pullOutExtractions") + .internal() + .doc("When true, extend variant field pushdown to extractions used inside aggregate " + + "functions, join conditions, sort keys, and projections above joins (including through " + + "chained joins), so that only the requested fields are read from the scan for these " + + "cases too instead of the whole variant column. Has no effect unless " + + "spark.sql.variant.pushVariantIntoScan is also true.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .booleanConf + .createWithDefault(true) + val PUSH_VARIANT_INTO_SCAN_DEFER_CAST_ERROR = buildConf("spark.sql.variant.pushVariantIntoScan.deferCastError") .internal() diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala index 54158d5bb4a94..0994220afe0c1 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala @@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.optimizer._ import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.connector.catalog.CatalogManager -import org.apache.spark.sql.execution.datasources.{MarkSingleTaskExecution, PruneFileSourcePartitions, PushVariantIntoScan, SchemaPruning, V1Writes} +import org.apache.spark.sql.execution.datasources.{MarkSingleTaskExecution, PruneFileSourcePartitions, PullOutVariantExtractions, PushVariantIntoScan, SchemaPruning, V1Writes} import org.apache.spark.sql.execution.datasources.v2.{GroupBasedRowLevelOperationScanPlanning, OptimizeMetadataOnlyDeleteFromTable, V2ScanPartitioningAndOrdering, V2ScanRelationPushDown, V2Writes} import org.apache.spark.sql.execution.dynamicpruning.{CleanupDynamicPruningFilters, PartitionPruning, RowLevelOperationRuntimeGroupFiltering} import org.apache.spark.sql.execution.python.{ExtractGroupingPythonUDFFromAggregate, ExtractPythonUDFFromAggregate, ExtractPythonUDFs, ExtractPythonUDTFs} @@ -37,6 +37,10 @@ class SparkOptimizer( override def earlyScanPushDownRules: Seq[Rule[LogicalPlan]] = // TODO: move SchemaPruning into catalyst Seq( + // Hoist variant extractions out of operators that variant-into-scan pushdown cannot see + // through (aggregate function arguments, etc.) into a Project above the scan, so the + // extractions below become visible to V2ScanRelationPushDown and PushVariantIntoScan. + PullOutVariantExtractions, SchemaPruning, GroupBasedRowLevelOperationScanPlanning, V1Writes, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PullOutVariantExtractions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PullOutVariantExtractions.scala new file mode 100644 index 0000000000000..94a1cb7b95ffd --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PullOutVariantExtractions.scala @@ -0,0 +1,468 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.datasources + +import scala.collection.mutable + +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeSet, Cast, Expression, + NamedExpression, SortOrder} +import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression +import org.apache.spark.sql.catalyst.expressions.variant.VariantGet +import org.apache.spark.sql.catalyst.plans.LeftExistence +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Join, LogicalPlan, Project, Sort} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.VariantType + +/** + * Hoists variant extractions out of operators that variant-into-scan pushdown cannot see through, + * into a [[Project]] directly below that operator, so the extraction becomes visible to the + * pushdown rules ([[PushVariantIntoScan]] for v1 and + * [[org.apache.spark.sql.execution.datasources.v2.V2ScanRelationPushDown]] for v2). The rule + * applies uniformly to all three table read paths: DS1 (e.g. `spark.read.parquet`), DS2 (e.g. + * `spark.read.parquet` with `spark.sql.sources.useV1SourceList` cleared), and Spark native + * catalog tables (managed or external tables created with `CREATE TABLE ... USING PARQUET`). In + * every case the downstream pushdown rule sees the hoisted extraction in the + * `Project`/`Filter` chain that [[org.apache.spark.sql.catalyst.planning.PhysicalOperation]] + * collapses above the scan; an extraction embedded in an `Aggregate`/`Join`/`Sort` is invisible + * to them and the whole variant column is read raw (or, when the column is also read raw for + * pass-through, shredded with a wasteful full-variant slot). + * + * This complements the rules that already relocate variant extractions into such a `Project`: + * `PhysicalOperation` (Project/Filter), `PullOutGroupingExpressions` (GROUP BY / DISTINCT keys), + * and `ExtractWindowExpressions` (window partition/order keys and window function arguments). The + * residual, un-relocated extractions live in aggregate function arguments, join conditions, and + * sort orders. This rule handles those three: + * + * - Aggregate: `Aggregate([name], [name, max(variant_get(v, '$.p'))], child)` becomes + * `Aggregate([name], [name, max(_ve0)], Project([name, variant_get(v, '$.p') AS _ve0], child))`. + * The Aggregate defines its own output, so the raw `v` is simply not passed through. + * + * - Sort / Join: these operators pass their child's columns through, so a naive hoist would keep + * a bare `v` alongside the hoisted extraction and the pushdown would still request the full + * variant. To avoid that we match the `Project` sitting above the barrier -- its `references` + * are exactly the columns still live above the barrier -- and drop a variant column that is no + * longer referenced once its extraction has been hoisted. A variant genuinely needed raw above + * the barrier stays in `references` and is preserved (keeping its full-variant slot). + * + * This handling REQUIRES a `Project` directly above the `Sort`/`Join`; that Project's + * `references` is what tells us which columns remain live. Any other shape above the barrier -- + * a bare/root `Sort`/`Join`, or one under a `Filter`/`Aggregate`/another operator -- is left + * unchanged: the extraction stays in the `Sort`/`Join`, the variant column is read raw, and the + * plan is correct (just not shredded). We deliberately do not hoist without a narrowing Project, + * because without knowing the live-above set we cannot safely decide whether the raw column can + * be dropped, and hoisting while blindly passing every column through would reintroduce the + * redundant full-variant slot this rule exists to remove. When the `Sort` itself sits over a + * `Join`, the hoisted order-key aliases are additionally pushed *through* the join to the scan + * side (see the through-a-`Join` note below), not merely parked in a `Project` above it. + * + * Pushing extractions *through* a `Join`: hoisting an extraction into a `Project`/child directly + * above a `Join` is not enough on its own, because `PhysicalOperation` (which both pushdown rules + * rely on) collapses only a contiguous `Project`/`Filter` chain and stops at the `Join`. A variant + * column `v` that reaches the barrier only via a hoisted `variant_get` would still flow up through + * the join as a bare attribute and be read raw at the scan. So once an extraction is hoisted above + * a `Join` (by the `Aggregate` or `Sort` case, or by the projection/condition handling of the + * `Project`-over-`Join` case), `pushSideAliases` pushes the resulting `_ve` aliases *down* through + * the join tree, routing each to the side whose output owns the referenced attribute, until it + * lands in a `Project` directly above a non-join child (the scan side). This descends any depth of + * chained joins in a single pass, so a variant fact-table column joined to several dimensions + * shreds to just the requested typed fields. Correctness is per join type: `Inner`/`Cross` and the + * outer joins (`LeftOuter`/`RightOuter`/`FullOuter`) push to either side -- the outer nullable side + * is value-preserving because `variant_get(NULL) = NULL` and null-padding commutes with the + * extraction. `LeftSemi`/`LeftAnti`/`ExistenceJoin` push only left-side aliases (their right side + * is not in the join output); a right-side alias on such a join is not pushed and the parent keeps + * a `Project` referencing it above the join. A variant still needed raw above the join stays live + * and keeps its full-variant slot, exactly as in the non-join cases. + * + * Every rewrite is a plain alias/`Project` introduction and is semantics-preserving; the pushdown + * shred/rewrite machinery is unchanged. The rule is gated on + * [[SQLConf.PUSH_VARIANT_INTO_SCAN_PULL_OUT_EXTRACTIONS]] (and is a no-op unless + * [[SQLConf.PUSH_VARIANT_INTO_SCAN]] is also enabled) and only fires when a hoistable extraction is + * present, so non-variant plans are untouched. + */ +object PullOutVariantExtractions extends Rule[LogicalPlan] { + + override def apply(plan: LogicalPlan): LogicalPlan = { + // Hoisting is only useful when variant-into-scan pushdown is enabled to consume the relocated + // extractions, and is independently gated so it can be turned off on its own. + if (!SQLConf.get.getConf(SQLConf.PUSH_VARIANT_INTO_SCAN) || + !SQLConf.get.getConf(SQLConf.PUSH_VARIANT_INTO_SCAN_PULL_OUT_EXTRACTIONS)) { + plan + } else { + // `transformUp` so inner operators are rewritten before their parents. + plan.transformUp { + case agg: Aggregate => rewriteAggregate(agg) + // Sort/Join are matched together with the Project above them so we know which columns + // remain live above the barrier (see class doc). A Sort/Join in any other shape -- + // bare/root or under a Filter/Aggregate/other operator -- is left unchanged (v read raw). + case p @ Project(projectList, s: Sort) => rewriteSortUnderProject(p, projectList, s) + case p @ Project(projectList, j: Join) => rewriteJoinUnderProject(p, projectList, j) + } + } + } + + // A variant column path: an `Attribute` (possibly through `GetStructField`) typed `VariantType`. + // Used by `isHoistable` to verify the extraction's source is a scan column, not a computed + // expression. Non-column variants (e.g. `variant_get(parse_json(x), ...)`) are excluded. + private def isVariantColumnPath(e: Expression): Boolean = + e.dataType.isInstanceOf[VariantType] && StructPath.unapply(e).isDefined + + private def isHoistable(e: Expression): Boolean = e match { + // Exclude the two-arg form variant_get(v, path) whose targetType is VariantType: it produces + // a variant-typed alias that the pushdown drops (see SPARK-57499), so hoisting adds a Project + // for no benefit and would also re-hoist on a second pass (breaking Once-batch idempotence). + case g: VariantGet => + !g.targetType.isInstanceOf[VariantType] && g.path.foldable && isVariantColumnPath(g.child) + // Only hoist `Cast` when its child is itself a `VariantGet` on a variant column path -- i.e. + // `cast(variant_get(v, '$.p', 'variant') as T)`. A `cast(v as T)` where `v` is a plain + // variant-typed attribute is handled directly by PushVariantIntoScan; hoisting it is + // unnecessary. More importantly, after a VariantGet is hoisted its alias `_ve0` has VariantType + // and would match `isVariantColumnPath`, causing `cast(_ve0 as string)` to look hoistable on + // a second pass and breaking the Once-batch idempotence check. + // A variant-to-variant Cast is also excluded: its sole requested field targets VariantType, + // which the pushdown drops (see SPARK-57499), so hoisting adds a Project for no benefit. + case c: Cast => + !c.dataType.isInstanceOf[VariantType] && + (c.child match { + case g: VariantGet => g.path.foldable && isVariantColumnPath(g.child) + case _ => false + }) + case _ => false + } + + /** + * Collects hoisted extractions as aliases, de-duplicated by canonical form so a repeated + * extraction maps to a single output slot. + */ + private class ExtractionHoister { + private val extracted = mutable.LinkedHashMap.empty[Expression, Alias] + + def aliases: Seq[NamedExpression] = extracted.values.toSeq + + def isEmpty: Boolean = extracted.isEmpty + + /** Returns (creating if needed) the alias attribute for a single hoistable extraction. */ + def aliasFor(ex: Expression): Attribute = + extracted.getOrElseUpdate(ex.canonicalized, Alias(ex, s"_ve${extracted.size}")()).toAttribute + + /** Replaces every hoistable extraction in `e` with a reference to its (new) alias. */ + def hoist(e: Expression): Expression = e.transformDown { + case ex if isHoistable(ex) => aliasFor(ex) + } + } + + // Recursively pushes hoisted extraction aliases down through a join tree until each lands in a + // `Project` directly above a non-join child (the scan side, or the `Filter`/`Project` chain above + // it that `PhysicalOperation` collapses). Hoisting an extraction into a `Project` above a `Join` + // is not sufficient on its own, because `PhysicalOperation` stops at the `Join`: the variant + // would still be read raw at the scan. This descends any depth of chained joins in a single call. + // + // `aliases` are the extraction aliases still to be placed. Each is routed to the join side whose + // output owns its referenced attribute. `needed` is the set of attributes that must remain in the + // pushed-into child's output -- the columns referenced above the push point plus the + // join-condition references accumulated while recursing, so join keys are never dropped. A + // variant consumed solely via a hoisted extraction is thereby not passed through (no redundant + // full-variant slot); one used raw elsewhere stays in `needed` and keeps its full-variant slot. + private def pushSideAliases( + child: LogicalPlan, + aliases: Seq[NamedExpression], + needed: AttributeSet): LogicalPlan = { + if (aliases.isEmpty) { + child + } else { + child match { + case join: Join => + val leftOutput = join.left.outputSet + val rightOutput = join.right.outputSet + // For LeftSemi/LeftAnti the right side is not in the join output, so a right-side alias + // could never be referenced above the join. Route only left-side aliases there; leave any + // right-side alias in place (it will not have been produced for such joins in practice). + val rightEligible = join.joinType match { + case LeftExistence(_) => false + case _ => true + } + val (leftAliases, rest) = + aliases.partition(_.references.subsetOf(leftOutput)) + val (rightAliases, unrouted) = + if (rightEligible) rest.partition(_.references.subsetOf(rightOutput)) + else (Seq.empty[NamedExpression], rest) + if (leftAliases.isEmpty && rightAliases.isEmpty) { + // Nothing cleanly routable (e.g. an ExistenceJoin, or an alias not confined to one + // side). Fall back to a Project above the join keeping the live columns. + val keep = join.output.filter(needed.contains) + Project(keep ++ aliases, join) + } else { + val neededHere = + needed ++ join.condition.map(_.references).getOrElse(AttributeSet.empty) + val newLeft = pushSideAliases(join.left, leftAliases, neededHere) + val newRight = pushSideAliases(join.right, rightAliases, neededHere) + val newJoin = join.copy(left = newLeft, right = newRight) + // An alias that could not be routed to a single side stays above the join. + if (unrouted.isEmpty) { + newJoin + } else { + val keep = newJoin.output.filter(needed.contains) + Project(keep ++ unrouted, newJoin) + } + } + // Descend through a pass-through `Project` (as `PhysicalOperation` does) to reach the + // join/scan beneath the intermediate Projects the optimizer leaves between chained joins. + // A `pushable` alias resolves against the grandchild and is pushed there; a `stay` alias + // references a column this Project introduces and cannot be pushed. + case project @ Project(projectList, grandChild) => + val (pushable, stay) = + aliases.partition(_.references.subsetOf(grandChild.outputSet)) + if (pushable.isEmpty) { + val keep = project.output.filter(needed.contains) + Project(keep ++ aliases, project) + } else { + // Prune to the still-live columns before re-adding the pushed alias references, else a + // bare variant consumed only via a now-pushed extraction would stay live and be read + // raw. Retain any projection a `stay` alias references so it stays resolvable. + val stayRefs = AttributeSet(stay.flatMap(_.references)) + val keptProjections = projectList.filter { e => + needed.contains(e.toAttribute) || stayRefs.contains(e.toAttribute) + } + val newGrandChild = pushSideAliases( + grandChild, pushable, needed ++ AttributeSet(keptProjections.flatMap(_.references))) + Project(keptProjections ++ pushable.map(_.toAttribute) ++ stay, newGrandChild) + } + // Any other operator -- notably a `Filter` above the scan: wrap a `Project` above it. We do + // NOT descend through a `Filter`, as that would force its variant predicate reference into + // `needed` and pass the raw variant through. The Filter and this Project collapse into the + // same `PhysicalOperation` chain, so the hoisted extraction is still seen at the scan. + case other => + val keep = other.output.filter(needed.contains) + Project(keep ++ aliases, other) + } + } + } + + // Pushes hoisted extraction aliases into a `Join` child: routes each alias to the join side whose + // output owns its referenced attribute and recurses via `pushSideAliases` so it lands in a + // `Project` directly above that side's scan (descending nested joins). The join propagates the + // `_ve` aliases up through its output, where the parent operator (Aggregate/Sort) references + // them. + // A `Project` wrapper above the join would leave the aliases invisible to the pushdown + // (`PhysicalOperation` stops at the join) and would not be revisited by the outer `transformUp`, + // so we push down instead. + // + // `needed` is the set of attributes that must remain live above the join (what the parent + // references). The join condition's references are added so join keys are never dropped. If any + // alias cannot be routed to exactly one pushable side -- a right-side alias on a `LeftExistence` + // join, or one that straddles both sides -- we fall back to wrapping a `Project` of the live + // columns plus every alias above the join, which keeps the parent's `_ve` references resolved (at + // the cost of reading that side's variant raw). + private def pushAliasesIntoJoin( + join: Join, + aliases: Seq[NamedExpression], + needed: AttributeSet): LogicalPlan = { + val leftOutput = join.left.outputSet + val rightOutput = join.right.outputSet + // For LeftSemi/LeftAnti/ExistenceJoin the right side is not in the join output, so a right-side + // alias could never be referenced above the join; only the left side is pushable. + val rightEligible = join.joinType match { + case LeftExistence(_) => false + case _ => true + } + val leftAliases = aliases.filter(_.references.subsetOf(leftOutput)) + val rightAliases = + if (rightEligible) { + aliases.filter(a => + !a.references.subsetOf(leftOutput) && a.references.subsetOf(rightOutput)) + } else { + Seq.empty[NamedExpression] + } + if (leftAliases.size + rightAliases.size != aliases.size) { + val keep = join.output.filter(needed.contains) + Project(keep ++ aliases, join) + } else { + val neededHere = needed ++ join.condition.map(_.references).getOrElse(AttributeSet.empty) + join.copy( + left = pushSideAliases(join.left, leftAliases, neededHere), + right = pushSideAliases(join.right, rightAliases, neededHere)) + } + } + + // Routes hoistable extractions found in `projectList` to the owning-side hoister and returns the + // projectList with those extractions replaced by references to their (new) alias attributes. + // Mirrors the join-condition routing in `rewriteJoinUnderProject`: a single variant extraction + // references exactly one attribute, hence one join side. Anything not cleanly on one side is left + // in place. + private def hoistProjectListExtractions( + projectList: Seq[NamedExpression], + leftOutput: AttributeSet, + rightOutput: AttributeSet, + leftHoister: ExtractionHoister, + rightHoister: ExtractionHoister): Seq[NamedExpression] = { + projectList.map { e => + e.transformDown { + case ex if isHoistable(ex) => + if (ex.references.subsetOf(leftOutput)) { + leftHoister.aliasFor(ex) + } else if (ex.references.subsetOf(rightOutput)) { + rightHoister.aliasFor(ex) + } else { + ex + } + }.asInstanceOf[NamedExpression] + } + } + + private def rewriteAggregate(agg: Aggregate): LogicalPlan = { + val hoister = new ExtractionHoister + // Only hoist extractions that sit inside an aggregate function's arguments (or filter). A + // top-level extraction in `aggregateExpressions` is a grouping-key reference (grouping keys are + // already pulled out by `PullOutGroupingExpressions`); hoisting it would leave the `Aggregate` + // referencing a column that is neither grouped nor aggregated. + val newAggExprs = agg.aggregateExpressions.map { + _.transform { + case ae: AggregateExpression => ae.withNewChildren(ae.children.map(hoister.hoist)) + }.asInstanceOf[NamedExpression] + } + + if (hoister.isEmpty) { + agg + } else { + // `referenced` is what the rewritten Aggregate still needs directly (aggregate-function args + // after hoisting, plus grouping keys). A variant consumed solely via a hoisted extraction is + // not in it and so is dropped rather than passed through raw -- see the class doc. + val referenced = AttributeSet( + newAggExprs.flatMap(_.references) ++ agg.groupingExpressions.flatMap(_.references)) + val newChild = agg.child match { + // Fuse into the Project that PullOutGroupingExpressions/PhysicalOperation already placed + // below the Aggregate; the Once-batch earlyScanPushDownRules has no CollapseProject to + // flatten a second stacked Project. Sound only when the fused Project still resolves + // against `grandChild.output` (see `canFuseIntoChildProject`); else keep the nested one. + case Project(projectList, grandChild) + if canFuseIntoChildProject(projectList, referenced, hoister, grandChild) => + val keptProjections = projectList.filter(e => referenced.contains(e.toAttribute)) + val fused = Project(keptProjections ++ hoister.aliases, grandChild) + grandChild match { + // The fused Project sits directly above a Join, which `PhysicalOperation` cannot see + // through -- and this newly created node will not be revisited by the outer + // `transformUp` in this pass. Push the hoisted aliases the rest of the way down through + // the join immediately, reusing the Project-over-Join handling. + case grandChildJoin: Join => + rewriteJoinUnderProject(fused, fused.projectList, grandChildJoin) + case _ => fused + } + // No fusable child Project. Only the still-referenced columns plus the hoisted aliases + // reach the child, so a variant used solely via a hoisted extraction is dropped. + case join: Join => + // The Aggregate's child is a Join. Push the hoisted aliases down through the join tree + // onto the owning side(s), where they land directly above the scan; the Aggregate + // references the `_ve` aliases, which the join propagates up through its output. Any + // variant still needed raw (in `referenced`) stays live. See `pushAliasesIntoJoin`. + pushAliasesIntoJoin(join, hoister.aliases, referenced) + case other => + val passthrough = other.output.filter(referenced.contains) + Project(passthrough ++ hoister.aliases, other) + } + agg.copy(aggregateExpressions = newAggExprs, child = newChild) + } + } + + // Fusing the hoisted aliases into an existing child `Project` (collapsing it onto its own child) + // is sound only if everything the fused Project carries still resolves against `grandChild`: + // - the retained projections (those whose output is still referenced), and + // - the hoisted alias expressions. + // If any of these references an attribute the Project introduces itself (not present in + // `grandChild.output`) -- e.g. a nested-variant extraction materialized as an intermediate alias + // -- collapsing would leave that reference unresolved, so we must not fuse. + private def canFuseIntoChildProject( + projectList: Seq[NamedExpression], + referenced: AttributeSet, + hoister: ExtractionHoister, + grandChild: LogicalPlan): Boolean = { + val grandChildOutput = grandChild.outputSet + val keptProjections = projectList.filter(e => referenced.contains(e.toAttribute)) + val carried = keptProjections ++ hoister.aliases + carried.forall(_.references.subsetOf(grandChildOutput)) + } + + private def rewriteSortUnderProject( + project: Project, projectList: Seq[NamedExpression], sort: Sort): LogicalPlan = { + val hoister = new ExtractionHoister + val newOrder = sort.order.map(hoister.hoist(_).asInstanceOf[SortOrder]) + if (hoister.isEmpty) { + project + } else { + // Columns still needed from the Sort's child: what the Project above consumes, plus what the + // rewritten order references directly (the alias attributes are supplied below, not from the + // child). A variant used only by the (now hoisted) order key is thereby dropped. + val needed = + AttributeSet(projectList.flatMap(_.references) ++ newOrder.flatMap(_.references)) + // `pushSideAliases` places the aliases in a Project directly above the child's non-join + // descendant, descending through any Join/Project chain between the Sort and the scan. For a + // flat child (scan / Filter above a scan) this is a single Project directly below the Sort, + // matching the previous behavior; when the Sort sits over a Join the aliases are pushed + // through it so the extraction reaches the scan instead of being read raw. + val newChild = pushSideAliases(sort.child, hoister.aliases, needed) + // Reproduce the original projectList so the alias columns do not leak into the output. + project.copy(child = sort.copy(order = newOrder, child = newChild)) + } + } + + private def rewriteJoinUnderProject( + project: Project, projectList: Seq[NamedExpression], join: Join): LogicalPlan = { + val leftOutput = join.left.outputSet + val rightOutput = join.right.outputSet + val leftHoister = new ExtractionHoister + val rightHoister = new ExtractionHoister + + // A single variant extraction references exactly one attribute, hence one join side; route it + // to that side's hoister. Anything not cleanly on one side is left in place. We hoist from both + // the join condition (its extractions feed the equi-key comparison) and the Project above the + // join (its extractions -- e.g. aggregate arguments hoisted here by `rewriteAggregate`, or a + // user's `SELECT variant_get(...)`), so the pushdown sees them below the join. + val newCondition = join.condition.map(_.transformDown { + case ex if isHoistable(ex) => + if (ex.references.subsetOf(leftOutput)) { + leftHoister.aliasFor(ex) + } else if (ex.references.subsetOf(rightOutput)) { + rightHoister.aliasFor(ex) + } else { + ex + } + }) + val newProjectList = + hoistProjectListExtractions(projectList, leftOutput, rightOutput, leftHoister, rightHoister) + + if (leftHoister.isEmpty && rightHoister.isEmpty) { + project + } else { + // Columns still live above the Join: what the (rewritten) Project consumes and what the + // rewritten condition references directly (alias attributes are supplied by the side Projects + // below). A side's variant used only by a hoisted extraction is dropped rather than passed + // through, so no redundant full-variant slot is requested. + val needed = AttributeSet( + newProjectList.flatMap(_.references) ++ + newCondition.map(_.references).getOrElse(AttributeSet.empty)) + // `pushSideAliases` places the aliases in a Project directly above the side's non-join child + // (recursing through nested joins), so the extractions reach the scan even when the side is + // itself a chain of joins. For a flat side (scan / Filter above a scan) this is a single + // Project, matching the previous behavior. + val newJoin = join.copy( + left = pushSideAliases(join.left, leftHoister.aliases, needed), + right = pushSideAliases(join.right, rightHoister.aliases, needed), + condition = newCondition) + project.copy(projectList = newProjectList, child = newJoin) + } + } +} 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 67bee55c22756..7f5a27c78e3c7 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 @@ -55,6 +55,63 @@ trait PushVariantIntoScanSuiteBase extends SharedSparkSession { } } + // A single parquet-backed table for `withVariantParquetTables`: the temp view name to expose, its + // column schema, and the row literals to insert. + protected case class VariantTable(view: String, schema: String, inserts: Seq[String]) + + // Like `withVariantParquetData` but sets up several parquet tables at once, each exposed as a + // temp view under the current read path (V1 or V2, per `useV2`). Used by the join tests, which + // need a variant fact table plus dimension tables. Each table is written via V1 parquet to an + // external location, then re-read as a view so the V2 suites exercise the DSv2 scan. + protected def withVariantParquetTables(tables: VariantTable*)(body: => Unit): Unit = { + withTempPath { dir => + val setupNames = tables.map(t => s"temp_setup_${t.view}") + withTable(setupNames: _*) { + tables.zip(setupNames).foreach { case (t, setupName) => + val path = s"${dir.getCanonicalPath}/${t.view}" + sql(s"create table $setupName (${t.schema}) using PARQUET location '$path'") + t.inserts.foreach(values => sql(s"insert into $setupName values $values")) + } + } + val sourceListConf: Seq[(String, String)] = + if (useV2) Seq(SQLConf.USE_V1_SOURCE_LIST.key -> "") else Nil + withSQLConf(sourceListConf: _*) { + tables.foreach { t => + spark.read.parquet(s"${dir.getCanonicalPath}/${t.view}").createOrReplaceTempView(t.view) + } + try body finally tables.foreach(t => spark.catalog.dropTempView(t.view)) + } + } + } + + // Find the scan relation (V1 or V2) whose output contains a column named `columnName`, and return + // that column's data type. Used by multi-scan join tests to inspect the fact table's variant + // column. Fails if zero or more than one matching scan is found. + protected def scanColumnType(plan: LogicalPlan, columnName: String): DataType = { + val matching = plan.collect { + case s: DataSourceV2ScanRelation if s.output.exists(_.name == columnName) => s.output + case l: LogicalRelation if l.output.exists(_.name == columnName) => l.output + } + assert(matching.length == 1, + s"Expected exactly one scan with a column named '$columnName' but found " + + s"${matching.length}:\n$plan") + matching.head.find(_.name == columnName).get.dataType + } + + // Assert that a variant column shredded to a struct and (optionally) that it has no full-variant + // slot -- i.e. the whole blob is not read. + protected def assertShreddedStruct( + dataType: DataType, expectFullVariantSlot: Boolean = false): StructType = { + val struct = dataType match { + case s: StructType => s + case other => fail(s"Expected the variant column shredded to a struct, but got $other") + } + val hasFullVariantSlot = struct.fields.exists(_.dataType.isInstanceOf[VariantType]) + assert(hasFullVariantSlot == expectFullVariantSlot, + s"Expected full-variant slot present=$expectFullVariantSlot, but got $struct") + struct + } + protected def localTimeZone = spark.sessionState.conf.sessionLocalTimeZone // Return a `StructField` with the expected `VariantMetadata`. @@ -73,6 +130,217 @@ trait PushVariantIntoScanSuiteBase extends SharedSparkSession { } } + // Locate the single scan relation in an optimized plan and return its output, regardless of the + // read path (V1 `LogicalRelation` or V2 `DataSourceV2ScanRelation`). Lets a shared test assert + // shredding for both the V1 and V2 pushdown rules. + protected def scanRelationOutput(plan: LogicalPlan): Seq[Attribute] = { + val outputs = plan.collect { + case s: DataSourceV2ScanRelation => s.output + case l: LogicalRelation => l.output + } + assert(outputs.length == 1, + s"Expected exactly one scan relation but found ${outputs.length}:\n$plan") + outputs.head + } + + test("aggregate function argument variant_get is hoisted below the aggregate and shredded") { + // `max(variant_get(v, '$.price'))` lives inside the aggregate function, above the aggregate + // barrier that PhysicalOperation-based pushdown cannot see through. PullOutVariantExtractions + // hoists it into a Project directly below the Aggregate so both the V1 and V2 pushdown rules + // shred `v` to just the `$.price` slot, with no full-variant slot (the raw column is never + // read). Runs under both the V1 and V2 pushdown paths. + withVariantParquetData( + "v variant, v2 variant, name string", + "(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')") { + val query = "select name, max(variant_get(v, '$.price', 'int')) as mx from T group by name" + val expected = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expected) // a -> 30, b -> 20 + val output = scanRelationOutput(sql(query).queryExecution.optimizedPlan) + val v = output.find(_.name == "v").getOrElse(fail(s"v missing from ${output.map(_.name)}")) + val vStruct = v.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 no full-variant slot, but got $vStruct") + } + } + + test("sort order variant_get is hoisted below the sort and shredded") { + // The sort key variant_get(v, '$.price') lives in Sort.order. PullOutVariantExtractions hoists + // it into a Project below the Sort; since only `name` is selected (v is not otherwise live + // above the Sort), v shreds to the $.price slot with no full-variant slot. Runs under V1/V2. + withVariantParquetData( + "v variant, name string", + "(parse_json('{\"price\": 3}'), 'x')", + "(parse_json('{\"price\": 1}'), 'z')", + "(parse_json('{\"price\": 2}'), 'y')") { + val query = "select name from T 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: $expectedOrder") + assert(sql(query).collect().map(_.getString(0)).toList == expectedOrder, + "ORDER BY variant_get produced the wrong row order") + val output = scanRelationOutput(sql(query).queryExecution.optimizedPlan) + val vStruct = output.find(_.name == "v").get.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 no full-variant slot, but got $vStruct") + } + } + + test("cast(v as variant) in aggregate function arg is not hoisted") { + // The aggregate argument is `cast(v as variant)`: a whole-variant read whose sole requested + // field targets VariantType. PullOutVariantExtractions must NOT hoist it (isHoistable excludes + // VariantType-target casts) -- hoisting would add a Project + `_ve` alias for no shredding + // benefit (the pushdown drops a lone VariantType slot anyway, see SPARK-57499). Assert the rule + // introduces no hoist alias, so the plan matches what the pull-out rule being off would give. + // Plan-shape only (no execution): a lone-VariantType shred is a separate, pre-existing concern. + // Runs under both the V1 and V2 pushdown paths. + withVariantParquetData( + "v variant, name string", + "(parse_json('{\"price\": 10}'), 'a')", + "(parse_json('{\"price\": 20}'), 'b')") { + val query = "select name, count(cast(v as variant)) as c from T group by name" + val plan = sql(query).queryExecution.optimizedPlan + val hoistAlias = plan.exists(_.expressions.exists(_.exists { + case a: Alias => a.name.startsWith("_ve") + case _ => false + })) + assert(!hoistAlias, + s"Expected no `_ve` hoist alias for a VariantType-target cast, but plan was:\n$plan") + } + } + + test("two-arg variant_get (returns VariantType) in aggregate function arg is not hoisted") { + // variant_get(v, '$.price') with no target type returns VariantType. isHoistable must NOT hoist + // it: the pushdown drops a lone VariantType requested field (SPARK-57499), so hoisting + // produces a `_ve` alias for no shredding benefit. It would also re-match on the second + // idempotence pass, breaking the Once-batch check. Assert no `_ve` alias is introduced. + withVariantParquetData( + "v variant, name string", + "(parse_json('{\"price\": 10}'), 'a')", + "(parse_json('{\"price\": 20}'), 'b')") { + val query = "select name, count(variant_get(v, '$.price')) as c from T group by name" + val plan = sql(query).queryExecution.optimizedPlan + val hoistAlias = plan.exists(_.expressions.exists(_.exists { + case a: Alias => a.name.startsWith("_ve") + case _ => false + })) + assert(!hoistAlias, + s"Expected no `_ve` hoist alias for a two-arg variant_get (VariantType result), " + + s"but plan was:\n$plan") + } + } + + test("two-arg variant_get (returns VariantType) in sort order is not hoisted") { + // Same exclusion as above but for a Sort order key. The rule must not introduce a `_ve` alias + // for a VariantType-returning extraction in Sort.order. Assert no `_ve` alias and that the + // query executes correctly (no idempotence error). + withVariantParquetData( + "v variant, name string", + "(parse_json('{\"price\": 3}'), 'x')", + "(parse_json('{\"price\": 1}'), 'z')", + "(parse_json('{\"price\": 2}'), 'y')") { + val query = "select name from T order by variant_get(v, '$.price')" + val plan = sql(query).queryExecution.optimizedPlan + val hoistAlias = plan.exists(_.expressions.exists(_.exists { + case a: Alias => a.name.startsWith("_ve") + case _ => false + })) + assert(!hoistAlias, + s"Expected no `_ve` hoist alias for a two-arg variant_get sort key, but plan was:\n$plan") + // Must execute without Once-batch idempotence error. + assert(sql(query).collect().length == 3) + } + } + + test("listagg(DISTINCT v:a::string) WITHIN GROUP (ORDER BY v:a::string) does not break " + + "Once-batch idempotence") { + // Regression test for the idempotence bug: `v:a::string` desugars to + // `cast(variant_get(v, '$.a') as string)`. The VariantGet is hoisted to a `_ve0: VariantType` + // alias. Without the fix, `cast(_ve0 as string)` would match isHoistable's Cast branch on the + // second pass (since `_ve0` is a VariantType-typed Attribute), producing a different plan and + // crashing with "Once strategy's idempotence is broken". With the fix, the Cast branch requires + // the child to itself be a VariantGet, so the already-hoisted alias is not re-hoisted. + withVariantParquetData( + "v variant", + "(parse_json('{\"a\": \"x\"}'))", + "(parse_json('{\"a\": \"y\"}'))", + "(parse_json('{\"a\": \"x\"}'))") { + val query = + "select listagg(distinct variant_get(v, '$.a', 'string'), ',') " + + "within group (order by variant_get(v, '$.a', 'string')) from T" + // Must execute without an idempotence exception and return the expected deduplicated result. + checkAnswer(sql(query), Seq(org.apache.spark.sql.Row("x,y"))) + } + } + + test("sort key hoisted while a different projected extraction stays above the sort: correct") { + // The projected extraction variant_get(v, '$.a') sits in the select list (above the Sort, + // lifted by PhysicalOperation), while the sort key variant_get(v, '$.b') is hoisted below the + // Sort. Because the projectList still references v raw, the hoist must NOT drop v -- it stays + // live above the Sort so the projected extraction can be computed. This guards the liveness + // computation in rewriteSortUnderProject: results and row order must both be correct. Runs + // under both the V1 and V2 pushdown paths. + withVariantParquetData( + "v variant, name string", + "(parse_json('{\"a\": 100, \"b\": 3}'), 'x')", + "(parse_json('{\"a\": 200, \"b\": 1}'), 'z')", + "(parse_json('{\"a\": 300, \"b\": 2}'), 'y')") { + val query = + "select name, variant_get(v, '$.a', 'int') as a " + + "from T order by variant_get(v, '$.b', 'int')" + val expected = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect().toSeq + } + // Baseline sanity: ordered by $.b ascending -> z(1), y(2), x(3) with their $.a values. + assert(expected.map(r => (r.getString(0), r.getInt(1))) == + Seq(("z", 200), ("y", 300), ("x", 100)), s"baseline sanity: $expected") + // Order-sensitive comparison: a mis-hoist that dropped v or mis-shredded would change the + // projected $.a value or the row order, which an order-insensitive checkAnswer would miss. + assert(sql(query).collect().toSeq.map(r => (r.getString(0), r.getInt(1))) == + expected.map(r => (r.getString(0), r.getInt(1))), + "projected extraction / order incorrect after hoisting the sort key") + } + } + + test("group by + aggregate hoist fuses into one Project, no stacked Project") { + // GROUP BY variant_get(v, '$.k') is relocated below the Aggregate by PullOutGroupingExpressions + // (creating a Project), and the aggregated variant_get(v, '$.price') is hoisted by + // PullOutVariantExtractions. The hoist must FUSE its aliases into that existing Project rather + // than stack a second one -- the earlyScanPushDownRules batch runs Once with no CollapseProject + // to flatten it afterward. Assert no two directly-stacked Projects survive in the optimized + // plan. Runs under both the V1 and V2 pushdown paths. + withVariantParquetData( + "v variant", + "(parse_json('{\"k\": \"a\", \"price\": 10}'))", + "(parse_json('{\"k\": \"a\", \"price\": 30}'))", + "(parse_json('{\"k\": \"b\", \"price\": 20}'))") { + val query = "select variant_get(v, '$.k', 'string') as k, " + + "max(variant_get(v, '$.price', 'int')) as mx " + + "from T group by variant_get(v, '$.k', 'string')" + val expected = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expected) // a -> 30, b -> 20 + val plan = sql(query).queryExecution.optimizedPlan + val stacked = plan.exists { + case Project(_, _: Project) => true + case _ => false + } + assert(!stacked, + s"Expected no directly-stacked Project nodes after hoist fusion, but got:\n$plan") + } + } + // Returns true iff `t` or any of its causes is an INVALID_VARIANT_CAST error. The failure may // surface directly or be wrapped in a task failure. protected def hasCastCondition(t: Throwable): Boolean = t match { @@ -275,6 +543,280 @@ trait PushVariantIntoScanSuiteBase extends SharedSparkSession { } } } + + // A variant fact table `SS(k int, data variant)` used by the push-through-join tests below. + private def factTable(view: String = "SS"): VariantTable = VariantTable( + view, "k int, data variant", + Seq( + "(1, parse_json('{\"qty\": 10, \"id\": \"a\"}'))", + "(1, parse_json('{\"qty\": 30, \"id\": \"a\"}'))", + "(2, parse_json('{\"qty\": 20, \"id\": \"b\"}'))")) + + test("aggregate over a single inner join: variant hoisted through the join and shredded") { + // AVG(variant_get(ss.data, '$.qty')) sits in an Aggregate whose child is a Join. The extraction + // is hoisted, then pushed down through the join onto the ss side so it lands directly above the + // ss scan. `data` is used ONLY via extractions (aggregate arg + GROUP BY key), so it shreds + // with no full-variant slot -- the whole blob is never read. Runs under both V1 and V2. + withVariantParquetTables( + factTable(), + VariantTable("DIM", "k int, name string", Seq("(1, 'x')", "(2, 'y')"))) { + val query = + "select variant_get(ss.data, '$.id', 'string') as id, " + + "avg(variant_get(ss.data, '$.qty', 'double')) as avg_qty " + + "from SS ss join DIM d on ss.k = d.k " + + "group by variant_get(ss.data, '$.id', 'string')" + val expected = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expected) // a -> 20.0, b -> 20.0 + assertShreddedStruct(scanColumnType(sql(query).queryExecution.optimizedPlan, "data")) + } + } + + test("aggregate over chained joins: variant pushed through all join levels and shredded") { + // The ss variant is joined to three dimensions. The aggregate-argument extractions must be + // pushed down through EVERY join level to reach the ss scan. This guards the recursion in + // pushSideAliases: a single non-recursive hoist would leave the extraction above the outermost + // join and `data` would still be read raw. Runs under both V1 and V2. + withVariantParquetTables( + factTable(), + VariantTable("D1", "k int, a string", Seq("(1, 'p')", "(2, 'q')")), + VariantTable("D2", "k int, b string", Seq("(1, 'r')", "(2, 's')")), + VariantTable("D3", "k int, c string", Seq("(1, 't')", "(2, 'u')"))) { + val query = + "select variant_get(ss.data, '$.id', 'string') as id, " + + "max(variant_get(ss.data, '$.qty', 'double')) as max_qty " + + "from SS ss " + + "join D1 on ss.k = D1.k join D2 on ss.k = D2.k join D3 on ss.k = D3.k " + + "group by variant_get(ss.data, '$.id', 'string')" + val expected = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expected) // a -> 30.0, b -> 20.0 + assertShreddedStruct(scanColumnType(sql(query).queryExecution.optimizedPlan, "data")) + } + } + + test("variant used raw and extracted across a join: keeps a full-variant slot") { + // ss.data is both selected raw (needs the whole blob) and extracted. After the extraction is + // pushed to the ss scan side, `data` stays live (still referenced raw), so it shreds to the + // typed slot PLUS a full-variant slot -- the raw value is preserved. Runs under both V1 and V2. + withVariantParquetTables( + factTable(), + VariantTable("DIM", "k int, name string", Seq("(1, 'x')", "(2, 'y')"))) { + val query = + "select ss.data as raw, variant_get(ss.data, '$.qty', 'double') as qty " + + "from SS ss join DIM d on ss.k = d.k" + val expected = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expected) + val struct = + assertShreddedStruct( + scanColumnType(sql(query).queryExecution.optimizedPlan, "data"), + expectFullVariantSlot = true) + // One typed slot ($.qty) plus the full-variant slot. + assert(struct.fields.length == 2, s"Expected two slots, but got $struct") + } + } + + test("aggregate over a left outer join: extraction on the preserved side is shredded") { + // ss is the preserved (left) side of a LEFT JOIN. Its variant extraction pushes onto the left + // side; unmatched rows (null-padded on the right) do not affect the ss-side extraction. Assert + // no full-variant slot and correct results including the unmatched row. Runs under V1 and V2. + withVariantParquetTables( + factTable(), + VariantTable("DIM", "k int, name string", Seq("(1, 'x')"))) { + val query = + "select variant_get(ss.data, '$.id', 'string') as id, " + + "max(variant_get(ss.data, '$.qty', 'double')) as max_qty " + + "from SS ss left join DIM d on ss.k = d.k " + + "group by variant_get(ss.data, '$.id', 'string')" + val expected = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expected) // a -> 30.0 (k=1 matched), b -> 20.0 (k=2 unmatched) + assertShreddedStruct(scanColumnType(sql(query).queryExecution.optimizedPlan, "data")) + } + } + + test("pull-out flag off reads the whole variant across a join") { + // With PUSH_VARIANT_INTO_SCAN_PULL_OUT_EXTRACTIONS disabled, the aggregate-argument extractions + // are not hoisted or pushed through the join. `ss.data` flows up through the join as a bare + // attribute, so the scan reads the whole blob -- either left raw as VariantType, or shredded to + // a single full-variant slot (struct<0:variant>). Either way there is no shredding benefit. + // Guards the config gate for the new join push-through path. Runs under both V1 and V2. + withVariantParquetTables( + factTable(), + VariantTable("DIM", "k int, name string", Seq("(1, 'x')", "(2, 'y')"))) { + val query = + "select variant_get(ss.data, '$.id', 'string') as id, " + + "avg(variant_get(ss.data, '$.qty', 'double')) as avg_qty " + + "from SS ss join DIM d on ss.k = d.k " + + "group by variant_get(ss.data, '$.id', 'string')" + withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN_PULL_OUT_EXTRACTIONS.key -> "false") { + val dataType = scanColumnType(sql(query).queryExecution.optimizedPlan, "data") + val readsWholeBlob = dataType match { + case _: VariantType => true + case s: StructType => s.fields.exists(_.dataType.isInstanceOf[VariantType]) + case _ => false + } + assert(readsWholeBlob, + s"Expected data read as a whole blob with the pull-out flag off, but got $dataType") + } + } + } + + test("aggregate over a left semi join: left-side extraction pushed and shredded") { + // A LEFT SEMI join's output is only the left (ss) columns. The ss-side aggregate extraction is + // pushed onto the left side and shreds with no full-variant slot; the semi join's existence + // semantics are unchanged. Runs under both V1 and V2. + withVariantParquetTables( + factTable(), + VariantTable("DIM", "k int", Seq("(1)"))) { + val query = + "select variant_get(ss.data, '$.id', 'string') as id, " + + "max(variant_get(ss.data, '$.qty', 'double')) as max_qty " + + "from SS ss left semi join DIM d on ss.k = d.k " + + "group by variant_get(ss.data, '$.id', 'string')" + val expected = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expected) // only k=1 rows survive -> a -> 30.0 + assertShreddedStruct(scanColumnType(sql(query).queryExecution.optimizedPlan, "data")) + } + } + + test("left outer join: extraction on the NULLABLE (right) side is pushed and value-preserving") { + // The extraction reads the DIM (right) variant, which is the NULLABLE side of a LEFT OUTER + // join: unmatched left rows are null-padded on the right. Pushing the extraction below the join + // computes variant_get BEFORE null-padding; the join then null-pads the `_ve` alias for + // unmatched rows. This must equal the baseline, where variant_get(NULL) is evaluated ABOVE the + // join for those rows -- i.e. null-padding commutes with the extraction (variant_get(NULL) = + // NULL). Asserting correctness here is the real point; also assert DIM.data shreds. The `ss` + // (left) side has its own unrelated extraction so both scans are exercised. Runs under V1/V2. + withVariantParquetTables( + factTable(), + VariantTable( + "DIM", "k int, data variant", + // Only k=1 matches SS's k in {1, 2}; SS rows with k=2 are unmatched -> right side null. + Seq("(1, parse_json('{\"label\": \"L1\"}'))"))) { + val query = + "select variant_get(ss.data, '$.id', 'string') as id, " + + "variant_get(d.data, '$.label', 'string') as label " + + "from SS ss left join DIM d on ss.k = d.k " + + "order by id, label" + val expected = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + // Order-sensitive comparison so a wrong null for unmatched rows cannot be masked. + assert(sql(query).collect().toSeq.map(r => (r.get(0), r.get(1))) == + expected.toSeq.map(r => (r.get(0), r.get(1))), + "Nullable-side extraction over a left outer join produced wrong values") + // Both SS.data and DIM.data are extracted only via variant_get; each shreds with no + // full-variant slot on its own scan (both scans expose a `data` column, so assert per scan). + val plan = sql(query).queryExecution.optimizedPlan + val dataTypes = plan.collect { + case s: DataSourceV2ScanRelation => s.output + case l: LogicalRelation => l.output + }.flatMap(_.filter(_.name == "data")).map(_.dataType) + assert(dataTypes.length == 2, s"Expected two `data` scan columns, but got:\n$plan") + dataTypes.foreach(dt => assertShreddedStruct(dt)) + } + } + + test("full outer join: extractions on both (nullable) sides are pushed and value-preserving") { + // In a FULL OUTER join both sides are nullable: unmatched rows are null-padded on the opposite + // side. Extractions on each side push onto that side, computed before null-padding; the join + // null-pads the resulting `_ve` aliases for the unmatched rows. Result must match the baseline + // that evaluates variant_get above the join (variant_get(NULL) = NULL). Runs under V1 and V2. + withVariantParquetTables( + VariantTable( + "L", "k int, data variant", + Seq("(1, parse_json('{\"lv\": \"a\"}'))", "(2, parse_json('{\"lv\": \"b\"}'))")), + VariantTable( + "R", "k int, data variant", + Seq("(2, parse_json('{\"rv\": \"y\"}'))", "(3, parse_json('{\"rv\": \"z\"}'))"))) { + val query = + "select variant_get(l.data, '$.lv', 'string') as lv, " + + "variant_get(r.data, '$.rv', 'string') as rv " + + "from L l full outer join R r on l.k = r.k " + + "order by lv, rv" + val expected = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + assert(sql(query).collect().toSeq.map(r => (r.get(0), r.get(1))) == + expected.toSeq.map(r => (r.get(0), r.get(1))), + "Both-nullable-side extractions over a full outer join produced wrong values") + // Both L.data and R.data are extracted only via variant_get, so each shreds with no + // full-variant slot on its own scan. Both scans expose a column named `data`, so assert on + // every `data` scan output rather than via the single-match scanColumnType helper. + val plan = sql(query).queryExecution.optimizedPlan + val dataTypes = plan.collect { + case s: DataSourceV2ScanRelation => s.output + case l: LogicalRelation => l.output + }.flatMap(_.filter(_.name == "data")).map(_.dataType) + assert(dataTypes.length == 2, s"Expected two `data` scan columns, but got:\n$plan") + dataTypes.foreach(dt => assertShreddedStruct(dt)) + } + } + + test("left semi join with a right-side variant_get in the condition: right side shredded") { + // The join condition extracts from the RIGHT (DIM) variant. For a LEFT SEMI join the right side + // is not in the join output, but the condition is still evaluated over both children, so the + // right-side extraction is hoisted onto the right side and its `_ve` alias feeds the condition. + // Exercises the right-side routing in rewriteJoinUnderProject for a LeftExistence join. Assert + // correct existence semantics and that DIM.data shreds with no full-variant slot. V1 and V2. + withVariantParquetTables( + factTable(), + VariantTable( + "DIM", "dk int, data variant", + Seq("(1, parse_json('{\"jk\": 1}'))", "(2, parse_json('{\"jk\": 9}'))"))) { + val query = + "select variant_get(ss.data, '$.id', 'string') as id " + + "from SS ss left semi join DIM d " + + "on ss.k = variant_get(d.data, '$.jk', 'int') " + + "order by id" + val expected = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expected) // ss.k in {1,2}; DIM jk in {1,9}; only k=1 rows survive + // DIM.data (right side) is used only via the condition extraction, so it shreds to the $.jk + // slot with no full-variant slot. Both SS and DIM expose a `data` column; assert both scans + // shred (the SS-side `$.id` extraction and the DIM-side `$.jk` condition extraction). + val plan = sql(query).queryExecution.optimizedPlan + val dataTypes = plan.collect { + case s: DataSourceV2ScanRelation => s.output + case l: LogicalRelation => l.output + }.flatMap(_.filter(_.name == "data")).map(_.dataType) + assert(dataTypes.length == 2, s"Expected two `data` scan columns, but got:\n$plan") + dataTypes.foreach(dt => assertShreddedStruct(dt)) + } + } + + test("sort over a join: order-key extraction is pushed through the join and shredded") { + // ORDER BY variant_get(ss.data, ...) sits in a Sort whose child is a Join (via the intermediate + // Project the optimizer leaves above the join). The hoisted order-key alias must be pushed DOWN + // through the join onto the ss side to reach the ss scan; parking it in a Project above the + // join would leave `data` read raw (PhysicalOperation stops at the join). Only ss.k is + // selected, so `data` is used only via the order key and shreds with no full-variant slot. + // Runs under V1/V2. + withVariantParquetTables( + factTable(), + VariantTable("DIM", "k int, name string", Seq("(1, 'x')", "(2, 'y')"))) { + val query = + "select ss.k from SS ss join DIM d on ss.k = d.k " + + "order by variant_get(ss.data, '$.qty', 'double')" + // Order-sensitive: compare the exact row sequence against the pushdown-disabled baseline. + val expectedOrder = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect().map(_.getInt(0)).toList + } + assert(sql(query).collect().map(_.getInt(0)).toList == expectedOrder, + s"ORDER BY variant_get over a join produced the wrong row order: got " + + s"${sql(query).collect().map(_.getInt(0)).toList}, expected $expectedOrder") + assertShreddedStruct(scanColumnType(sql(query).queryExecution.optimizedPlan, "data")) + } + } } // V1 DataSource tests with parameterized reader type @@ -651,6 +1193,72 @@ abstract class PushVariantIntoScanV2SuiteBase extends QueryTest with PushVariant } } + test(s"V2 test - nested variant extraction in aggregate function arg is hoisted and shredded " + + s"($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (vs struct, name string) " + + s"using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(named_struct('v1', parse_json('{\"price\": 10}'), 'i', 1), 'a'), " + + "(named_struct('v1', parse_json('{\"price\": 30}'), 'i', 2), 'a'), " + + "(named_struct('v1', parse_json('{\"price\": 20}'), 'i', 3), 'b')") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // The extraction is on a variant nested in a struct (vs.v1). Its child is a GetStructField + // chain rooted at an attribute, which isVariantColumnPath accepts, so it is hoisted and the + // nested variant shreds to the $.price slot. + val query = + "select name, max(variant_get(vs.v1, '$.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 vsStruct = scanRelation.output.find(_.name == "vs").get.dataType match { + case s: StructType => s + case other => fail(s"Expected vs to stay a struct, but got $other") + } + val v1Struct = vsStruct.fields.find(_.name == "v1").get.dataType match { + case s: StructType => s + case other => fail(s"Expected nested v1 shredded to struct, but got $other") + } + assert(!v1Struct.fields.exists(_.dataType.isInstanceOf[VariantType]), + s"Expected no full-variant slot in nested v1, but got $v1Struct") + } + } + } + + test(s"V2 test - cast(variant) in aggregate function arg is hoisted and shredded ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, name string) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('\"10\"'), 'a'), (parse_json('\"30\"'), 'a'), (parse_json('\"20\"'), 'b')") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // The aggregate argument is `cast(v as string)`, not variant_get. isHoistable's Cast branch + // (matching collectRequestedFields) makes it eligible, so v shreds to the cast slot. + val query = "select name, max(cast(v as string)) 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) + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + val vStruct = scanRelation.output.find(_.name == "v").get.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 no full-variant slot, but got $vStruct") + } + } + } + test(s"V2 test - no pushdown when struct is used ($readerName)") { withTempPath { dir => val path = dir.getCanonicalPath @@ -1031,7 +1639,7 @@ abstract class PushVariantIntoScanV2SuiteBase extends QueryTest with PushVariant } } - test(s"V2 test - order by variant_get: correct ordering, variant read raw, sibling pruned " + + test(s"V2 test - order by variant_get: hoisted and shredded, correct ordering, sibling pruned " + s"($readerName)") { withTempPath { dir => val path = dir.getCanonicalPath @@ -1045,11 +1653,11 @@ abstract class PushVariantIntoScanV2SuiteBase extends QueryTest with PushVariant } 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. + // The sort key variant_get(v, '$.price') lives in Sort.order. PullOutVariantExtractions + // hoists it into a Project below the Sort, and since v is not otherwise referenced above + // the Sort (only `name` is selected), the raw v is dropped and v shreds to the $.price slot + // (no full-variant slot). Compare ORDER-SENSITIVELY: a placeholder slot would 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 @@ -1061,15 +1669,18 @@ abstract class PushVariantIntoScanV2SuiteBase extends QueryTest with PushVariant 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}") + val vStruct = scanRelation.output.find(_.name == "v").get.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 no full-variant slot, but got $vStruct") } } } - test(s"V2 test - aggregate max(variant_get) with no local filter: no codegen crash, variant " + - s"read raw ($readerName)") { + test(s"V2 test - aggregate max(variant_get) with no local filter: hoisted and shredded " + + s"($readerName)") { withTempPath { dir => val path = dir.getCanonicalPath withTable("temp_v1") { @@ -1082,11 +1693,10 @@ abstract class PushVariantIntoScanV2SuiteBase extends QueryTest with PushVariant } 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(, - // ...)) would fail to codegen. Keeping v raw yields correct aggregates with a valid plan. + // variant_get(v, '$.price') is inside an aggregate function, above the aggregate barrier. + // PullOutVariantExtractions hoists it into a Project directly below the Aggregate, so the + // pushdown shreds v to just the $.price slot: no full-variant slot (so no boolean + // placeholder to break codegen), and the raw column is never read. 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") { @@ -1095,15 +1705,152 @@ abstract class PushVariantIntoScanV2SuiteBase extends QueryTest with PushVariant 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}") + val vStruct = vAttr.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 no full-variant slot, but got $vStruct") 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 " + + test(s"V2 test - aggregate over the whole variant is left raw, no hoist ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, name string) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"price\": 10}'), 'a'), (parse_json('{\"price\": 20}'), 'b')") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // count(v) references the whole variant, not a field extraction. PullOutVariantExtractions + // has nothing to hoist, and the pushdown keeps v raw (shredding a sole full-variant request + // saves no I/O). + val query = "select name, count(v) as c from T_V2 group by name" + 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 + assert(vAttr.dataType.isInstanceOf[VariantType], + s"Expected v left raw, but got ${vAttr.dataType}") + } + } + } + + test(s"V2 test - GROUP BY key and aggregate function argument both shredded, dedup " + + s"($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"k\": \"a\", \"price\": 10}')), " + + "(parse_json('{\"k\": \"a\", \"price\": 30}')), " + + "(parse_json('{\"k\": \"b\", \"price\": 20}'))") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // $.k is the GROUP BY key (relocated below the Aggregate by PullOutGroupingExpressions); + // $.price appears twice inside aggregate functions (relocated by the pull-out rule). + // The two $.price extractions must dedup to a single struct slot, and there must be no + // full-variant slot. + val query = "select variant_get(v, '$.k', 'string') as k, " + + "max(variant_get(v, '$.price', 'int')) as mx, " + + "min(variant_get(v, '$.price', 'int')) as mn " + + "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) // a -> (30, 10), b -> (20, 20) + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + val vStruct = scanRelation.output.find(_.name == "v").get.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 no full-variant slot, but got $vStruct") + // Exactly two data slots: one for $.k, one shared for the two $.price extractions. + assert(vStruct.fields.length == 2, + s"Expected 2 shredded fields ($$.k and a deduped $$.price), but got $vStruct") + } + } + } + + test(s"V2 test - aggregate function arg with a non-literal path is not hoisted ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, p string, name string) " + + s"using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"price\": 10}'), '$.price', 'a'), " + + "(parse_json('{\"price\": 20}'), '$.price', 'b')") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // The extraction path is a column (non-foldable), which the pushdown cannot handle, so the + // rule must not hoist it and v is read raw. Guards the `path.foldable` predicate. + val query = "select name, max(variant_get(v, p, '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) + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + val vAttr = scanRelation.output.find(_.name == "v").get + assert(vAttr.dataType.isInstanceOf[VariantType], + s"Expected v left raw for a non-literal path, but got ${vAttr.dataType}") + } + } + } + + test(s"V2 test - join selecting the whole variant keeps v raw ($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, name string) using PARQUET location '$itemsPath'") + sql("insert into temp_items values " + + "(parse_json('{\"country_code\": \"CN\"}'), 'widget'), " + + "(parse_json('{\"country_code\": \"JP\"}'), '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") + // i.v is selected raw (needed above the Join) AND used in the join key. The hoist must NOT + // drop it: v shreds with a full-variant slot plus the $.country_code slot. Guards the join + // liveness computation. + val query = + "select i.v, 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() + } + 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}")) + val vStruct = itemsScan.output.find(_.name == "v").get.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 a full-variant slot because i.v is selected raw, but got $vStruct") + } + } + } + + test(s"V2 test - join on variant_get key: hoisted and shredded, sibling pruned " + s"($readerName)") { withTempPath { dir => val itemsPath = dir.getCanonicalPath + "/items" @@ -1127,21 +1874,125 @@ abstract class PushVariantIntoScanV2SuiteBase extends QueryTest with PushVariant 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. + // v is the join key (referenced only via variant_get in the join condition). + // PullOutVariantExtractions hoists the extraction onto the items side, dropping the raw v + // so it shreds to just the $.country_code slot (no full-variant slot). v2 is pruned. assert(itemsScan.output.map(_.name).toSet == Set("v", "name"), s"Expected items scan output {v, name} but got ${itemsScan.output.map(_.name)}") + val vStruct = itemsScan.output.find(_.name == "v").get.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 no full-variant slot, but got $vStruct") + } + } + } + + test(s"V2 test - self join on variant_get key: both sides shredded ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, name string) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"id\": 1}'), 'a'), (parse_json('{\"id\": 2}'), 'b')") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // Both sides' join keys are variant_get(_, '$.id'); each is hoisted onto its own side and + // both variant columns shred to the $.id slot with no full-variant slot. + val query = "select a.name, b.name from T_V2 a join T_V2 b " + + "on variant_get(a.v, '$.id', 'int') = variant_get(b.v, '$.id', 'int')" + val expectedRows = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") { + sql(query).collect() + } + checkAnswer(sql(query), expectedRows) + val scans = sql(query).queryExecution.optimizedPlan.collect { + case s: DataSourceV2ScanRelation => s + } + assert(scans.length == 2, s"Expected two scans but got ${scans.length}") + scans.foreach { s => + val vStruct = s.output.find(_.name == "v").get.dataType match { + case st: StructType => st + case other => fail(s"Expected v shredded to struct, but got $other") + } + assert(!vStruct.fields.exists(_.dataType.isInstanceOf[VariantType]), + s"Expected no full-variant slot, but got $vStruct") + } + } + } + } + + test(s"V2 test - join on variant_get key with no narrowing Project: left raw, correct " + + 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, name string) using PARQUET location '$itemsPath'") + sql("insert into temp_items values " + + "(parse_json('{\"country_code\": \"CN\"}'), 'widget'), " + + "(parse_json('{\"country_code\": \"JP\"}'), '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") + // `select *` leaves the Join without a narrowing Project directly above it, so + // PullOutVariantExtractions does not hoist the join-key extraction (see class doc: hoisting + // requires a Project above the barrier to know the live-above set). The plan must still be + // correct with v read raw -- this locks in the documented fallback behavior. + val query = + "select * 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() + } + 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}")) val vAttr = itemsScan.output.find(_.name == "v").get assert(vAttr.dataType.isInstanceOf[VariantType], - s"Expected v left as raw VariantType, but got ${vAttr.dataType}") + s"Expected v left raw with no narrowing Project above the join, " + + s"but got ${vAttr.dataType}") + } + } + } + + test(s"V2 test - pull-out flag off leaves aggregate function arg raw ($readerName)") { + withTempPath { dir => + val path = dir.getCanonicalPath + withTable("temp_v1") { + sql(s"create table temp_v1 (v variant, name string) using PARQUET location '$path'") + sql("insert into temp_v1 values " + + "(parse_json('{\"price\": 10}'), 'a'), (parse_json('{\"price\": 20}'), 'b')") + } + withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "", + SQLConf.PUSH_VARIANT_INTO_SCAN_PULL_OUT_EXTRACTIONS.key -> "false") { + spark.read.parquet(path).createOrReplaceTempView("T_V2") + // With the pull-out rule disabled, the aggregate function argument is not hoisted and v is + // read raw (pushdown still enabled). Confirms the flag gates the rule. + 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) + val scanRelation = findScanRelation(sql(query).queryExecution.optimizedPlan) + val vAttr = scanRelation.output.find(_.name == "v").get + assert(vAttr.dataType.isInstanceOf[VariantType], + s"Expected v left raw with pull-out disabled, but got ${vAttr.dataType}") } } } @@ -1177,8 +2028,8 @@ abstract class PushVariantIntoScanV2SuiteBase extends QueryTest with PushVariant } } - test(s"V2 test - join on variant_get key through an aliasing projection: no crash, no " + - s"mis-shred ($readerName)") { + test(s"V2 test - join on variant_get key through an aliasing projection: hoisted and shredded " + + s"($readerName)") { withTempPath { dir => val itemsPath = dir.getCanonicalPath + "/items" val countriesPath = dir.getCanonicalPath + "/countries" @@ -1196,9 +2047,9 @@ abstract class PushVariantIntoScanV2SuiteBase extends QueryTest with PushVariant 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). + // it via variant_get. The optimizer inlines the alias, PullOutVariantExtractions hoists the + // extraction onto the items side, and the join key shreds to the $.country_code slot with + // correct results (checkAnswer) and no wrong data. v2 is pruned. val query = "select c.country_name from " + "(select v as vw, name as nm from ITEMS_V2) i join COUNTRIES_V2 c " + @@ -1217,10 +2068,14 @@ abstract class PushVariantIntoScanV2SuiteBase extends QueryTest with PushVariant // 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}") + // The variant join key is hoisted and shredded to the $.country_code slot, no full-variant. + val keyAttr = itemsScan.output.find(a => a.name == "v" || a.name == "vw").get + val vStruct = keyAttr.dataType match { + case s: StructType => s + case other => fail(s"Expected the variant join key shredded to struct, but got $other") + } + assert(!vStruct.fields.exists(_.dataType.isInstanceOf[VariantType]), + s"Expected no full-variant slot, but got $vStruct") } } } @@ -1692,6 +2547,36 @@ abstract class PushVariantIntoScanV2SuiteBase extends QueryTest with PushVariant } } + test(s"Spark native managed table: aggregate hoist shreds without full-variant slot " + + s"($readerName)") { + // A Spark native (managed catalog) Parquet table created with CREATE TABLE ... USING PARQUET + // resolves to a V1 `LogicalRelation` regardless of USE_V1_SOURCE_LIST -- that config only + // affects the `spark.read`/DataSource resolution of external file reads, not managed catalog + // tables. PullOutVariantExtractions must still hoist the aggregate-function extraction just + // like it does for external-file reads, and the pushdown must shred v to the $.price slot + // with no full-variant slot. `scanRelationOutput` handles whichever scan node the plan uses. + withTable("T_native") { + sql("create table T_native (v variant, name string) using parquet") + sql("insert into T_native values " + + "(parse_json('{\"price\": 10}'), 'a'), " + + "(parse_json('{\"price\": 30}'), 'a'), " + + "(parse_json('{\"price\": 20}'), 'b')") + val query = + "select name, max(variant_get(v, '$.price', 'int')) as mx from T_native 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 output = scanRelationOutput(sql(query).queryExecution.optimizedPlan) + val vStruct = output.find(_.name == "v").get.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 no full-variant slot for a native managed table, but got $vStruct") + } + } + test(s"V2 No push down for JSON ($readerName)") { withTempPath { dir => val path = dir.getCanonicalPath