diff --git a/src/commands.zig b/src/commands.zig index cb9c67d..c41fe2c 100644 --- a/src/commands.zig +++ b/src/commands.zig @@ -577,12 +577,13 @@ fn cmd_find(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { const skip_usize = std.math.cast(usize, skip) orelse break :blk 0; break :blk skip_usize +| limit; }; - // 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); + // An index whose order already is the requested one lets the scan stop + // at the page boundary and skip sorting entirely. Otherwise a sort has + // to see every match before it can tell which ones the page contains. + var index_sorted = false; + _ = try scan_sorted(ctx, db_name, coll_name, filter, page_end, &matched, sort_keys, &index_sorted); - if (sort_keys.len > 0) { + if (sort_keys.len > 0 and !index_sorted) { // 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. @@ -616,6 +617,29 @@ fn scan_matching( limit: usize, out: ?*std.ArrayListUnmanaged(*const bson.Document), ) !usize { + return scan_sorted(ctx, db_name, coll_name, filter, limit, out, &.{}, null); +} + +/// `scan_matching` plus the option of having an index produce the ordering. +/// When `sorted` is given it reports whether the candidates came out in +/// `sort` order, in which case the caller must not sort them again — and +/// `limit` is then a genuine early stop rather than an arbitrary subset. +fn scan_sorted( + ctx: *Context, + db_name: []const u8, + coll_name: []const u8, + filter: []const bson.Pair, + limit: usize, + out: ?*std.ArrayListUnmanaged(*const bson.Document), + sort: []const query.SortKey, + sorted: ?*bool, +) !usize { + if (sorted) |flag| flag.* = false; + // Stopping early is only meaningful when the candidates come out in the + // order the caller asked for. Without a sort any subset of that size is + // a valid page; with one, the limit is honoured only if an index turns + // out to supply the ordering. + var lim: usize = if (sort.len == 0) limit else 0; const coll = ctx.engine.get_collection(db_name, coll_name) orelse return 0; const filter_doc = bson.Document{ .arena = undefined, .pairs = filter }; var n: usize = 0; @@ -637,25 +661,27 @@ fn scan_matching( if (!try query.matches(ctx.gpa, &filter_doc, doc)) continue; if (out) |list| try list.append(ctx.gpa, doc); n += 1; - if (limit != 0 and n >= limit) break; + if (lim != 0 and n >= lim) break; } return n; } // Secondary-index plan: candidates in index order, re-filtered. The // returned ids alias the docs map keys, valid under the read lock. - if (try index.plan(ctx.gpa, coll.indexes.items, filter)) |p| { + if (try index.plan(ctx.gpa, coll.indexes.items, filter, sort)) |p| { var plan = p; defer plan.deinit(ctx.gpa); var ids: std.ArrayListUnmanaged([]const u8) = .empty; defer ids.deinit(ctx.gpa); try plan.search(ctx.gpa, &ids); + if (sorted) |flag| flag.* = plan.provides_sort; + if (plan.provides_sort) lim = limit; for (ids.items) |id| { const doc = coll.docs.get(id) orelse continue; if (!try query.matches(ctx.gpa, &filter_doc, doc)) continue; if (out) |list| try list.append(ctx.gpa, doc); n += 1; - if (limit != 0 and n >= limit) break; + if (lim != 0 and n >= lim) break; } return n; } @@ -665,7 +691,7 @@ fn scan_matching( if (!try query.matches(ctx.gpa, &filter_doc, entry.value_ptr.*)) continue; if (out) |list| try list.append(ctx.gpa, entry.value_ptr.*); n += 1; - if (limit != 0 and n >= limit) break; + if (lim != 0 and n >= lim) break; } return n; } diff --git a/src/index.zig b/src/index.zig index 50d3a92..d883499 100644 --- a/src/index.zig +++ b/src/index.zig @@ -748,6 +748,11 @@ pub const Plan = struct { lo_incl: bool, hi: ?bson.Value, hi_incl: bool, + /// The candidates come out already in the requested sort order, so the + /// caller can skip sorting and stop as soon as the page is full. + provides_sort: bool = false, + /// That order is the reverse of the index's. + backward: bool = false, pub fn deinit(self: *Plan, gpa: std.mem.Allocator) void { for (self.lookup_keys.items) |k| gpa.free(k); @@ -778,6 +783,13 @@ pub const Plan = struct { try self.index.lookup_range(gpa, key, self.lo, self.lo_incl, self.hi, self.hi_incl, out); } } + // Entries come out of the array in key order, so a forward scan is + // already sorted; a backward one just reads it the other way. This + // must happen before the dedupe pass below, which sorts by id and + // would destroy the order — the conditions that set provides_sort + // are exactly the ones under which that pass is skipped. + if (self.provides_sort and self.backward) std.mem.reverse([]const u8, out.items); + const may_repeat = self.index.multikey or self.lookup_keys.items.len > 1; if (may_repeat and out.items.len > 1) { std.mem.sort([]const u8, out.items, {}, less_ids); @@ -797,7 +809,7 @@ pub const Plan = struct { /// documents: the one covering the longest leading run of equality/$in /// predicates, optionally with a range on the next key. Returns null when /// nothing usable remains — the caller scans. -pub fn plan(gpa: std.mem.Allocator, indexes: []const Index, filter: []const bson.Pair) !?Plan { +pub fn plan(gpa: std.mem.Allocator, indexes: []const Index, filter: []const bson.Pair, sort: []const query.SortKey) !?Plan { if (indexes.len == 0) return null; var clauses: std.ArrayListUnmanaged(Clause) = .empty; defer clauses.deinit(gpa); @@ -805,7 +817,7 @@ pub fn plan(gpa: std.mem.Allocator, indexes: []const Index, filter: []const bson var best: ?Plan = null; for (indexes) |*ix| { - var cand = (try evaluate_index(gpa, ix, clauses.items)) orelse continue; + var cand = (try evaluate_index(gpa, ix, clauses.items, sort)) orelse continue; if (best) |b| { if (plan_better(&cand, &b)) { best.?.deinit(gpa); @@ -822,13 +834,39 @@ pub fn plan(gpa: std.mem.Allocator, indexes: []const Index, filter: []const bson fn plan_better(a: *const Plan, b: *const Plan) bool { if (a.key_len() != b.key_len()) return a.key_len() > b.key_len(); + // Same selectivity: providing the sort saves ordering the whole result. + if (a.provides_sort != b.provides_sort) return a.provides_sort; const a_range = a.lo != null or a.hi != null; const b_range = b.lo != null or b.hi != null; if (a_range != b_range) return a_range; return a.lookup_keys.items.len < b.lookup_keys.items.len; } -fn evaluate_index(gpa: std.mem.Allocator, ix: *const Index, clauses: []const Clause) !?Plan { +/// Whether scanning `ix` in order satisfies `sort`, and if so whether that +/// means reading the array backwards. +/// +/// The sort keys must line up with the index components that follow the +/// equality-pinned prefix: those components are fixed to one value each, so +/// they do not affect the order of what remains. Directions must agree +/// uniformly — every key the same way round, or every key opposite — since +/// the array can only be read forwards or backwards. +/// +/// A multikey index is excluded: it emits a document once per indexed +/// value, so its order is not an order on documents. +fn index_provides_sort(ix: *const Index, run: usize, sort: []const query.SortKey) ?bool { + if (sort.len == 0 or ix.multikey) return null; + if (run + sort.len > ix.keys.len) return null; + + const backward = sort[0].descending != ix.keys[run].descending; + for (sort, 0..) |sk, i| { + const k = ix.keys[run + i]; + if (!std.mem.eql(u8, sk.path, k.path)) return null; + if ((sk.descending != k.descending) != backward) return null; + } + return backward; +} + +fn evaluate_index(gpa: std.mem.Allocator, ix: *const Index, clauses: []const Clause, sort: []const query.SortKey) !?Plan { const n = ix.keys.len; var infos: [max_index_keys]CompInfo = undefined; for (0..n) |i| { @@ -852,7 +890,11 @@ fn evaluate_index(gpa: std.mem.Allocator, ix: *const Index, clauses: []const Cla hi = infos[run].hi; hi_incl = infos[run].hi_incl; } - if (run == 0 and lo == null and hi == null) return null; + const sort_dir = index_provides_sort(ix, run, sort); + // With no filter to narrow anything down, a full index scan is only + // worth it when it is what produces the ordering — otherwise scanning + // the docs map directly is strictly cheaper. + if (run == 0 and lo == null and hi == null and sort_dir == null) return null; // A two-sided range on a multikey index can under-approximate: a doc // like {a: [1, 2]} satisfies {a: {$gt: 5, $lt: 25}} with the array for @@ -898,6 +940,10 @@ fn evaluate_index(gpa: std.mem.Allocator, ix: *const Index, clauses: []const Cla .hi = hi, .hi_incl = hi_incl, }; + // Several lookup keys ($in) concatenate disjoint ranges, whose + // concatenation is not ordered. + pl.provides_sort = sort_dir != null and combos == 1; + pl.backward = sort_dir orelse false; errdefer pl.deinit(gpa); if (run == 0) { const empty = try gpa.alloc(bson.Value, 0); @@ -1609,7 +1655,7 @@ test "planner picks eq run, ranges, and bails on sparse null" { .{ .key = "a", .value = .{ .int32 = 1 } }, .{ .key = "b", .value = .{ .int32 = 2 } }, }; - var p = (try plan(gpa, &.{ix}, &f)).?; + var p = (try plan(gpa, &.{ix}, &f, &.{})).?; defer p.deinit(gpa); try testing.expectEqual(@as(usize, 2), p.key_len()); try testing.expect(p.lo == null and p.hi == null); @@ -1620,7 +1666,7 @@ test "planner picks eq run, ranges, and bails on sparse null" { .{ .key = "a", .value = .{ .int32 = 1 } }, .{ .key = "b", .value = .{ .doc = &.{.{ .key = "$gt", .value = .{ .int32 = 2 } }} } }, }; - var p = (try plan(gpa, &.{ix}, &f)).?; + var p = (try plan(gpa, &.{ix}, &f, &.{})).?; defer p.deinit(gpa); try testing.expectEqual(@as(usize, 1), p.key_len()); try testing.expect(p.hi == null and p.lo != null and !p.lo_incl); @@ -1628,14 +1674,14 @@ test "planner picks eq run, ranges, and bails on sparse null" { // {a: 1} only → prefix run of 1. { const f = [_]bson.Pair{.{ .key = "a", .value = .{ .int32 = 1 } }}; - var p = (try plan(gpa, &.{ix}, &f)).?; + var p = (try plan(gpa, &.{ix}, &f, &.{})).?; defer p.deinit(gpa); try testing.expectEqual(@as(usize, 1), p.key_len()); } // Pure range on the first key → key_len 0 with a bound. { const f = [_]bson.Pair{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$gte", .value = .{ .int32 = 1 } }} } }}; - var p = (try plan(gpa, &.{ix}, &f)).?; + var p = (try plan(gpa, &.{ix}, &f, &.{})).?; defer p.deinit(gpa); try testing.expectEqual(@as(usize, 0), p.key_len()); try testing.expect(p.lo != null and p.lo_incl); @@ -1643,23 +1689,23 @@ test "planner picks eq run, ranges, and bails on sparse null" { // Unusable filter → no plan. { const f = [_]bson.Pair{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$regex", .value = .{ .string = "^x" } }} } }}; - try testing.expect((try plan(gpa, &.{ix}, &f)) == null); + try testing.expect((try plan(gpa, &.{ix}, &f, &.{})) == null); const or_f = [_]bson.Pair{.{ .key = "$or", .value = .{ .array = &.{ .{ .doc = &.{.{ .key = "a", .value = .{ .int32 = 1 } }} }, } } }}; - try testing.expect((try plan(gpa, &.{ix}, &or_f)) == null); + try testing.expect((try plan(gpa, &.{ix}, &or_f, &.{})) == null); } // Sparse index bails on a null component. { const f = [_]bson.Pair{.{ .key = "a", .value = .null }}; - try testing.expect((try plan(gpa, &.{sp}, &f)) == null); + try testing.expect((try plan(gpa, &.{sp}, &f, &.{})) == null); // Non-sparse is fine with null. - var p = (try plan(gpa, &.{ix}, &f)).?; + var p = (try plan(gpa, &.{ix}, &f, &.{})).?; defer p.deinit(gpa); try testing.expect(p.key_len() == 1); // A null inside $in bails too. const fin = [_]bson.Pair{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$in", .value = .{ .array = &.{ .{ .int32 = 1 }, .null } } }} } }}; - try testing.expect((try plan(gpa, &.{sp}, &fin)) == null); + try testing.expect((try plan(gpa, &.{sp}, &fin, &.{})) == null); } // $in cartesian product is capped. { @@ -1670,6 +1716,6 @@ test "planner picks eq run, ranges, and bails on sparse null" { .{ .key = "b", .value = .{ .doc = &.{.{ .key = "$in", .value = .{ .array = &members } }} } }, }; // 20 * 20 = 400 > 100 → fall back to a scan. - try testing.expect((try plan(gpa, &.{ix}, &f)) == null); + try testing.expect((try plan(gpa, &.{ix}, &f, &.{})) == null); } }