query/commands: top-k sort selection and an allocation-free decorate pass
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:
@@ -564,18 +564,26 @@ fn cmd_find(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
|
|||||||
|
|
||||||
var matched: std.ArrayListUnmanaged(*const bson.Document) = .empty;
|
var matched: std.ArrayListUnmanaged(*const bson.Document) = .empty;
|
||||||
defer matched.deinit(ctx.gpa);
|
defer matched.deinit(ctx.gpa);
|
||||||
// Stop scanning once the page is filled. Only sound without a sort,
|
// Documents needed to fill the page, counting the skipped prefix; 0
|
||||||
// which has to see every match before it can tell which ones the page
|
// means unbounded.
|
||||||
// contains, and the bound has to cover the skipped prefix too because
|
const page_end: usize = if (limit == 0) 0 else blk: {
|
||||||
// scan_matching counts matches rather than returned documents.
|
|
||||||
const need: usize = if (sort_keys.len > 0 or limit == 0) 0 else blk: {
|
|
||||||
const skip_usize = std.math.cast(usize, skip) orelse break :blk 0;
|
const skip_usize = std.math.cast(usize, skip) orelse break :blk 0;
|
||||||
break :blk skip_usize +| limit;
|
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) {
|
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 rest = if (skip < matched.items.len) matched.items[skip..] else &.{};
|
||||||
const page = if (limit > 0 and limit < rest.len) rest[0..limit] else rest;
|
const page = if (limit > 0 and limit < rest.len) rest[0..limit] else rest;
|
||||||
|
|||||||
288
src/query.zig
288
src/query.zig
@@ -552,42 +552,155 @@ pub const SortKey = struct {
|
|||||||
descending: bool,
|
descending: bool,
|
||||||
};
|
};
|
||||||
|
|
||||||
/// Sort `docs` in place by `keys`. Candidate values are collected up front
|
/// The first value `collect_values` would append at `path`, or null when it
|
||||||
/// (allocations happen before the sort), so the comparator itself is pure
|
/// would append none. Same traversal in the same order, same depth cutoff —
|
||||||
/// and cannot fail — OOM during collection propagates as QueryError.
|
/// 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 {
|
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;
|
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 {
|
/// Place the `k` smallest documents by `keys`, in order, at the front of
|
||||||
doc: *const bson.Document,
|
/// `docs`. **`docs[k..]` is left in an unspecified order** — callers must
|
||||||
values: [][]const bson.Value,
|
/// only read the first `k`.
|
||||||
};
|
///
|
||||||
const entries = try arena.alloc(SortedDoc, docs.len);
|
/// A query that sorts a whole collection to return one page pays
|
||||||
for (docs, 0..) |d, i| {
|
/// n log n comparisons to discard almost all of the result. This keeps a
|
||||||
const values = try arena.alloc([]const bson.Value, keys.len);
|
/// k-element max-heap instead: one comparison against the heap root per
|
||||||
for (keys, 0..) |k, ki| {
|
/// document, and only the survivors are ever ordered.
|
||||||
var list: std.ArrayListUnmanaged(bson.Value) = .empty;
|
pub fn sort_docs_top_k(arena: std.mem.Allocator, docs: []*const bson.Document, keys: []const SortKey, k: usize) QueryError!void {
|
||||||
try collect_values(arena, d.pairs, k.path, &list, 0);
|
if (keys.len == 0 or docs.len < 2) return;
|
||||||
values[ki] = list.items;
|
if (k == 0) return;
|
||||||
}
|
if (k >= docs.len) return sort_docs(arena, docs, keys);
|
||||||
entries[i] = .{ .doc = d, .values = values };
|
|
||||||
|
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 {
|
for (entries[0..k], 0..) |e, j| docs[j] = e.doc;
|
||||||
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..) |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
|
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" {
|
test "projection inclusion and exclusion" {
|
||||||
const d = doc_of(&.{
|
const d = doc_of(&.{
|
||||||
.{ .key = "_id", .value = .{ .int32 = 1 } },
|
.{ .key = "_id", .value = .{ .int32 = 1 } },
|
||||||
|
|||||||
Reference in New Issue
Block a user