index/commands: let an index supply the sort order

A sort ordered every match before discarding all but one page, even when
an index already held the candidates in exactly that order. The planner
now recognizes that case and the scan streams the page straight out.

An index provides the sort when the sort keys line up with the components
after the equality-pinned prefix (those are fixed to one value each, so
they do not affect the order of what follows) and every direction agrees
uniformly -- all the same way round, or all opposite, since the array can
only be read forwards or backwards. Multikey indexes are excluded: they
emit a document once per indexed value, so their order is not an order on
documents. So is an $in, whose disjoint ranges concatenate unordered.

Entries already come out of the array in key order, so forward scans were
sorted all along; what destroyed it was the dedupe pass sorting by id.
The conditions above are exactly the ones under which that pass is
skipped, so ordered output needs only reversing for a backward scan.

evaluate_index gains a plan shape it did not have: a full index scan when
the ordering is the reason to use it. Without that, find({}).sort(...)
was unreachable -- an empty filter yields no clauses and the planner bailed
before looking at any index. It is guarded to only appear when the sort is
satisfied, since otherwise scanning the docs map directly is cheaper.

The early stop had to move into the scan, which is the only place that
knows whether the order came from an index: a limit is a valid page
boundary without a sort, or with one an index supplies, and otherwise
means nothing. Getting this wrong the other way -- limiting first and
re-scanning -- would have doubled the work for every unindexed sort.

  find({}).sort({k: 1}).limit(20) over 65,536 x 16 KB documents:
    4.0ms -> 1.0ms
  sort({_id: -1}) is unchanged at 4.3ms: nothing ordered covers _id yet.

Checked against the definition rather than by example: eleven query shapes
-- forward, backward, skip, unlimited, filtered, equality-prefixed both
directions, compound, directions the index cannot serve, and a range --
each compared through the real driver against the page of the equivalent
full sort.

Verified: 78 unit tests under ReleaseFast and ReleaseSafe, e2e
29/16/17/3/2, the crash pair, e2e6 72/72.

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,
+    // An index whose order already is the requested one lets the scan stop
-    // which has to see every match before it can tell which ones the page
+    // at the page boundary and skip sorting entirely. Otherwise a sort has
-    // contains.
+    // 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);
+    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;
 }

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);
     }
 }