diff --git a/src/aggregate.zig b/src/aggregate.zig new file mode 100644 index 0000000..72f6267 --- /dev/null +++ b/src/aggregate.zig @@ -0,0 +1,163 @@ +const std = @import("std"); +const rowmod = @import("row.zig"); +const querymod = @import("query.zig"); + +const Allocator = std.mem.Allocator; +const Value = rowmod.Value; +const Row = rowmod.Row; +const Page = querymod.Page; +const Predicate = querymod.Predicate; + +/// GROUP BY / aggregate engine over already-materialized `Row` slices (the +/// same currency `join.zig` uses): callers filter and fetch rows through +/// `Table.select`/`Table.scanAll` first, then feed them here one at a time. +/// Keeping this decoupled from `query.zig`'s `Table` (rather than adding a +/// method on it) mirrors the existing split between the storage-facing +/// `Table` and the purely row-shaped `join.zig`. +pub const AggKind = enum { + /// Counts rows. `field.len == 0` counts every observed row + /// (`COUNT(*)`); otherwise only rows where `field` is present + /// (`COUNT(field)`). + count, + sum, + avg, + min, + max, + /// Counts distinct non-none values of `field` (SQL's + /// `COUNT(DISTINCT field)`). Always requires a `field` - there is no + /// bare `COUNT(DISTINCT *)` form, mirroring `sum`/`avg`/`min`/`max`. + count_distinct, +}; + +pub const AggSpec = struct { + /// Ignored for `.count` with `COUNT(*)` semantics (pass `""`). + field: []const u8, + kind: AggKind, +}; + +pub const AggregateError = error{ + /// `.sum`/`.avg` observed a field holding a non-numeric `Value` + /// (anything but `.int`/`.float`). + NonNumericField, + /// The running integer sum of a `.sum`/`.avg` accumulator would not fit + /// in `i64`. Never silently wraps or drops precision by switching to + /// float behind the caller's back. + Overflow, + /// `sortGroupsBy` was asked to order by a field that names neither the + /// GROUP BY key nor any `specLabel` of the aggregate results - the + /// ORDER BY counterpart of `expr.parse`'s "malformed input is an error, + /// not a guess" convention `filterHaving` already applies to HAVING. + UnknownOrderField, +} || Allocator.Error; + +/// One row of `GroupBy.finish()`'s result: `group` is the distinct group-key +/// value, `values` holds one result `Value` per input `AggSpec`, same order. +pub const GroupRow = struct { + group: Value, + values: []Value, + + pub fn deinit(self: *GroupRow, allocator: Allocator) void { + self.group.free(allocator); + for (self.values) |v| v.free(allocator); + allocator.free(self.values); + } +}; + +fn freeGroupRows(allocator: Allocator, rows: []GroupRow) void { + for (rows) |*r| r.deinit(allocator); + allocator.free(rows); +} + +/// Trims a `GroupBy.finish()`/`run()` result to a `query.zig` `Page` window +/// (0 = unlimited, the same convention `Table.selectPage`/`join.applyPage` +/// use) - the GROUP BY counterpart of pagination for `table_aggregate`/ +/// `remote-aggregate`. Grouping has already had to materialize every +/// distinct group to know how many there are, so - exactly like +/// `join.applyPage` - there is no cheaper way to honor `offset`/`limit` +/// than windowing the already-complete result; this always runs after +/// `run()` returns, never instead of it. Not a generic +/// `join.applyPage(GroupRow, ...)` call because `GroupRow.deinit` takes an +/// `allocator` argument (it owns heap payloads on `.group`/`.values`) +/// where `join.applyPage`'s `T.deinit()` takes none. +/// +/// Takes ownership of `rows`: frees the backing slice, `.deinit()`s every +/// row that falls outside `[offset, offset+limit)`, and returns a freshly +/// allocated slice holding only the kept rows in their original (first-seen) +/// order. Callers own the returned slice and the rows in it exactly as they +/// would have owned `rows`. +pub fn applyPageGroups(allocator: Allocator, rows: []GroupRow, page: Page) ![]GroupRow { + const total = rows.len; + const start = @min(page.offset, total); + const end = if (page.limit == 0) total else @min(start + page.limit, total); + + for (rows[0..start]) |*r| r.deinit(allocator); + for (rows[end..total]) |*r| r.deinit(allocator); + + const kept = end - start; + const result = allocator.alloc(GroupRow, kept) catch |err| { + for (rows[start..end]) |*r| r.deinit(allocator); + allocator.free(rows); + return err; + }; + @memcpy(result, rows[start..end]); + allocator.free(rows); + return result; +} + +/// The row field name a `GroupRow`/global aggregate result is exposed under +/// for HAVING and CLI output purposes: bare `AggKind` name for `field.len == +/// 0` (`COUNT(*)`), otherwise `"kind_field"`. Must stay in sync with +/// `main.zig`'s `printAggLabel`, which prints this same format directly to a +/// writer without allocating - `filterHaving`/`globalMatchesHaving` need an +/// owned string instead because it becomes a `Row` field name, so this is a +/// separate (but format-identical) helper rather than a shared one. +pub fn specLabel(allocator: Allocator, spec: AggSpec) Allocator.Error![]u8 { + if (spec.field.len == 0) return allocator.dupe(u8, @tagName(spec.kind)); + return std.fmt.allocPrint(allocator, "{s}_{s}", .{ @tagName(spec.kind), spec.field }); +} + +/// Builds the synthetic probe `Row` a HAVING predicate is evaluated against: +/// one field per `spec` (named via `specLabel`) holding that spec's result +/// `Value` for this group, plus (when `groupField.len > 0`) one field named +/// `groupField` holding the group's own key - so a HAVING clause can +/// reference an aggregate result (`count`, `sum_score`, ...) or the group key +/// itself through the exact same `expr.zig` grammar `WHERE` already uses. +/// Does not take ownership of `groupValue`/`values` - `Row.set` dupes. +fn buildHavingProbe(allocator: Allocator, groupField: []const u8, groupValue: ?Value, specs: []const AggSpec, values: []const Value) !Row { + var probe = Row.init(allocator); + errdefer probe.deinit(); + if (groupValue) |gv| try probe.set(groupField, gv); + for (specs, values) |spec, v| { + const label = try specLabel(allocator, spec); + defer allocator.free(label); + try probe.set(label, v); + } + return probe; +} + +/// Filters an already-grouped `run()` result by a HAVING predicate - the +/// GROUP BY counterpart of `Predicate.evaluate` on plain rows, applied AFTER +/// aggregation instead of before it (the same WHERE/HAVING split SQL makes). +/// Takes ownership of `groups` (same convention as `applyPageGroups`): every +/// group that fails `having` is `.deinit()`-ed here, the rest are moved +/// (first-seen order preserved) into a freshly allocated result slice. +/// Callers typically run this between `run()` and `applyPageGroups` - HAVING +/// narrows which groups exist before LIMIT/OFFSET windows them, exactly like +/// SQL evaluates `HAVING` before `LIMIT`. +pub fn filterHaving(allocator: Allocator, groups: []GroupRow, groupField: []const u8, specs: []const AggSpec, having: Predicate) AggregateError![]GroupRow { + var kept: std.ArrayListUnmanaged(GroupRow) = .{}; + var i: usize = 0; + // `groups[0..i)` is always fully accounted for by this point: each one + // either moved into `kept` (accepted) or was already `.deinit()`-ed + // in-loop (rejected) - touching it again here would double-free. Only + // `groups[i..]` (the current index, mid-evaluation when this errdefer + // fires, plus everything not yet reached) is still owned by `groups` + // and needs disposing, the same "always dispose of the input on every + // path" contract `applyPageGroups` upholds for its own OOM case. + errdefer { + for (kept.items) |*g| g.deinit(allocator); + kept.deinit(allocator); + for (groups[i..]) |*g| g.deinit(allocator); + allocator.free(groups); + } +