Module late_materialize

Expand description

Late materialisation of a bounded top-N aggregate.

A GROUP BY that lists a key and the columns that key determines carries those columns through every join and every shuffle beneath it, only to display a handful of them at the end. This rule groups on the key alone, takes the top N, and re-fetches the wide columns for the survivors.

§The query that motivated this, and the number that justifies it

TPC-H q10 groups by seven columns — c_custkey, c_name, c_acctbal, c_phone, n_name, c_address, c_comment — and returns twenty rows. Six of the seven are functionally determined by c_custkey; c_comment alone averages ~73 B/row.

Measured at SF100 on the three-node cluster by hand-writing the narrowed query (same joins, same filters, same LIMIT, but GROUP BY c_custkey):

                      real q10 (wide)      narrowed
  s1  customer scan   52 task-s / 3.40 GB  16.5 / 240 MB   14x fewer bytes
  s2  orders⋈customer 8,968 task-s         38.9            230x
  s5  final agg+TopK  1,510                6.3             240x
  wall clock          1784.57 s            120.88 s        14.8x

s2 fell 230x while its input bytes fell only ~10x, so the cost is superlinear in the wide columns — per-row string handling in the hash join, not wire volume. That is why nothing about transport fixed it (see the q10-dist-s2-is-the-whole-query record: neither a cross-stage runtime filter nor a deeper shuffle prefetch moved it).

§Two rewrites that do not work, so they are not attempted again

Declaring the key alone is not enough. ParquetTableSpec::with_primary_key (shipped separately) gives DataFusion the functional dependency, and optimize_projections will happily shrink a GROUP BY with it — but only (columns the parent requires) ∪ (minimal FD subset). q10 selects all seven grouped columns, so the parent requires them and no key declaration can prune them.

Narrowing the group key alone is not enough either, and this one was measured rather than reasoned. Rewriting q10 as GROUP BY c_custkey plus first_value(...) per determined column — exactly what a local composite-key rule would emit — ran at SF100 in 1955.65 s against the 1784.57 s baseline, ~10% slower, and not one stage improved: s1 shipped the identical 3.40 GB. first_value still takes the wide columns as aggregate inputs, so they cross every join and shuffle exactly as before. Narrowing the group key only saves hashing, and hashing was never the cost.

The cost is the wide columns flowing through the joins. Only a join-back removes them, which is why this rule is non-local: the non-locality is the optimization, not an inconvenience around it.

§The rewrite

  Sort: revenue DESC, fetch=20
    Projection: c_custkey, c_name, revenue, c_acctbal, n_name, ...
      Aggregate: groupBy=[c_custkey, c_name, c_acctbal, c_phone,
                          n_name, c_address, c_comment]
                 aggr=[sum(...)]
        <customer ⋈ orders ⋈ lineitem ⋈ nation>

becomes

  Sort: revenue DESC, fetch=20                          (unchanged)
    Projection: ...                                     (unchanged)
      Projection: <exactly the aggregate's old schema>
        Inner Join: customer.c_nationkey = nation.n_nationkey
          Inner Join: __krishiv_lm.c_custkey = customer.c_custkey
            SubqueryAlias: __krishiv_lm
              Sort: sum(...) DESC, fetch=20
                Aggregate: groupBy=[c_custkey], aggr=[sum(...)]
                  <customer ⋈ orders ⋈ lineitem ⋈ nation>
            TableScan: customer
          TableScan: nation

Only the Aggregate node is replaced; everything above it keeps the exact same schema, so the enclosing Projection and Sort are untouched. The inner Sort carries the same fetch, which is what bounds the join-back to twenty rows — and is why the rule refuses without one.

Nothing prunes the wide columns from the narrow branch directly: once the aggregate stops referencing them, DataFusion’s own optimize_projections does it on the next pass, and the customer scan under the aggregate drops to [c_custkey, c_nationkey].

§Reaching a column through more than one table

n_name lives in nation, not in customer, so no direct join-back on c_custkey can fetch it. It is still determined: c_custkey → (customer’s key) → c_nationkey, c_nationkey = n_nationkey is an equality of the original join, and n_nationkey is nation’s key → n_name. DataFusion’s FD machinery does not compose dependencies across join equalities, so this rule computes its own closure:

a table’s columns become available when every column of its declared primary key is either already available or equated — by an inner-join ON pair in the aggregate’s input — to a column that is.

The same closure decides which group columns may be dropped and, run forwards, emits the join-back chain, so the two can never disagree about what is recoverable.

§Why it is safe

The key really determines the columns. Constraint::PrimaryKey is unverified here, exactly as Spark/Databricks RELY is. That single declaration is doing two jobs — uniqueness (so the join-back returns one row per key) and non-nullness (so the key equi-joins at all) — which is precisely what a primary key means. Constraint::Unique is refused: DataFusion marks it nullable, and a null key would silently fetch the wrong row or none.
Inner joins only, everywhere. Every node between the aggregate and its base scans must preserve column values: an outer join null-pads its non-preserved side, so a column re-fetched from the base table would come back non-null where the original plan had a null. Any node this rule does not understand — a SubqueryAlias, a nested aggregate, a union — makes it decline rather than guess.
The join-back is Inner, deliberately. A Left join is the instinct, and it is wrong here for a mechanical reason: PartitionMode::CollectLeft with a join type that emits unmatched build rows cannot be split across distributed tasks, so redistribute_unsplittable_broadcast_joins would convert it to a hash-partitioned join and shuffle the very columns this rule exists to keep off the wire. Inner is also exactly right semantically: the surviving key came from a row that already joined.
The ordering is computable before the columns are. Every ORDER BY expression must resolve to a retained key column or an aggregate output; a sort on a deferred column would need the column it is deferring. Group columns the sort names are added back to the key rather than refused.
Bounded output only. Without a fetch the join-back re-joins every group and the rewrite is a pure loss. See MAX_LATE_MATERIALIZE_FETCH.
Names cannot be crossed. Columns are matched by fully-qualified name, and a projection that reuses an input’s qualified name for a different expression makes the rule decline; duplicate qualified names across the collected scans (a self-join) do too.

Set KRISHIV_LATE_MATERIALIZATION=off to disable.

Structs§

LateMaterializeTopKAggregate: Replace a bounded top-N aggregate’s determined grouping columns with a join-back, so they never enter the joins beneath it.

Constants§

LATE_MATERIALIZATION_ENV: Environment switch for late materialisation.

Functions§

late_materialization_enabled: Whether late materialisation is enabled (default: yes).