commands: the stages that rewrite a document

`$addFields`, `$set`, `$unset`, `$replaceRoot` and `$unwind`. The corpus goes
0 pass / 24 fail to 21 / 3, and the three left are `$project`'s computed
fields, renames and nested inclusions, which want the `query.project` that
`find` shares and are their own change.

**The design review was wrong about what this needed, and the corpus is what
settled it.** Tier 2 was scoped as a per-stage iterator on the grounds that
`$unwind` is 1->N and "there is no way to express that in a window over the
input". That was true of the window as it stood, and stopped being true the
moment `$project` was made to rebuild the stream instead of moving bounds over
it -- a stage that rebuilds can emit as many documents as it likes, or none. So
all five share one shape: read the window, build a new list, replace the
stream. No iterator, no rewrite.

What the recording settled, and what a hand-written test would have got wrong:

  - `$addFields` whose expression resolves to nothing leaves the field out
    entirely rather than setting it to null -- so `set_path` is only reached
    when there is a value, and `eval_expr`'s absent/null distinction earns its
    keep a second time.
  - `$addFields: {"n.z": 1}` sets the nested path and keeps its siblings, and
    an existing field is replaced *where it stands*, which is what makes the
    stage "add or overwrite" rather than "append".
  - `$unwind` drops a document whose field is missing or an empty array, keeps
    one whose field is not an array *whole*, and numbers `includeArrayIndex`
    from zero. Three separate behaviours where one guess would have covered
    them all wrongly.
  - `$replaceRoot` of a missing path and of a non-document are the same error,
    40228.

`ReplaceRootNotDocument` joins `EvalError` rather than being reported at the
stage: it is a failure only a document can produce, which is the line that set
already draws.

191/191 unit tests in ReleaseFast and ReleaseSafe, 83/83 fuzz, e2e 49, e2e3 16,
e2e4 17, e2e6 72, e2e7 86, crud corpus unchanged at 201/90/196.
This commit is contained in:
A.Shakhmatov
2026-08-09 22:37:18 +03:00
parent 88ac010c3c
commit 1a5386ff00
2 changed files with 302 additions and 6 deletions

View File

@@ -116,6 +116,8 @@ pub const ErrorCode = enum(i32) {
location_switch_no_default = 40069,
location_divide_by_zero = 4848401,
location_non_numeric_arithmetic = 7157723,
location_replace_root_not_document = 40228,
location_unwind_bad_path = 28818,
location_unknown_group_operator = 15952,
location_group_needs_id = 15955,
location_accumulator_not_object = 40234,
@@ -2339,6 +2341,41 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
in_trees = true;
start = 0;
end = trees.items.len;
} 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, "$unwind"))
{
// The stages that rewrite a document, all one shape: read the
// window, build a new list, replace the stream. The design review
// expected these to need a per-stage iterator because `$unwind` is
// 1->N -- true of the window as it stood, and no longer true once a
// stage rebuilds the list rather than moving bounds over it.
const arena = reply.arena_alloc();
var built: std.ArrayListUnmanaged(*const bson.Document) = .empty;
errdefer built.deinit(ctx.gpa);
const rewrite = try compile_rewrite(ctx, reply, arena, stage_name, stage[0].value) orelse return;
var w: usize = start;
while (w < end) : (w += 1) {
const doc = if (in_trees) trees.items[w] else try doc_tree(arena, coll, offs.items[w]);
const ec: EvalCtx = .{
.arena = arena,
.coll = coll,
.src = .{ .docs = (&doc)[0..1] },
.i = 0,
};
apply_rewrite(ec, rewrite, doc, &built, ctx.gpa) catch |err| {
try report_eval_error(reply, err);
return;
};
}
offs.deinit(ctx.gpa);
offs = .empty;
trees.deinit(ctx.gpa);
trees = built;
built = .empty;
in_trees = true;
start = 0;
end = trees.items.len;
} else if (std.mem.eql(u8, stage_name, "$out") or std.mem.eql(u8, stage_name, "$merge")) {
const is_out = std.mem.eql(u8, stage_name, "$out");
// MongoDB requires either to be last, and this server needs it too:
@@ -2962,9 +2999,262 @@ fn compile_switch(reply: *wire.Reply, arena: std.mem.Allocator, spec: bson.Value
/// A failure only a document can produce, so it cannot be caught at compile
/// time. Both codes measured against mongod.
const EvalError = error{ DivideByZero, SwitchNoDefault, NonNumericArithmetic } ||
const EvalError = error{ DivideByZero, SwitchNoDefault, NonNumericArithmetic, ReplaceRootNotDocument } ||
std.mem.Allocator.Error || error{ EndOfStream, Overflow, InvalidBson };
/// A compiled document-rewriting stage.
const Rewrite = union(enum) {
/// `$addFields` and `$set`, which are the same stage under two names.
add_fields: []const Expr.Field,
unset: []const []const u8,
replace_root: Expr,
unwind: UnwindSpec,
};
fn compile_rewrite(
ctx: *Context,
reply: *wire.Reply,
arena: std.mem.Allocator,
name: []const u8,
spec: bson.Value,
) !?Rewrite {
_ = ctx;
if (std.mem.eql(u8, name, "$unwind")) {
return Rewrite{ .unwind = (try unwind_spec(reply, spec)) orelse return null };
}
if (std.mem.eql(u8, name, "$unset")) {
return Rewrite{ .unset = (try unset_paths(reply, arena, spec)) orelse return null };
}
if (std.mem.eql(u8, name, "$replaceRoot")) {
const d = doc_arg(spec) orelse {
try bad_value(reply, "$replaceRoot requires a document");
return null;
};
const new_root = bson.get_pair(d, "newRoot") orelse {
try bad_value(reply, "$replaceRoot requires newRoot");
return null;
};
return Rewrite{ .replace_root = (try compile_expr(reply, arena, new_root, "$replaceRoot's newRoot")) orelse return null };
}
const d = doc_arg(spec) orelse {
try bad_value(reply, "$addFields requires a document");
return null;
};
const fields = try arena.alloc(Expr.Field, d.len);
for (d, 0..) |p, i| {
const sub = (try compile_expr(reply, arena, p.value, "an $addFields value")) orelse return null;
fields[i] = .{ .key = p.key, .value = sub };
}
return Rewrite{ .add_fields = fields };
}
/// Put one document through a rewrite, appending whatever it produces. A stage
/// may emit no documents (`$unwind` of an empty array) or several.
fn apply_rewrite(
ec: EvalCtx,
rewrite: Rewrite,
doc: *const bson.Document,
out: *std.ArrayListUnmanaged(*const bson.Document),
gpa: std.mem.Allocator,
) EvalError!void {
switch (rewrite) {
.add_fields => |fields| {
var pairs: []const bson.Pair = doc.pairs;
for (fields) |f| {
// A value that resolves to nothing leaves the field out
// entirely, rather than setting it to 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));
},
.unset => |paths| {
var pairs: []const bson.Pair = doc.pairs;
for (paths) |path| pairs = try unset_path(ec.arena, pairs, path);
try out.append(gpa, try tree_of(ec.arena, pairs));
},
.replace_root => |e| {
const v = (try eval_expr(ec, e)) orelse return error.ReplaceRootNotDocument;
const pairs = switch (v) {
.doc => |p| p,
else => return error.ReplaceRootNotDocument,
};
try out.append(gpa, try tree_of(ec.arena, pairs));
},
.unwind => |spec| {
const found = try eval_expr(ec, .{ .path = spec.path });
const items: []const bson.Value = switch (found orelse bson.Value.null) {
.array => |a| a,
// A non-array is kept whole, and a missing or null field is
// dropped unless the caller asked to keep it.
.null => if (spec.keep_empty) &.{} else return,
else => {
try out.append(gpa, doc);
return;
},
};
if (items.len == 0) {
if (!spec.keep_empty) return;
var pairs = try unset_path(ec.arena, doc.pairs, spec.path);
if (spec.index_field) |f| pairs = try set_path(ec.arena, pairs, f, .null);
try out.append(gpa, try tree_of(ec.arena, pairs));
return;
}
for (items, 0..) |item, i| {
var pairs = try set_path(ec.arena, doc.pairs, spec.path, item);
if (spec.index_field) |f| {
pairs = try set_path(ec.arena, pairs, f, .{ .int64 = @intCast(i) });
}
try out.append(gpa, try tree_of(ec.arena, pairs));
}
},
}
}
fn tree_of(arena: std.mem.Allocator, pairs: []const bson.Pair) EvalError!*const bson.Document {
const d = try arena.create(bson.Document);
d.* = .{ .arena = undefined, .pairs = pairs };
return d;
}
/// A document with `path` set to `v`, creating the intermediate documents a
/// dotted path names. Siblings are kept and an existing field is replaced where
/// it stands, which is what makes `$addFields` "add or overwrite" rather than
/// "append".
fn set_path(
arena: std.mem.Allocator,
pairs: []const bson.Pair,
path: []const u8,
v: bson.Value,
) EvalError![]const bson.Pair {
const dot = std.mem.indexOfScalar(u8, path, '.');
const head = if (dot) |d| path[0..d] else path;
const value: bson.Value = if (dot) |d| blk: {
const existing = switch (bson.get_pair(pairs, head) orelse bson.Value.null) {
.doc => |sub| sub,
// A non-document in the way is replaced by one, as mongod does.
else => &.{},
};
break :blk .{ .doc = try set_path(arena, existing, path[d + 1 ..], v) };
} else v;
for (pairs, 0..) |p, i| {
if (!std.mem.eql(u8, p.key, head)) continue;
const out = try arena.alloc(bson.Pair, pairs.len);
@memcpy(out, pairs);
out[i] = .{ .key = head, .value = value };
return out;
}
const out = try arena.alloc(bson.Pair, pairs.len + 1);
@memcpy(out[0..pairs.len], pairs);
out[pairs.len] = .{ .key = head, .value = value };
return out;
}
/// The same document without `path`. A path naming nothing changes nothing.
fn unset_path(
arena: std.mem.Allocator,
pairs: []const bson.Pair,
path: []const u8,
) EvalError![]const bson.Pair {
const dot = std.mem.indexOfScalar(u8, path, '.');
const head = if (dot) |d| path[0..d] else path;
var out: std.ArrayListUnmanaged(bson.Pair) = .empty;
try out.ensureTotalCapacity(arena, pairs.len);
for (pairs) |p| {
if (!std.mem.eql(u8, p.key, head)) {
out.appendAssumeCapacity(p);
continue;
}
const d = dot orelse continue; // a leaf: drop it
const sub = switch (p.value) {
.doc => |x| x,
else => {
out.appendAssumeCapacity(p);
continue;
},
};
out.appendAssumeCapacity(.{
.key = p.key,
.value = .{ .doc = try unset_path(arena, sub, path[d + 1 ..]) },
});
}
return out.items;
}
/// The paths a `$unset` names: one string, or an array of them.
fn unset_paths(reply: *wire.Reply, arena: std.mem.Allocator, v: bson.Value) !?[]const []const u8 {
switch (v) {
.string => |str| {
const one = try arena.alloc([]const u8, 1);
one[0] = str;
return one;
},
.array => |items| {
const out = try arena.alloc([]const u8, items.len);
for (items, 0..) |item, i| {
out[i] = switch (item) {
.string => |str| str,
else => {
try bad_value(reply, "$unset takes field names");
return null;
},
};
}
return out;
},
else => {
try bad_value(reply, "$unset takes a field name or an array of them");
return null;
},
}
}
/// What a `$unwind` was asked to do. Both spellings land here: the bare
/// `"$path"` string and the document form.
const UnwindSpec = struct {
path: []const u8,
keep_empty: bool = false,
index_field: ?[]const u8 = null,
};
fn unwind_spec(reply: *wire.Reply, v: bson.Value) !?UnwindSpec {
const raw: bson.Value = switch (v) {
.doc => |d| bson.get_pair(d, "path") orelse {
try bad_value(reply, "$unwind requires a path");
return null;
},
else => v,
};
const str = switch (raw) {
.string => |x| x,
else => {
try bad_value(reply, "$unwind requires a string path");
return null;
},
};
if (str.len == 0 or str[0] != '$') {
const detail = try std.fmt.allocPrint(
reply.arena_alloc(),
"path option to $unwind stage should be prefixed with a '$': {s}",
.{str},
);
try reply.put_error(@intFromEnum(ErrorCode.location_unwind_bad_path), "Location28818", detail);
return null;
}
var spec: UnwindSpec = .{ .path = str[1..] };
if (v == .doc) {
if (bson.get_pair(v.doc, "preserveNullAndEmptyArrays")) |b| spec.keep_empty = query.truthy(b);
if (bson.get_pair(v.doc, "includeArrayIndex")) |f| {
spec.index_field = switch (f) {
.string => |x| x,
else => null,
};
}
}
return spec;
}
/// Turn a failure only a document could produce into the reply mongod gives.
fn report_eval_error(reply: *wire.Reply, err: EvalError) !void {
switch (err) {
@@ -2978,6 +3268,11 @@ fn report_eval_error(reply: *wire.Reply, err: EvalError) !void {
"Location40069",
"$switch could not find a matching branch for an input, and no default was specified.",
),
error.ReplaceRootNotDocument => try reply.put_error(
@intFromEnum(ErrorCode.location_replace_root_not_document),
"Location40228",
"$replaceRoot requires a document as its newRoot",
),
error.NonNumericArithmetic => try reply.put_error(
@intFromEnum(ErrorCode.location_non_numeric_arithmetic),
"Location7157723",

View File

@@ -64,7 +64,7 @@ accumulators in:
```
group-accumulators.json 19 pass 0 fail 0 skip
expressions.json 27 pass 0 fail 0 skip
document-stages.json 0 pass 24 fail 0 skip
document-stages.json 21 pass 3 fail 0 skip
```
`group-accumulators` found its first real disagreement on the way to 18:
@@ -73,10 +73,11 @@ that counted documents rather than numbers would have passed every test
anybody would think to write by hand.
`expressions.json` was recorded before the evaluator was written and read
1 pass / 26 fail against it; it is green now. `document-stages.json` is the
next tier's spec, recorded and not implemented -- `$addFields`/`$set`,
`$unset`, `$replaceRoot`, `$unwind` and `$project`'s computed fields, none of
which this server has, so it starts at zero.
1 pass / 26 fail against it; it is green now. `document-stages.json` was recorded at
0 pass / 24 fail; `$addFields`/`$set`, `$unset`, `$replaceRoot` and `$unwind`
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: