diff --git a/src/commands.zig b/src/commands.zig index 7e05ecd..f1e2474 100644 --- a/src/commands.zig +++ b/src/commands.zig @@ -564,7 +564,15 @@ 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); - _ = try scan_matching(ctx, db_name, coll_name, filter, 0, &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, 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: { + 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); if (sort_keys.len > 0) { try query.sort_docs(reply.arena_alloc(), matched.items, sort_keys); diff --git a/src/db.zig b/src/db.zig index 658d91e..90d053e 100644 --- a/src/db.zig +++ b/src/db.zig @@ -352,15 +352,17 @@ pub const Engine = struct { return error.IndexOptionsConflict; } - // Build entries over the existing documents, checking uniqueness as - // we go (the index is not exposed until the end, so mutating it is - // safe). Each document's batch is inserted immediately, so on any - // later failure the deferred ix.deinit frees every inserted entry - // key. Nothing is persisted on failure. + // Build entries over the existing documents (the index is not + // exposed until the end, so mutating it is safe). Entries are + // appended unsorted and ordered once at the end — inserting each + // document into a sorted array memmoves the tail every time, which + // is what made this quadratic. On any failure the deferred + // ix.deinit frees every appended key. Nothing is persisted. var doc_it = coll.docs.iterator(); while (doc_it.next()) |entry| { - _ = try ix.add_doc(self.gpa, entry.value_ptr.*, entry.key_ptr.*, true); + try ix.append_doc_entries(self.gpa, entry.value_ptr.*, entry.key_ptr.*); } + _ = try ix.finish_bulk(true); // Reserve the collection slot, then persist and publish. try coll.indexes.ensureUnusedCapacity(self.gpa, 1); @@ -619,16 +621,17 @@ pub const Engine = struct { if (ix.entries.items.len > 0) continue; // defensive var doc_it = coll_entry.value_ptr.docs.iterator(); while (doc_it.next()) |doc_entry| { - const duplicate = ix.add_doc(self.gpa, doc_entry.value_ptr.*, doc_entry.key_ptr.*, false) catch |err| switch (err) { + ix.append_doc_entries(self.gpa, doc_entry.value_ptr.*, doc_entry.key_ptr.*) catch |err| switch (err) { error.ParallelArrays => { std.debug.print("mongo-lite: WARNING: index '{s}' cannot index an existing document; entry skipped\n", .{ix.name}); continue; }, else => return err, }; - if (duplicate) { - std.debug.print("mongo-lite: WARNING: unique index '{s}' has duplicate keys in existing data; duplicates not enforced for existing documents\n", .{ix.name}); - } + } + // Tolerated, not enforced: the database must always open. + if (try ix.finish_bulk(false)) { + std.debug.print("mongo-lite: WARNING: unique index '{s}' has duplicate keys in existing data; duplicates not enforced for existing documents\n", .{ix.name}); } } } diff --git a/src/index.zig b/src/index.zig index aa51a59..ab92809 100644 --- a/src/index.zig +++ b/src/index.zig @@ -231,6 +231,53 @@ pub const Index = struct { return duplicate; } + /// Append one document's entries without maintaining sort order. Pairs + /// with `finish_bulk`, which sorts the whole array once at the end. + /// + /// This is the bulk counterpart of `add_doc`. Inserting documents one at + /// a time keeps the array sorted by memmoving the tail on every entry, + /// so building an index over n documents moves O(n²) bytes — that is the + /// whole cost of createIndex on a large collection. Appending and + /// sorting once is O(n log n) comparisons and no memmove. + pub fn append_doc_entries(self: *Index, gpa: std.mem.Allocator, doc: *const bson.Document, id: []const u8) !void { + var built = try self.build_entries(gpa, doc, id); + // Runs on success too: the append below drains the keys, leaving + // only the (now empty) ArrayList buffer to free. + defer built.deinit(gpa); + try self.entries.appendSlice(gpa, built.entries.items); + if (built.multikey) self.multikey = true; + // Key slices now belong to the index. + built.entries.items.len = 0; + } + + /// Sort the entries collected by `append_doc_entries` into the index's + /// total order and, for a unique index, look for duplicate keys — one + /// adjacent-pair scan instead of a binary search per document. + /// + /// With `enforce_unique` false a duplicate is tolerated rather than + /// rejected, matching `add_doc`; the return value reports whether that + /// happened. + pub fn finish_bulk(self: *Index, enforce_unique: bool) error{DuplicateKeyIndex}!bool { + std.mem.sort(Entry, self.entries.items, {}, entry_less); + if (!self.unique or self.entries.items.len < 2) return false; + + var duplicate = false; + for (self.entries.items[1..], 0..) |cur, prev_i| { + const prev = self.entries.items[prev_i]; + // Every key in one index has the same component count, so a + // prefix comparison over the previous key is a full key + // comparison. Sorting puts equal keys next to each other. + if (prefix_order(prev.key, cur) != .eq) continue; + // A document's own entries were deduped at build time, so equal + // keys under one id are not a conflict — same rule as + // check_unique's exclude_id. + if (std.mem.eql(u8, prev.id, cur.id)) continue; + if (enforce_unique) return error.DuplicateKeyIndex; + duplicate = true; + } + return duplicate; + } + /// Remove every entry for `id` and free its key slices. Infallible. /// One compaction pass: removing in place would memmove the tail per hit. pub fn remove_id(self: *Index, gpa: std.mem.Allocator, id: []const u8) void { @@ -287,21 +334,36 @@ pub const Index = struct { hi_incl: bool, out: *std.ArrayListUnmanaged([]const u8), ) !void { - const start = self.lower_bound_prefix(prefix); - const end = self.upper_bound_prefix(prefix); - var i = start; - while (i < end) : (i += 1) { - const v = self.entries.items[i].key[prefix.len]; - if (lo) |l| { - const o = bson.compare(v, l); - if (o == .lt or (o == .eq and !lo_incl)) continue; - } - if (hi) |h| { - const o = bson.compare(v, h); - if (o == .gt or (o == .eq and !hi_incl)) continue; - } - try out.append(gpa, self.entries.items[i].id); + // The array is sorted on this component too, so both bounds are + // binary searches. Scanning the whole equality band and filtering + // made a range on the first component touch every entry in the + // index — O(n) for what is O(log n + result). + var ext: [max_index_keys]bson.Value = undefined; + @memcpy(ext[0..prefix.len], prefix); + + // Clamp into the equality band: lo/hi constrain only the component + // at prefix.len, so the bounds alone would reach past the entries + // that share the prefix. + var start = self.lower_bound_prefix(prefix); + var end = self.upper_bound_prefix(prefix); + + if (lo) |l| { + ext[prefix.len] = l; + const key = ext[0 .. prefix.len + 1]; + // Inclusive wants the first entry not less than lo; exclusive + // wants the first one strictly greater. + const bound = if (lo_incl) self.lower_bound_prefix(key) else self.upper_bound_prefix(key); + start = @max(start, bound); } + if (hi) |h| { + ext[prefix.len] = h; + const key = ext[0 .. prefix.len + 1]; + const bound = if (hi_incl) self.upper_bound_prefix(key) else self.lower_bound_prefix(key); + end = @min(end, bound); + } + if (start >= end) return; + + for (self.entries.items[start..end]) |e| try out.append(gpa, e.id); } /// First entry whose first `prefix.len` components are not less than @@ -1131,6 +1193,91 @@ test "compound index prefix search and range on the next key" { try testing.expect(std.mem.eql(u8, out.items[0], "c1") or std.mem.eql(u8, out.items[0], "c2")); } +test "lookup_range matches a brute-force filter over random data" { + const gpa = testing.allocator; + + // lookup_range binary-searches both ends instead of scanning the + // equality band. Bounds like that are easy to get subtly wrong at the + // inclusive/exclusive edges and where the band ends, so check the whole + // result set against the definition rather than spot-checking. + var prng = std.Random.DefaultPrng.init(0x5eed_1234); + const rand = prng.random(); + + var ids: std.ArrayListUnmanaged([]u8) = .empty; + defer { + for (ids.items) |s| gpa.free(s); + ids.deinit(gpa); + } + + var ix = try simple_index(gpa, &.{ "a", "b" }, false, false); + defer ix.deinit(gpa); + + // Deliberately few distinct values so equal keys, and therefore the + // boundaries between them, come up constantly. + const n = 400; + var facts: [n]struct { a: i32, b: i32 } = undefined; + for (0..n) |i| { + const a = rand.intRangeAtMost(i32, 0, 4); + const b = rand.intRangeAtMost(i32, 0, 9); + facts[i] = .{ .a = a, .b = b }; + const id = try std.fmt.allocPrint(gpa, "id{d}", .{i}); + try ids.append(gpa, id); + const d = doc_of(&.{ + .{ .key = "_id", .value = .{ .int32 = @intCast(i) } }, + .{ .key = "a", .value = .{ .int32 = a } }, + .{ .key = "b", .value = .{ .int32 = b } }, + }); + try ix.append_doc_entries(gpa, &d, id); + } + _ = try ix.finish_bulk(false); + + var out: std.ArrayListUnmanaged([]const u8) = .empty; + defer out.deinit(gpa); + + for (0..600) |case| { + const a = rand.intRangeAtMost(i32, 0, 4); + const lo_v = rand.intRangeAtMost(i32, -1, 10); + const hi_v = rand.intRangeAtMost(i32, -1, 10); + const lo_incl = rand.boolean(); + const hi_incl = rand.boolean(); + const use_lo = rand.boolean(); + const use_hi = rand.boolean(); + + out.clearRetainingCapacity(); + try ix.lookup_range( + gpa, + &.{.{ .int32 = a }}, + if (use_lo) .{ .int32 = lo_v } else null, + lo_incl, + if (use_hi) .{ .int32 = hi_v } else null, + hi_incl, + &out, + ); + + var expected: usize = 0; + for (facts, 0..) |f, i| { + _ = i; + if (f.a != a) continue; + if (use_lo) { + if (f.b < lo_v) continue; + if (f.b == lo_v and !lo_incl) continue; + } + if (use_hi) { + if (f.b > hi_v) continue; + if (f.b == hi_v and !hi_incl) continue; + } + expected += 1; + } + testing.expectEqual(expected, out.items.len) catch |err| { + std.debug.print( + "case {d}: a={d} lo={?d} incl={} hi={?d} incl={}\n", + .{ case, a, if (use_lo) lo_v else null, lo_incl, if (use_hi) hi_v else null, hi_incl }, + ); + return err; + }; + } +} + test "id fast path guards and $in" { // plan_id may heap-copy the single value; run it through the allocator // and free. The helper asserts on whether a plan was produced.