apache / apache/datafusion

Tighten `ScalarSubqueryToJoin` Projection Rewrite State and Coverage

Open
#22,670 0 comments 0 reactions 1 assignee Claimed by @kosiew View on GitHub
Dominant language
Rust
Stars
9.3k
Forks
2.4k
Avg merge
3d 7h
Merged PRs (30d)
344

Description

## Summary
`ScalarSubqueryToJoin` projection rewrite logic already uses slot-indexed ownership (`alias_to_index` + `rewrite_exprs`) and no longer depends on `Expr -> Expr` rewrite maps.

However, state is still distributed across parallel containers and tested mostly through a small set of projection scenarios. This issue proposes a follow-up cleanup to make ownership/rewrite progression more explicit and add targeted regression coverage.

Goal: reduce maintenance risk in a historically fragile path without changing behavior.

## Context
In `datafusion/optimizer/src/scalar_subquery_to_join.rs` (projection branch), correlated scalar subqueries are extracted, then each subquery is converted to a join and projection expressions are rewritten iteratively.

Current shape (conceptually):

- collect all extracted `(subquery, alias)` pairs
- map each alias to a projection slot index
- keep per-slot rewritten expression state in a vector
- for each built join, locate owning expression via alias, then rewrite it

This is functionally correct today, but ownership and rewrite progression are still spread across loosely coupled containers (`all_subqueries`, `alias_to_index`, `rewrite_exprs`), which makes future changes harder to reason about.

## Problem
The current implementation has maintainability risk:

1. Ownership of subquery aliases and per-slot rewrite state is represented indirectly via parallel vectors/maps.
2. The projection branch performs iterative in-place rewrites where subtle ordering/name-preservation regressions are possible.
3. Test coverage does not clearly stress multi-subquery-per-slot and repeated structurally similar projection expressions.

## Proposed Design
Keep behavior unchanged, but make projection rewrite state explicit in one slot-owned model.

Example model:

```rust
struct ProjectionRewriteState {
rewritten_expr: Expr,
subquery_aliases: Vec,
}
```

Drive rewrite logic from:

- `Vec` indexed by projection slot
- `HashMap` mapping extracted alias to owning slot
- extraction-order queue of `(Subquery, String)` (or equivalent)

High-level flow:

1. Initialize per-slot state from projection expressions.
2. During extraction, append aliases into the owning slot and populate `alias -> slot_index`.
3. For each `(subquery, alias)` join conversion:
- resolve owning slot index from `alias -> slot_index`
- rewrite `states[slot_index].rewritten_expr` in place using `expr_check_map`
4. Build final projection expressions from per-slot rewritten expressions, preserving output names/aliases exactly as today.

## Why This Helps
1. Encodes ownership and rewrite progression directly in slot state.
2. Reduces cognitive overhead from parallel-container coupling.
3. Makes it easier to audit extraction/rewrite order guarantees.
4. Improves readability and lowers regression risk in future optimizer work.

## Scope
In scope:

- Refactor projection arm of `ScalarSubqueryToJoin` to explicit slot-owned state (without behavioral change).
- Keep existing behavior and output schema/name preservation unchanged.
- Update/add tests as needed.

Out of scope:

- Broader redesign of scalar-subquery rewrite outside projection branch.
- Behavioral changes to which queries are optimized.

## Acceptance Criteria
1. No behavior regression for correlated scalar subquery rewrites in projection context.
2. Projection rewrite bookkeeping is modeled through explicit slot-owned rewrite state.
3. Alias ownership remains represented by slot index (or equivalent direct ownership model).
4. Existing tests pass and targeted projection tests are added for this refactor-sensitive area.

## Test Plan
Add/adjust tests to cover:

1. Multiple projection expressions, each with correlated scalar subqueries.
2. A projection expression containing multiple correlated scalar subqueries.
3. Repeated/structurally-similar expressions in projection list (ensure correct slot ownership rewrite).
4. Existing regression scenario from the recent fix.

Prefer SQL-visible behavior checks via sqllogictest where appropriate, plus focused optimizer/unit tests near the rule implementation.

## Risks and Mitigations
Risk: subtle rewrite-order or alias-preservation regression.

Mitigation:

- Preserve extraction and rewrite order semantics.
- Validate output expression naming behavior against current logic.
- Add targeted tests for mixed projection/subquery layouts.

## Related PR
- #22313

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.