Consider the following python script:
ht1 = ... ht2 = ... ht2.aggregate(hl.agg.count() + ht1.aggregate(hl.agg.count(), _localize=False)
This will generate an IR that looks like:
(TableAggregate ht2 (ApplyBinOp + (ApplyAggOp count) (TableAggregate ht1 (ApplyAggOp count))))
My initial design for lowering TableAggregate was roughly the following:
- Compute requiredness over the whole IR to be lowered
- Recursively lower with the RequirednessAnalysis as auxillary information
- When hitting a TableAggregate, lower the child by extracting with the computed requiredness.
I expected that we’d only need to compute requiredness once at the beginning. However, Arcturus pointed out that the above strategy is insufficient for the above IR, which has a nested TableAggregate inside the query.
I investigated further, and none of the current IRs support a pattern that has a relational IR inside a value IR (the
newRow for TableMapRows, the filter for TableFilter, etc).
Proposal 1: restructure IR before running LowerToCDA
One option to fix this is to introduce a rewriting pass that runs before LowerToCDA, which extracts these nested patterns as
RelationalLet IRs near the root of the computation.
This has an additional benefit, which is that we gain the same property guaranteed by
InterpretNonCompilable in the TableIR.execute runtime: a TableAggregate IR that appears with nesting depth > 0 will still only be run once.
The complexity here is that we need to deal with figuring out how these RelationalLet nodes are lowered to normal lets, with bindings inserted as broadcast values in CDA where necessary.
Proposal 2: recursively lower value IR children
Another option is to recursively lower value IR children of nodes like TableFilter/TableAggregate in
LowerTableIR. This is simpler, but has a few negative properties:
- We need to recompute requiredness for each of these nodes. This scales quadratically with the size of an IR.
- We can emit CDAs inside of loops, and we currently don’t have a code motion pass to clean this up.