table valued fn that returns segments in a table

src/Table.zig

@@ -314,11 +314,13 @@ pub fn update(
     @panic("delete and update are not supported");
 }
 
+const BestIndexError = error{} || Allocator.Error || vtab.BestIndexError;
+
 pub fn bestIndex(
     self: *Self,
     cb_ctx: *vtab.CallbackContext,
     best_index_info: vtab.BestIndexInfo,
-) !void {
+) BestIndexError!void {
     try index.chooseBestIndex(cb_ctx.arena, self.schema.sort_key, best_index_info);
 }
 

src/functions/Segments.zig

@@ -0,0 +1,205 @@
+//! Table valued function that returns all the segments in a table
+
+const std = @import("std");
+const fmt = std.fmt;
+const mem = std.mem;
+const Allocator = std.mem.Allocator;
+const ArenaAllocator = std.heap.ArenaAllocator;
+
+const sqlite = @import("../sqlite3.zig");
+const Conn = sqlite.Conn;
+const vtab = sqlite.vtab;
+const BestIndexError = vtab.BestIndexError;
+const BestIndexInfo = vtab.BestIndexInfo;
+const Result = vtab.Result;
+
+const schema_mod = @import("../schema.zig");
+const Schema = schema_mod.Schema;
+const SchemaManager = schema_mod.Manager;
+
+const row_group = @import("../row_group.zig");
+const RowGroupIndex = row_group.Index;
+
+allocator: Allocator,
+conn: Conn,
+
+const Self = @This();
+
+pub fn connect(
+    self: *Self,
+    allocator: Allocator,
+    conn: Conn,
+    _: *vtab.CallbackContext,
+    _: []const []const u8,
+) !void {
+    self.allocator = allocator;
+    self.conn = conn;
+}
+
+pub fn disconnect(_: *Self) void {}
+
+pub fn ddl(_: *Self, allocator: Allocator) ![:0]const u8 {
+    return fmt.allocPrintZ(
+        allocator,
+        \\CREATE TABLE stanchion_segments(
+        \\  row_group_index INTEGER NOT NULL,
+        \\  column_rank INTEGER NOT NULL,
+        \\  column_name TEXT NOT NULL,
+        \\  segment_id INTEGER NULL,
+        \\  table_name TEXT HIDDEN
+        \\)
+    ,
+        .{},
+    );
+}
+
+pub fn bestIndex(
+    _: *Self,
+    _: *vtab.CallbackContext,
+    best_index_info: vtab.BestIndexInfo,
+) BestIndexError!void {
+    // Find the equality constraint on `table_name`
+    for (0..best_index_info.constraintsLen()) |idx| {
+        const c = best_index_info.constraint(idx);
+        // column index 4 is the `table_name` hidden column
+        if (c.usable() and c.columnIndex() == 4 and c.op() == .eq) {
+            c.setOmitTest(true);
+            c.includeArgInFilter(1);
+            return;
+        }
+    }
+
+    return BestIndexError.QueryImpossible;
+}
+
+pub fn open(
+    self: *Self,
+    _: *vtab.CallbackContext,
+) !Cursor {
+    return Cursor.init(self.allocator, self.conn);
+}
+
+pub const Cursor = struct {
+    lifetime_arena: ArenaAllocator,
+    conn: Conn,
+
+    table_name: []const u8,
+    schema: Schema,
+    rg_index: RowGroupIndex,
+    rg_index_cursor: RowGroupIndex.EntriesCursor,
+    rg_entry: RowGroupIndex.Entry,
+
+    rg_idx: i64,
+    col_idx: i64,
+
+    pub fn init(allocator: Allocator, conn: Conn) Cursor {
+        return .{
+            .lifetime_arena = ArenaAllocator.init(allocator),
+            .conn = conn,
+
+            .table_name = undefined,
+            .schema = undefined,
+            .rg_index = undefined,
+            .rg_index_cursor = undefined,
+            .rg_entry = undefined,
+
+            .rg_idx = 0,
+            .col_idx = -1,
+        };
+    }
+
+    pub fn deinit(self: *Cursor) void {
+        self.rg_index_cursor.deinit();
+        self.rg_index.deinit();
+        self.lifetime_arena.deinit();
+    }
+
+    pub fn begin(
+        self: *Cursor,
+        cb_ctx: *vtab.CallbackContext,
+        _: i32,
+        _: [:0]const u8,
+        filter_args: vtab.FilterArgs,
+    ) !void {
+        const table_name_ref = (try filter_args.readValue(0)).asText();
+        self.table_name = try self.lifetime_arena.allocator()
+            .dupe(u8, table_name_ref);
+
+        var schema_manager = try SchemaManager.init(cb_ctx.arena, self.conn, self.table_name);
+        defer schema_manager.deinit();
+        self.schema = try schema_manager.load(&self.lifetime_arena, cb_ctx.arena);
+
+        self.rg_index = RowGroupIndex.open(self.conn, self.table_name, &self.schema);
+        errdefer self.rg_index.deinit();
+        self.rg_index_cursor = try self.rg_index.cursor(cb_ctx.arena);
+        errdefer self.rg_index_cursor.deinit();
+
+        if (!self.rg_index_cursor.eof()) {
+            self.rg_entry = .{
+                .rowid_segment_id = 0,
+                .column_segment_ids = try self.lifetime_arena.allocator()
+                    .alloc(i64, self.schema.columns.len),
+                .record_count = 0,
+            };
+            self.rg_index_cursor.readEntry(&self.rg_entry);
+        }
+    }
+
+    fn readTableName(allocator: Allocator, filter_args: vtab.FilterArgs) ![]const u8 {
+        const table_name_ref = (try filter_args.readValue(0)).asText();
+        return allocator.dupe(u8, table_name_ref);
+    }
+
+    pub fn eof(self: *Cursor) bool {
+        return self.rg_index_cursor.eof();
+    }
+
+    pub fn next(self: *Cursor) !void {
+        self.col_idx += 1;
+
+        if (self.col_idx >= self.schema.columns.len) {
+            try self.rg_index_cursor.next();
+            if (!self.rg_index_cursor.eof()) {
+                self.rg_index_cursor.readEntry(&self.rg_entry);
+            }
+            self.rg_idx += 1;
+            self.col_idx = -1;
+        }
+    }
+
+    pub fn rowid(self: *Cursor) !i64 {
+        const cols_len: i64 = @intCast(self.schema.columns.len);
+        return (self.rg_idx * (cols_len + 1)) + self.col_idx + 2;
+    }
+
+    pub fn column(self: *Cursor, result: Result, col_idx: usize) !void {
+        switch (col_idx) {
+            // row_group_index
+            0 => result.setI64(self.rg_idx),
+            // column_rank
+            1 => result.setI64(self.col_idx),
+            // column_name
+            2 => {
+                if (self.col_idx == -1) {
+                    result.setText("rowid");
+                    return;
+                }
+                const col_name = self.schema.columns[@intCast(self.col_idx)].name;
+                result.setText(col_name);
+            },
+            // segment_id
+            3 => {
+                var seg_id: i64 = undefined;
+                if (self.col_idx == -1) {
+                    seg_id = self.rg_entry.rowid_segment_id;
+                } else {
+                    seg_id = self.rg_entry.column_segment_ids[@intCast(self.col_idx)];
+                }
+                result.setI64(seg_id);
+            },
+            // table_name
+            4 => result.setText(self.table_name),
+            else => result.setNull(),
+        }
+    }
+};

src/main.zig

@@ -2,15 +2,18 @@ const std = @import("std");
 const Allocator = std.mem.Allocator;
 const GeneralPurposeAllocator = std.heap.GeneralPurposeAllocator;
 
-const c = @import("./sqlite3/c.zig").c;
+const c = @import("sqlite3/c.zig").c;
 const sqlite = @import("sqlite3.zig");
 const sqlite_errors = sqlite.errors;
 const vtab = sqlite.vtab;
 const Conn = sqlite.Conn;
 
-const Table = @import("./Table.zig");
+const Table = @import("Table.zig");
 const StanchionVTab = vtab.VirtualTable(Table);
 
+const SegmentsFn = @import("functions/Segments.zig");
+const SegmentsTabValFn = vtab.VirtualTable(SegmentsFn);
+
 var allocator: GeneralPurposeAllocator(.{}) = undefined;
 
 pub export fn sqlite3_stanchion_init(
@@ -24,15 +27,27 @@ pub export fn sqlite3_stanchion_init(
     // with a struct containing the common data (see `ModuleContext` in `zig-sqlite`)
     allocator = GeneralPurposeAllocator(.{}){};
 
-    const res = c.sqlite3_create_module_v2(
+    var res = c.sqlite3_create_module_v2(
         db,
         "stanchion",
         &StanchionVTab.module,
         &allocator,
         null,
     );
     if (res != c.SQLITE_OK) {
-        err_msg.* = @constCast(@ptrCast("error creating module"));
+        err_msg.* = @constCast(@ptrCast("error creating stanchion module"));
+        return res;
+    }
+
+    res = c.sqlite3_create_module_v2(
+        db,
+        "stanchion_segments",
+        &SegmentsTabValFn.module,
+        &allocator,
+        null,
+    );
+    if (res != c.SQLITE_OK) {
+        err_msg.* = @constCast(@ptrCast("error creating stanchion_segments module"));
         return res;
     }