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

query/commands: top-k sort selection and an allocation-free decorate pass

Browse Source 
sort+limit ordered the entire result set to return one page: 65,536
documents sorted to hand back 20. Two independent costs.

The decorate pass built, per document per sort key, an ArrayList of every
value at the path -- but the comparator only ever reads element 0. Added
first_value_at, which mirrors collect_values' traversal exactly (same
order, same depth cutoff) and stops at the first hit, and moved the
decorated values into one flat allocation. That equivalence is the whole
correctness argument, so it is pinned by a test covering dotted paths,
arrays of documents, numeric element addressing, repeated keys, missing
paths and the depth cutoff.

sort_docs_top_k keeps a k-element max-heap instead of ordering
everything: one comparison against the heap root per document, and only
the survivors are ever sorted. cmd_find uses it when the page is at most
a quarter of the matches, where the heap's bookkeeping still pays for
itself, and falls back to a full sort otherwise. It leaves docs[k..]
unordered, which is safe because the page is a prefix of the first k.

cmd_aggregate's $sort is deliberately untouched: a later stage can read
the whole stream, and top-k would silently corrupt the tail.

  find({}).sort({_id:-1}).limit(20) over 65,536 x 16 KB documents:
    baseline                        40.0ms
    decorate only (top-k disabled)  23.9ms
    decorate + top-k                 4.3ms

Correctness checked end to end as well: the limited page is identical to
the prefix of the equivalent full sort. The top-k test compares against a
full sort across ascending, descending and compound keys, for k of 1, 2,
20, n-1, n and n+1, over data with heavy ties; verified it fails when the
heap's child comparison is inverted.
This commit is contained in:
Alexey Shakhmatov
2026-08-02 18:33:04 +03:00
parent ff37e6c813
commit 9eecb5092c
2 changed files with 273 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
src/commands.zig
View File

@@ -564,19 +564,27 @@ 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) {
// 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;
try emit_docs(reply, db_name, coll_name, proj_pairs, page);

278
src/query.zig
View File

@@ -552,44 +552,157 @@ 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.
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;
/// 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();
const SortedDoc = struct {
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,
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 };
}
/// `keys.len` values; a path that yields nothing sorts as null.
vals: []const bson.Value,
};
const Ctx = struct {
const SortCtx = struct {
keys: []const SortKey,
fn lessThan(ctx: @This(), a: SortedDoc, b: SortedDoc) bool {
fn less(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);
const o = bson.compare(a.vals[ki], b.vals[ki]);
if (o != .eq) return if (k.descending) o == .gt else o == .lt;
}
return false;
}
};
std.mem.sort(SortedDoc, entries, Ctx{ .keys = keys }, Ctx.lessThan);
};
/// 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;
}
/// 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);
}
for (entries[0..k], 0..) |e, j| docs[j] = 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;
}
}
// ---------------------------------------------------------------------------
// Projection
// ---------------------------------------------------------------------------
@@ -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 } },