Track max_document_id

Author: lalinsky

src/FileSegment.zig

@@ -22,6 +22,7 @@ dir: std.fs.Dir,
 info: SegmentInfo = .{},
 attributes: std.StringHashMapUnmanaged(u64) = .{},
 docs: std.AutoHashMapUnmanaged(u32, bool) = .{},
+max_doc_id: u32 = 0,
 index: std.ArrayListUnmanaged(u32) = .{},
 block_size: usize = 0,
 blocks: []const u8,

src/IndexReader.zig

@@ -58,6 +58,15 @@ pub fn getDocInfo(self: *Self, doc_id: u32) !?DocInfo {
     return null;
 }
 
+pub fn getMaxDocId(self: *Self) u32 {
+    var result: u32 = 0;
+    inline for (segment_lists) |n| {
+        const segments = @field(self, n);
+        result = @max(result, segments.value.getMaxDocId());
+    }
+    return result;
+}
+
 pub fn getVersion(self: *Self) u64 {
     if (self.memory_segments.value.getLast()) |node| {
         return node.value.info.version;
@@ -86,5 +95,8 @@ pub fn getAttributes(self: *Self, allocator: std.mem.Allocator) !std.StringHashM
         }
     }
 
+    // builtin attributes
+    try attributes.put(allocator, "max_document_id", self.getMaxDocId());
+
     return attributes;
 }

src/MemorySegment.zig

@@ -21,6 +21,7 @@ allocator: std.mem.Allocator,
 info: SegmentInfo = .{},
 attributes: std.StringHashMapUnmanaged(u64) = .{},
 docs: std.AutoHashMapUnmanaged(u32, bool) = .{},
+max_doc_id: u32 = 0,
 items: std.ArrayListUnmanaged(Item) = .{},
 frozen: bool = false,
 
@@ -81,6 +82,7 @@ pub fn build(self: *Self, changes: []const Change) !void {
     try self.docs.ensureTotalCapacity(self.allocator, num_docs);
     try self.items.ensureTotalCapacity(self.allocator, num_items);
 
+    self.max_doc_id = 0;
     var i = changes.len;
     while (i > 0) {
         i -= 1;
@@ -94,6 +96,7 @@ pub fn build(self: *Self, changes: []const Change) !void {
                     for (op.hashes, 0..) |hash, j| {
                         items[j] = .{ .hash = hash, .id = op.id };
                     }
+                    self.max_doc_id = @max(self.max_doc_id, op.id);
                 }
             },
             .delete => |op| {
@@ -131,6 +134,8 @@ pub fn merge(self: *Self, merger: *SegmentMerger(Self)) !void {
     self.docs.deinit(self.allocator);
     self.docs = merger.segment.docs.move();
 
+    self.max_doc_id = merger.segment.max_doc_id;
+
     self.items.clearRetainingCapacity();
     try self.items.ensureTotalCapacity(self.allocator, merger.estimated_size);
     while (true) {

src/filefmt.zig

@@ -435,6 +435,12 @@ pub fn readSegmentFile(dir: fs.Dir, info: SegmentInfo, segment: *FileSegment) !v
         try unpacker.readMapInto(&docs);
         segment.docs.deinit(segment.allocator);
         segment.docs = docs.unmanaged.move();
+
+        var iter = segment.docs.keyIterator();
+        segment.max_doc_id = 0;
+        while (iter.next()) |key_ptr| {
+            segment.max_doc_id = @max(segment.max_doc_id, key_ptr.*);
+        }
     }
 
     const block_size = header.block_size;

src/segment_list.zig

@@ -110,6 +110,14 @@ pub fn SegmentList(Segment: type) type {
             }
         }
 
+        pub fn getMaxDocId(self: Self) u32 {
+            var result: u32 = 0;
+            for (self.nodes.items) |node| {
+                result = @max(result, node.value.max_doc_id);
+            }
+            return result;
+        }
+
         pub fn getDocInfo(self: Self, doc_id: u32) ?DocInfo {
             var result: ?DocInfo = null;
             for (self.nodes.items) |node| {

src/segment_merger.zig

@@ -9,6 +9,7 @@ pub const MergedSegmentInfo = struct {
     info: SegmentInfo = .{},
     attributes: std.StringHashMapUnmanaged(u64) = .{},
     docs: std.AutoHashMapUnmanaged(u32, bool) = .{},
+    max_doc_id: u32 = 0,
 
     pub fn deinit(self: *MergedSegmentInfo, allocator: std.mem.Allocator) void {
         var iter = self.attributes.iterator();
@@ -115,6 +116,7 @@ pub fn SegmentMerger(comptime Segment: type) type {
             }
 
             try self.segment.docs.ensureTotalCapacity(self.allocator, total_docs);
+            self.segment.max_doc_id = 0;
             for (sources) |*source| {
                 const segment = source.reader.segment;
                 var docs_added: usize = 0;
@@ -127,6 +129,7 @@ pub fn SegmentMerger(comptime Segment: type) type {
                     if (!self.collection.hasNewerVersion(doc_id, segment.info.version)) {
                         try self.segment.docs.put(self.allocator, doc_id, doc_status);
                         docs_added += 1;
+                        self.segment.max_doc_id = @max(self.segment.max_doc_id, doc_id);
                     } else {
                         try source.skip_docs.put(self.allocator, doc_id, {});
                     }

tests/test_index_api.py

@@ -25,6 +25,7 @@ def test_get_index_not_found(client, index_name):
 def test_get_index(client, index_name, create_index, fmt):
     req = client.post(f'/{index_name}/_update', json={
         'changes': [
+            {'insert': {'id': 1, 'hashes': [101, 201, 301]}},
             {'set_attribute': {'name': 'foo', 'value': 1234}},
         ],
     })
@@ -33,9 +34,9 @@ def test_get_index(client, index_name, create_index, fmt):
     req = client.get(f'/{index_name}', headers=headers(fmt))
     assert req.status_code == 200, req.content
     if fmt == 'json':
-        expected = {'version': 1, 'segments': 1, 'attributes': {'foo': 1234}}
+        expected = {'version': 1, 'segments': 1, 'attributes': {'foo': 1234, 'max_document_id': 1}}
     else:
-        expected = {'v': 1, 's': 1, 'a': {'foo': 1234}}
+        expected = {'v': 1, 's': 1, 'a': {'foo': 1234, 'max_document_id': 1}}
     assert decode(fmt, req.content) == expected