diff --git a/src/commands.zig b/src/commands.zig index f1e2474..44793e1 100644 --- a/src/commands.zig +++ b/src/commands.zig @@ -564,18 +564,26 @@ fn cmd_find(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { var matched: std.ArrayListUnmanaged(*const bson.Document) = .empty; defer matched.deinit(ctx.gpa); - // Stop scanning once the page is filled. Only sound without a sort, - // which has to see every match before it can tell which ones the page - // contains, and the bound has to cover the skipped prefix too because - // scan_matching counts matches rather than returned documents. - const need: usize = if (sort_keys.len > 0 or limit == 0) 0 else blk: { + // Documents needed to fill the page, counting the skipped prefix; 0 + // means unbounded. + const page_end: usize = if (limit == 0) 0 else blk: { const skip_usize = std.math.cast(usize, skip) orelse break :blk 0; break :blk skip_usize +| limit; }; - _ = try scan_matching(ctx, db_name, coll_name, filter, need, &matched); + // Stop scanning once the page is filled. Only sound without a sort, + // which has to see every match before it can tell which ones the page + // contains. + _ = try scan_matching(ctx, db_name, coll_name, filter, if (sort_keys.len > 0) 0 else page_end, &matched); if (sort_keys.len > 0) { - try query.sort_docs(reply.arena_alloc(), matched.items, sort_keys); + // Selecting the page is much cheaper than ordering everything when + // the page is a small fraction of the matches. Above that fraction + // the heap's bookkeeping stops paying for itself. + if (page_end > 0 and page_end *| 4 <= matched.items.len) { + try query.sort_docs_top_k(reply.arena_alloc(), matched.items, sort_keys, page_end); + } else { + try query.sort_docs(reply.arena_alloc(), matched.items, sort_keys); + } } const rest = if (skip < matched.items.len) matched.items[skip..] else &.{}; const page = if (limit > 0 and limit < rest.len) rest[0..limit] else rest; diff --git a/src/query.zig b/src/query.zig index 5ac16ab..faea98d 100644 --- a/src/query.zig +++ b/src/query.zig @@ -552,42 +552,155 @@ pub const SortKey = struct { descending: bool, }; -/// Sort `docs` in place by `keys`. Candidate values are collected up front -/// (allocations happen before the sort), so the comparator itself is pure -/// and cannot fail — OOM during collection propagates as QueryError. +/// The first value `collect_values` would append at `path`, or null when it +/// would append none. Same traversal in the same order, same depth cutoff — +/// it just stops at the first hit instead of building a list. +/// +/// Sorting only ever reads element 0 of the collected list, so materializing +/// the rest cost one allocation per document per sort key. +pub fn first_value_at(pairs: []const bson.Pair, path: []const u8, depth: usize) ?bson.Value { + var it = std.mem.splitScalar(u8, path, '.'); + const first = it.next() orelse return null; + const rest = it.rest(); + + for (pairs) |p| { + if (!std.mem.eql(u8, p.key, first)) continue; + if (rest.len == 0) { + if (depth < 8) return p.value; + } else { + if (first_from_value(p.value, rest, depth + 1)) |v| return v; + } + } + return null; +} + +fn first_from_value(v: bson.Value, path: []const u8, depth: usize) ?bson.Value { + if (depth > 8) return null; + switch (v) { + .doc => |pairs| return first_value_at(pairs, path, depth), + .array => |items| { + var pit = std.mem.splitScalar(u8, path, '.'); + const seg = pit.next() orelse return null; + if (std.fmt.parseInt(usize, seg, 10)) |idx| { + if (idx >= items.len) return null; + const rest = pit.rest(); + if (rest.len == 0) { + if (depth < 8) return items[idx]; + return null; + } + return first_from_value(items[idx], rest, depth + 1); + } else |_| {} + for (items) |item| { + switch (item) { + .doc => if (first_value_at(item.doc, path, depth)) |x| return x, + else => {}, + } + } + return null; + }, + else => return null, + } +} + +/// A document paired with the one value per sort key the comparator reads. +const SortedDoc = struct { + doc: *const bson.Document, + /// `keys.len` values; a path that yields nothing sorts as null. + vals: []const bson.Value, +}; + +const SortCtx = struct { + keys: []const SortKey, + fn less(ctx: @This(), a: SortedDoc, b: SortedDoc) bool { + for (ctx.keys, 0..) |k, ki| { + const o = bson.compare(a.vals[ki], b.vals[ki]); + if (o != .eq) return if (k.descending) o == .gt else o == .lt; + } + return false; + } +}; + +/// Pull each document's sort-key values into one flat allocation, so the +/// comparator is pure and cannot fail. +fn decorate(arena: std.mem.Allocator, docs: []*const bson.Document, keys: []const SortKey) QueryError![]SortedDoc { + const entries = try arena.alloc(SortedDoc, docs.len); + const flat = try arena.alloc(bson.Value, docs.len * keys.len); + for (docs, 0..) |d, i| { + const vals = flat[i * keys.len ..][0..keys.len]; + for (keys, 0..) |k, ki| { + vals[ki] = first_value_at(d.pairs, k.path, 0) orelse .null; + } + entries[i] = .{ .doc = d, .vals = vals }; + } + return entries; +} + +/// Sort `docs` in place by `keys`. pub fn sort_docs(arena: std.mem.Allocator, docs: []*const bson.Document, keys: []const SortKey) QueryError!void { if (keys.len == 0 or docs.len < 2) return; + const entries = try decorate(arena, docs, keys); + std.mem.sort(SortedDoc, entries, SortCtx{ .keys = keys }, SortCtx.less); + for (entries, 0..) |e, i| docs[i] = e.doc; +} - const SortedDoc = struct { - doc: *const bson.Document, - values: [][]const bson.Value, - }; - const entries = try arena.alloc(SortedDoc, docs.len); - for (docs, 0..) |d, i| { - const values = try arena.alloc([]const bson.Value, keys.len); - for (keys, 0..) |k, ki| { - var list: std.ArrayListUnmanaged(bson.Value) = .empty; - try collect_values(arena, d.pairs, k.path, &list, 0); - values[ki] = list.items; - } - entries[i] = .{ .doc = d, .values = values }; +/// Place the `k` smallest documents by `keys`, in order, at the front of +/// `docs`. **`docs[k..]` is left in an unspecified order** — callers must +/// only read the first `k`. +/// +/// A query that sorts a whole collection to return one page pays +/// n log n comparisons to discard almost all of the result. This keeps a +/// k-element max-heap instead: one comparison against the heap root per +/// document, and only the survivors are ever ordered. +pub fn sort_docs_top_k(arena: std.mem.Allocator, docs: []*const bson.Document, keys: []const SortKey, k: usize) QueryError!void { + if (keys.len == 0 or docs.len < 2) return; + if (k == 0) return; + if (k >= docs.len) return sort_docs(arena, docs, keys); + + const entries = try decorate(arena, docs, keys); + const ctx = SortCtx{ .keys = keys }; + + // Max-heap over the first k: the root is the worst of the best-so-far. + var i = k / 2; + while (i > 0) { + i -= 1; + sift_down(entries[0..k], i, ctx); + } + // Anything better than the root replaces it; anything worse is dropped + // after a single comparison. + for (entries[k..]) |e| { + if (!ctx.less(e, entries[0])) continue; + entries[0] = e; + sift_down(entries[0..k], 0, ctx); + } + // Drain the heap back-to-front, which leaves entries[0..k] ascending. + var end = k; + while (end > 1) { + end -= 1; + const tmp = entries[0]; + entries[0] = entries[end]; + entries[end] = tmp; + sift_down(entries[0..end], 0, ctx); } - const Ctx = struct { - keys: []const SortKey, - fn lessThan(ctx: @This(), a: SortedDoc, b: SortedDoc) bool { - for (ctx.keys, 0..) |k, ki| { - const aval: bson.Value = if (a.values[ki].len > 0) a.values[ki][0] else .null; - const bval: bson.Value = if (b.values[ki].len > 0) b.values[ki][0] else .null; - const o = bson.compare(aval, bval); - if (o != .eq) return if (k.descending) o == .gt else o == .lt; - } - return false; - } - }; - std.mem.sort(SortedDoc, entries, Ctx{ .keys = keys }, Ctx.lessThan); + for (entries[0..k], 0..) |e, j| docs[j] = e.doc; +} - for (entries, 0..) |e, i| docs[i] = e.doc; +/// Restore the max-heap property at `root` over `heap`. +fn sift_down(heap: []SortedDoc, root: usize, ctx: SortCtx) void { + var parent = root; + while (true) { + const left = parent * 2 + 1; + if (left >= heap.len) return; + const right = left + 1; + // The larger child under `less`, i.e. the one that must rise. + var largest = left; + if (right < heap.len and ctx.less(heap[left], heap[right])) largest = right; + if (!ctx.less(heap[parent], heap[largest])) return; + const tmp = heap[parent]; + heap[parent] = heap[largest]; + heap[largest] = tmp; + parent = largest; + } } // --------------------------------------------------------------------------- @@ -892,6 +1005,121 @@ test "sort compares by BSON order" { try testing.expect(docs2[0] == &b); // stable-ish: order untouched by missing key } +test "first_value_at agrees with collect_values on its first element" { + // This equivalence is the entire correctness argument for the sort + // decorate pass, so exercise the traversal shapes that differ: + // dotted paths, arrays of documents (multikey), numeric element + // addressing, repeated keys, and the depth cutoff. + const gpa = testing.allocator; + + const inner = [_]bson.Pair{ + .{ .key = "x", .value = .{ .int32 = 7 } }, + .{ .key = "y", .value = .{ .string = "deep" } }, + }; + const arr_docs = [_]bson.Value{ + .{ .doc = &[_]bson.Pair{.{ .key = "v", .value = .{ .int32 = 1 } }} }, + .{ .doc = &[_]bson.Pair{.{ .key = "v", .value = .{ .int32 = 2 } }} }, + }; + const plain_arr = [_]bson.Value{ .{ .int32 = 10 }, .{ .int32 = 20 } }; + + const d = doc_of(&.{ + .{ .key = "n", .value = .{ .int32 = 5 } }, + .{ .key = "sub", .value = .{ .doc = &inner } }, + .{ .key = "items", .value = .{ .array = &arr_docs } }, + .{ .key = "nums", .value = .{ .array = &plain_arr } }, + // A repeated key: collect_values appends both, sorting takes the first. + .{ .key = "dup", .value = .{ .int32 = 100 } }, + .{ .key = "dup", .value = .{ .int32 = 200 } }, + }); + + const paths = [_][]const u8{ + "n", "sub", "sub.x", "sub.y", "sub.missing", + "items", "items.v", "items.0", "items.1.v", "items.9", + "nums", "nums.0", "nums.1", "nums.5", + "dup", "missing", "n.deeper", "", "sub.x.y", + }; + + for (paths) |path| { + var list: std.ArrayListUnmanaged(bson.Value) = .empty; + defer list.deinit(gpa); + try collect_values(gpa, d.pairs, path, &list, 0); + + const first = first_value_at(d.pairs, path, 0); + if (list.items.len == 0) { + testing.expect(first == null) catch |e| { + std.debug.print("path '{s}': collect empty but first_value_at returned a value\n", .{path}); + return e; + }; + } else { + testing.expect(first != null) catch |e| { + std.debug.print("path '{s}': collect got {d} values but first_value_at returned null\n", .{ path, list.items.len }); + return e; + }; + testing.expectEqual(std.math.Order.eq, bson.compare(list.items[0], first.?)) catch |e| { + std.debug.print("path '{s}': first value mismatch\n", .{path}); + return e; + }; + } + } +} + +test "top-k selection matches a full sort on the leading page" { + const gpa = testing.allocator; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + + // Values with many ties, so the heap's boundary behaviour is exercised. + var prng = std.Random.DefaultPrng.init(0xA11CE); + const rand = prng.random(); + const n = 300; + + var storage: [n]bson.Document = undefined; + var pairs: [n][2]bson.Pair = undefined; + for (0..n) |i| { + pairs[i] = .{ + .{ .key = "a", .value = .{ .int32 = rand.intRangeAtMost(i32, 0, 9) } }, + .{ .key = "b", .value = .{ .int32 = @intCast(i) } }, + }; + storage[i] = doc_of(pairs[i][0..]); + } + + const key_sets = [_][]const SortKey{ + &.{.{ .path = "a", .descending = false }}, + &.{.{ .path = "a", .descending = true }}, + // Second key breaks every tie, so the page is fully determined. + &.{ .{ .path = "a", .descending = false }, .{ .path = "b", .descending = false } }, + &.{ .{ .path = "a", .descending = true }, .{ .path = "b", .descending = false } }, + }; + + for (key_sets) |keys| { + for ([_]usize{ 1, 2, 20, 299, 300, 301 }) |k| { + var full: [n]*const bson.Document = undefined; + var topk: [n]*const bson.Document = undefined; + for (0..n) |i| { + full[i] = &storage[i]; + topk[i] = &storage[i]; + } + + try sort_docs(arena.allocator(), &full, keys); + try sort_docs_top_k(arena.allocator(), &topk, keys, k); + + const page = @min(k, n); + for (0..page) |i| { + // Ties make document identity ambiguous, so compare the + // sort keys rather than the pointers. + for (keys) |sk| { + const want = first_value_at(full[i].pairs, sk.path, 0) orelse bson.Value.null; + const got = first_value_at(topk[i].pairs, sk.path, 0) orelse bson.Value.null; + testing.expectEqual(std.math.Order.eq, bson.compare(want, got)) catch |e| { + std.debug.print("k={d} pos={d} key='{s}' diverged from the full sort\n", .{ k, i, sk.path }); + return e; + }; + } + } + } + } +} + test "projection inclusion and exclusion" { const d = doc_of(&.{ .{ .key = "_id", .value = .{ .int32 = 1 } },