You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2026/07/27 16:04:41 UTC
(spark) branch branch-4.x updated: [SPARK-58089][SQL] Push variant extractions through Aggregate/Sort/Join
This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 4ab9e91c5f81 [SPARK-58089][SQL] Push variant extractions through Aggregate/Sort/Join
4ab9e91c5f81 is described below
commit 4ab9e91c5f8165fd01f3914ac773939c27f00a79
Author: Qiegang Long <ql...@users.noreply.github.com>
AuthorDate: Tue Jul 28 00:04:08 2026 +0800
[SPARK-58089][SQL] Push variant extractions through Aggregate/Sort/Join
### What changes were proposed in this pull request?
Introduce a new optimizer rule `PullOutVariantExtractions` that hoists
`variant_get` / `Cast(variant)` extractions out of three operator types
that the existing `PushVariantIntoScan` / `V2ScanRelationPushDown` rules
cannot see through:
- **Aggregate function arguments** – e.g. `max(variant_get(v, '$.price', 'int'))`:
the extraction is moved into a `Project` directly below the `Aggregate` and the
aggregate references the resulting alias. The bare variant column is suppressed
unless it is also needed raw (e.g. as a `GROUP BY` key), so no redundant
full-variant slot is generated.
- **Sort order keys** – e.g. `ORDER BY variant_get(v, '$.price', 'int')`:
matched as `Project → Sort`; the extraction is hoisted below the `Sort` and
the original `Project` is reproduced to prevent the alias from leaking into
the output.
- **Join conditions and projections above joins** – matched as `Project → Join`;
extractions in both the join condition and the outer `Project` are routed to
the owning join side. A `pushSideAliases` helper then pushes the aliases
*through* any depth of chained joins so they land in a `Project` directly
above the scan (where `PhysicalOperation` collapses them with the scan, making
them visible to the pushdown). This is necessary because `PhysicalOperation`
stops at a `Join` node.
A `Sort` sitting over a `Join` is handled by fusing the two cases: the
order-key aliases are pushed through the join tree, not left in a `Project`
above it.
The rule is gated by a new internal config
`spark.sql.variant.pushVariantIntoScan.pullOutExtractions` (default `true`)
and is a no-op unless `spark.sql.variant.pushVariantIntoScan` is also enabled.
Non-variant plans are untouched.
The rule is registered as the first rule in `SparkOptimizer.earlyScanPushDownRules`,
before `SchemaPruning` and the V2 scan pushdown rules.
### Why are the changes needed?
Before this change, a `variant_get` inside an aggregate function argument, sort key,
or join condition caused the whole variant column to be read raw (or shredded with a
redundant full-variant slot). For example:
```sql
SELECT name, max(variant_get(v, '$.price', 'int')) FROM T GROUP BY name
```
read the entire v column even though only the price field was needed.
After this change, Spark shreds only the requested typed fields, avoiding the
full-variant I/O.
The change improves query performance for variant referenced in aggregrate, join, sort.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
Added new units. Also run correctness tests against some known workload.
#### Performance result
Test framework: https://github.com/cloudera-labs/variant-conformance-benchmark
Dataset: tpc-ds dataset (SF=5), spark native parquet table, payload in variant or json.
Run setup: pre-warm jvm, three runs, median query timing reported by spark.
Run A with pullout rule enabled vs B with the rule disabled.
```
Compare: tpcds-flat-pullout-20260714 vs tpcds-flat-no-pullout-20260714 (metric: query_median)
Run A: /Users/qlong/opensources/variant-conformance-benchmark/results/tpcds-flat-pullout-20260714/tpcds-flat/timings-variant.csv
Run B: /Users/qlong/opensources/variant-conformance-benchmark/results/tpcds-flat-no-pullout-20260714/tpcds-flat/timings-variant.csv
query median_A median_B delta_s delta_%
------ -------- -------- ------- -------
q07 1.40 3.28 -1.88 -57.2%
q12 0.07 0.07 -0.00 -2.8%
q19 0.06 0.07 -0.01 -9.9%
q26 1.14 3.40 -2.26 -66.4%
q42 0.56 2.59 -2.03 -78.2%
q52 0.61 3.04 -2.43 -80.0%
q55 0.54 2.68 -2.14 -79.9%
q63 0.58 2.95 -2.37 -80.4%
q68 1.00 3.42 -2.42 -70.8%
q73 0.57 2.88 -2.31 -80.1%
q79 0.85 3.26 -2.41 -73.9%
q98 0.65 3.32 -2.67 -80.5%
Summary: 12 queries, 12 comparable
Geo-mean delta: -69.5% (Run A faster)
Total (query_median): 8.0s vs 31.0s
```
Run A with pullout rule enabled vs B using json for payload, this test shows the **performance advantage of variant over json**.
```
Compare: tpcds-flat-pullout-20260714 vs tpcds-flat-pullout-20260714 (metric: query_median)
Run A: /Users/qlong/opensources/variant-conformance-benchmark/results/tpcds-flat-pullout-20260714/tpcds-flat/timings-variant.csv
Run B: /Users/qlong/opensources/variant-conformance-benchmark/results/tpcds-flat-pullout-20260714/tpcds-flat/timings-string_json.csv
query median_A median_B delta_s delta_%
------ -------- -------- ------- -------
q07 1.40 3.26 -1.85 -56.9%
q12 0.07 0.13 -0.06 -47.0%
q19 0.06 0.23 -0.17 -72.2%
q26 1.14 2.17 -1.03 -47.5%
q42 0.56 0.56 +0.01 +1.1%
q52 0.61 1.09 -0.48 -44.1%
q55 0.54 0.47 +0.07 +15.2%
q63 0.58 0.88 -0.30 -34.2%
q68 1.00 1.34 -0.34 -25.1%
q73 0.57 0.70 -0.13 -18.0%
q79 0.85 1.62 -0.77 -47.4%
q98 0.65 1.19 -0.54 -45.4%
Summary: 12 queries, 12 comparable
Geo-mean delta: -39.3% (Run A faster)
Total (query_median): 8.0s vs 13.6s
```
### Was this patch authored or co-authored using generative AI tooling?
Co-authored with Claude Code
Closes #57190 from qlong/variant-pushout-aggregate.
Authored-by: Qiegang Long <ql...@users.noreply.github.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
(cherry picked from commit adad83831f22076a855456aeae7ea62b3895e851)
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../org/apache/spark/sql/internal/SQLConf.scala | 13 +
.../spark/sql/execution/SparkOptimizer.scala | 6 +-
.../datasources/PullOutVariantExtractions.scala | 487 +++++++++
.../datasources/PushVariantIntoScanSuite.scala | 1039 +++++++++++++++++++-
4 files changed, 1511 insertions(+), 34 deletions(-)
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 c6e0c9bc60bf..97594802e258 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
@@ -6697,6 +6697,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 54158d5bb4a9..0994220afe0c 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 000000000000..6695e24ad618
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PullOutVariantExtractions.scala
@@ -0,0 +1,487 @@
+/*
+ * 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.catalyst.trees.TreePattern.{AGGREGATE, JOIN, SORT}
+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).
+ *
+ * A `Join`, and a `Sort` directly under a narrowing `Project`, are handled with that Project's
+ * `references` telling us which columns remain live: a variant used only via the hoisted
+ * extraction is dropped (no redundant full-variant slot). This is the common `Sort` shape when
+ * only a subset of columns is selected (e.g. `SELECT name ... ORDER BY variant_get(v, ...)`),
+ * where the analyzer adds a `Project` above the `Sort` to project away the order-only columns.
+ *
+ * A bare/root `Sort` -- one with no narrowing `Project` above it, as in
+ * `SELECT * FROM t ORDER BY variant_get(v, '$.b')` (all order columns are already selected, so
+ * the analyzer adds no narrowing `Project`) -- is also handled, by a second `transformUp` pass
+ * (see `apply`). Its output must equal its child's output exactly, so ALL child columns stay
+ * live: a variant present in the output stays raw AND gains an extra `_ve` slot, producing a
+ * multi-slot shredded struct with a full-variant slot -- the same shape as a `Sort`-under-
+ * `Project` whose `v` is also selected. `Sort`s under a `GlobalLimit`/`LocalLimit` (a
+ * `SELECT * ... ORDER BY ... LIMIT n`) are covered by the same second pass. A `Sort`/`Join` in
+ * any other shape -- under a `Filter`/`Aggregate`/another non-`Project` operator that is not a
+ * limit wrapper -- is left unchanged: the extraction stays in place, the variant is read raw,
+ * and the plan is correct (just not shredded). 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.
+ *
+ * Cast-error surface: relocating a strict extraction (`variant_get(..., failOnError = true)` or a
+ * strict `Cast`) below a `Join` means it is evaluated at the scan on rows the join later eliminates
+ * -- so a cast failure can surface for a row the un-hoisted plan would never have cast (here the
+ * eliminating rows come from the *other* joined table). This is the same pre-existing trade-off as
+ * [[PushVariantIntoScan]] pushing casts below a `Filter`, not a new error class: in both, the
+ * strict cast runs before the operator that would have discarded the failing row. This rule only
+ * relocates the extraction into a `Project`; [[PushVariantIntoScan]] still does the scan-level
+ * materialization and, when [[SQLConf.PUSH_VARIANT_INTO_SCAN_DEFER_CAST_ERROR]] is set (default
+ * false), wraps the cast with a per-row cast-error companion slot so the error is only raised when
+ * the original expression consumes the failing row. That deferral is provenance-agnostic -- it acts
+ * on the relocated extraction regardless of how it reached the `Project` -- so enabling the flag
+ * suppresses the join-eliminated-row error exactly as it does the filter-eliminated-row case. With
+ * the flag off (the default), the strict cast raises immediately on any failing scanned row, as
+ * documented for below-`Filter` pushdown.
+ */
+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 {
+ // First pass: handle operators where a narrowing Project above the barrier tells us which
+ // columns remain live. `transformUp` so inner operators are rewritten before their parents.
+ // Sort/Join are matched together with the Project above them (see class doc). Prune to
+ // subtrees that actually contain one of the matched operators so plans with none (and the
+ // Once-batch idempotence re-run) skip the traversal.
+ val afterProjectCases = plan.transformUpWithPruning(
+ _.containsAnyPattern(AGGREGATE, SORT, JOIN)) {
+ case agg: Aggregate => rewriteAggregate(agg)
+ case p @ Project(projectList, s: Sort) => rewriteSortUnderProject(p, projectList, s)
+ case p @ Project(projectList, j: Join) => rewriteJoinUnderProject(p, projectList, j)
+ }
+ // Second pass: handle a Sort NOT under a narrowing Project -- a bare/root Sort (as in
+ // `SELECT * ORDER BY variant_get(...)`), or one under a GlobalLimit/LocalLimit. This must be
+ // a separate pass, not another arm in the first `transformUp`: that traversal is bottom-up,
+ // so a bare-Sort arm would fire on the inner Sort of a `Project(_, Sort)` shape before the
+ // parent Project is visited, keeping v raw and pre-empting `rewriteSortUnderProject`. After
+ // the first pass, any Sort that was under a Project has its order rewritten to non-hoistable
+ // `_ve` attribute references, so `rewriteBareSort` is a no-op on it (idempotent). Prune to
+ // Sort-bearing subtrees so Sort-free plans skip this pass entirely.
+ afterProjectCases.transformUpWithPruning(_.containsPattern(SORT)) {
+ case s: Sort => rewriteBareSort(s)
+ }
+ }
+ }
+
+ // 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. It would also break Once-batch idempotence: the hoisted `_ve0` is a VariantType
+ // attribute, so a wrapping `cast(_ve0 as string)` (e.g. from `v:a::string`, which desugars to
+ // `cast(variant_get(v, '$.a') as string)`) would match the Cast branch below on the next pass.
+ case g: VariantGet =>
+ !g.targetType.isInstanceOf[VariantType] && g.path.foldable && isVariantColumnPath(g.child)
+ // A variant-to-variant `Cast` (e.g. `cast(v as variant)`) is a whole-variant read: its sole
+ // requested field targets `VariantType`, which the pushdown drops (see SPARK-57499) because a
+ // lone variant-typed slot saves no I/O and is mishandled by the reader. Hoisting it would add a
+ // Project + alias for no shredding benefit and leave the column raw anyway, so exclude it.
+ // A hoistable `Cast` yields a non-variant-typed alias, so it is never re-hoisted (idempotent).
+ case c: Cast => !c.dataType.isInstanceOf[VariantType] && isVariantColumnPath(c.child)
+ 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/ExistenceJoin (all matched by `LeftExistence`) 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)
+ }
+ }
+ }
+
+ // 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. `pushSideAliases` pushes 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. 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.
+ pushSideAliases(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))
+ }
+ }
+
+ // Handles a `Sort` that is NOT directly under a narrowing `Project` -- a bare/root Sort such as
+ // `SELECT * FROM t ORDER BY variant_get(v, '$.b')` (all order columns are already selected, so
+ // the analyzer adds no narrowing Project above the Sort), or a Sort under a
+ // GlobalLimit/LocalLimit.
+ // Unlike `rewriteSortUnderProject`, there is no Project above to tell us which columns remain
+ // live; a bare Sort's output IS its child's output, so ALL child columns must stay live. A
+ // variant present in the output (SELECT *) therefore stays raw AND gains an extra `_ve` slot,
+ // yielding a multi-slot shredded struct with a full-variant slot -- the same shape as a
+ // Sort-under-Project whose `v` is also selected. See the class doc.
+ private def rewriteBareSort(sort: Sort): LogicalPlan = {
+ val hoister = new ExtractionHoister
+ val newOrder = sort.order.map(hoister.hoist(_).asInstanceOf[SortOrder])
+ if (hoister.isEmpty) {
+ sort
+ } else {
+ // All child output columns stay live: the bare Sort's output equals its child's output. The
+ // hoisted `_ve` aliases are supplied to `pushSideAliases` separately (not via `needed`).
+ val needed = sort.child.outputSet
+ // `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.
+ val newChild = pushSideAliases(sort.child, hoister.aliases, needed)
+ // Wrap in a Project restoring the original output so the `_ve` alias columns do not leak.
+ // `sort.child.output` (original attributes, in order) is the correct projectList -- NOT
+ // `newSort.output`, which would include the `_ve` columns `pushSideAliases` added below.
+ // `sort.copy` preserves `sort.global` and `sort.hint`.
+ Project(sort.child.output, 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 67bee55c2275..d6ed274327a6 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,268 @@ 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("bare/root sort (SELECT *): order-key variant_get is hoisted and shredded") {
+ // SELECT * keeps v in the query output, so the sort key variant_get(v, '$.price') lives in a
+ // bare/root Sort with no narrowing Project above it (the analyzer adds no Project since every
+ // order column is already selected). PullOutVariantExtractions handles this via rewriteBareSort
+ // in its second pass. Because v is in the SELECT * output it stays live, so v shreds to a
+ // multi-slot struct WITH a full-variant slot (the raw blob is still materialized) plus the
+ // typed $.price slot the sort reads. 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 * from T order by variant_get(v, '$.price', 'int')"
+ // Order-sensitive: compare the non-variant `name` column sequence.
+ val expectedOrder = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") {
+ sql(query).collect().map(_.getString(1)).toList
+ }
+ assert(expectedOrder == List("z", "y", "x"), s"baseline sanity: $expectedOrder")
+ assert(sql(query).collect().map(_.getString(1)).toList == expectedOrder,
+ "SELECT * ORDER BY variant_get produced the wrong row order")
+ val output = scanRelationOutput(sql(query).queryExecution.optimizedPlan)
+ assertShreddedStruct(
+ output.find(_.name == "v").get.dataType, expectFullVariantSlot = true)
+ }
+ }
+
+ test("bare/root sort (SELECT id, v): order-key variant_get is hoisted and shredded") {
+ // Explicitly selecting the variant column (SELECT id, v) produces the same bare/root Sort shape
+ // as SELECT *: no narrowing Project above the Sort. Exercises rewriteBareSort's second pass;
+ // v stays live (it is selected) so it shreds with a full-variant slot plus the $.b slot.
+ withVariantParquetData(
+ "id int, v variant",
+ "(1, parse_json('{\"b\": 3}'))",
+ "(2, parse_json('{\"b\": 1}'))",
+ "(3, parse_json('{\"b\": 2}'))") {
+ val query = "select id, v from T order by variant_get(v, '$.b', 'int')"
+ val expectedOrder = withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN.key -> "false") {
+ sql(query).collect().map(_.getInt(0)).toList
+ }
+ assert(expectedOrder == List(2, 3, 1), s"baseline sanity: $expectedOrder")
+ assert(sql(query).collect().map(_.getInt(0)).toList == expectedOrder,
+ "SELECT id, v ORDER BY variant_get produced the wrong row order")
+ val output = scanRelationOutput(sql(query).queryExecution.optimizedPlan)
+ assertShreddedStruct(
+ output.find(_.name == "v").get.dataType, expectFullVariantSlot = true)
+ }
+ }
+
+ test("bare sort over a join (SELECT *): order-key extraction pushed through the join") {
+ // A bare/root Sort whose child is a Join, with all join columns selected (no narrowing Project
+ // above the Sort). The hoisted order-key alias must be pushed DOWN through the join to the ss
+ // scan side. `data` is in the query output, so it shreds with a full-variant slot plus the
+ // $.qty slot. Runs under V1/V2.
+ withVariantParquetTables(
+ factTable(),
+ VariantTable("DIM", "k int, name string", Seq("(1, 'x')", "(2, 'y')"))) {
+ val query =
+ "select ss.k, ss.data, d.name from SS ss join DIM d on ss.k = d.k " +
+ "order by variant_get(ss.data, '$.qty', 'double')"
+ 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"SELECT with join columns ORDER BY variant_get produced the wrong row order: " +
+ s"got ${sql(query).collect().map(_.getInt(0)).toList}, expected $expectedOrder")
+ assertShreddedStruct(
+ scanColumnType(sql(query).queryExecution.optimizedPlan, "data"),
+ expectFullVariantSlot = true)
+ }
+ }
+
+ 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("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)`, where the inner two-arg variant_get returns
+ // VariantType. Excluding that two-arg form from hoisting (isHoistable's VariantGet branch)
+ // leaves the whole `cast(variant_get(...))` untouched -- so nothing is hoisted here and there
+ // is no VariantType `_ve` alias for the Cast branch to re-hoist on the next pass. Without the
+ // exclusion the VariantGet would hoist to `_ve0: VariantType`, then `cast(_ve0 as string)`
+ // would re-match, producing a different plan and crashing with
+ // "Once strategy's idempotence is broken".
+ 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 +594,317 @@ 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("strict cast pushed through a join: cast-error deferral covers join-eliminated rows") {
+ // SS has a row whose data is a non-castable type; that row is eliminated by the join (its k
+ // has no match in DIM). The aggregate-argument extraction `cast(ss.data as int)` is hoisted
+ // onto the SS side and pushed through the join to the scan.
+ // - With deferral OFF: the strict cast is evaluated at the scan on the eliminated row and
+ // raises INVALID_VARIANT_CAST (the same trade-off as below-Filter pushdown).
+ // - With deferral ON: the cast-error companion slot carries the error instead; the join
+ // eliminates the row before its error is consumed; the query returns correct values.
+ // Data: SS has k=1 (string data, bad for cast) and k=2 (int data, good); DIM has k=2 only.
+ // Result: only the k=2 row survives the join and is included in the aggregate.
+ withVariantParquetTables(
+ VariantTable(
+ "SS", "k int, data variant",
+ Seq("(1, parse_json('\"hello\"'))", "(2, parse_json('42'))")),
+ VariantTable(
+ "DIM", "k int, name string",
+ Seq("(2, 'match')"))) {
+ val query = "select avg(cast(ss.data as int)) as avg_val from SS ss join DIM d on ss.k = d.k"
+
+ // Without deferral, the strict cast pushed into the scan raises on the k=1 row (malformed)
+ // even though that row is eliminated by the join condition ss.k = d.k (no k=1 in DIM).
+ withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN_DEFER_CAST_ERROR.key -> "false") {
+ val ex = intercept[Exception](sql(query).collect())
+ assert(hasCastCondition(ex), s"Expected INVALID_VARIANT_CAST, got $ex")
+ }
+
+ // With deferral, the cast error is deferred and not raised because the failing row is
+ // eliminated before its error is consumed. The query returns the correct average.
+ withSQLConf(SQLConf.PUSH_VARIANT_INTO_SCAN_DEFER_CAST_ERROR.key -> "true") {
+ val rows = sql(query).collect()
+ assert(rows.length == 1, s"Expected 1 row, got ${rows.length}")
+ val avgVal = rows(0).getDouble(0)
+ assert(avgVal == 42.0, s"Expected avg=42.0 (only k=2 row survives), got $avgVal")
+ }
+ }
+ }
+
+ 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 +1281,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<v1 variant, i int>, 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 +1727,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 +1741,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 +1757,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 +1781,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(<boolean>,
- // ...)) 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 +1793,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 +1962,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 +2116,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 +2135,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 +2156,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 +2635,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
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org