diff --git a/src/aggregate.zig b/src/aggregate.zig index 1f7c82d..3756c94 100644 --- a/src/aggregate.zig +++ b/src/aggregate.zig @@ -344,3 +344,161 @@ const AccState = struct { } } + /// Produces the final result `Value`, transferring ownership of any + /// heap payload (`.min`/`.max` on text/bytes fields) out of this + /// accumulator - a subsequent `deinit` sees `null` and frees nothing, + /// so the same payload is never freed twice. `.count_distinct` instead + /// frees its whole tracking set right here (the result is a plain + /// `.int`, nothing to hand ownership of), which is why this needs + /// `allocator` - both call sites (`GroupBy.finish`, `runGlobal`) skip a + /// separate `deinit` call after `take()` entirely, so this is the only + /// place that set is ever freed on the success path. + fn take(self: *AccState, allocator: Allocator) Value { + return switch (self.kind) { + .count => .{ .int = @intCast(self.n) }, + .sum => if (self.allInt) .{ .int = self.sumInt } else .{ .float = self.sumFloat }, + .avg => blk: { + if (self.n == 0) break :blk .{ .none = {} }; + const total: f64 = if (self.allInt) @floatFromInt(self.sumInt) else self.sumFloat; + break :blk .{ .float = total / @as(f64, @floatFromInt(self.n)) }; + }, + .min => blk: { + const v = self.minv orelse Value{ .none = {} }; + self.minv = null; + break :blk v; + }, + .max => blk: { + const v = self.maxv orelse Value{ .none = {} }; + self.maxv = null; + break :blk v; + }, + .count_distinct => blk: { + const result = Value{ .int = @intCast(self.distinct.count()) }; + var it = self.distinct.keyIterator(); + while (it.next()) |k| allocator.free(k.*); + self.distinct.deinit(allocator); + self.distinct = .{}; + break :blk result; + }, + }; + } + + fn deinit(self: *AccState, allocator: Allocator) void { + if (self.minv) |v| v.free(allocator); + if (self.maxv) |v| v.free(allocator); + var it = self.distinct.keyIterator(); + while (it.next()) |k| allocator.free(k.*); + self.distinct.deinit(allocator); + } +}; + +const Bundle = struct { + groupValue: Value, + accs: []AccState, + + fn deinit(self: *Bundle, allocator: Allocator) void { + self.groupValue.free(allocator); + for (self.accs) |*a| a.deinit(allocator); + allocator.free(self.accs); + } +}; + +/// Streaming GROUP BY accumulator: call `observe` once per candidate row +/// (already filtered by whatever `Predicate` the caller wants applied), then +/// `finish` to materialize one `GroupRow` per distinct group-key value seen, +/// in first-seen order (stable, not sorted - callers that want sorted +/// output can sort the returned slice by `.group` using `row.compare`). +pub const GroupBy = struct { + allocator: Allocator, + groupField: []const u8, + specs: []const AggSpec, + map: std.StringHashMapUnmanaged(*Bundle) = .{}, + order: std.ArrayListUnmanaged([]const u8) = .{}, + + pub fn init(allocator: Allocator, groupField: []const u8, specs: []const AggSpec) GroupBy { + return .{ .allocator = allocator, .groupField = groupField, .specs = specs }; + } + + /// Safe to call after `finish()` (which empties the internal maps) or on + /// a `GroupBy` that was never `finish()`-ed (e.g. `observe` returned an + /// error partway through a batch) - either way, exactly what is still + /// owned gets freed once. + pub fn deinit(self: *GroupBy) void { + var it = self.map.iterator(); + while (it.next()) |entry| { + self.allocator.free(entry.key_ptr.*); + entry.value_ptr.*.deinit(self.allocator); + self.allocator.destroy(entry.value_ptr.*); + } + self.map.deinit(self.allocator); + self.order.deinit(self.allocator); + self.map = .{}; + self.order = .{}; + } + + /// Rows missing `groupField` entirely are excluded from every group + /// (no "null group" bucket) - the same omission `Table.createIndex`'s + /// backfill applies to rows missing an indexed field. + pub fn observe(self: *GroupBy, row: *const Row) AggregateError!void { + const groupVal = row.get(self.groupField) orelse return; + const key = try rowmod.encodeValue(self.allocator, groupVal); + + if (self.map.get(key)) |bundle| { + self.allocator.free(key); + for (bundle.accs) |*a| try a.observe(self.allocator, row); + return; + } + + // Scoped to this block so the errdefers below only guard the + // allocation itself: once `bundle` is handed to `self.map.put` + // just after, ownership has moved there, and a later failure (the + // final observe loop) must NOT free it out from under the map - + // an errdefer spanning the whole function would do exactly that + // and leave a dangling pointer in `self.map`. + const bundle = blk: { + const b = try self.allocator.create(Bundle); + errdefer self.allocator.destroy(b); + const accs = try self.allocator.alloc(AccState, self.specs.len); + errdefer self.allocator.free(accs); + for (self.specs, 0..) |spec, i| accs[i] = .{ .kind = spec.kind, .field = spec.field }; + const dupedGroup = try groupVal.dupe(self.allocator); + b.* = .{ .groupValue = dupedGroup, .accs = accs }; + break :blk b; + }; + + try self.map.put(self.allocator, key, bundle); + try self.order.append(self.allocator, key); + for (bundle.accs) |*a| try a.observe(self.allocator, row); + } + + /// Consumes the accumulated groups into an owned `[]GroupRow` and + /// leaves `self` empty (a subsequent `deinit()` is then a no-op). + pub fn finish(self: *GroupBy) AggregateError![]GroupRow { + var results = try self.allocator.alloc(GroupRow, self.order.items.len); + var filled: usize = 0; + errdefer freeGroupRows(self.allocator, results[0..filled]); + + for (self.order.items) |key| { + const bundle = self.map.get(key).?; + var values = try self.allocator.alloc(Value, bundle.accs.len); + errdefer self.allocator.free(values); + for (bundle.accs, 0..) |*a, j| values[j] = a.take(self.allocator); + results[filled] = .{ .group = bundle.groupValue, .values = values }; + filled += 1; + } + + var it = self.map.iterator(); + while (it.next()) |entry| { + self.allocator.free(entry.key_ptr.*); + self.allocator.free(entry.value_ptr.*.accs); + self.allocator.destroy(entry.value_ptr.*); + } + self.map.deinit(self.allocator); + self.order.deinit(self.allocator); + self.map = .{}; + self.order = .{}; + + return results; + } +}; +