Skip to content

feat(substrait): serialize correlated subqueries with OuterReference field references#23488

Open
clflushopt wants to merge 2 commits into
apache:mainfrom
clflushopt:feat/substrait-producer-outer-references
Open

feat(substrait): serialize correlated subqueries with OuterReference field references#23488
clflushopt wants to merge 2 commits into
apache:mainfrom
clflushopt:feat/substrait-producer-outer-references

Conversation

@clflushopt

@clflushopt clflushopt commented Jul 12, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

SingleDistinctToGroupBy rewrites a grouped AGG(DISTINCT col) into a cheaper nested double aggregate. Today the rule only recognizes a bare Expr::AggregateFunction in Aggregate.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:

  • SQL (count(DISTINCT v) AS n): the planner hoists AS n into an outer Projection, leaving a bare Expr::AggregateFunction in aggr_expr.
  • DataFrame API (count(col("v")).distinct().alias("n")): .alias(..) wraps the aggregate directly, so aggr_expr holds Expr::Alias(AggregateFunction).

is_single_distinct_agg matched only the bare shape and returned false for Alias(AggregateFunction), so the rule never fired for DataFrame-API distinct aggregates.

What changes are included in this PR?

  • Unwrap a single Expr::Alias layer in the matcher (is_single_distinct_agg) and in the rewrite, mirroring the one-layer unalias() idiom already used in the optimizer (e.g. decorrelate.rs), so both front-end shapes are recognized identically.
  • The user-facing alias is preserved without extra work: the rewrite already rebuilds the final output Projection from the original Aggregate schema (schema.qualified_field(idx)), so AS n survives.
  • Bare-aggregate behavior is unchanged (byte-identical): every pre-existing snapshot passes untouched.

Are these changes tested?

Yes — new snapshot tests in single_distinct_to_groupby.rs:

  • single_distinct_aliased — the DataFrame-API Alias(AggregateFunction) shape is now rewritten, and the output column stays named n.
  • single_distinct_aliased_and_groupby — same, alongside a real GROUP BY column.
  • non_distinct_aliased and two_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 in joins.slt fail in --substrait-round-trip mode.

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 FieldReference with an OuterReference root type and a steps_out depth. The consumer side of was implemented in #20439, which resolves OuterReference field 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?

  • SubstraitProducer gains outer-schema-stack methods mirroring SubstraitConsumer: push_outer_schema / pop_outer_schema (default no-ops) and get_outer_schema(steps_out) (default None), plus a handle_outer_reference_column method 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 of not_impl_err.
  • DefaultSubstraitProducer maintains the stack in a Vec<DFSchemaRef>.
  • The four subquery producers (from_in_subquery, from_scalar_subquery, from_exists, from_set_comparison) push the enclosing query's schema around the subquery plan conversion (via a shared produce_subquery_rel, analogous to the consumer's consume_subquery_rel).
  • from_outer_reference_column (previously unused and emitting an incorrect plain RootReference) now resolves the column against the outer-schema stack, innermost first, and emits a FieldReference with an OuterReference root and the corresponding steps_out.
  • to_substrait_rex dispatches Expr::OuterReferenceColumn to the new handler instead of erroring.

Are these changes tested?

Yes:

  • New round-trip tests in roundtrip_logical_plan.rs covering correlated EXISTS, correlated IN subquery, correlated scalar subquery, and a nested correlated subquery that crosses two subquery boundaries (steps_out = 2). Each test asserts the produced plan contains an OuterReference at the expected depth and that the plan round-trips through the existing consumer with its schema intact.
  • Consumer-side resolution was already covered by the tests added in fix(substrait): Correctly parse field references in subqueries #20439; these tests now exercise both halves together.
  • joins.slt in --substrait-round-trip mode goes from 38 failures to 12; the remaining failures are pre-existing gaps unrelated to outer references (USING join constraint, plan-level lateral LogicalPlan::Subquery, duplicate unqualified field names).

Are there any user-facing changes?

  • Plans containing correlated subqueries now serialize instead of returning "not implemented", emitting spec-standard OuterReference field references.
  • SubstraitProducer has three new provided methods and handle_outer_reference_column; all have defaults, so existing implementations continue to compile.
  • The signature of the public helper from_outer_reference_column changed (it now takes the producer and the outer field) — its previous form resolved against the wrong schema and emitted a plain RootReference, and it was not called from anywhere in the crate.

…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.
@github-actions github-actions Bot added the substrait Changes to the substrait crate label Jul 12, 2026
@clflushopt

Copy link
Copy Markdown
Contributor Author

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.

@github-actions

Copy link
Copy Markdown

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
     Cloning apache/main
    Building datafusion-substrait v54.0.0 (current)
       Built [ 346.688s] (current)
     Parsing datafusion-substrait v54.0.0 (current)
      Parsed [   0.017s] (current)
    Building datafusion-substrait v54.0.0 (baseline)
       Built [ 345.153s] (baseline)
     Parsing datafusion-substrait v54.0.0 (baseline)
      Parsed [   0.018s] (baseline)
    Checking datafusion-substrait v54.0.0 -> v54.0.0 (no change; assume patch)
     Checked [   0.223s] 223 checks: 222 pass, 1 fail, 0 warn, 30 skip

--- failure function_parameter_count_changed: pub fn parameter count changed ---

Description:
A publicly-visible function now takes a different number of parameters.
        ref: https://doc.rust-lang.org/cargo/reference/semver.html#fn-change-arity
       impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.48.0/src/lints/function_parameter_count_changed.ron

Failed in:
  datafusion_substrait::logical_plan::producer::from_outer_reference_column now takes 4 parameters instead of 2, in /home/runner/work/datafusion/datafusion/datafusion/substrait/src/logical_plan/producer/expr/field_reference.rs:96

     Summary semver requires new major version: 1 major and 0 minor checks failed
    Finished [ 694.039s] datafusion-substrait

@github-actions github-actions Bot added the auto detected api change Auto detected API change label Jul 12, 2026

@bvolpato bvolpato 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.

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?

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

Labels

auto detected api change Auto detected API change substrait Changes to the substrait crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[substrait] [sqllogictest] Cannot convert Exists { subquery: <subquery>, negated: true } to Substrait

2 participants