M2.5: the aggregation engine, gated by a corpus recorded from mongod #3

Merged
dev merged 7 commits from m2.5-aggregation-engine into main 2026-08-09 19:52:38 +00:00
2 changed files with 126 additions and 90 deletions
Showing only changes of commit fc611a2c64 - Show all commits

View File

@@ -2296,34 +2296,6 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
} else if (std.mem.eql(u8, stage_name, "$limit")) { } else if (std.mem.eql(u8, stage_name, "$limit")) {
const n = try stage_count(reply, stage[0].value, "$limit") orelse return; const n = try stage_count(reply, stage[0].value, "$limit") orelse return;
end = @min(end, start + n); end = @min(end, start + n);
} else if (std.mem.eql(u8, stage_name, "$project")) {
const pp = doc_arg(stage[0].value) orelse return bad_value(reply, "$project requires a document");
if (try refuse_unprojectable(reply, pp)) return;
// Applied here rather than remembered for the emit. It used to set
// a variable that only the last `$project` in a pipeline could win
// and that no later stage could see -- so a `$match` after a
// `$project` still matched on a field the projection had removed.
const arena = reply.arena_alloc();
var projected: std.ArrayListUnmanaged(*const bson.Document) = .empty;
errdefer projected.deinit(ctx.gpa);
try projected.ensureTotalCapacity(ctx.gpa, end - start);
if (in_trees) {
for (trees.items[start..end]) |d| {
projected.appendAssumeCapacity(try projected_tree(arena, d, pp));
}
} else {
for (offs.items[start..end]) |off| {
projected.appendAssumeCapacity(try projected_tree(arena, try doc_tree(arena, coll, off), pp));
}
}
offs.deinit(ctx.gpa);
offs = .empty;
trees.deinit(ctx.gpa);
trees = projected;
projected = .empty;
in_trees = true;
start = 0;
end = trees.items.len;
} else if (std.mem.eql(u8, stage_name, "$group")) { } else if (std.mem.eql(u8, stage_name, "$group")) {
const gp = doc_arg(stage[0].value) orelse return bad_value(reply, "$group requires a document"); const gp = doc_arg(stage[0].value) orelse return bad_value(reply, "$group requires a document");
const src: Stream = if (in_trees) const src: Stream = if (in_trees)
@@ -2343,7 +2315,7 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
end = trees.items.len; end = trees.items.len;
} else if (std.mem.eql(u8, stage_name, "$addFields") or std.mem.eql(u8, stage_name, "$set") or } else if (std.mem.eql(u8, stage_name, "$addFields") or std.mem.eql(u8, stage_name, "$set") or
std.mem.eql(u8, stage_name, "$unset") or std.mem.eql(u8, stage_name, "$replaceRoot") or std.mem.eql(u8, stage_name, "$unset") or std.mem.eql(u8, stage_name, "$replaceRoot") or
std.mem.eql(u8, stage_name, "$unwind")) std.mem.eql(u8, stage_name, "$unwind") or std.mem.eql(u8, stage_name, "$project"))
{ {
// The stages that rewrite a document, all one shape: read the // The stages that rewrite a document, all one shape: read the
// window, build a new list, replace the stream. The design review // window, build a new list, replace the stream. The design review
@@ -2512,55 +2484,18 @@ fn count_only_pipeline(reply: *wire.Reply, stages: []const bson.Value) !?CountSh
return .{ .filter = filter, .id_value = id_value, .accs = accs.items }; return .{ .filter = filter, .id_value = id_value, .accs = accs.items };
} }
/// Refuse a `$project` this engine cannot carry out, answering the client. /// Refuse a `$project` that mixes inclusion with exclusion, which mongod
/// Returns true when it did. /// refuses too -- so this is parity rather than a limitation of this server.
/// Judged on the *flattened* flags, so a nested spec is treated the same way a
/// dotted one is.
/// ///
/// The stage takes inclusion and exclusion flags and nothing else. A computed /// The rest of what this used to refuse -- computed fields and nested specs --
/// field (`{y: {$literal: 5}}`) needs the expression evaluator that does not /// is implemented now. They were refused because both read as falsy, which put
/// exist yet, and a nested spec (`{a: {b: 1}}`) needs a narrowing /// the whole projection into its exclusion branch and returned the entire
/// `query.project` does not do. Both used to be read as *falsy*, which put the /// document minus that field.
/// whole projection into its exclusion branch: `{$project: {y: {$literal: 5}}}` fn refuse_mixed_projection(reply: *wire.Reply, flags: []const bson.Pair) !bool {
/// returned every document with `y` removed, where mongod adds a computed `y`.
/// Measured on a live server, not inferred.
///
/// Codes read off mongod 8.3.7. It refuses the empty and the mixed forms too,
/// so those two are parity rather than a limitation of this server.
fn refuse_unprojectable(reply: *wire.Reply, pp: []const bson.Pair) !bool {
if (pp.len == 0) {
try reply.put_error(
@intFromEnum(ErrorCode.location_project_empty),
"Location51272",
"Invalid $project :: caused by :: projection specification must have at least one field",
);
return true;
}
var include: ?bool = null; var include: ?bool = null;
for (pp) |p| { for (flags) |p| {
switch (p.value) {
.bool, .int32, .int64, .double => {},
else => {
// A `$`-led document is an expression mongod evaluates and
// names in its own message; anything else is a nested spec,
// where the honest message is the one this server can stand
// behind.
const detail = if (p.value == .doc and p.value.doc.len > 0 and
p.value.doc[0].key.len > 0 and p.value.doc[0].key[0] == '$')
try std.fmt.allocPrint(
reply.arena_alloc(),
"Invalid $project :: caused by :: Unknown expression {s}",
.{p.value.doc[0].key},
)
else
try std.fmt.allocPrint(
reply.arena_alloc(),
"Invalid $project :: caused by :: field '{s}' must be an inclusion or " ++
"exclusion flag: this server projects no computed or nested fields",
.{p.key},
);
try reply.put_error(@intFromEnum(ErrorCode.location_project_unknown_expression), "Location31325", detail);
return true;
},
}
// `_id` is the one field that may be excluded from an inclusion // `_id` is the one field that may be excluded from an inclusion
// projection, so it never decides which kind this is. // projection, so it never decides which kind this is.
if (std.mem.eql(u8, p.key, "_id")) continue; if (std.mem.eql(u8, p.key, "_id")) continue;
@@ -3002,6 +2937,83 @@ fn compile_switch(reply: *wire.Reply, arena: std.mem.Allocator, spec: bson.Value
const EvalError = error{ DivideByZero, SwitchNoDefault, NonNumericArithmetic, ReplaceRootNotDocument } || const EvalError = error{ DivideByZero, SwitchNoDefault, NonNumericArithmetic, ReplaceRootNotDocument } ||
std.mem.Allocator.Error || error{ EndOfStream, Overflow, InvalidBson }; std.mem.Allocator.Error || error{ EndOfStream, Overflow, InvalidBson };
/// A `$project` taken apart: the inclusion/exclusion flags flattened to dotted
/// paths, and the computed fields as expressions.
///
/// Flattening is what lets a nested spec work without touching
/// `query.project`, which `find` shares: `{n: {x: 1}}` *is* `{"n.x": 1}`, and a
/// dotted path is something the projection already narrows correctly. The
/// nested form used to read as *falsy*, which flipped the whole projection into
/// its exclusion branch and returned every document minus that field.
const ProjectSpec = struct {
flags: []const bson.Pair,
computed: []const Expr.Field,
/// A projection that only computes -- `{$project: {b: <expr>}}` -- keeps
/// `_id` and nothing else. `query.project` cannot say that: with no
/// non-`_id` flag it reads the spec as an *exclusion* and returns the whole
/// document, so this case is built directly instead.
id_only: bool = false,
};
fn compile_project(
reply: *wire.Reply,
arena: std.mem.Allocator,
pp: []const bson.Pair,
) !?ProjectSpec {
var flags: std.ArrayListUnmanaged(bson.Pair) = .empty;
var computed: std.ArrayListUnmanaged(Expr.Field) = .empty;
if (try flatten_project(reply, arena, pp, "", &flags, &computed)) return null;
if (flags.items.len == 0 and computed.items.len == 0) {
try reply.put_error(
@intFromEnum(ErrorCode.location_project_empty),
"Location51272",
"Invalid $project :: caused by :: projection specification must have at least one field",
);
return null;
}
return ProjectSpec{
.flags = flags.items,
.computed = computed.items,
.id_only = flags.items.len == 0,
};
}
/// Walk a projection spec into flags and computed fields. Returns true when it
/// has already answered the client.
fn flatten_project(
reply: *wire.Reply,
arena: std.mem.Allocator,
pp: []const bson.Pair,
prefix: []const u8,
flags: *std.ArrayListUnmanaged(bson.Pair),
computed: *std.ArrayListUnmanaged(Expr.Field),
) CompileError!bool {
for (pp) |p| {
const key = if (prefix.len == 0)
p.key
else
try std.fmt.allocPrint(arena, "{s}.{s}", .{ prefix, p.key });
switch (p.value) {
.bool, .int32, .int64, .double => try flags.append(arena, .{ .key = key, .value = p.value }),
.doc => |d| {
// `{a: {$op: ...}}` computes; `{a: {b: 1}}` narrows.
if (d.len > 0 and d[0].key.len > 0 and d[0].key[0] == '$') {
const e = (try compile_expr(reply, arena, p.value, "a $project value")) orelse return true;
try computed.append(arena, .{ .key = key, .value = e });
} else if (try flatten_project(reply, arena, d, key, flags, computed)) {
return true;
}
},
// A bare path is a rename, which is a computed field like any other.
else => {
const e = (try compile_expr(reply, arena, p.value, "a $project value")) orelse return true;
try computed.append(arena, .{ .key = key, .value = e });
},
}
}
return false;
}
/// A compiled document-rewriting stage. /// A compiled document-rewriting stage.
const Rewrite = union(enum) { const Rewrite = union(enum) {
/// `$addFields` and `$set`, which are the same stage under two names. /// `$addFields` and `$set`, which are the same stage under two names.
@@ -3009,6 +3021,7 @@ const Rewrite = union(enum) {
unset: []const []const u8, unset: []const []const u8,
replace_root: Expr, replace_root: Expr,
unwind: UnwindSpec, unwind: UnwindSpec,
project: ProjectSpec,
}; };
fn compile_rewrite( fn compile_rewrite(
@@ -3019,6 +3032,17 @@ fn compile_rewrite(
spec: bson.Value, spec: bson.Value,
) !?Rewrite { ) !?Rewrite {
_ = ctx; _ = ctx;
if (std.mem.eql(u8, name, "$project")) {
const pp = doc_arg(spec) orelse {
try bad_value(reply, "$project requires a document");
return null;
};
const parsed = (try compile_project(reply, arena, pp)) orelse return null;
// Mixed inclusion and exclusion is judged on the *flattened* flags, so
// a nested spec is treated the same way a dotted one is.
if (try refuse_mixed_projection(reply, parsed.flags)) return null;
return Rewrite{ .project = parsed };
}
if (std.mem.eql(u8, name, "$unwind")) { if (std.mem.eql(u8, name, "$unwind")) {
return Rewrite{ .unwind = (try unwind_spec(reply, spec)) orelse return null }; return Rewrite{ .unwind = (try unwind_spec(reply, spec)) orelse return null };
} }
@@ -3081,6 +3105,23 @@ fn apply_rewrite(
}; };
try out.append(gpa, try tree_of(ec.arena, pairs)); try out.append(gpa, try tree_of(ec.arena, pairs));
}, },
.project => |spec| {
var kept: std.ArrayListUnmanaged(bson.Pair) = .empty;
if (spec.id_only) {
if (bson.get_pair(doc.pairs, "_id")) |id| try kept.append(ec.arena, .{ .key = "_id", .value = id });
} else {
query.project(ec.arena, doc, &.{ .arena = undefined, .pairs = spec.flags }, &kept) catch
return error.OutOfMemory;
}
var pairs: []const bson.Pair = kept.items;
for (spec.computed) |f| {
// Same rule as `$addFields`: an expression that resolves to
// nothing leaves the field out rather than setting it null.
const v = (try eval_expr(ec, f.value)) orelse continue;
pairs = try set_path(ec.arena, pairs, f.key, v);
}
try out.append(gpa, try tree_of(ec.arena, pairs));
},
.unwind => |spec| { .unwind => |spec| {
const found = try eval_expr(ec, .{ .path = spec.path }); const found = try eval_expr(ec, .{ .path = spec.path });
const items: []const bson.Value = switch (found orelse bson.Value.null) { const items: []const bson.Value = switch (found orelse bson.Value.null) {
@@ -5301,15 +5342,13 @@ test "$project is a stage, not a note about how to print the answer" {
}, },
.code = 31254, .code = 31254,
}, },
// A computed field and a nested spec stood here until the expression
// evaluator and the flattening landed; both work now. What is left is
// the one shape mongod refuses too.
.{ .{
.name = "a computed field", .name = "an unknown expression",
.spec = &.{.{ .key = "y", .value = .{ .doc = &.{.{ .key = "$literal", .value = .{ .int32 = 5 } }} } }}, .spec = &.{.{ .key = "y", .value = .{ .doc = &.{.{ .key = "$bogusExpr", .value = .{ .int32 = 5 } }} } }},
.code = 31325, .code = 168,
},
.{
.name = "a nested spec",
.spec = &.{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "b", .value = .{ .int32 = 1 } }} } }},
.code = 31325,
}, },
}; };
for (cases) |c| { for (cases) |c| {

View File

@@ -64,7 +64,7 @@ accumulators in:
``` ```
group-accumulators.json 19 pass 0 fail 0 skip group-accumulators.json 19 pass 0 fail 0 skip
expressions.json 27 pass 0 fail 0 skip expressions.json 27 pass 0 fail 0 skip
document-stages.json 21 pass 3 fail 0 skip document-stages.json 24 pass 0 fail 0 skip
``` ```
`group-accumulators` found its first real disagreement on the way to 18: `group-accumulators` found its first real disagreement on the way to 18:
@@ -73,11 +73,8 @@ that counted documents rather than numbers would have passed every test
anybody would think to write by hand. anybody would think to write by hand.
`expressions.json` was recorded before the evaluator was written and read `expressions.json` was recorded before the evaluator was written and read
1 pass / 26 fail against it; it is green now. `document-stages.json` was recorded at 1 pass / 26 fail against it; it is green now. `document-stages.json` was recorded at 0 pass / 24 fail and is green. The whole
0 pass / 24 fail; `$addFields`/`$set`, `$unset`, `$replaceRoot` and `$unwind` corpus is: 70 cases, every answer byte-identical to mongod 8.3.7.
took it to 21. The three left are `$project`'s computed fields, its renames and
its nested inclusions, which want the shared `query.project` and are their own
change.
What recording *that* settled: What recording *that* settled: