///|
/// Cap on how many hash partitions one aggregation may create. The join caps
/// at 256 and splits an over-budget partition again; aggregation cannot
/// recurse (a group's state must stay in one partition), so it stays lower.
const AGG_PARTITION_CAP : Int32 = 64
// Grace hash aggregation: when a GROUP BY's input
// exceeds the memory budget, the input is partitioned by the hash of
// its group keys into per-partition spill runs (the same SpillStore the
// sort and the join use), and the existing in-memory aggregate() runs
// one partition at a time. Equal group keys always land in the same
// partition, so each partition's group rows are FINAL — median,
// string_agg and count-distinct state never spans partitions, and no
// combine phase exists. The engine still never touches the filesystem.
///|
/// Dispatch a GROUP BY aggregation under the memory budget: within
/// budget (or budget-less, or a global aggregate with no GROUP BY) the
/// historical in-memory aggregate() runs; over budget the aggregation
/// spills (grace hash) when a store is provided, and is a clean error
/// without one.
fn aggregate_within_budget(
joined : Array[@types.DataBatch],
group_exprs : Array[PhysExpr],
group_types : Array[@types.DataType],
aggs : Array[AggSpec],
opts : ExecOptions,
stats : Array[(String, Int64)],
) -> @types.DataBatch raise @types.SqlError {
if opts.budget <= 0 || group_exprs.length() == 0 {
return aggregate(joined, group_exprs, group_types, aggs)
}
if batches_bytes(joined) <= opts.budget {
return aggregate(joined, group_exprs, group_types, aggs)
}
match opts.spill {
Some(store) =>
grace_agg(
joined,
group_exprs,
group_types,
aggs,
opts.budget,
store,
stats,
)
None =>
raise @types.SqlError::Execute(
"aggregate input exceeds the memory budget (\{batches_bytes(joined)} bytes > \{opts.budget}); provide a spill store (--spill-dir) or raise the budget (--memory)",
)
}
}
///|
/// Grace hash aggregation: partition the input by group-key hash into
/// spill runs, aggregate each partition in memory (its group rows are
/// final), and concatenate. P is derived from the estimated input bytes
/// so a partition's input fits the budget; a partition that still
/// overflows (few distinct keys — every row of a group shares its
/// partition) aggregates in memory, exceeding the budget for that
/// partition only — group state cannot be split further without a
/// combine phase (future work).
fn grace_agg(
joined : Array[@types.DataBatch],
group_exprs : Array[PhysExpr],
group_types : Array[@types.DataType],
aggs : Array[AggSpec],
budget : Int32,
store : SpillStore,
stats : Array[(String, Int64)],
) -> @types.DataBatch raise @types.SqlError {
let p = partition_count(batches_bytes(joined), budget, AGG_PARTITION_CAP)
let writers : Array[RunWriter] = []
for _ in 0.. 0 {
stats.push(("spill partitions", p.to_int64()))
}
grace_finish(store, writers, group_exprs, group_types, aggs)
}
///|
/// Aggregate every partition the writers hold, one run at a time, remove their
/// files and concatenate the groups. Shared by the collected and the pulled
/// grace aggregation: both end with every partition's runs on the store.
fn grace_finish(
store : SpillStore,
writers : Array[RunWriter],
group_exprs : Array[PhysExpr],
group_types : Array[@types.DataType],
aggs : Array[AggSpec],
) -> @types.DataBatch raise @types.SqlError {
let out : Array[@types.DataBatch] = []
let dtypes : Array[@types.DataType] = []
for k in 0.. Int32 {
let mut p = budget / 4 / AGG_MIN_RUN
if p < 2 {
p = 2
}
if p > AGG_PULL_PARTITION_CAP {
p = AGG_PULL_PARTITION_CAP
}
p
}
///|
/// The partition writers of one pulled aggregation and the single buffer they
/// share. Each writer used to hold its own `budget / 4`, so up to 64 of them
/// could keep 16 times the budget; here the buffers add up to `limit`, and when
/// they do not fit the largest one is written out as a run.
priv struct PartitionPool {
writers : Array[RunWriter]
limit : Int32
mut pooled : Int32
// the most the shared buffers held at once, before the flush that brought
// them back under `limit`
mut peak : Int32
// estimated in-memory bytes each partition has written out as runs
spilled : Array[Int64]
}
///|
/// Route each row of `batch` into its key-hash partition buffer, then write out
/// the largest buffers until the pool fits again.
fn PartitionPool::add(
self : PartitionPool,
batch : @types.DataBatch,
key_exprs : Array[PhysExpr],
) -> Unit raise @types.SqlError {
let p = self.writers.length()
let key_vecs : Array[@types.ColumnVector] = []
for e in key_exprs {
key_vecs.push(eval_vec(e, batch))
}
let idx : Array[Array[Int32]] = []
for _ in 0.. 0 {
let part = take_rows(batch, idx[k])
let bytes = estimate_batch_bytes(part)
self.writers[k].stage(part, bytes)
self.pooled += bytes
}
}
if self.pooled > self.peak {
self.peak = self.pooled
}
while self.pooled > self.limit {
self.flush_largest()
}
}
///|
/// Write the largest buffer as one run and release what it held.
fn PartitionPool::flush_largest(
self : PartitionPool,
) -> Unit raise @types.SqlError {
let mut best = 0
for k in 1.. self.writers[best].acc_bytes {
best = k
}
}
let w = self.writers[best]
let held = w.acc_bytes
if held == 0 {
self.pooled = 0
return
}
w.flush()
self.pooled -= held
self.spilled[best] += held.to_int64()
}
///|
/// What a pulled aggregation hands back besides its groups: how many partitions
/// it spilled into (0 when it fit its budget), the most input bytes it held at
/// once, and the most one partition wrote out. Both byte counts are the
/// engine's size estimate (`estimate_batch_bytes`), not file sizes.
priv struct PulledAgg {
batch : @types.DataBatch
partitions : Int32
peak_buffered : Int64
largest_spilled : Int64
// how many segments a mergeable fold split into; 0 when it did not fold that
// way, so a run that fits its budget reports nothing
segments : Int32
}
///|
/// The specs that re-fold a batch of partial *state*: the keys are the batch's
/// leading columns and every state column is folded by the operation it is the
/// partial of — a sum of sums, a count of counts, a min of mins, a max of maxes.
/// `avg` is the reason the state exists at all: its two columns are the part-sum
/// and the part-count, so an average of averages never happens.
fn merge_state_specs(
aggs : Array[AggSpec],
group_types : Array[@types.DataType],
) -> (Array[PhysExpr], Array[AggSpec]) {
let keys : Array[PhysExpr] = []
for i, t in group_types {
keys.push(PhysExpr::ColRef(i, t))
}
let mut at = group_types.length()
let specs : Array[AggSpec] = []
for spec in aggs {
match spec.kind {
Avg => {
specs.push({
kind: Sum,
input: Some(PhysExpr::ColRef(at, @types.Float64)),
out_type: @types.Float64,
sep: None,
q: None,
})
specs.push({
kind: Sum,
input: Some(PhysExpr::ColRef(at + 1, @types.Int64)),
out_type: @types.Int64,
sep: None,
q: None,
})
at = at + 2
}
Count | CountStar => {
specs.push({
kind: Sum,
input: Some(PhysExpr::ColRef(at, @types.Int64)),
out_type: @types.Int64,
sep: None,
q: None,
})
at = at + 1
}
Sum | Min | Max => {
specs.push({ ..spec, input: Some(PhysExpr::ColRef(at, spec.out_type)), })
at = at + 1
}
_ => ()
}
}
(keys, specs)
}
///|
/// Turn a merged state batch into the grouped result: the keys pass through, and
/// each aggregate's state becomes its value — an average is its sum over its
/// count, with no value in the group at all giving NULL rather than a division.
fn finalize_state(
state : @types.DataBatch,
aggs : Array[AggSpec],
group_types : Array[@types.DataType],
) -> @types.DataBatch {
let g = group_types.length()
let n = state.row_count()
let cols = state.columns()
let columns : Array[@types.ColumnVector] = []
for c in 0.. {
let sums = cols[at]
let counts = cols[at + 1]
let vec = @types.ColumnVector::make(@types.Float64, n)
for r in 0.. c
_ => 0L
}
vec.set(
r,
if count == 0L {
@types.Null
} else {
match sums.get(r) {
@types.Float64(s) => @types.Float64(s / count.to_double())
@types.Int64(s) =>
@types.Float64(s.to_double() / count.to_double())
_ => @types.Null
}
},
)
}
columns.push(vec)
at = at + 2
}
Sum | Count | CountStar | Min | Max => {
columns.push(cols[at])
at = at + 1
}
_ => ()
}
}
@types.DataBatch::new(columns, n)
}
///|
/// Whether the fold's keys are sparse enough for the state to be worth carrying
/// as rows: at most one distinct group key per eight rows of the first batch. A
/// shape whose keys are nearly as many as its rows gains nothing from segments —
/// the merged partials would be as large as the input it replaced — and keeps the
/// partitioned path, which is the shape q20 is.
fn fold_keys_are_sparse(
batch : @types.DataBatch,
group_exprs : Array[PhysExpr],
) -> Bool raise @types.SqlError {
let n = batch.row_count()
// a batch too small to say anything about the key density splits freely: the
// state of a few hundred rows is cheap either way
if n < 1024 || group_exprs.length() == 0 {
return true
}
let cols : Array[@types.ColumnVector] = []
for e in group_exprs {
cols.push(eval_vec(e, batch))
}
let seen : Map[String, Bool] = Map([])
for r in 0.. n {
return false
}
}
true
}
///|
/// The fold for an aggregation whose partials merge: read the pull in
/// `budget`-sized segments, fold each one into its own grouped batch, and merge
/// the partials as they arrive — re-folding them together whenever they outgrow
/// the budget, so the merge itself stays bounded. q1 folds 5.9M rows of lineitem
/// into four groups, and the partitioned path wrote and read 199 MB to do it;
/// this path touches the spill store never (§94).
fn aggregate_pulled_segments(
next : () -> @types.DataBatch? raise @types.SqlError,
group_exprs : Array[PhysExpr],
group_types : Array[@types.DataType],
aggs : Array[AggSpec],
limit : Int64,
) -> PulledAgg raise @types.SqlError {
let (keys, mspecs) = merge_state_specs(aggs, group_types)
let mut partials : Array[@types.DataBatch] = []
let mut segments : Int32 = 0
let mut partial_bytes = 0L
let mut peak = 0L
let mut ended = false
while !ended {
let buffered : Array[@types.DataBatch] = []
let mut bytes = 0L
let mut segment_done = false
while !segment_done {
match next() {
Some(batch) => {
buffered.push(batch)
bytes += estimate_batch_bytes(batch).to_int64()
segment_done = bytes > limit
}
None => {
ended = true
segment_done = true
}
}
}
if buffered.length() > 0 {
segments = segments + 1
let partial = aggregate_seq(
BatchSeq::of_batches(buffered),
group_exprs,
group_types,
aggs,
state=true,
)
partials.push(partial)
partial_bytes += estimate_batch_bytes(partial).to_int64()
if partials.length() > 1 && partial_bytes > limit {
let merged = aggregate_seq(
BatchSeq::of_batches(partials),
keys,
group_types,
mspecs,
state=true,
)
partials = [merged]
partial_bytes = estimate_batch_bytes(merged).to_int64()
}
}
let held = bytes + partial_bytes
if held > peak {
peak = held
}
}
let state = match partials.length() {
0 =>
aggregate_seq(
BatchSeq::of_batches([]),
keys,
group_types,
mspecs,
state=true,
)
1 => partials[0]
_ =>
aggregate_seq(
BatchSeq::of_batches(partials),
keys,
group_types,
mspecs,
state=true,
)
}
let batch = finalize_state(state, aggs, group_types)
{ batch, partitions: 0, peak_buffered: peak, largest_spilled: 0L, segments, }
}
///|
/// A GROUP BY under `--memory`, fed from a pull instead of a collected input.
/// Batches are buffered while their estimated bytes (summed in 64 bits) stay
/// within the budget; if the pull ends first, the buffered batches are aggregated
/// exactly as the within-budget path does and no spill store is needed. The first
/// overflow needs one: the buffer is then routed into a fixed number of hash
/// partitions that share one `budget / 4` pool, the rest of the pull follows, and
/// each partition is aggregated run by run. Retained input is the budget plus one
/// scan batch, whatever the total row count; group state and the grouped result
/// are not bounded by this. The result carries the counters `PulledAgg` names.
fn aggregate_pulled(
next : () -> @types.DataBatch? raise @types.SqlError,
group_exprs : Array[PhysExpr],
group_types : Array[@types.DataType],
aggs : Array[AggSpec],
opts : ExecOptions,
) -> PulledAgg raise @types.SqlError {
let limit = opts.budget.to_int64()
// the shape of the first batch decides whether the fold's state is carried as
// rows: only a sparse key set makes the segments cheaper than partitioning,
// and the decision is made before any segment is folded
let mut first : @types.DataBatch? = next()
let split = match first {
Some(batch) =>
aggs_state_mergeable(aggs) && fold_keys_are_sparse(batch, group_exprs)
None => false
}
let pull = fn() -> @types.DataBatch? raise @types.SqlError {
match first {
Some(batch) => {
first = None
Some(batch)
}
None => next()
}
}
if split {
return aggregate_pulled_segments(
pull, group_exprs, group_types, aggs, limit,
)
}
let buffered : Array[@types.DataBatch] = []
let mut bytes = 0L
let mut exhausted = false
while !exhausted && bytes <= limit {
match pull() {
Some(batch) => {
buffered.push(batch)
bytes += estimate_batch_bytes(batch).to_int64()
}
None => exhausted = true
}
}
if exhausted {
return {
batch: aggregate(buffered, group_exprs, group_types, aggs),
partitions: 0,
peak_buffered: bytes,
largest_spilled: 0L,
segments: 0,
}
}
let store = match opts.spill {
Some(store) => store
None =>
raise @types.SqlError::Execute(
"aggregate input exceeds the memory budget (\{bytes} bytes > \{opts.budget}); provide a spill store (--spill-dir) or raise the budget (--memory)",
)
}
let p = agg_pull_partitions(opts.budget)
let writers : Array[RunWriter] = []
for _ in 0.. largest {
largest = s
}
}
let peak = pool.peak.to_int64()
{
batch: grace_finish(store, writers, group_exprs, group_types, aggs),
partitions: p,
// the buffer drains into the pool, so the most held at once is whichever
// of the two phases was larger
peak_buffered: if peak > bytes {
peak
} else {
bytes
},
largest_spilled: largest,
segments: 0,
}
}