feat: add take and deinit to aggregate
This commit is contained in:
1 parent
dd8d35b953
commit
9b639f1592
1 file changed
+158
@@ -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;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
Reference in new issue
Block a user