support bounded range when using sort key index

dgllghr · Jan 17, 2024 · 99c1b83 · 99c1b83
1 parent e91f933
commit 99c1b83
Show file tree

Hide file tree

Showing 22 changed files with 702 additions and 246 deletions.
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -18,8 +18,16 @@ jobs:
       matrix:
         os: [ "ubuntu-latest", "macos-latest", "windows-latest" ]
         sqlite:
-          - version: 3.21.0
-            year: 2017
+          - version: 3.23.1
+            year: 2018
+          - version: 3.29.0
+            year: 2019
+          - version: 3.35.5
+            year: 2021
+          - version: 3.38.5
+            year: 2022
+          - version: 3.41.2
+            year: 2023
           - version: 3.44.2
             year: 2023
     runs-on: ${{ matrix.os }}

diff --git a/build.zig b/build.zig
@@ -82,7 +82,7 @@ pub fn build(b: *Build) void {
         "bash", "test/runtest.sh",
     });
     // Build the loadable extension and then run the integration test script
-    const loadable_ext_itest = buildLoadableExt(b, target, .ReleaseFast);
+    const loadable_ext_itest = buildLoadableExt(b, target, .Debug);
     const install_loadable_ext_itest = b.addInstallArtifact(loadable_ext_itest, .{});
     const integration_test_step = b.step("itest", "Run integration tests");
     run_integration_tests.step.dependOn(&install_loadable_ext_itest.step);

diff --git a/ci/linux-x86_64-test.sh b/ci/linux-x86_64-test.sh
@@ -101,8 +101,10 @@ zig build test --summary all
 
 ## Integration tests
 if [ $sqlite32bit -ne 0 ]; then
+    echo ">> Running with x86 sqlite"
     zig build itest --summary all -Dtarget=x86-linux
 else
+    echo ">> Running with x86_64 sqlite"
     zig build itest --summary all
 fi
 

diff --git a/src/BlobManager.zig b/src/BlobManager.zig
@@ -106,8 +106,9 @@ pub fn create(self: *Self, tmp_arena: *ArenaAllocator, size: u32) !Handle {
     defer self.insert_stmt.release();
 
     try stmt.bind(.Int64, 1, @as(i64, @intCast(size)));
-    _ = try stmt.next();
-    const id = stmt.read(.Int64, false, 0);
+
+    try stmt.exec();
+    const id = self.conn.lastInsertRowid();
     errdefer self.delete(tmp_arena, id) catch |e| {
         log.err("unable to destroy segment {}: {}", .{ id, e });
     };
@@ -124,7 +125,6 @@ fn insertDml(self: *const Self, arena: *ArenaAllocator) ![]const u8 {
     return fmt.allocPrintZ(arena.allocator(),
         \\INSERT INTO "{s}" (blob)
         \\VALUES (ZEROBLOB(?))
-        \\RETURNING id
     , .{self.table_name});
 }
 

diff --git a/src/CursorRange.zig b/src/CursorRange.zig
diff --git a/src/PendingInserts.zig b/src/PendingInserts.zig
@@ -19,7 +19,10 @@ const DataType = schema_mod.ColumnType.DataType;
 const Schema = schema_mod.Schema;
 
 const TableData = @import("TableData.zig");
-const CursorRange = @import("CursorRange.zig");
+
+const index_mod = @import("index.zig");
+const RangeOp = index_mod.RangeOp;
+const SortKeyRange = index_mod.SortKeyRange;
 
 conn: Conn,
 vtab_table_name: []const u8,
@@ -256,7 +259,7 @@ pub const Cursor = struct {
 
     pub fn rewind(self: *Cursor) !void {
         try self.stmt.resetExec();
-        self.is_eof = !(try self.stmt.next());
+        try self.next();
     }
 
     pub fn readRowid(self: Cursor) !ValueRef {
@@ -307,7 +310,6 @@ fn cursorFromStartQuery(self: *const Self, arena: *ArenaAllocator) ![]const u8 {
 }
 
 /// Only 1 pending inserts cursor can be open at a time
-/// TODO use a pool of stmts to allow for multiple cursors at once
 pub fn cursorFrom(
     self: *Self,
     tmp_arena: *ArenaAllocator,
@@ -346,12 +348,12 @@ fn cursorFromKeyQuery(self: *const Self, arena: *ArenaAllocator) ![]const u8 {
     });
 }
 
-pub fn cursorPartial(self: *Self, tmp_arena: *ArenaAllocator, range: CursorRange) !Cursor {
+pub fn cursorPartial(self: *Self, tmp_arena: *ArenaAllocator, range: SortKeyRange) !Cursor {
     const query = try self.cursorPartialQuery(tmp_arena, range);
     const stmt = try self.conn.prepare(query);
 
-    for (0..range.key.valuesLen()) |idx| {
-        const value = try range.key.readValue(idx);
+    for (0..range.args.valuesLen()) |idx| {
+        const value = try range.args.readValue(idx);
         try value.bind(stmt, idx + 1);
     }
 
@@ -364,30 +366,69 @@ pub fn cursorPartial(self: *Self, tmp_arena: *ArenaAllocator, range: CursorRange
     };
 }
 
-fn cursorPartialQuery(self: *const Self, arena: *ArenaAllocator, range: CursorRange) ![]const u8 {
-    const key_len = range.key.valuesLen();
-    return fmt.allocPrintZ(arena.allocator(),
-        \\SELECT rowid, {}
-        \\FROM "{s}_pendinginserts"
-        \\WHERE ({}) {s} ({})
-        \\ORDER BY {}, rowid
-    , .{
-        sql_fmt.ColumnListLenFormatter("col_{d}"){ .len = self.columns_len },
-        self.vtab_table_name,
-        sql_fmt.ColumnListIndicesFormatter("col_{d}"){ .indices = self.sort_key[0..key_len] },
-        lastOpSymbol(range.last_op),
-        sql_fmt.ParameterListFormatter{ .len = key_len },
-        sql_fmt.ColumnListIndicesFormatter("col_{d}"){ .indices = self.sort_key },
-    });
+fn cursorPartialQuery(self: *const Self, arena: *ArenaAllocator, range: SortKeyRange) ![]const u8 {
+    const args_len = range.args.valuesLen();
+    switch (range.last_op) {
+        .between => |between| {
+            const upper_bound_indices = try arena.allocator()
+                .alloc(usize, args_len - 1);
+            for (0..(args_len - 2)) |i| {
+                upper_bound_indices[i] = i;
+            }
+            upper_bound_indices[args_len - 2] = args_len - 1;
+
+            return fmt.allocPrintZ(arena.allocator(),
+                \\WITH args AS (SELECT {})
+                \\SELECT rowid, {}
+                \\FROM "{s}_pendinginserts"
+                \\JOIN args a
+                \\WHERE ({}) {s} ({}) AND ({}) {s} ({})
+                \\ORDER BY {}, rowid
+            , .{
+                sql_fmt.ColumnListLenFormatter("? AS arg_{d}"){ .len = args_len },
+                sql_fmt.ColumnListLenFormatter("col_{d}"){ .len = self.columns_len },
+                self.vtab_table_name,
+                sql_fmt.ColumnListIndicesFormatter("col_{d}"){
+                    .indices = self.sort_key[0..(args_len - 1)],
+                },
+                if (between.lower_inc) ">=" else ">",
+                sql_fmt.ColumnListLenFormatter("a.arg_{d}"){ .len = args_len - 1 },
+                sql_fmt.ColumnListIndicesFormatter("col_{d}"){
+                    .indices = self.sort_key[0..(args_len - 1)],
+                },
+                if (between.upper_inc) "<=" else "<",
+                sql_fmt.ColumnListIndicesFormatter("a.arg_{d}"){ .indices = upper_bound_indices },
+                sql_fmt.ColumnListIndicesFormatter("col_{d}"){ .indices = self.sort_key },
+            });
+        },
+        else => |op| {
+            return fmt.allocPrintZ(arena.allocator(),
+                \\SELECT rowid, {}
+                \\FROM "{s}_pendinginserts"
+                \\WHERE ({}) {s} ({})
+                \\ORDER BY {}, rowid
+            , .{
+                sql_fmt.ColumnListLenFormatter("col_{d}"){ .len = self.columns_len },
+                self.vtab_table_name,
+                sql_fmt.ColumnListIndicesFormatter("col_{d}"){
+                    .indices = self.sort_key[0..args_len],
+                },
+                lastOpSymbol(op),
+                sql_fmt.ParameterListFormatter{ .len = args_len },
+                sql_fmt.ColumnListIndicesFormatter("col_{d}"){ .indices = self.sort_key },
+            });
+        },
+    }
 }
 
-fn lastOpSymbol(last_op: CursorRange.LastOp) []const u8 {
+fn lastOpSymbol(last_op: RangeOp) []const u8 {
     return switch (last_op) {
-        .lt => "<",
-        .le => "<=",
+        .to_exc => "<",
+        .to_inc => "<=",
         .eq => "=",
-        .gt => ">",
-        .ge => ">=",
+        .from_exc => ">",
+        .from_inc => ">=",
+        else => unreachable,
     };
 }