Optimizer: pull a row-local operator chain (UNNEST / projections) above an exchange to shuffle less (cost-based)

[kaikalur]

Summary

Add an optimizer rule that pushes an UnnestNode below an adjacent repartitioning ExchangeNode, valid when none of the exchange's partitioning keys is an unnested output variable (or the ordinality variable) — i.e. the exchange partitions only on columns that UNNEST passes through unchanged (its replicate columns). Under that condition UNNEST is partition-preserving on those keys, so Exchange[K](Unnest(x)) ≡ Unnest(Exchange[K](x)) when K ∩ unnestedVariables = ∅ (equivalently K ⊆ replicate(UNNEST)). Choosing between the two placements is then a pure cost decision.

Motivation

When an UNNEST sits just below a repartitioning exchange (e.g. the build side of a downstream join that keys on a replicate column carried through the unnest), the exchange shuffles the post-unnest, fan-out-expanded row set. Whether that's good or bad depends on the shape:

• If the unnested array element is narrow and the source row is wide (the wide payload is consumed by the unnest), the post-unnest rows are more rows but fewer bytes, so shuffling above the unnest (current behavior) is cheaper — the rule should NOT fire.
• If the fan-out is large relative to element width (post-unnest is many more bytes), shuffling the pre-unnest input moves far fewer bytes — the rule should fire.

Today the placement is fixed; there is no rule to move UNNEST across the exchange even when doing so reduces shuffle bytes. This also helps stacked self-join / late-materialization plans keep the repartition on the cheaper side.

Proposal

• New rule (default OFF, session-property gated) matching Exchange (repartition on K) -> [pass-through Project?] -> Unnest where K contains no unnested output variable and no ordinality variable (K ⊆ UNNEST.replicateVariables).
• Rewrite to [Project?] -> Unnest -> Exchange (repartition on K) (push the exchange below the unnest), preserving semantics because K is replicate.
• Cost-based trigger: apply only when the estimated pre-unnest shuffle size (rows × width) is smaller than the post-unnest shuffle size — i.e. when fan-out × element-width does not dominate the consumed source width.

Acceptance

• Unit (plan-shape) test: rule fires / does not fire per the "no exchange key is an unnested/ordinality variable" guard and the cost threshold; semantics-preserving.
• e2e: enabled-vs-disabled result equality over ... CROSS JOIN UNNEST(...) ... shapes (TPC-H, e.g. lineitem / orders).

(Follow-up; relates to existing UNNEST pushdown rules such as pushdown_through_unnest.)

[kaikalur]

Generalization: pull a maximal row-local chain above the exchange (don't shuffle expanded data)

While scoping the UNNEST-below-exchange rule above, a production INSERT … SELECT … CROSS JOIN UNNEST(…) made it clear the same idea generalizes well beyond "UNNEST under a hash-repartition exchange". The unifying principle:

A repartitioning/scaling Exchange should shuffle data at the point in the plan where it is smallest. If a chain of row-local, partition-key-preserving operators sits below the exchange and expands the data (UNNEST fan-out, width-increasing projections), pull the whole chain above the exchange so the shuffle happens on the smaller pre-expansion input.

Invariant

For an operator Op and exchange Exchange[K]: Op(Exchange[K](x)) ≡ Exchange[K](Op(x)) iff

1. Op keeps K pass-through (does not compute or drop the partition / hash / ordering columns), and
2. Op is row-local / partition-independent (processes each row, or each K-partition, independently) and deterministic.

When K = ∅ (a scaling / round-robin writer exchange) condition (1) is vacuous, so any row-local deterministic operator may cross it.

Motivating shape (generic)

INSERT INTO target
SELECT  <wide replicate columns>, t.elem, <derived columns>
FROM    wide_fact f
JOIN    (SELECT DISTINCT key FROM …) d ON f.key = d.key
CROSS JOIN UNNEST(f.<array_or_map_column>) AS t(elem)

Planned today (top = consumer, data flows upward):

TableWriter
  └─ Exchange[SCALED, K=∅]            ← shuffles POST-unnest, post-projection rows
       └─ Project (derive output columns)
            └─ Unnest (fan-out; wide replicate columns duplicated per array element)
                 └─ Project (extract array/map to unnest)
                      └─ InnerJoin (key)

The UNNEST replicates the wide carry-through columns once per array element, so the SCALED writer exchange shuffles the fully fanned-out, widened row set. In one production query this exchange moved ~1 PB, while the join output feeding the chain was ~316 TB — i.e. the write shuffled ~3.2× more bytes than necessary, a large contributor to the query exceeding its execution-time limit.

Proposed rewrite — pull everything except the join above the exchange

TableWriter
  └─ Project (derive output columns)
       └─ Unnest
            └─ Project (extract array/map)
                 └─ Exchange[SCALED, K=∅]   ← now shuffles the JOIN OUTPUT (~316 TB)
                      └─ InnerJoin (key)

The join stays put (multi-input, partition-consuming barrier). The expansion happens on the consumer side, after the shuffle. Same result, ~3.2× less shuffle (and ~700 TB less network) on this query.

Note: this reduces shuffle/network, not the table-writer CPU — the writer still emits the same expanded bytes. It directly targets the shuffle and wall-time, not write cost.

Rule: PullRowLocalChainAboveExchange (default OFF, cost-gated)

• Rule<ExchangeNode>, pattern exchange().matching(scope.isRemote()) — same entry point as PullConstantProjectionAboveExchange (feat(optimizer): Pull constant projections above remote exchanges #27499).
• From the exchange's (single) source, walk down a maximal chain of allowlisted row-local nodes — start with { deterministic Project, Unnest } — until a barrier (Join / Aggregation / TableScan / Exchange / any non-allowlisted node).
• Legality: K (partitioning ∪ hash ∪ ordering vars) must be pass-through to the barrier's outputs (not produced by an in-chain Unnest or a non-identity Project); every Project assignment deterministic (RowExpressionDeterminismEvaluator); single-source exchange; not replicateNullsAndAny.
• Cost gate (isCostBased, bail on LOW confidence): fire iff
  stats(barrier).getOutputSizeForVariables(varsChainNeeds) * threshold < stats(chainTop).getOutputSizeInBytes()
  (pre-expansion shuffle strictly smaller than the current post-expansion shuffle).
• Rewrite (reparent — no symbol surgery): build E' over the barrier (copy type / scope / scaleWriters / bucketToPartition; partition on K mapped to barrier-side vars; output layout = the vars the chain consumes), then reparent the chain nodes on top of E' bottom→top via replaceChildren, allocating fresh PlanNodeIds. Because each chain node already references the barrier's outputs and E' passes them through unchanged, all variable references stay valid. Add a top rename Project only if the original exchange renamed its inputs (otherwise bail in v1).

Relationship to existing rules

• Subsumes the original ask (push UNNEST below a repartition exchange) as the case chain = {Unnest}, K ⊆ replicate variables.
• Generalizes PullConstantProjectionAboveExchange (feat(optimizer): Pull constant projections above remote exchanges #27499): that rule is the degenerate case chain = {single Project}, no Unnest, lifting only constant assignments.
• Loop-prevention vs PushProjectionThroughExchange (which unconditionally pushes computed projections below exchanges — PRs feat(optimizer): Add option to not push down remote project below exchange #26943 and feat(optimizer): Pull constant projections above remote exchanges #27499 already add guards against its over-eager pushdown): register this rule as a standalone IterativeOptimizer after the final projectionPushDown pass (same slot as feat(optimizer): Pull constant projections above remote exchanges #27499), and add a gated skip in PushProjectionThroughExchange for projections belonging to a hoisted chain (gated on the new session property, so default plans are unchanged when OFF). PushdownThroughUnnest cannot re-push the above-unnest projection (it references unnested outputs).

Acceptance

• Unit (plan-shape): chain pulled above the exchange when K is pass-through and cost says pre < post; does not fire when K is produced in-chain, below the cost threshold, with a non-deterministic projection, or when disabled. Cover {Unnest}, {Project → Unnest → Project}, and the SCALED (K = ∅) writer-exchange shape.
• e2e: assertQuery(enabled, sql, disabled, sql) over … CROSS JOIN UNNEST(…) and INSERT … SELECT … CROSS JOIN UNNEST(…) (± WITH ORDINALITY) on TPC-H tables — result equality with the optimization enabled vs disabled.

Session property: optimizer.pull-row-local-chain-above-exchange / pull_row_local_chain_above_exchange (default false), plus a byte-reduction threshold property mirroring partial_aggregation_byte_reduction_threshold.