dev/MultiforaDB
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

query/commands/wire: trim the scan and request paths

Browse Source 
Matching allocated an ArrayList per filter field per candidate document,
on the process-wide allocator, to hold what is almost always a single
value. Candidates now collect into a stack buffer that spills to the heap
only for arrays: measured 15.7 -> 12.0ms on a 65,536-document range scan.

The OOM-propagation test moves with it. Its point is that a failed
collection must surface as an error rather than an empty candidate list,
which would make $ne and $exists:false report a match -- a wrong answer
rather than a failed one. That invariant still holds on the spill path, so
the test now uses an array long enough to reach the allocator, and a new
test pins the flip side: the common single-value match now completes
correctly even when the allocator always fails, because it never calls it.

Query operators were dispatched by a chain of up to fourteen mem.eql per
value per document, with $gt/$gte/$lt/$lte re-comparing the operator name
inside the loop over candidate values. Names resolve to an enum once per
filter field. Command dispatch likewise walked a 30-entry table comparing
strings; it is a comptime StaticStringMap now.

Each request built a fresh reply arena and handed its pages straight back.
One reply per connection, reset between requests, keeps them.

countDocuments() arrives as [{$match: F}?, {$group: {_id: <literal>,
n: {$sum: 1}}}], which the general path answered by materializing every
matching document and discarding them all. It is now recognized and
answered from a counting scan: countDocuments({}) 2.3 -> 1.5ms.

The detector is deliberately conservative -- grouping by "$field", summing
a field, an unmodelled accumulator or any extra stage all fall through to
the general path, since those need the documents themselves. A unit test
pins each accept and reject, and the whole count path was checked against
the general one through the real driver, including the shapes that must
not take it.

The filtered range-scan row does not move: it is bound by walking 65,536
documents that each live in their own arena, not by the matcher. That is
Phase 4 work.

Verified: 77 unit tests under ReleaseFast and ReleaseSafe, e2e
29/16/17/3/2, the crash pair, e2e6 72/72.
This commit is contained in:
Alexey Shakhmatov
2026-08-02 19:13:29 +03:00
parent 552b916833
commit 75e412a4af
5 changed files with 327 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

159
src/commands.zig
View File

@@ -86,11 +86,18 @@ const command_table = [_]Command{
.{ .name = "listIndexes", .kind = .read, .handler = cmd_list_indexes },
};
/// Command name to its index in `command_table`, resolved at comptime so a
/// request costs one hash instead of a walk down the whole table comparing
/// strings.
const command_index = blk: {
var kvs: [command_table.len]struct { []const u8, usize } = undefined;
for (&command_table, 0..) |c, i| kvs[i] = .{ c.name, i };
break :blk std.StaticStringMap(usize).initComptime(kvs);
};
pub fn dispatch(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
const name = msg.command_name();
const cmd = for (&command_table) |*c| {
if (std.mem.eql(u8, c.name, name)) break c;
} else {
const cmd = if (command_index.get(name)) |i| &command_table[i] else {
var buf: [256]u8 = undefined;
const errmsg = try std.fmt.bufPrint(&buf, "no such command: '{s}'", .{name});
return reply.put_error(@intFromEnum(ErrorCode.command_not_found), "CommandNotFound", errmsg);
@@ -884,6 +891,41 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
else => return bad_value(reply, "pipeline must be an array"),
};
// countDocuments() reaches us as [{$match: F}?, {$group: {_id: <literal>,
// n: {$sum: 1}}}]. The general path answers that by materializing every
// matching document and then throwing them all away, so recognize the
// shape and answer it from a counting scan instead.
if (try count_only_pipeline(reply, stages)) |shape| {
const n = try scan_matching(ctx, db_name, coll_name, shape.filter, 0, null);
// No documents means no groups at all, not a group holding zero —
// same as the general path, which builds groups per document.
var docs: []const *const bson.Document = &.{};
if (n > 0) {
const arena = reply.arena_alloc();
const pairs = try arena.alloc(bson.Pair, 1 + shape.accs.len);
pairs[0] = .{ .key = "_id", .value = try bson.copy_value(arena, shape.id_value) };
for (shape.accs, 0..) |acc, i| {
// Mirrors run_group's coercion exactly: an integral sum in
// int32 range comes back as int32, otherwise a double.
const sum: f64 = @as(f64, @floatFromInt(n)) * acc.term;
pairs[1 + i] = .{
.key = acc.key,
.value = if (sum == @floor(sum) and sum <= 2_147_483_647 and sum >= -2_147_483_648)
.{ .int32 = @intFromFloat(sum) }
else
.{ .double = sum },
};
}
const doc = try arena.create(bson.Document);
doc.* = bson.Document{ .arena = undefined, .pairs = pairs };
const one = try arena.alloc(*const bson.Document, 1);
one[0] = doc;
docs = one;
}
try emit_docs(reply, db_name, coll_name, null, docs);
return reply.put_ok();
}
// The pipeline operates on a stream of documents; each stage transforms
// the current window [start, end) of `stream`, and $group replaces the
// stream entirely (so $sort/$limit after it apply to the groups).
@@ -973,6 +1015,71 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
try reply.put_ok();
}
/// A pipeline whose whole answer is the number of matching documents.
const CountShape = struct {
filter: []const bson.Pair,
/// The literal every document groups under.
id_value: bson.Value,
accs: []const Acc,
const Acc = struct { key: []const u8, term: f64 };
};
/// Recognize `[{$match: F}?, {$group: {_id: <literal>, k: {$sum: <number>}}}]`
/// — the shape a driver sends for countDocuments().
///
/// Deliberately conservative: a `_id` of `"$field"`, an accumulator over a
/// field, or any other stage needs the documents themselves, so anything
/// that is not exactly this shape returns null and takes the general path.
fn count_only_pipeline(reply: *wire.Reply, stages: []const bson.Value) !?CountShape {
if (stages.len == 0 or stages.len > 2) return null;
var filter: []const bson.Pair = &.{};
if (stages.len == 2) {
const first = switch (stages[0]) {
.doc => |p| p,
else => return null,
};
if (first.len != 1 or !std.mem.eql(u8, first[0].key, "$match")) return null;
filter = doc_arg(first[0].value) orelse return null;
}
const last = switch (stages[stages.len - 1]) {
.doc => |p| p,
else => return null,
};
if (last.len != 1 or !std.mem.eql(u8, last[0].key, "$group")) return null;
const gp = doc_arg(last[0].value) orelse return null;
const id_value = bson.get_pair(gp, "_id") orelse return null;
switch (id_value) {
// A field path or a computed id groups per document.
.string => |s| if (s.len > 0 and s[0] == '$') return null,
.doc, .array => return null,
else => {},
}
var accs: std.ArrayListUnmanaged(CountShape.Acc) = .empty;
for (gp) |p| {
if (std.mem.eql(u8, p.key, "_id")) continue;
const spec = switch (p.value) {
.doc => |d| d,
else => return null,
};
if (spec.len != 1 or !std.mem.eql(u8, spec[0].key, "$sum")) return null;
const term: f64 = switch (spec[0].value) {
.int32 => |i| @floatFromInt(i),
.int64 => |i| @floatFromInt(i),
.double => |d| d,
// $sum over a field depends on the documents.
else => return null,
};
try accs.append(reply.arena_alloc(), .{ .key = p.key, .term = term });
}
return .{ .filter = filter, .id_value = id_value, .accs = accs.items };
}
/// Minimal $group: supports `_id` of null/literal/"$field" and `$sum`
/// accumulators (constant or "$field").
fn run_group(ctx: *Context, reply: *wire.Reply, group_pairs: []const bson.Pair, docs: []const *const bson.Document) !?std.ArrayListUnmanaged(*const bson.Document) {
@@ -1877,6 +1984,52 @@ fn clear_id_list(gpa: std.mem.Allocator, list: *std.ArrayListUnmanaged([]u8)) vo
list.clearRetainingCapacity();
}
test "count_only_pipeline accepts only shapes a count can answer" {
// The fast path skips materializing documents, so mis-accepting a
// pipeline would silently return a wrong aggregate rather than a slow
// one. Pin exactly which shapes it claims.
var reply = wire.Reply.init(testing.allocator);
defer reply.deinit();
const group_count = bson.Value{ .doc = &.{.{ .key = "$group", .value = .{ .doc = &.{
.{ .key = "_id", .value = .null },
.{ .key = "n", .value = .{ .doc = &.{.{ .key = "$sum", .value = .{ .int32 = 1 } }} } },
} } }} };
const match_k = bson.Value{ .doc = &.{.{ .key = "$match", .value = .{ .doc = &.{.{ .key = "k", .value = .{ .int32 = 2 } }} } }} };
// Accepted: the two shapes countDocuments() produces.
try testing.expect(try count_only_pipeline(&reply, &.{group_count}) != null);
const with_match = try count_only_pipeline(&reply, &.{ match_k, group_count });
try testing.expect(with_match != null);
try testing.expectEqual(@as(usize, 1), with_match.?.filter.len);
try testing.expectEqualStrings("k", with_match.?.filter[0].key);
// Rejected: grouping by a field value needs the documents.
try testing.expect(try count_only_pipeline(&reply, &.{.{ .doc = &.{.{ .key = "$group", .value = .{ .doc = &.{
.{ .key = "_id", .value = .{ .string = "$k" } },
.{ .key = "n", .value = .{ .doc = &.{.{ .key = "$sum", .value = .{ .int32 = 1 } }} } },
} } }} }}) == null);
// Rejected: summing a field, not a constant.
try testing.expect(try count_only_pipeline(&reply, &.{.{ .doc = &.{.{ .key = "$group", .value = .{ .doc = &.{
.{ .key = "_id", .value = .null },
.{ .key = "t", .value = .{ .doc = &.{.{ .key = "$sum", .value = .{ .string = "$x" } }} } },
} } }} }}) == null);
// Rejected: an accumulator we do not model at all.
try testing.expect(try count_only_pipeline(&reply, &.{.{ .doc = &.{.{ .key = "$group", .value = .{ .doc = &.{
.{ .key = "_id", .value = .null },
.{ .key = "m", .value = .{ .doc = &.{.{ .key = "$max", .value = .{ .string = "$x" } }} } },
} } }} }}) == null);
// Rejected: any extra stage, since it could reshape the result.
try testing.expect(try count_only_pipeline(&reply, &.{ match_k, group_count, .{ .doc = &.{.{ .key = "$limit", .value = .{ .int32 = 1 } }} } }) == null);
// Rejected: a leading stage that is not $match.
try testing.expect(try count_only_pipeline(&reply, &.{ .{ .doc = &.{.{ .key = "$sort", .value = .{ .doc = &.{.{ .key = "k", .value = .{ .int32 = 1 } }} } }} }, group_count }) == null);
// Rejected: empty pipeline.
try testing.expect(try count_only_pipeline(&reply, &.{}) == null);
}
test "indexed queries are equivalent to scans over a mixed corpus" {
var threaded: std.Io.Threaded = .init_single_threaded;
defer threaded.deinit();