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;
|
||||
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;
|
||||
|
||||
Reference in New Issue
Block a user