feat(substrait): serialize correlated subqueries with OuterReference field references#23488
feat(substrait): serialize correlated subqueries with OuterReference field references#23488clflushopt wants to merge 2 commits into
Conversation
…field references The Substrait producer previously errored on Expr::OuterReferenceColumn, making it impossible to serialize unoptimized plans containing correlated subqueries. The consumer already resolves OuterReference field references against a stack of outer schemas; this adds the producing half: - SubstraitProducer gains push_outer_schema/pop_outer_schema/get_outer_schema (mirroring SubstraitConsumer) and a handle_outer_reference_column hook. - DefaultSubstraitProducer maintains the outer-schema stack. - The subquery producers (InSubquery, ScalarSubquery, Exists, SetComparison) push the enclosing schema around the subquery plan. - from_outer_reference_column resolves the column against the stack, innermost first, and emits a FieldReference with an OuterReference root and the corresponding steps_out, per the Substrait spec.
|
The inception for this PR came from a discussion I had with @vbarua on Datafusion's Discord; do note that initially this was marked as a bit tricky to do and would warrant a discussion so looking back at the history specifically #18987 and its consumer half in #20439 I found a good middle ground to attempt an approach. Open to change or even close and discuss. |
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
bvolpato
left a comment
There was a problem hiding this comment.
Direction and implementation look good. Standard OuterReference encoding matches #20439, and local Substrait suite passed.
PR description starts with unrelated SingleDistinctToGroupBy content. Could you remove that copied section?
Which issue does this PR close?
Rationale for this change
SingleDistinctToGroupByrewrites a groupedAGG(DISTINCT col)into a cheaper nested double aggregate. Today the rule only recognizes a bareExpr::AggregateFunctioninAggregate.aggr_expr, so it fires for plans from the SQL planner but silently skips the logically-identical plan built through the DataFrame API — that path runs materially slower for the same query.The divergence is purely in where each front-end places the output alias:
count(DISTINCT v) AS n): the planner hoistsAS ninto an outerProjection, leaving a bareExpr::AggregateFunctioninaggr_expr.count(col("v")).distinct().alias("n")):.alias(..)wraps the aggregate directly, soaggr_exprholdsExpr::Alias(AggregateFunction).is_single_distinct_aggmatched only the bare shape and returnedfalseforAlias(AggregateFunction), so the rule never fired for DataFrame-API distinct aggregates.What changes are included in this PR?
Expr::Aliaslayer in the matcher (is_single_distinct_agg) and in the rewrite, mirroring the one-layerunalias()idiom already used in the optimizer (e.g.decorrelate.rs), so both front-end shapes are recognized identically.Projectionfrom the originalAggregateschema (schema.qualified_field(idx)), soAS nsurvives.Are these changes tested?
Yes — new snapshot tests in
single_distinct_to_groupby.rs:single_distinct_aliased— the DataFrame-APIAlias(AggregateFunction)shape is now rewritten, and the output column stays namedn.single_distinct_aliased_and_groupby— same, alongside a realGROUP BYcolumn.non_distinct_aliasedandtwo_distinct_aliased_and_groupby— negative cases: an aliased non-distinct aggregate, and two aliased distinct aggregates on different fields, are still (correctly) left untouched.Reverting the production change makes the two positive tests fail (the rule no longer fires) while every other test stays green, confirming the tests are wired to the fix.
Are there any user-facing changes?
No API changes. DataFrame-API queries using
count(distinct …)(and other single-distinct aggregates) now receive the same optimized plan as their SQL equivalents — a performance improvement, with no change to results or output column names.Rationale for this change
The Substrait producer errors on
Expr::OuterReferenceColumn, so any plan containing a correlated subquery cannot be serialized currently. DataFusion's round-trip tests don't hit this because the optimizer decorrelates subqueries into joins before serialization, but any workflow that serializes unoptimized plans (e.g. sending raw plans between systems for later optimization in my case) fails on queries like TPC-H q2/q4/q17/q20/q21/q22, and ~26 cases injoins.sltfail in--substrait-round-tripmode.A previous attempt (#18987) was closed because it introduced a non-standard mechanism for outer references, producing plans only DataFusion would be able consume but Substrait already represents correlated references natively via a
FieldReferencewith anOuterReferenceroot type and asteps_outdepth. The consumer side of was implemented in #20439, which resolvesOuterReferencefield references against a stack of outer schemas. This PR implements the producing half, symmetric with that design, so correlated plans round-trip using only standard Substrait.What changes are included in this PR?
SubstraitProducergains outer-schema-stack methods mirroringSubstraitConsumer:push_outer_schema/pop_outer_schema(default no-ops) andget_outer_schema(steps_out)(defaultNone), plus ahandle_outer_reference_columnmethod so custom producers can override the behaviour like every other expression kind. Defaults are backward compatible: existing custom producers are unaffected unless a plan actually contains an outer reference, in which case they now get an actionable error instead ofnot_impl_err.DefaultSubstraitProducermaintains the stack in aVec<DFSchemaRef>.from_in_subquery,from_scalar_subquery,from_exists,from_set_comparison) push the enclosing query's schema around the subquery plan conversion (via a sharedproduce_subquery_rel, analogous to the consumer'sconsume_subquery_rel).from_outer_reference_column(previously unused and emitting an incorrect plainRootReference) now resolves the column against the outer-schema stack, innermost first, and emits aFieldReferencewith anOuterReferenceroot and the correspondingsteps_out.to_substrait_rexdispatchesExpr::OuterReferenceColumnto the new handler instead of erroring.Are these changes tested?
Yes:
roundtrip_logical_plan.rscovering correlatedEXISTS, correlatedINsubquery, correlated scalar subquery, and a nested correlated subquery that crosses two subquery boundaries (steps_out = 2). Each test asserts the produced plan contains anOuterReferenceat the expected depth and that the plan round-trips through the existing consumer with its schema intact.joins.sltin--substrait-round-tripmode goes from 38 failures to 12; the remaining failures are pre-existing gaps unrelated to outer references (USINGjoin constraint, plan-level lateralLogicalPlan::Subquery, duplicate unqualified field names).Are there any user-facing changes?
OuterReferencefield references.SubstraitProducerhas three new provided methods andhandle_outer_reference_column; all have defaults, so existing implementations continue to compile.from_outer_reference_columnchanged (it now takes the producer and the outer field) — its previous form resolved against the wrong schema and emitted a plainRootReference, and it was not called from anywhere in the crate.