[SPARK-58089][SQL] Push variant extractions through Aggregate/Sort/Join - #57190
[SPARK-58089][SQL] Push variant extractions through Aggregate/Sort/Join#57190qlong wants to merge 2 commits into
Conversation
ef2a92e to
7cace8c
Compare
`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
7cace8c to
bc27fe1
Compare
|
@cloud-fan Can you take a look for this change to extend scope of variant extraction pushdown? |
cloud-fan
left a comment
There was a problem hiding this comment.
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:
pushAliasesIntoJoinduplicates theJoinbranch ofpushSideAliases— 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( |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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 => |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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 |
There was a problem hiding this comment.
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.
…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`.
22f282a to
a8f418e
Compare
|
To verify performance, use https://github.com/cloudera-labs/variant-conformance-benchmark |
|
@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
left a comment
There was a problem hiding this comment.
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.
### 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 Thanks for reviewing the PR |
What changes were proposed in this pull request?
Introduce a new optimizer rule
PullOutVariantExtractionsthat hoistsvariant_get/Cast(variant)extractions out of three operator typesthat the existing
PushVariantIntoScan/V2ScanRelationPushDownrulescannot see through:
Aggregate function arguments – e.g.
max(variant_get(v, '$.price', 'int')):the extraction is moved into a
Projectdirectly below theAggregateand theaggregate references the resulting alias. The bare variant column is suppressed
unless it is also needed raw (e.g. as a
GROUP BYkey), so no redundantfull-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 theSortandthe original
Projectis reproduced to prevent the alias from leaking intothe output.
Join conditions and projections above joins – matched as
Project → Join;extractions in both the join condition and the outer
Projectare routed tothe owning join side. A
pushSideAliaseshelper then pushes the aliasesthrough any depth of chained joins so they land in a
Projectdirectlyabove the scan (where
PhysicalOperationcollapses them with the scan, makingthem visible to the pushdown). This is necessary because
PhysicalOperationstops at a
Joinnode.A
Sortsitting over aJoinis handled by fusing the two cases: theorder-key aliases are pushed through the join tree, not left in a
Projectabove it.
The rule is gated by a new internal config
spark.sql.variant.pushVariantIntoScan.pullOutExtractions(defaulttrue)and is a no-op unless
spark.sql.variant.pushVariantIntoScanis also enabled.Non-variant plans are untouched.
The rule is registered as the first rule in
SparkOptimizer.earlyScanPushDownRules,before
SchemaPruningand the V2 scan pushdown rules.Why are the changes needed?
Before this change, a
variant_getinside 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:
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.
Run A with pullout rule enabled vs B using json for payload, this test shows the performance advantage of variant over json.
Was this patch authored or co-authored using generative AI tooling?
Co-authored with Claude Code