feat(sstable): add reader with bloom filter and block index

src/sstable.zig

@@ -155,3 +155,257 @@ pub const Builder = struct {
         self.offset += bytes.len + 4;
     }
 };
+
+const IndexEntry = struct {
+    offset: u64,
+    blockLen: u32,
+    firstKey: []const u8,
+};
+
+const blockPrefix = 1 + 4;
+
+pub const Reader = struct {
+    allocator: Allocator,
+    buffer: []u8,
+    dataBuf: []u8,
+    index: []IndexEntry,
+    bloom: Bloom,
+    fileId: u64,
+    maxSeq: u64,
+    numEntries: u64,
+    physicalBytes: u64,
+    level: usize,
+    minKey: []const u8,
+    maxKey: []const u8,
+
+    pub fn open(allocator: Allocator, dir: std.fs.Dir, name: []const u8, fileId: u64) !Reader {
+        const file = try dir.openFile(name, .{});
+        defer file.close();
+        const buffer = try file.readToEndAlloc(allocator, std.math.maxInt(usize));
+        errdefer allocator.free(buffer);
+        const physicalBytes = buffer.len;
+
+        if (buffer.len < footerLen) return error.CorruptSstable;
+        const footerStart = buffer.len - footerLen;
+        if (readU32(buffer, footerStart) != magic) return error.CorruptSstable;
+        const indexOffset = readU64(buffer, footerStart + 8);
+        const indexLen = readU32(buffer, footerStart + 16);
+        const bloomOffset = readU64(buffer, footerStart + 20);
+        const bloomLen = readU32(buffer, footerStart + 28);
+        const numEntries = readU64(buffer, footerStart + 32);
+        const maxSeq = readU64(buffer, footerStart + 40);
+
+        const bloomData = try verifiedBlock(buffer, bloomOffset, bloomLen);
+        var bloom = try Bloom.deserialize(allocator, bloomData);
+        errdefer bloom.deinit();
+
+        const indexData = try verifiedBlock(buffer, indexOffset, indexLen);
+
+        var decoded = std.ArrayList(u8).init(allocator);
+        errdefer decoded.deinit();
+        var entries = std.ArrayList(IndexEntry).init(allocator);
+        errdefer entries.deinit();
+
+        var pos: usize = 0;
+        while (pos < indexData.len) {
+            const physOffset = readU64(indexData, pos);
+            const physLen = readU32(indexData, pos + 8);
+            const firstKeyLen = readU32(indexData, pos + 12);
+            const keyStart = pos + 16;
+            const firstKey = indexData[keyStart .. keyStart + firstKeyLen];
+
+            const physical = try verifiedBlock(buffer, physOffset, physLen);
+            if (physical.len < blockPrefix) return error.CorruptSstable;
+            const codec: Codec = @enumFromInt(physical[0]);
+            const originalLen = readU32(physical, 1);
+            const payload = physical[blockPrefix..];
+            const plainBlock = try compress.decode(allocator, codec, payload, originalLen);
+            defer allocator.free(plainBlock);
+
+            const decOffset = decoded.items.len;
+            try decoded.appendSlice(plainBlock);
+
+            try entries.append(.{ .offset = decOffset, .blockLen = @intCast(plainBlock.len), .firstKey = firstKey });
+            pos = keyStart + firstKeyLen;
+        }
+
+        const dataBuf = try decoded.toOwnedSlice();
+        errdefer allocator.free(dataBuf);
+        const indexSlice = try entries.toOwnedSlice();
+        errdefer allocator.free(indexSlice);
+
+        var minKey: []const u8 = buffer[0..0];
+        var maxKey: []const u8 = buffer[0..0];
+        if (indexSlice.len > 0) {
+            minKey = indexSlice[0].firstKey;
+            const lastEntry = indexSlice[indexSlice.len - 1];
+            const lastBlock = dataBuf[lastEntry.offset .. lastEntry.offset + lastEntry.blockLen];
+            var blockPos: usize = 0;
+            while (blockPos < lastBlock.len) {
+                const rec = decodeRecord(lastBlock, blockPos);
+                maxKey = rec.record.key;
+                blockPos = rec.next;
+            }
+        }
+
+        return .{
+            .allocator = allocator,
+            .buffer = buffer,
+            .dataBuf = dataBuf,
+            .index = indexSlice,
+            .bloom = bloom,
+            .fileId = fileId,
+            .maxSeq = maxSeq,
+            .numEntries = numEntries,
+            .physicalBytes = physicalBytes,
+            .level = 0,
+            .minKey = minKey,
+            .maxKey = maxKey,
+        };
+    }
+
+    pub fn deinit(self: *Reader) void {
+        self.allocator.free(self.index);
+        self.bloom.deinit();
+        self.allocator.free(self.dataBuf);
+        self.allocator.free(self.buffer);
+    }
+
+    pub fn mayContain(self: *const Reader, key: []const u8) bool {
+        return self.bloom.mayContain(key);
+    }
+
+    pub fn find(self: *const Reader, key: []const u8) ?Record {
+        if (!self.bloom.mayContain(key)) return null;
+        const blockIdx = self.findBlock(key) orelse return null;
+        const block = self.blockBytes(blockIdx);
+        var pos: usize = 0;
+        while (pos < block.len) {
+            const decoded = decodeRecord(block, pos);
+            switch (order(decoded.record.key, key)) {
+                .eq => return decoded.record,
+                .gt => return null,
+                .lt => pos = decoded.next,
+            }
+        }
+        return null;
+    }
+
+    pub fn findAsOf(self: *const Reader, key: []const u8, maxSeq: u64) ?Record {
+        const rec = self.find(key) orelse return null;
+        if (rec.seq > maxSeq) return null;
+        return rec;
+    }
+
+    pub fn findCachedAsOf(self: *const Reader, blockCache: *BlockCache, allocator: Allocator, key: []const u8, maxSeq: u64) !?OwnedHit {
+        if (!self.bloom.mayContain(key)) return null;
+        const blockIdx = self.findBlock(key) orelse return null;
+        const entry = self.index[blockIdx];
+        const blockKey = BlockKey{ .fileId = self.fileId, .offset = entry.offset };
+
+        const cached = blockCache.get(blockKey);
+        const block = cached orelse self.dataBuf[entry.offset .. entry.offset + entry.blockLen];
+        if (cached == null) try blockCache.put(blockKey, block);
+
+        var pos: usize = 0;
+        while (pos < block.len) {
+            const decoded = decodeRecord(block, pos);
+            switch (order(decoded.record.key, key)) {
+                .eq => {
+                    if (decoded.record.seq > maxSeq) return null;
+                    return OwnedHit{
+                        .seq = decoded.record.seq,
+                        .kind = decoded.record.kind,
+                        .value = try allocator.dupe(u8, decoded.record.value),
+                    };
+                },
+                .gt => return null,
+                .lt => pos = decoded.next,
+            }
+        }
+        return null;
+    }
+
+    pub fn cursor(self: *const Reader, start: ?[]const u8) Cursor {
+        var c = Cursor{ .reader = self, .blockIdx = 0, .pos = 0, .nextPos = 0, .cur = null };
+        if (self.index.len == 0) return c;
+        c.blockIdx = if (start) |s| (self.findBlock(s) orelse 0) else 0;
+        c.refresh();
+        if (start) |s| {
+            while (c.valid() and order(c.cur.?.key, s) == .lt) c.advance();
+        }
+        return c;
+    }
+
+    pub fn reverseCursor(self: *const Reader, allocator: Allocator, end: ?[]const u8, endInclusive: bool) !ReverseCursor {
+        var c = ReverseCursor{
+            .reader = self,
+            .allocator = allocator,
+            .positions = std.ArrayList(usize).init(allocator),
+            .blockIdx = 0,
+            .idx = 0,
+            .exhausted = false,
+            .cur = null,
+        };
+        if (self.index.len == 0) {
+            c.exhausted = true;
+            return c;
+        }
+
+        var blockIdx: usize = if (end) |e| (self.findBlock(e) orelse {
+            c.exhausted = true;
+            return c;
+        }) else self.index.len - 1;
+
+        try c.loadBlock(blockIdx);
+        if (end) |e| {
+            var chosen: ?usize = null;
+            for (c.positions.items, 0..) |p, i| {
+                const rec = decodeRecord(self.blockBytes(blockIdx), p);
+                const within = if (endInclusive) order(rec.record.key, e) != .gt else order(rec.record.key, e) == .lt;
+                if (within) chosen = i else break;
+            }
+            if (chosen) |i| {
+                c.idx = i;
+            } else {
+                if (blockIdx == 0) {
+                    c.exhausted = true;
+                    c.cur = null;
+                    return c;
+                }
+                blockIdx -= 1;
+                try c.loadBlock(blockIdx);
+                c.idx = c.positions.items.len - 1;
+            }
+        } else {
+            c.idx = c.positions.items.len - 1;
+        }
+        c.setCurrent();
+        return c;
+    }
+
+    fn findBlock(self: *const Reader, key: []const u8) ?usize {
+        var lo: usize = 0;
+        var hi: usize = self.index.len;
+        var ans: ?usize = null;
+        while (lo < hi) {
+            const mid = lo + (hi - lo) / 2;
+            if (order(self.index[mid].firstKey, key) != .gt) {
+                ans = mid;
+                lo = mid + 1;
+            } else {
+                hi = mid;
+            }
+        }
+        return ans;
+    }
+
+    fn blockBytes(self: *const Reader, idx: usize) []const u8 {
+        const entry = self.index[idx];
+        return self.dataBuf[entry.offset .. entry.offset + entry.blockLen];
+    }
+
+    pub fn blockCount(self: *const Reader) usize {
+        return self.index.len;
+    }
+};