Skip to content

[SPARK-58089][SQL] Push variant extractions through Aggregate/Sort/Join - #57190

Closed
qlong wants to merge 2 commits into
apache:masterfrom
qlong:variant-pushout-aggregate
Closed

[SPARK-58089][SQL] Push variant extractions through Aggregate/Sort/Join#57190
qlong wants to merge 2 commits into
apache:masterfrom
qlong:variant-pushout-aggregate

Conversation

@qlong

@qlong qlong commented Jul 10, 2026

Copy link
Copy Markdown
Contributor

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:

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

`PushVariantIntoScan` (v1) and `V2ScanRelationPushDown` (v2)
both rely on `PhysicalOperation`, which collapses only a
contiguous `Project`/`Filter` chain and stops at `Aggregate`,
`Sort`, and `Join`. As a result, `variant_get` expressions
embedded in aggregate function arguments, sort-order
expressions, or join conditions are invisible to the pushdown
rules, and the full variant column is read raw instead of
being shredded to the requested typed fields.

When an `Aggregate` or `Sort` sits above a `Join`, the
barrier compounds: even a `Project` hoisted above the join
tree is still unreachable from the scan side.

Example queries that fail to shred without this fix:

  -- Aggregate: extraction in agg arg / GROUP BY
  SELECT AVG(variant_get(data, '$.qty', 'double'))
  FROM   store_sales
  GROUP  BY variant_get(data, '$.id', 'string')

  -- Join: extraction in ON condition
  SELECT ss.k
  FROM   store_sales ss
  JOIN   date_dim d
    ON ss.date_sk = variant_get(d.data, '$.sk', 'int')

  -- Aggregate over Join (TPC-DS Q26 shape)
  SELECT AVG(variant_get(ss.data, '$.qty', 'double'))
  FROM   store_sales ss
  JOIN   date_dim  d  ON ss.date_sk  = d.date_sk
  JOIN   store     st ON ss.store_sk = st.store_sk
  GROUP  BY variant_get(ss.data, '$.id', 'string')

In all cases the scan emits `data:struct<0:variant>` (full
blob) instead of `data:struct<0:double, 1:string>`.

Introduce a new optimizer rule `PullOutVariantExtractions`,
registered as the first rule in `earlyScanPushDownRules`
(before `SchemaPruning`), which hoists `variant_get` calls
out of `Aggregate`, `Sort`, and `Join` into a `Project`
directly below the operator so the downstream pushdown rules
can see them.

- **Aggregate**: hoist extractions from aggregate function
  arguments into a `Project` below the `Aggregate`. The
  `Aggregate` defines its own output so the raw variant
  column is not passed through.

- **Sort / Join**: match the `Project` sitting directly above
  the `Sort`/`Join` (its `references` give the live-above
  set) and drop any variant column no longer referenced once
  its extraction is hoisted.

- **Push-through-Join**: hoisting above a `Join` is not
  enough because `PhysicalOperation` stops there. The new
  `pushSideAliases` helper recursively routes each hoisted
  `_ve` alias down through the join tree to the side whose
  output owns the referenced attribute, landing it in a
  `Project` directly above the scan. Handles any depth of
  chained joins in one pass.for review
@qlong
qlong force-pushed the variant-pushout-aggregate branch from 7cace8c to bc27fe1 Compare July 13, 2026 23:25
@qlong

qlong commented Jul 14, 2026

Copy link
Copy Markdown
Contributor Author

@cloud-fan Can you take a look for this change to extend scope of variant extraction pushdown?

@cloud-fan cloud-fan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

0 blocking, 2 non-blocking, 1 nit.
Well-scoped, well-documented, and exceptionally well-tested. The hoist is value-preserving and the synthesized _ve aliases are consistently nullable, so pushing to an outer-join nullable side never narrows declared nullability. The items below are optional / confirmation-only.

Design / architecture (1)

  • PullOutVariantExtractions.scala:283: pushAliasesIntoJoin duplicates the Join branch of pushSideAliases — see inline

Correctness (1)

  • PullOutVariantExtractions.scala:149: strict-cast error can raise on join-eliminated rows — confirm intended — see inline

Nits: 1 minor item (see inline comments).

Verification

Traced the hoist X(...variant_get(v)...) -> X(..._ve...) over Project([_ve = variant_get(v)], child) below the barrier (and through joins to the scan). Equivalent for empty / single / many-row input, under duplicates, and for NULLs (the extraction is nullIntolerant). For an extraction pushed to an outer-join nullable side, it is computed before null-padding and the join null-pads _ve — equivalent because null-padding commutes with the extraction (variant_get(NULL)=NULL). Output-attribute check: VariantGet.nullable and Cast(VariantType->_) are always nullable=true, and Join.computeOutput widens the null-supplying side to nullable, so every parent reference to a pushed _ve is nullable-over-nullable (AGREES) — no plan-integrity nullability mismatch.

// join, or one that straddles both sides -- we fall back to wrapping a `Project` of the live
// columns plus every alias above the join, which keeps the parent's `_ve` references resolved (at
// the cost of reading that side's variant raw).
private def pushAliasesIntoJoin(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a near-duplicate of the case join: Join branch of pushSideAliases (lines 201-234). Since a single variant extraction references exactly one attribute (hence one join side), aliases always route cleanly and the leftAliases.size + rightAliases.size != aliases.size fallback here never fires in practice — so rewriteAggregate (line 384) could call pushSideAliases(join, hoister.aliases, referenced) directly and this method could go away. Intentional split for readability, or worth collapsing?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for review. Refactored to reduce duplicated code

// 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 =>

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pushing a strict variant_get(..., failOnError=true) / strict Cast below a Join evaluates it at the scan on rows the join later eliminates, so a cast failure can surface for a row the un-hoisted plan would never cast. This is the same behavior class as the existing PushVariantIntoScan pushing casts below Filters (the reason PUSH_VARIANT_INTO_SCAN_DEFER_CAST_ERROR exists), but here the eliminated-row surface is driven by the other joined table. Can you confirm this widened error surface is intended and consistent with the "No user-facing change" claim?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added new test that confirms PUSH_VARIANT_INTO_SCAN_DEFER_CAST_ERROR covers the hoisted-through-join case. No user-facing change.

case join: Join =>
val leftOutput = join.left.outputSet
val rightOutput = join.right.outputSet
// For LeftSemi/LeftAnti the right side is not in the join output, so a right-side alias

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The guard just below is case LeftExistence(_) => false, which also matches ExistenceJoin, but this comment only mentions LeftSemi/LeftAnti. The parallel comment in pushAliasesIntoJoin (lines 289-290) lists all three — worth aligning here for completeness.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

updated comment

…versals

Address PR review feedback.
- Collapse `pushAliasesIntoJoin` into `pushSideAliases`: the
  aggregate-over-join path now calls `pushSideAliases(join, ...)`
  directly and the near-duplicate method is removed. A hoisted `_ve`
  alias always references exactly one attribute (its source is a single
  variant column path), so it routes cleanly to one join side and the
  partial-routing fallback that distinguished the two methods is
  unreachable -- making the two equivalent for this call site.
- Document the widened cast-error surface: relocating a strict
  extraction below a `Join` can surface a cast failure on a
  join-eliminated row -- the same pre-existing trade-off as
  `PushVariantIntoScan` pushing casts below a `Filter` (see
  `PUSH_VARIANT_INTO_SCAN_DEFER_CAST_ERROR`, default false). This rule
  only relocates the extraction; `PushVariantIntoScan` still does the
  scan-level materialization and applies the same deferral
  companion-slot logic downstream, so enabling that flag suppresses the
  join-eliminated-row error exactly as it does the filter case.
- Prune both `transformUp` passes with `transformUpWithPruning`: gate
  the first on `AGGREGATE`/`SORT`/`JOIN` and the second on `SORT`, so
  plans lacking those operators (and the Once-batch idempotence re-run)
  skip the traversal.
- Align the `pushSideAliases` join comment to note `LeftExistence` also
  matches `ExistenceJoin`, not just `LeftSemi`/`LeftAnti`.
@qlong
qlong force-pushed the variant-pushout-aggregate branch from 22f282a to a8f418e Compare July 16, 2026 00:18
@qlong
qlong requested a review from cloud-fan July 20, 2026 03:29
@qlong

qlong commented Jul 24, 2026

Copy link
Copy Markdown
Contributor Author

@qlong

qlong commented Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

@cloud-fan Can you take another look at this? It would be nice if this can get into 4.3 release. I added link to the benchmark which shows about 4 times speed up for tpc-ds like workload.

@cloud-fan cloud-fan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

3 addressed, 0 remaining, 0 new.
0 blocking, 0 non-blocking, 0 nits.
The prior review's three observations are addressed, and the current patch has no remaining design, correctness, completeness, or documentation findings.

Verification

Traced every rewrite from the original Aggregate/Sort/Join expression to a single child alias and verified that retained attributes include parent references plus join-condition keys. Checked empty/many-row cardinality, duplicate extraction deduplication, sort ordering, NULL padding across outer joins, alias metadata, configuration-off behavior, and the Project/Filter topology consumed by both scan-pushdown paths. Also re-adjudicated all prior threads and the mechanical link, text, and contract scans; no tests were run as part of this review.

@cloud-fan cloud-fan closed this in adad838 Jul 27, 2026
cloud-fan pushed a commit that referenced this pull request Jul 27, 2026
### 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 <qlong@users.noreply.github.com>
Signed-off-by: Wenchen Fan <wenchen@databricks.com>
(cherry picked from commit adad838)
Signed-off-by: Wenchen Fan <wenchen@databricks.com>
@cloud-fan

Copy link
Copy Markdown
Contributor

Merge Summary:

Posted by merge_spark_pr.py

@qlong

qlong commented Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

@cloud-fan Thanks for reviewing the PR

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants