diff --git a/REDIS_COMPATIBILITY.md b/REDIS_COMPATIBILITY.md index 6e54449..e7fb5fd 100644 --- a/REDIS_COMPATIBILITY.md +++ b/REDIS_COMPATIBILITY.md @@ -176,18 +176,35 @@ Redis modules extend Redis functionality with custom data types and commands. Ze | TS.DELETERULE | No | Delete a downsampling rule | | TS.ALTER | Yes | Alter the configuration of a time series | -### Popular Redis Modules - -| Module | Supported | Notes | -| --------------- | --------- | --------------------------------------------------- | -| RedisJSON | No | JSON data type and operations | -| RediSearch | No | Full-text search and indexing | -| RedisGraph | No | Graph database functionality | -| RedisTimeSeries | No | Time series data structures | -| RedisBloom | No | Probabilistic data structures (Bloom filters, etc.) | -| RedisGears | No | Programmable data processing engine | - -### Module System +## RediSearch (FT) + +Search indexes are self-contained (documents stored inside the index via the +legacy `FT.ADD` API). Vectors use the `FLAT` (exact KNN) algorithm with +`L2`, `IP`, and `COSINE` distance metrics. Indexes are rebuilt from AOF on boot. + +| Command | Supported | Notes | +| ------------ | --------- | --------------------------------------------------------------------------------------- | +| FT.CREATE | Yes | Schema with TEXT, NUMERIC, TAG, and VECTOR (FLAT) fields | +| FT.DROPINDEX | Yes | Delete an index | +| FT._LIST | Yes | List all search indexes | +| FT.INFO | Yes | Index schema and document count | +| FT.ADD | Yes | Add a document (supports REPLACE) | +| FT.DEL | Yes | Delete a document | +| FT.GET | Yes | Fetch a document's fields | +| FT.SEARCH | Yes | Text/numeric/tag filters, BM25 scoring, SORTBY, LIMIT, RETURN, PARAMS, and `=>[KNN]` | + +## Popular Redis Modules + +| Module | Supported | Notes | +| --------------- | --------- | ------------------------------------------------------------------ | +| RedisJSON | No | JSON data type and operations | +| RediSearch | Partial | FT commands implemented natively (see RediSearch (FT) section) | +| RedisGraph | No | Graph database functionality | +| RedisTimeSeries | Partial | Time series data structures (see Time Series section) | +| RedisBloom | Partial | Probabilistic data structures (see Bloom Filter commands) | +| RedisGears | No | Programmable data processing engine | + +## Module System | Feature | Supported | Notes | | ----------------- | --------- | -------------------------------- | @@ -198,11 +215,11 @@ Redis modules extend Redis functionality with custom data types and commands. Ze ## Summary -**Total Commands**: 77 -- **Fully Implemented**: 28 commands +**Total Commands**: 85 +- **Fully Implemented**: 36 commands - **Partially Implemented**: 0 commands - **Not Implemented**: 49 commands -**Implementation Coverage**: ~36% +**Implementation Coverage**: ~42% This compatibility matrix will be updated as new commands are implemented in Zedis. diff --git a/src/aof/aof.zig b/src/aof/aof.zig index 9c56a7b..81a5a48 100644 --- a/src/aof/aof.zig +++ b/src/aof/aof.zig @@ -140,7 +140,7 @@ pub const Reader = struct { const testing = std.testing; test "aof reading test" { - const reg_init = @import("../commands/init.zig"); + const reg_init = @import("../commands/init_registry.zig"); // Read a command and test that the value is stored as expected const test_file_data = "*3\r\n$3\r\nset\r\n$1\r\nt\r\n$4\r\ntest\r\n"; @@ -169,7 +169,7 @@ test "aof reading test" { try testing.expect(std.mem.eql(u8, store.get("t").?.value.short_string.asSlice(), "test")); } test "aof writing test" { - const reg_init = @import("../commands/init.zig"); + const reg_init = @import("../commands/init_registry.zig"); // Execute a command and test that it writes it correctly const test_file_name = "aof_writing_test.aof"; diff --git a/src/benchmarks/bench_commands.zig b/src/benchmarks/bench_commands.zig index 60e4638..54d3a4a 100644 --- a/src/benchmarks/bench_commands.zig +++ b/src/benchmarks/bench_commands.zig @@ -1,7 +1,7 @@ const std = @import("std"); const Store = @import("../store.zig").Store; const CommandRegistry = @import("../commands/registry.zig").CommandRegistry; -const initRegistry = @import("../commands/init.zig").initRegistry; +const initRegistry = @import("../commands/init_registry.zig").initRegistry; const Parser = @import("../parser.zig"); const Value = Parser.Value; const bench_runner = @import("bench_runner.zig"); diff --git a/src/commands/init.zig b/src/commands/init.zig deleted file mode 100644 index 4362efc..0000000 --- a/src/commands/init.zig +++ /dev/null @@ -1,501 +0,0 @@ -const std = @import("std"); -const Allocator = std.mem.Allocator; -const CommandRegistry = @import("registry.zig").CommandRegistry; -const connection_commands = @import("connection.zig"); -const string = @import("string.zig"); -const list = @import("list.zig"); -const rdb = @import("../commands/rdb.zig"); -const pubsub = @import("../commands/pubsub.zig"); -const ts = @import("../commands/time_series.zig"); -const key = @import("../commands/keys.zig"); -const server_commands = @import("../commands/server.zig"); - -pub fn initRegistry(allocator: Allocator) !CommandRegistry { - var registry = CommandRegistry.init(allocator); - - try registry.register(.{ - .name = "PING", - .handler = .{ .default = connection_commands.ping }, - .min_args = 1, - .max_args = 2, - .description = "Ping the server", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "ECHO", - .handler = .{ .default = connection_commands.echo }, - .min_args = 2, - .max_args = 2, - .description = "Echo the given string", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "QUIT", - .handler = .{ .client_handler = connection_commands.quit }, - .min_args = 1, - .max_args = 1, - .description = "Close the connection", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "SET", - .handler = .{ .store_handler = string.set }, - .min_args = 3, - .max_args = 3, - .description = "Set string value of a key", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "GET", - .handler = .{ .store_handler = string.get }, - .min_args = 2, - .max_args = 2, - .description = "Get string value of a key", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "INCR", - .handler = .{ .store_handler = string.incr }, - .min_args = 2, - .max_args = 2, - .description = "Increment the value of a key", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "DECR", - .handler = .{ .store_handler = string.decr }, - .min_args = 2, - .max_args = 2, - .description = "Decrement the value of a key", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "HELP", - .handler = .{ .default = connection_commands.help }, - .min_args = 1, - .max_args = 1, - .description = "Show help message", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "DEL", - .handler = .{ .store_handler = string.del }, - .min_args = 2, - .max_args = null, - .description = "Delete key", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "SAVE", - .handler = .{ .client_handler = rdb.save }, - .min_args = 1, - .max_args = 1, - .description = "The SAVE commands performs a synchronous save of the dataset producing a point in time snapshot of all the data inside the Redis instance, in the form of an RDB file.", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "PUBLISH", - .handler = .{ .client_handler = pubsub.publish }, - .min_args = 3, - .max_args = 3, - .description = "Publish message", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "SUBSCRIBE", - .handler = .{ .client_handler = pubsub.subscribe }, - .min_args = 2, - .max_args = null, - .description = "Subscribe to channels", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "EXPIRE", - .handler = .{ .store_handler = string.expire }, - .min_args = 3, - .max_args = null, - .description = "Expire key", - // TODO: convert to expireat - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "EXPIREAT", - .handler = .{ .store_handler = string.expireAt }, - .min_args = 3, - .max_args = null, - .description = "Expire key", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "AUTH", - .handler = .{ .client_handler = connection_commands.auth }, - .min_args = 2, - .max_args = 2, - .description = "Authenticate to the server", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "CONFIG", - .handler = .{ .client_handler = connection_commands.config }, - .min_args = 2, - .max_args = null, - .description = "Inspect and update server configuration", - .write_to_aof = false, - }); - - // List commands: LPUSH, RPUSH, LPOP, RPOP, LLEN, LRANGE - - try registry.register(.{ - .name = "LPUSH", - .handler = .{ .store_handler = list.lpush }, - .min_args = 3, - .max_args = null, - .description = "Prepend one or multiple values to a list", - // TODO: test - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "RPUSH", - .handler = .{ .store_handler = list.rpush }, - .min_args = 3, - .max_args = null, - .description = "Append one or multiple values to a list", - // TODO: test - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "LPOP", - .handler = .{ .store_handler = list.lpop }, - .min_args = 2, - .max_args = 3, - .description = "Remove and return the first element of a list", - // TODO: test - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "RPOP", - .handler = .{ .store_handler = list.rpop }, - .min_args = 2, - .max_args = 3, - .description = "Remove and return the last element of a list", - // TODO: test - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "LLEN", - .handler = .{ .store_handler = list.llen }, - .min_args = 2, - .max_args = 2, - .description = "Get the length of a list", - // TODO: test - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "LINDEX", - .handler = .{ .store_handler = list.lindex }, - .min_args = 3, - .max_args = 3, - .description = "Get an element from a list by its index", - // TODO: test - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "LSET", - .handler = .{ .store_handler = list.lset }, - .min_args = 4, - .max_args = 4, - .description = "Set the value of an element in a list by its index", - // TODO: test - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "LRANGE", - .handler = .{ .store_handler = list.lrange }, - .min_args = 4, - .max_args = 4, - .description = "Get a range of elements from a list", - // TODO: test - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "APPEND", - .handler = .{ .store_handler = string.append }, - .min_args = 3, - .max_args = 3, - .description = "Append a value to a key", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "STRLEN", - .handler = .{ .store_handler = string.strlen }, - .min_args = 2, - .max_args = 2, - .description = "Get the length of the value stored in a key", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "GETSET", - .handler = .{ .store_handler = string.getset }, - .min_args = 3, - .max_args = 3, - .description = "Set a key and return its old value", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "MGET", - .handler = .{ .store_handler = string.mget }, - .min_args = 2, - .max_args = null, - .description = "Get the values of multiple keys", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "MSET", - .handler = .{ .store_handler = string.mset }, - .min_args = 3, - .max_args = null, - .description = "Set multiple key-value pairs", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "SETEX", - .handler = .{ .store_handler = string.setex }, - .min_args = 4, - .max_args = 4, - .description = "Set a key with expiration time", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "SETNX", - .handler = .{ .store_handler = string.setnx }, - .min_args = 3, - .max_args = 3, - .description = "Set a key only if it doesn't exist", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "INCRBY", - .handler = .{ .store_handler = string.incrby }, - .min_args = 3, - .max_args = 3, - .description = "Increment a key by a specific amount", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "DECRBY", - .handler = .{ .store_handler = string.decrby }, - .min_args = 3, - .max_args = 3, - .description = "Decrement a key by a specific amount", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "INCRBYFLOAT", - .handler = .{ .store_handler = string.incrbyfloat }, - .min_args = 3, - .max_args = 3, - .description = "Increment a key by a floating point number", - .write_to_aof = true, - }); - - // Key commands - - try registry.register(.{ - .name = "KEYS", - .handler = .{ .store_handler = key.keys }, - .min_args = 2, - .max_args = 2, - .description = "Find all keys matching a pattern", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "EXISTS", - .handler = .{ .store_handler = key.exists }, - .min_args = 2, - .max_args = null, - .description = "Check if key exists", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "TTL", - .handler = .{ .store_handler = key.ttl }, - .min_args = 2, - .max_args = 2, - .description = "Get remaining time to live of a key", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "PERSIST", - .handler = .{ .store_handler = key.persist }, - .min_args = 2, - .max_args = 2, - .description = "Remove expiration from a key", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "TYPE", - .handler = .{ .store_handler = key.typeCmd }, - .min_args = 2, - .max_args = 2, - .description = "Get the data type of a key", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "RENAME", - .handler = .{ .store_handler = key.rename }, - .min_args = 3, - .max_args = 3, - .description = "Rename a key", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "RANDOMKEY", - .handler = .{ .store_handler = key.randomkey }, - .min_args = 1, - .max_args = 1, - .description = "Return a random key", - .write_to_aof = false, - }); - - // Time series commands - try registry.register(.{ - .name = "TS.CREATE", - .handler = .{ .store_handler = ts.ts_create }, - .min_args = 2, - .max_args = null, - .description = "Create a new time series", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "TS.ADD", - .handler = .{ .store_handler = ts.ts_add }, - .min_args = 4, - .max_args = null, - .description = "Add a new sample to a time series", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "TS.GET", - .handler = .{ .store_handler = ts.ts_get }, - .min_args = 2, - .max_args = 2, - .description = "Get the last sample from a time series", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "TS.INCRBY", - .handler = .{ .store_handler = ts.ts_incrby }, - .min_args = 4, - .max_args = 4, - .description = "Increment the last value and add as a new sample", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "TS.DECRBY", - .handler = .{ .store_handler = ts.ts_decrby }, - .min_args = 4, - .max_args = 4, - .description = "Decrement the last value and add as a new sample", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "TS.ALTER", - .handler = .{ .store_handler = ts.ts_alter }, - .min_args = 2, - .max_args = null, - .description = "Alter time series properties", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "TS.ALTER", - .handler = .{ .store_handler = ts.ts_alter }, - .min_args = 2, - .max_args = null, - .description = "Alter time series properties", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "TS.RANGE", - .handler = .{ .store_handler = ts.ts_range }, - .min_args = 4, - .max_args = null, - .description = "Query a range of samples from a time series", - .write_to_aof = false, - }); - - // Server commands - - try registry.register(.{ - .name = "DBSIZE", - .handler = .{ .store_handler = server_commands.db_size }, - .min_args = 1, - .max_args = 1, - .description = "Get database size", - .write_to_aof = false, - }); - - try registry.register(.{ - .name = "FLUSHDB", - .handler = .{ .client_handler = server_commands.flush_db }, - .min_args = 1, - .max_args = 1, - .description = "Flush the store", - .write_to_aof = true, - }); - - try registry.register(.{ - .name = "FLUSHALL", - .handler = .{ .client_handler = server_commands.flush_all }, - .min_args = 1, - .max_args = 1, - .description = "Flush the store", - .write_to_aof = true, - }); - - return registry; -} diff --git a/src/commands/keys.zig b/src/commands/keys.zig index cc01658..3c5f67f 100644 --- a/src/commands/keys.zig +++ b/src/commands/keys.zig @@ -66,6 +66,7 @@ pub fn typeCmd(writer: *Writer, store: *Store, args: []const Value) !void { .list => "list", .time_series => "tseries-type", .bloom_filter => "bloom_filter", + .search_index => "search_index", } else "none"; try resp.writeBulkString(writer, type_str); diff --git a/src/commands/registry.zig b/src/commands/registry.zig index 3867b56..d54fe79 100644 --- a/src/commands/registry.zig +++ b/src/commands/registry.zig @@ -88,6 +88,15 @@ pub const CommandRegistry = struct { error.AuthInvalidPassword => "ERR invalid password", error.AlreadyExists => "ERR key already exists", error.TSDB_DuplicateTimestamp => "ERR duplicate timestamp", + error.DocumentExists => "ERR document already exists", + error.DocumentNotFound => "ERR no such document", + error.FieldNotFound => "ERR unknown field", + error.UnsupportedVectorAlgorithm => "ERR unsupported vector index algorithm", + error.SyntaxError => "ERR syntax error", + error.InvalidArgument => "ERR invalid argument", + error.QuerySyntax => "ERR invalid query syntax", + error.UnknownField => "ERR unknown field", + error.MissingParam => "ERR missing query parameter", else => blk: { std.log.err("Handler for command '{s}' failed with error: {s}", .{ command_name, diff --git a/src/commands/search.zig b/src/commands/search.zig new file mode 100644 index 0000000..0ab02cb --- /dev/null +++ b/src/commands/search.zig @@ -0,0 +1,805 @@ +const std = @import("std"); +const storeModule = @import("../store.zig"); +const Store = storeModule.Store; +const Value = @import("../parser.zig").Value; +const resp = @import("./resp.zig"); +const SearchIndex = @import("../search/index.zig").SearchIndex; +const FieldDef = @import("../search/index.zig").FieldDef; +const FieldValue = @import("../search/index.zig").FieldValue; +const FieldType = @import("../search/index.zig").FieldType; +const VectorParams = @import("../search/vector.zig").VectorParams; +const VectorType = @import("../search/vector.zig").VectorType; +const DistanceMetric = @import("../search/vector.zig").DistanceMetric; +const Clock = @import("../clock.zig"); +const search_engine = @import("../search/search.zig"); +const SearchOptions = search_engine.SearchOptions; + +const Io = std.Io; +const Writer = Io.Writer; +const Allocator = std.mem.Allocator; +const eqlIgnoreCase = std.ascii.eqlIgnoreCase; + +pub fn ft_create(writer: *Writer, store: *Store, args: []const Value) !void { + // FT.CREATE index [ON HASH] [PREFIX count prefix...] SCHEMA field type [type options]... + if (args.len < 2) return error.WrongNumberOfArguments; + const key = args[1].asSlice(); + + if (store.exists(key)) return error.AlreadyExists; + + // Scan forward for the SCHEMA marker. + var i: usize = 2; + var schema_start: ?usize = null; + while (i < args.len) : (i += 1) { + if (eqlIgnoreCase(args[i].asSlice(), "SCHEMA")) { + schema_start = i + 1; + break; + } + } + const start = schema_start orelse return error.SyntaxError; + if (start >= args.len) return error.SyntaxError; + + var fields: std.ArrayListUnmanaged(FieldDef) = .empty; + errdefer { + for (fields.items) |f| store.allocator.free(f.name); + fields.deinit(store.allocator); + } + + i = start; + while (i < args.len) { + const fname = args[i].asSlice(); + if (i + 1 >= args.len) return error.SyntaxError; + const ftype = args[i + 1].asSlice(); + + const name_owned = try store.allocator.dupe(u8, fname); + errdefer store.allocator.free(name_owned); + + if (eqlIgnoreCase(ftype, "TEXT")) { + try fields.append(store.allocator, .{ .name = name_owned, .field_type = .text }); + i += 2; + } else if (eqlIgnoreCase(ftype, "NUMERIC")) { + try fields.append(store.allocator, .{ .name = name_owned, .field_type = .numeric }); + i += 2; + } else if (eqlIgnoreCase(ftype, "TAG")) { + try fields.append(store.allocator, .{ .name = name_owned, .field_type = .tag }); + i += 2; + } else if (eqlIgnoreCase(ftype, "VECTOR")) { + // VECTOR ... (count = number of tokens) + if (i + 3 >= args.len) return error.SyntaxError; + const algo = args[i + 2].asSlice(); + if (!eqlIgnoreCase(algo, "FLAT")) return error.UnsupportedVectorAlgorithm; + const attr_tokens = try args[i + 3].asU64(); + if (attr_tokens % 2 != 0) return error.InvalidArgument; + if (i + 4 + attr_tokens > args.len) return error.SyntaxError; + + var params: VectorParams = undefined; + var dim_set = false; + var type_set = false; + var metric_set = false; + + var j: usize = i + 4; + var token: u64 = 0; + while (token < attr_tokens) : (token += 2) { + const attr = args[j].asSlice(); + const val = args[j + 1].asSlice(); + if (eqlIgnoreCase(attr, "TYPE")) { + params.typ = VectorType.fromSlice(val) orelse return error.InvalidArgument; + type_set = true; + } else if (eqlIgnoreCase(attr, "DIM")) { + params.dim = try args[j + 1].asU16(); + if (params.dim == 0) return error.InvalidArgument; + dim_set = true; + } else if (eqlIgnoreCase(attr, "DISTANCE_METRIC")) { + params.metric = DistanceMetric.fromSlice(val) orelse return error.InvalidArgument; + metric_set = true; + } + j += 2; + } + if (!(dim_set and type_set and metric_set)) return error.InvalidArgument; + + try fields.append(store.allocator, .{ .name = name_owned, .field_type = .{ .vector = params } }); + i = j; + } else { + return error.InvalidArgument; + } + } + + // Reject duplicate field names. + var seen: std.StringHashMapUnmanaged(void) = .empty; + defer seen.deinit(store.allocator); + for (fields.items) |f| { + if (seen.contains(f.name)) return error.InvalidArgument; + try seen.put(store.allocator, f.name, {}); + } + + const owned_fields = try fields.toOwnedSlice(store.allocator); + const si = SearchIndex.init(store.allocator, owned_fields); + try store.createSearchIndex(key, si); + + try resp.writeOK(writer); +} + +pub fn ft_dropindex(writer: *Writer, store: *Store, args: []const Value) !void { + // FT.DROPINDEX idx [DD] + const key = args[1].asSlice(); + if (!store.delete(key)) return error.KeyNotFound; + try resp.writeOK(writer); +} + +pub fn ft_list(writer: *Writer, store: *Store, args: []const Value) !void { + _ = args; + const all_keys = try store.keys(store.allocator, ""); + defer store.allocator.free(all_keys); + + var count: usize = 0; + for (all_keys) |k| { + if (store.getType(k) == .search_index) count += 1; + } + + try resp.writeListLen(writer, count); + for (all_keys) |k| { + if (store.getType(k) == .search_index) { + try resp.writeBulkString(writer, k); + } + } +} + +fn typeString(field_type: FieldType) []const u8 { + return switch (field_type) { + .text => "TEXT", + .numeric => "NUMERIC", + .tag => "TAG", + .vector => "VECTOR", + }; +} + +pub fn ft_info(writer: *Writer, store: *Store, args: []const Value) !void { + // FT.INFO idx + const key = args[1].asSlice(); + const index = try store.getSearchIndex(key) orelse return error.KeyNotFound; + + try resp.writeListLen(writer, 6); + try resp.writeBulkString(writer, "index_name"); + try resp.writeBulkString(writer, key); + try resp.writeBulkString(writer, "num_docs"); + try resp.writeInt(writer, @as(i64, @intCast(index.docCount()))); + try resp.writeBulkString(writer, "attributes"); + try resp.writeListLen(writer, index.fields.len * 2); + for (index.fields) |f| { + try resp.writeBulkString(writer, f.name); + try resp.writeBulkString(writer, typeString(f.field_type)); + } +} + +pub fn ft_add(writer: *Writer, store: *Store, args: []const Value) !void { + // FT.ADD idx docId score [NOSAVE] [REPLACE] [LANGUAGE lang] FIELDS n f1 v1 ... + if (args.len < 4) return error.WrongNumberOfArguments; + const key = args[1].asSlice(); + const doc_id = args[2].asSlice(); + _ = try args[3].asF64(); // legacy score; unused for ranking + + var replace = false; + var fields_start: ?usize = null; + var i: usize = 4; + while (i < args.len) { + const a = args[i].asSlice(); + if (eqlIgnoreCase(a, "REPLACE")) { + replace = true; + i += 1; + } else if (eqlIgnoreCase(a, "NOSAVE")) { + i += 1; + } else if (eqlIgnoreCase(a, "LANGUAGE")) { + i += 2; + } else if (eqlIgnoreCase(a, "FIELDS")) { + fields_start = i + 1; + break; + } else { + i += 1; + } + } + + const start = fields_start orelse return error.SyntaxError; + if (start >= args.len) return error.SyntaxError; + const pair_count = try args[start].asU64(); + if (start + 1 + 2 * pair_count > args.len) return error.SyntaxError; + + const index = try store.getSearchIndex(key) orelse return error.KeyNotFound; + if (index.getDocument(doc_id) != null and !replace) return error.DocumentExists; + + const values = try store.allocator.alloc(?FieldValue, index.fields.len); + defer store.allocator.free(values); + for (values) |*v| v.* = null; + + var p: usize = 0; + while (p < pair_count) : (p += 1) { + const fname = args[start + 1 + 2 * p].asSlice(); + const fval = args[start + 1 + 2 * p + 1].asSlice(); + const fidx = index.fieldIndex(fname) orelse return error.FieldNotFound; + switch (index.fields[fidx].field_type) { + .numeric => { + values[fidx] = .{ .numeric = std.fmt.parseFloat(f64, fval) catch return error.InvalidFloat }; + }, + else => values[fidx] = .{ .string = fval }, + } + } + + if (replace and index.getDocument(doc_id) != null) { + _ = index.removeDocument(doc_id); + } + try index.addDocument(doc_id, values); + try resp.writeOK(writer); +} + +pub fn ft_del(writer: *Writer, store: *Store, args: []const Value) !void { + // FT.DEL idx docId [DD] + const key = args[1].asSlice(); + const doc_id = args[2].asSlice(); + const index = try store.getSearchIndex(key) orelse return error.KeyNotFound; + const removed = index.removeDocument(doc_id); + try resp.writeInt(writer, @as(i64, @intFromBool(removed))); +} + +pub fn ft_get(writer: *Writer, store: *Store, args: []const Value) !void { + // FT.GET idx docId + const key = args[1].asSlice(); + const doc_id = args[2].asSlice(); + const index = try store.getSearchIndex(key) orelse return error.KeyNotFound; + const doc = index.getDocument(doc_id) orelse return error.DocumentNotFound; + + try resp.writeListLen(writer, index.fields.len * 2); + for (index.fields, 0..) |f, fi| { + try resp.writeBulkString(writer, f.name); + const fv = doc.fields[fi]; + if (fv) |v| switch (v) { + .string => |s| try resp.writeBulkString(writer, s), + .numeric => |n| try resp.writeDoubleBulkString(writer, n), + } else try resp.writeNull(writer); + } +} + +pub fn ft_search(writer: *Writer, store: *Store, args: []const Value) !void { + // FT.SEARCH idx query [NOCONTENT] [WITHSCORES] [SORTBY f [ASC|DESC]] + // [LIMIT off num] [RETURN n f...] [PARAMS n k v...] [DIALECT n] + if (args.len < 3) return error.WrongNumberOfArguments; + const key = args[1].asSlice(); + const query_str = args[2].asSlice(); + const index = try store.getSearchIndex(key) orelse return error.KeyNotFound; + + var opts = SearchOptions{}; + var params: std.StringHashMapUnmanaged([]const u8) = .empty; + defer params.deinit(store.allocator); + + var i: usize = 3; + while (i < args.len) { + const a = args[i].asSlice(); + if (eqlIgnoreCase(a, "NOCONTENT")) { + opts.no_content = true; + i += 1; + } else if (eqlIgnoreCase(a, "WITHSCORES")) { + opts.with_scores = true; + i += 1; + } else if (eqlIgnoreCase(a, "SORTBY")) { + if (i + 1 >= args.len) return error.SyntaxError; + opts.sort_by = args[i + 1].asSlice(); + i += 2; + if (i < args.len and eqlIgnoreCase(args[i].asSlice(), "ASC")) { + opts.sort_asc = true; + i += 1; + } else if (i < args.len and eqlIgnoreCase(args[i].asSlice(), "DESC")) { + opts.sort_asc = false; + i += 1; + } + } else if (eqlIgnoreCase(a, "LIMIT")) { + if (i + 2 >= args.len) return error.SyntaxError; + opts.limit_offset = try args[i + 1].asUsize(); + opts.limit_count = try args[i + 2].asUsize(); + i += 3; + } else if (eqlIgnoreCase(a, "RETURN")) { + if (i + 1 >= args.len) return error.SyntaxError; + const count = try args[i + 1].asUsize(); + if (i + 1 + count >= args.len + 1) return error.SyntaxError; + var fields: std.ArrayListUnmanaged([]const u8) = .empty; + defer fields.deinit(store.allocator); + var j: usize = 0; + while (j < count) : (j += 1) { + try fields.append(store.allocator, args[i + 2 + j].asSlice()); + } + opts.return_fields = try fields.toOwnedSlice(store.allocator); + defer store.allocator.free(opts.return_fields.?); + i += 2 + count; + } else if (eqlIgnoreCase(a, "PARAMS")) { + if (i + 1 >= args.len) return error.SyntaxError; + const count = try args[i + 1].asUsize(); + if (i + 1 + 2 * count > args.len) return error.SyntaxError; + var j: usize = 0; + while (j < count) : (j += 1) { + const name = args[i + 2 + 2 * j].asSlice(); + const value = args[i + 3 + 2 * j].asSlice(); + try params.put(store.allocator, name, value); + } + i += 2 + 2 * count; + } else if (eqlIgnoreCase(a, "DIALECT")) { + i += 2; + } else { + return error.SyntaxError; + } + } + + var results = try search_engine.search(store.allocator, index, query_str, ¶ms, &opts); + defer results.deinit(store.allocator); + + try resp.writeListLen(writer, results.items.len); + for (results.items) |r| { + try resp.writeBulkString(writer, r.doc_id); + if (opts.with_scores) try resp.writeDoubleBulkString(writer, r.score); + if (opts.no_content) continue; + + const doc = index.getDocument(r.doc_id) orelse continue; + const nfields = if (opts.return_fields) |rfs| rfs.len else index.fields.len; + try resp.writeListLen(writer, nfields * 2); + + if (opts.return_fields) |rfs| { + for (rfs) |fname| { + try resp.writeBulkString(writer, fname); + const field_idx = index.fieldIndex(fname) orelse { + try resp.writeNull(writer); + continue; + }; + const fv = doc.fields[field_idx]; + if (fv) |v| switch (v) { + .string => |s| try resp.writeBulkString(writer, s), + .numeric => |n| try resp.writeDoubleBulkString(writer, n), + } else try resp.writeNull(writer); + } + } else { + for (index.fields, 0..) |f, fi| { + try resp.writeBulkString(writer, f.name); + const fv = doc.fields[fi]; + if (fv) |v| switch (v) { + .string => |s| try resp.writeBulkString(writer, s), + .numeric => |n| try resp.writeDoubleBulkString(writer, n), + } else try resp.writeNull(writer); + } + } + } +} + +const testing = std.testing; + +test "FT.CREATE then FT.INFO" { + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + const allocator = arena.allocator(); + var clock = Clock.init(testing.io, 0); + var store = try Store.init(allocator, testing.io, &clock, .{ .initial_capacity = 16 }); + defer store.deinit(); + + var buf: [4096]u8 = undefined; + var writer = Writer.fixed(&buf); + + const create_args = [_]Value{ + .{ .data = "FT.CREATE" }, + .{ .data = "idx" }, + .{ .data = "SCHEMA" }, + .{ .data = "title" }, + .{ .data = "TEXT" }, + .{ .data = "price" }, + .{ .data = "NUMERIC" }, + }; + try ft_create(&writer, &store, &create_args); + try testing.expectEqualStrings("+OK\r\n", writer.buffered()); + + var buf2: [4096]u8 = undefined; + var writer2 = Writer.fixed(&buf2); + const info_args = [_]Value{ .{ .data = "FT.INFO" }, .{ .data = "idx" } }; + try ft_info(&writer2, &store, &info_args); + const out = writer2.buffered(); + try testing.expect(out[0] == '*'); +} + +test "FT.CREATE with VECTOR field" { + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + const allocator = arena.allocator(); + var clock = Clock.init(testing.io, 0); + var store = try Store.init(allocator, testing.io, &clock, .{ .initial_capacity = 16 }); + defer store.deinit(); + + var buf: [4096]u8 = undefined; + var writer = Writer.fixed(&buf); + + const create_args = [_]Value{ + .{ .data = "FT.CREATE" }, + .{ .data = "vecidx" }, + .{ .data = "SCHEMA" }, + .{ .data = "v" }, + .{ .data = "VECTOR" }, + .{ .data = "FLAT" }, + .{ .data = "6" }, + .{ .data = "TYPE" }, + .{ .data = "FLOAT32" }, + .{ .data = "DIM" }, + .{ .data = "2" }, + .{ .data = "DISTANCE_METRIC" }, + .{ .data = "L2" }, + }; + try ft_create(&writer, &store, &create_args); + try testing.expectEqualStrings("+OK\r\n", writer.buffered()); + try testing.expectEqual(@as(usize, 1), store.size()); +} + +test "FT.ADD document then FT.GET" { + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + const allocator = arena.allocator(); + var clock = Clock.init(testing.io, 0); + var store = try Store.init(allocator, testing.io, &clock, .{ .initial_capacity = 16 }); + defer store.deinit(); + + var buf: [4096]u8 = undefined; + var writer = Writer.fixed(&buf); + const create_args = [_]Value{ + .{ .data = "FT.CREATE" }, .{ .data = "idx" }, .{ .data = "SCHEMA" }, + .{ .data = "title" }, .{ .data = "TEXT" }, .{ .data = "price" }, + .{ .data = "NUMERIC" }, + }; + try ft_create(&writer, &store, &create_args); + + var buf2: [4096]u8 = undefined; + var writer2 = Writer.fixed(&buf2); + const add_args = [_]Value{ + .{ .data = "FT.ADD" }, + .{ .data = "idx" }, + .{ .data = "doc1" }, + .{ .data = "1.0" }, + .{ .data = "FIELDS" }, + .{ .data = "2" }, + .{ .data = "title" }, + .{ .data = "hello world" }, + .{ .data = "price" }, + .{ .data = "9.99" }, + }; + try ft_add(&writer2, &store, &add_args); + try testing.expectEqualStrings("+OK\r\n", writer2.buffered()); + + var buf3: [4096]u8 = undefined; + var writer3 = Writer.fixed(&buf3); + const get_args = [_]Value{ .{ .data = "FT.GET" }, .{ .data = "idx" }, .{ .data = "doc1" } }; + try ft_get(&writer3, &store, &get_args); + const out = writer3.buffered(); + try testing.expect(out[0] == '*'); + + var buf4: [4096]u8 = undefined; + var writer4 = Writer.fixed(&buf4); + const del_args = [_]Value{ .{ .data = "FT.DEL" }, .{ .data = "idx" }, .{ .data = "doc1" } }; + try ft_del(&writer4, &store, &del_args); + try testing.expectEqualStrings(":1\r\n", writer4.buffered()); +} + +test "FT.ADD rejects duplicate without REPLACE" { + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + const allocator = arena.allocator(); + var clock = Clock.init(testing.io, 0); + var store = try Store.init(allocator, testing.io, &clock, .{ .initial_capacity = 16 }); + defer store.deinit(); + + var buf: [4096]u8 = undefined; + var writer = Writer.fixed(&buf); + const create_args = [_]Value{ + .{ .data = "FT.CREATE" }, .{ .data = "idx" }, .{ .data = "SCHEMA" }, + .{ .data = "title" }, .{ .data = "TEXT" }, + }; + try ft_create(&writer, &store, &create_args); + + const add = [_]Value{ + .{ .data = "FT.ADD" }, .{ .data = "idx" }, .{ .data = "doc1" }, .{ .data = "1.0" }, + .{ .data = "FIELDS" }, .{ .data = "1" }, .{ .data = "title" }, .{ .data = "a" }, + }; + var buf2: [4096]u8 = undefined; + var writer2 = Writer.fixed(&buf2); + try ft_add(&writer2, &store, &add); + + var buf3: [4096]u8 = undefined; + var writer3 = Writer.fixed(&buf3); + try testing.expectError(error.DocumentExists, ft_add(&writer3, &store, &add)); +} + +test "FT._LIST returns created indexes" { + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + const allocator = arena.allocator(); + var clock = Clock.init(testing.io, 0); + var store = try Store.init(allocator, testing.io, &clock, .{ .initial_capacity = 16 }); + defer store.deinit(); + + var buf: [4096]u8 = undefined; + var writer = Writer.fixed(&buf); + const create_args = [_]Value{ + .{ .data = "FT.CREATE" }, .{ .data = "idx1" }, .{ .data = "SCHEMA" }, + .{ .data = "title" }, .{ .data = "TEXT" }, + }; + try ft_create(&writer, &store, &create_args); + try store.set("notanindex", "x"); + + var buf2: [4096]u8 = undefined; + var writer2 = Writer.fixed(&buf2); + const list_args = [_]Value{.{ .data = "FT._LIST" }}; + try ft_list(&writer2, &store, &list_args); + const out = writer2.buffered(); + try testing.expect(out[0] == '*'); + try testing.expect(std.mem.indexOf(u8, out, "idx1") != null); +} + +test "FT.DROPINDEX removes index" { + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + const allocator = arena.allocator(); + var clock = Clock.init(testing.io, 0); + var store = try Store.init(allocator, testing.io, &clock, .{ .initial_capacity = 16 }); + defer store.deinit(); + + var buf: [4096]u8 = undefined; + var writer = Writer.fixed(&buf); + const create_args = [_]Value{ + .{ .data = "FT.CREATE" }, .{ .data = "idx1" }, .{ .data = "SCHEMA" }, + .{ .data = "title" }, .{ .data = "TEXT" }, + }; + try ft_create(&writer, &store, &create_args); + + var buf2: [4096]u8 = undefined; + var writer2 = Writer.fixed(&buf2); + const drop_args = [_]Value{ .{ .data = "FT.DROPINDEX" }, .{ .data = "idx1" } }; + try ft_dropindex(&writer2, &store, &drop_args); + try testing.expectEqualStrings("+OK\r\n", writer2.buffered()); + try testing.expect(store.getSearchIndex("idx1") == null); +} + +fn vecF32Blob(values: []const f32) []const u8 { + return std.mem.sliceAsBytes(values); +} + +test "FT.ADD under kv allocator does not OOM" { + const KeyValueAllocator = @import("../kv_allocator.zig"); + var kv = try KeyValueAllocator.init(testing.allocator, 1024 * 1024, .allkeys_lru); + defer kv.deinit(); + + var clock = Clock.init(testing.io, 0); + var store = try Store.init(kv.allocator(), testing.io, &clock, .{ + .initial_capacity = 16, + .eviction_policy = .allkeys_lru, + }); + defer store.deinit(); + kv.attachStore(&store); + + var buf: [8192]u8 = undefined; + var writer = Writer.fixed(&buf); + const create_args = [_]Value{ + .{ .data = "FT.CREATE" }, + .{ .data = "vidx" }, + .{ .data = "SCHEMA" }, + .{ .data = "title" }, + .{ .data = "TEXT" }, + .{ .data = "price" }, + .{ .data = "NUMERIC" }, + .{ .data = "color" }, + .{ .data = "TAG" }, + .{ .data = "v" }, + .{ .data = "VECTOR" }, + .{ .data = "FLAT" }, + .{ .data = "6" }, + .{ .data = "TYPE" }, + .{ .data = "FLOAT32" }, + .{ .data = "DIM" }, + .{ .data = "2" }, + .{ .data = "DISTANCE_METRIC" }, + .{ .data = "L2" }, + }; + try ft_create(&writer, &store, &create_args); + + var buf2: [8192]u8 = undefined; + var writer2 = Writer.fixed(&buf2); + const add = [_]Value{ + .{ .data = "FT.ADD" }, .{ .data = "vidx" }, .{ .data = "doc1" }, .{ .data = "1.0" }, + .{ .data = "FIELDS" }, .{ .data = "4" }, .{ .data = "title" }, .{ .data = "the quick brown fox jumps over the lazy dog, repeatedly, with many words to grow the posting arrays well beyond any initial capacity" }, + .{ .data = "price" }, .{ .data = "9.99" }, .{ .data = "color" }, .{ .data = "red,blue" }, + .{ .data = "v" }, vecF32Blob(&[_]f32{ 0, 0 }), + }; + try ft_add(&writer2, &store, &add); + try testing.expectEqualStrings("+OK\r\n", writer2.buffered()); +} + +test "FT.SEARCH text query returns matching docs" { + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + const allocator = arena.allocator(); + var clock = Clock.init(testing.io, 0); + var store = try Store.init(allocator, testing.io, &clock, .{ .initial_capacity = 16 }); + defer store.deinit(); + + var buf: [4096]u8 = undefined; + var writer = Writer.fixed(&buf); + const create_args = [_]Value{ + .{ .data = "FT.CREATE" }, .{ .data = "idx" }, .{ .data = "SCHEMA" }, + .{ .data = "title" }, .{ .data = "TEXT" }, + }; + try ft_create(&writer, &store, &create_args); + + const add = [_]Value{ + .{ .data = "FT.ADD" }, .{ .data = "idx" }, .{ .data = "doc1" }, .{ .data = "1.0" }, + .{ .data = "FIELDS" }, .{ .data = "1" }, .{ .data = "title" }, .{ .data = "the quick fox" }, + }; + var buf2: [4096]u8 = undefined; + var writer2 = Writer.fixed(&buf2); + try ft_add(&writer2, &store, &add); + + const add2 = [_]Value{ + .{ .data = "FT.ADD" }, .{ .data = "idx" }, .{ .data = "doc2" }, .{ .data = "1.0" }, + .{ .data = "FIELDS" }, .{ .data = "1" }, .{ .data = "title" }, .{ .data = "the lazy dog" }, + }; + var buf3: [4096]u8 = undefined; + var writer3 = Writer.fixed(&buf3); + try ft_add(&writer3, &store, &add2); + + var buf4: [4096]u8 = undefined; + var writer4 = Writer.fixed(&buf4); + const search_args = [_]Value{ + .{ .data = "FT.SEARCH" }, .{ .data = "idx" }, .{ .data = "quick" }, + }; + try ft_search(&writer4, &store, &search_args); + const out = writer4.buffered(); + try testing.expect(out[0] == '*'); + try testing.expect(std.mem.indexOf(u8, out, "doc1") != null); + try testing.expect(std.mem.indexOf(u8, out, "doc2") == null); +} + +test "FT.SEARCH numeric range filter" { + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + const allocator = arena.allocator(); + var clock = Clock.init(testing.io, 0); + var store = try Store.init(allocator, testing.io, &clock, .{ .initial_capacity = 16 }); + defer store.deinit(); + + var buf: [4096]u8 = undefined; + var writer = Writer.fixed(&buf); + const create_args = [_]Value{ + .{ .data = "FT.CREATE" }, .{ .data = "idx" }, .{ .data = "SCHEMA" }, + .{ .data = "price" }, .{ .data = "NUMERIC" }, + }; + try ft_create(&writer, &store, &create_args); + + const add = [_]Value{ + .{ .data = "FT.ADD" }, .{ .data = "idx" }, .{ .data = "doc1" }, .{ .data = "1.0" }, + .{ .data = "FIELDS" }, .{ .data = "1" }, .{ .data = "price" }, .{ .data = "9.99" }, + }; + var buf2: [4096]u8 = undefined; + var writer2 = Writer.fixed(&buf2); + try ft_add(&writer2, &store, &add); + + const add2 = [_]Value{ + .{ .data = "FT.ADD" }, .{ .data = "idx" }, .{ .data = "doc2" }, .{ .data = "1.0" }, + .{ .data = "FIELDS" }, .{ .data = "1" }, .{ .data = "price" }, .{ .data = "80.00" }, + }; + var buf3: [4096]u8 = undefined; + var writer3 = Writer.fixed(&buf3); + try ft_add(&writer3, &store, &add2); + + var buf4: [4096]u8 = undefined; + var writer4 = Writer.fixed(&buf4); + const search_args = [_]Value{ + .{ .data = "FT.SEARCH" }, .{ .data = "idx" }, .{ .data = "@price:[0 50]" }, + }; + try ft_search(&writer4, &store, &search_args); + const out = writer4.buffered(); + try testing.expect(std.mem.indexOf(u8, out, "doc1") != null); + try testing.expect(std.mem.indexOf(u8, out, "doc2") == null); +} + +test "FT.SEARCH tag filter" { + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + const allocator = arena.allocator(); + var clock = Clock.init(testing.io, 0); + var store = try Store.init(allocator, testing.io, &clock, .{ .initial_capacity = 16 }); + defer store.deinit(); + + var buf: [4096]u8 = undefined; + var writer = Writer.fixed(&buf); + const create_args = [_]Value{ + .{ .data = "FT.CREATE" }, .{ .data = "idx" }, .{ .data = "SCHEMA" }, + .{ .data = "color" }, .{ .data = "TAG" }, + }; + try ft_create(&writer, &store, &create_args); + + const add = [_]Value{ + .{ .data = "FT.ADD" }, .{ .data = "idx" }, .{ .data = "doc1" }, .{ .data = "1.0" }, + .{ .data = "FIELDS" }, .{ .data = "1" }, .{ .data = "color" }, .{ .data = "red,blue" }, + }; + var buf2: [4096]u8 = undefined; + var writer2 = Writer.fixed(&buf2); + try ft_add(&writer2, &store, &add); + + const add2 = [_]Value{ + .{ .data = "FT.ADD" }, .{ .data = "idx" }, .{ .data = "doc2" }, .{ .data = "1.0" }, + .{ .data = "FIELDS" }, .{ .data = "1" }, .{ .data = "color" }, .{ .data = "green" }, + }; + var buf3: [4096]u8 = undefined; + var writer3 = Writer.fixed(&buf3); + try ft_add(&writer3, &store, &add2); + + var buf4: [4096]u8 = undefined; + var writer4 = Writer.fixed(&buf4); + const search_args = [_]Value{ + .{ .data = "FT.SEARCH" }, .{ .data = "idx" }, .{ .data = "@color:{red}" }, + }; + try ft_search(&writer4, &store, &search_args); + const out = writer4.buffered(); + try testing.expect(std.mem.indexOf(u8, out, "doc1") != null); + try testing.expect(std.mem.indexOf(u8, out, "doc2") == null); +} + +test "FT.SEARCH vector KNN returns nearest doc" { + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + const allocator = arena.allocator(); + var clock = Clock.init(testing.io, 0); + var store = try Store.init(allocator, testing.io, &clock, .{ .initial_capacity = 16 }); + defer store.deinit(); + + var buf: [4096]u8 = undefined; + var writer = Writer.fixed(&buf); + const create_args = [_]Value{ + .{ .data = "FT.CREATE" }, + .{ .data = "vidx" }, + .{ .data = "SCHEMA" }, + .{ .data = "title" }, + .{ .data = "TEXT" }, + .{ .data = "v" }, + .{ .data = "VECTOR" }, + .{ .data = "FLAT" }, + .{ .data = "6" }, + .{ .data = "TYPE" }, + .{ .data = "FLOAT32" }, + .{ .data = "DIM" }, + .{ .data = "2" }, + .{ .data = "DISTANCE_METRIC" }, + .{ .data = "L2" }, + }; + try ft_create(&writer, &store, &create_args); + + const add = [_]Value{ + .{ .data = "FT.ADD" }, .{ .data = "vidx" }, .{ .data = "doc1" }, .{ .data = "1.0" }, + .{ .data = "FIELDS" }, .{ .data = "2" }, .{ .data = "title" }, .{ .data = "first" }, + .{ .data = "v" }, vecF32Blob(&[_]f32{ 0, 0 }), + }; + var buf2: [8192]u8 = undefined; + var writer2 = Writer.fixed(&buf2); + try ft_add(&writer2, &store, &add); + + const add2 = [_]Value{ + .{ .data = "FT.ADD" }, .{ .data = "vidx" }, .{ .data = "doc2" }, .{ .data = "1.0" }, + .{ .data = "FIELDS" }, .{ .data = "2" }, .{ .data = "title" }, .{ .data = "second" }, + .{ .data = "v" }, vecF32Blob(&[_]f32{ 10, 10 }), + }; + var buf3: [8192]u8 = undefined; + var writer3 = Writer.fixed(&buf3); + try ft_add(&writer3, &store, &add2); + + var buf4: [8192]u8 = undefined; + var writer4 = Writer.fixed(&buf4); + const search_args = [_]Value{ + .{ .data = "FT.SEARCH" }, + .{ .data = "vidx" }, + .{ .data = "*=>[KNN 1 @v $B]" }, + .{ .data = "PARAMS" }, + .{ .data = "2" }, + .{ .data = "B" }, + vecF32Blob(&[_]f32{ 0, 0 }), + .{ .data = "DIALECT" }, + .{ .data = "2" }, + }; + try ft_search(&writer4, &store, &search_args); + const out = writer4.buffered(); + try testing.expect(std.mem.indexOf(u8, out, "doc1") != null); + try testing.expect(std.mem.indexOf(u8, out, "doc2") == null); +} diff --git a/src/commands/string.zig b/src/commands/string.zig index 6142af7..ec050a8 100644 --- a/src/commands/string.zig +++ b/src/commands/string.zig @@ -30,7 +30,7 @@ pub fn get(writer: *Writer, store: *Store, args: []const Value) !void { .int => |i| { try resp.writeIntBulkString(writer, i); }, - .list, .time_series, .bloom_filter => return error.WrongType, + .list, .time_series, .bloom_filter, .search_index => return error.WrongType, } } else { try resp.writeNull(writer); diff --git a/src/rdb/zdb.zig b/src/rdb/zdb.zig index b920137..52e17b6 100644 --- a/src/rdb/zdb.zig +++ b/src/rdb/zdb.zig @@ -49,6 +49,7 @@ pub const Writer = struct { .list => 0x01, .time_series => 0x0A, .bloom_filter => 0x0B, + .search_index => 0x0C, }; } diff --git a/src/search/index.zig b/src/search/index.zig new file mode 100644 index 0000000..f2feb49 --- /dev/null +++ b/src/search/index.zig @@ -0,0 +1,350 @@ +const std = @import("std"); +const InvertedIndex = @import("inverted.zig").InvertedIndex; +const NumericIndex = @import("numeric.zig").NumericIndex; +const TagIndex = @import("tag.zig").TagIndex; +const VectorIndex = @import("vector.zig").VectorIndex; +const VectorParams = @import("vector.zig").VectorParams; + +pub const Allocator = std.mem.Allocator; + +pub const FieldType = union(enum) { + text, + numeric, + tag, + vector: VectorParams, +}; + +pub const FieldDef = struct { + name: []const u8, + field_type: FieldType, +}; + +pub const FieldValue = union(enum) { + string: []const u8, + numeric: f64, +}; + +pub const Document = struct { + internal_id: u64, + fields: []?FieldValue, +}; + +pub const SearchIndex = struct { + allocator: Allocator, + fields: []FieldDef, + field_map: std.StringHashMapUnmanaged(u16) = .empty, + docs: std.StringHashMapUnmanaged(Document) = .empty, + doc_counter: u64 = 0, + + inverted: std.AutoHashMapUnmanaged(u16, InvertedIndex) = .empty, + numeric: std.AutoHashMapUnmanaged(u16, NumericIndex) = .empty, + tag: std.AutoHashMapUnmanaged(u16, TagIndex) = .empty, + vector: std.AutoHashMapUnmanaged(u16, VectorIndex) = .empty, + + pub fn init(allocator: Allocator, fields: []FieldDef) SearchIndex { + var field_map: std.StringHashMapUnmanaged(u16) = .empty; + field_map.ensureTotalCapacity(allocator, @intCast(fields.len)) catch {}; + for (fields, 0..) |field, i| { + field_map.putAssumeCapacity(field.name, @intCast(i)); + } + return .{ + .allocator = allocator, + .fields = fields, + .field_map = field_map, + }; + } + + pub fn deinit(self: *SearchIndex) void { + var doc_it = self.docs.iterator(); + while (doc_it.next()) |entry| { + const doc = entry.value_ptr; + self.freeFieldValues(doc); + self.allocator.free(doc.fields); + self.allocator.free(entry.key_ptr.*); + } + self.docs.deinit(self.allocator); + self.field_map.deinit(self.allocator); + + var inv_it = self.inverted.iterator(); + while (inv_it.next()) |e| e.value_ptr.deinit(); + self.inverted.deinit(self.allocator); + + var num_it = self.numeric.iterator(); + while (num_it.next()) |e| e.value_ptr.deinit(); + self.numeric.deinit(self.allocator); + + var tag_it = self.tag.iterator(); + while (tag_it.next()) |e| e.value_ptr.deinit(); + self.tag.deinit(self.allocator); + + var vec_it = self.vector.iterator(); + while (vec_it.next()) |e| e.value_ptr.deinit(); + self.vector.deinit(self.allocator); + + for (self.fields) |field| { + self.allocator.free(field.name); + } + self.allocator.free(self.fields); + } + + pub fn fieldIndex(self: *const SearchIndex, name: []const u8) ?u16 { + return self.field_map.get(name); + } + + pub fn docCount(self: *const SearchIndex) u64 { + return self.docs.count(); + } + + /// Adds a document. `values` must be aligned to `self.fields` (null for + /// absent fields). Returns error.DocumentExists if the doc id is taken. + pub fn addDocument(self: *SearchIndex, doc_id: []const u8, values: []const ?FieldValue) !void { + if (self.docs.contains(doc_id)) return error.DocumentExists; + if (values.len != self.fields.len) return error.FieldCountMismatch; + + const internal_id = self.doc_counter + 1; + self.doc_counter = internal_id; + + const owned_key = try self.allocator.dupe(u8, doc_id); + errdefer self.allocator.free(owned_key); + + const owned_fields = try self.allocator.alloc(?FieldValue, self.fields.len); + errdefer self.allocator.free(owned_fields); + for (values, 0..) |fv, i| { + owned_fields[i] = if (fv) |v| switch (v) { + .string => |s| .{ .string = try self.allocator.dupe(u8, s) }, + .numeric => v, + } else null; + } + + const doc = Document{ .internal_id = internal_id, .fields = owned_fields }; + try self.docs.put(self.allocator, owned_key, doc); + + for (self.fields, 0..) |field, i| { + const fv = owned_fields[i] orelse continue; + const field_idx: u16 = @intCast(i); + try self.indexValue(field.field_type, field_idx, internal_id, fv); + } + } + + /// Removes a document and de-indexes its fields. Returns false if absent. + pub fn removeDocument(self: *SearchIndex, doc_id: []const u8) bool { + const kv = self.docs.fetchRemove(doc_id) orelse return false; + const doc = kv.value; + + for (self.fields, 0..) |field, i| { + const fv = doc.fields[i] orelse continue; + const field_idx: u16 = @intCast(i); + self.deindexValue(field.field_type, field_idx, doc.internal_id, fv); + } + + self.allocator.free(kv.key); + self.freeFieldValues(&doc); + self.allocator.free(doc.fields); + return true; + } + + pub fn getDocument(self: *const SearchIndex, doc_id: []const u8) ?*const Document { + return self.docs.getPtr(doc_id); + } + + pub fn getInverted(self: *const SearchIndex, field_idx: u16) ?*const InvertedIndex { + return self.inverted.getPtr(field_idx); + } + + pub fn getNumeric(self: *const SearchIndex, field_idx: u16) ?*const NumericIndex { + return self.numeric.getPtr(field_idx); + } + + pub fn getTag(self: *const SearchIndex, field_idx: u16) ?*const TagIndex { + return self.tag.getPtr(field_idx); + } + + pub fn getVector(self: *const SearchIndex, field_idx: u16) ?*const VectorIndex { + return self.vector.getPtr(field_idx); + } + + fn indexValue(self: *SearchIndex, field_type: FieldType, field_idx: u16, doc_id: u64, fv: FieldValue) !void { + switch (field_type) { + .text => switch (fv) { + .string => |s| try self.indexText(field_idx, doc_id, s), + else => {}, + }, + .numeric => switch (fv) { + .numeric => |n| try self.indexNumeric(field_idx, doc_id, n), + else => {}, + }, + .tag => switch (fv) { + .string => |s| try self.indexTag(field_idx, doc_id, s), + else => {}, + }, + .vector => |vp| switch (fv) { + .string => |blob| try self.indexVector(field_idx, vp, doc_id, blob), + else => {}, + }, + } + } + + fn deindexValue(self: *SearchIndex, field_type: FieldType, field_idx: u16, doc_id: u64, fv: FieldValue) void { + switch (field_type) { + .text => switch (fv) { + .string => |s| if (self.inverted.getPtr(field_idx)) |ii| ii.remove(doc_id, s), + else => {}, + }, + .numeric => switch (fv) { + .numeric => |n| { + _ = n; + if (self.numeric.getPtr(field_idx)) |ni| ni.remove(doc_id); + }, + else => {}, + }, + .tag => switch (fv) { + .string => |s| if (self.tag.getPtr(field_idx)) |ti| ti.remove(doc_id, s), + else => {}, + }, + .vector => switch (fv) { + .string => |blob| { + _ = blob; + if (self.vector.getPtr(field_idx)) |vi| vi.remove(doc_id); + }, + else => {}, + }, + } + } + + fn indexText(self: *SearchIndex, field_idx: u16, doc_id: u64, text: []const u8) !void { + const gop = try self.inverted.getOrPut(self.allocator, field_idx); + if (!gop.found_existing) gop.value_ptr.* = InvertedIndex.init(self.allocator); + try gop.value_ptr.add(doc_id, text); + } + + fn indexNumeric(self: *SearchIndex, field_idx: u16, doc_id: u64, value: f64) !void { + const gop = try self.numeric.getOrPut(self.allocator, field_idx); + if (!gop.found_existing) gop.value_ptr.* = NumericIndex.init(self.allocator); + try gop.value_ptr.add(doc_id, value); + } + + fn indexTag(self: *SearchIndex, field_idx: u16, doc_id: u64, value: []const u8) !void { + const gop = try self.tag.getOrPut(self.allocator, field_idx); + if (!gop.found_existing) gop.value_ptr.* = TagIndex.init(self.allocator); + try gop.value_ptr.add(doc_id, value); + } + + fn indexVector(self: *SearchIndex, field_idx: u16, params: VectorParams, doc_id: u64, blob: []const u8) !void { + const gop = try self.vector.getOrPut(self.allocator, field_idx); + if (!gop.found_existing) gop.value_ptr.* = VectorIndex.init(self.allocator, params); + try gop.value_ptr.add(doc_id, blob); + } + + fn freeFieldValues(self: *SearchIndex, doc: *const Document) void { + for (doc.fields) |field_value| { + if (field_value) |fv| { + switch (fv) { + .string => |s| self.allocator.free(s), + .numeric => {}, + } + } + } + } +}; + +const testing = std.testing; + +fn makeFields() ![]FieldDef { + return try testing.allocator.dupe(FieldDef, &.{ + .{ .name = try testing.allocator.dupe(u8, "title"), .field_type = .text }, + .{ .name = try testing.allocator.dupe(u8, "price"), .field_type = .numeric }, + .{ .name = try testing.allocator.dupe(u8, "color"), .field_type = .tag }, + .{ .name = try testing.allocator.dupe(u8, "v"), .field_type = .{ .vector = .{ + .dim = 2, + .typ = .float32, + .metric = .l2, + } } }, + }); +} + +fn vecF32(values: []const f32) []const u8 { + return std.mem.sliceAsBytes(values); +} + +test "SearchIndex init and deinit" { + const fields = try makeFields(); + var index = SearchIndex.init(testing.allocator, fields); + defer index.deinit(); + + try testing.expectEqual(@as(usize, 4), index.fields.len); + try testing.expectEqual(@as(u16, 0), index.field_map.get("title").?); + try testing.expectEqual(@as(u16, 1), index.field_map.get("price").?); + try testing.expectEqual(@as(u16, 2), index.field_map.get("color").?); + try testing.expectEqual(@as(u16, 3), index.field_map.get("v").?); +} + +test "SearchIndex add document indexes all field types" { + const fields = try makeFields(); + var index = SearchIndex.init(testing.allocator, fields); + defer index.deinit(); + + try index.addDocument("doc1", &.{ + .{ .string = "the quick fox" }, + .{ .numeric = 9.99 }, + .{ .string = "red,blue" }, + .{ .string = vecF32(&[_]f32{ 0, 0 }) }, + }); + try index.addDocument("doc2", &.{ + .{ .string = "quick dog" }, + .{ .numeric = 25.0 }, + .{ .string = "green" }, + .{ .string = vecF32(&[_]f32{ 10, 10 }) }, + }); + + try testing.expectEqual(@as(u64, 2), index.docCount()); + + const ii = index.getInverted(0).?; + try testing.expectEqual(@as(usize, 2), ii.postings("quick").?.len); + + const ni = index.getNumeric(1).?; + var range: std.ArrayListUnmanaged(u64) = .empty; + defer range.deinit(testing.allocator); + try ni.range(10.0, 50.0, &range); + try testing.expectEqual(@as(usize, 1), range.items.len); + + const ti = index.getTag(2).?; + try testing.expect(ti.docIds("red").?.contains(index.docs.get("doc1").?.internal_id)); + + const vi = index.getVector(3).?; + try testing.expectEqual(@as(usize, 2), vi.count()); +} + +test "SearchIndex remove document de-indexes" { + const fields = try makeFields(); + var index = SearchIndex.init(testing.allocator, fields); + defer index.deinit(); + + try index.addDocument("doc1", &.{ + .{ .string = "alpha beta" }, + .{ .numeric = 1.0 }, + .{ .string = "red" }, + .{ .string = vecF32(&[_]f32{ 1, 1 }) }, + }); + + try testing.expect(index.removeDocument("doc1")); + try testing.expect(!index.removeDocument("doc1")); + try testing.expectEqual(@as(u64, 0), index.docCount()); + + const ii = index.getInverted(0).?; + try testing.expectEqual(@as(u32, 0), ii.docFreq("alpha")); + try testing.expectEqual(@as(usize, 0), index.getVector(3).?.count()); +} + +test "SearchIndex addDocument rejects duplicates" { + const fields = try makeFields(); + var index = SearchIndex.init(testing.allocator, fields); + defer index.deinit(); + + try index.addDocument("doc1", &.{ .{ .string = "a" }, null, null, null }); + try testing.expectError(error.DocumentExists, index.addDocument("doc1", &.{ + .{ .string = "b" }, + null, + null, + null, + })); +} diff --git a/src/search/inverted.zig b/src/search/inverted.zig new file mode 100644 index 0000000..ad86114 --- /dev/null +++ b/src/search/inverted.zig @@ -0,0 +1,195 @@ +const std = @import("std"); +const Tokenizer = @import("tokenize.zig").Tokenizer; + +pub const Allocator = std.mem.Allocator; + +const k1: f64 = 1.2; +const b: f64 = 0.75; + +pub const TermPosting = struct { + doc_id: u64, + tf: u32, +}; + +pub const TermInfo = struct { + ddff: u32 = 0, + posts: std.ArrayListUnmanaged(TermPosting) = .empty, +}; + +/// Per-field inverted index: term -> (df, postings), plus per-doc lengths +/// for BM25 scoring. +pub const InvertedIndex = struct { + allocator: Allocator, + terms: std.StringHashMapUnmanaged(TermInfo) = .empty, + doc_len: std.AutoHashMapUnmanaged(u64, u32) = .empty, + num_docs: u64 = 0, + total_len: u64 = 0, + + pub fn init(allocator: Allocator) InvertedIndex { + return .{ .allocator = allocator }; + } + + pub fn deinit(self: *InvertedIndex) void { + var it = self.terms.iterator(); + while (it.next()) |e| { + self.allocator.free(e.key_ptr.*); + e.value_ptr.posts.deinit(self.allocator); + } + self.terms.deinit(self.allocator); + self.doc_len.deinit(self.allocator); + } + + pub fn add(self: *InvertedIndex, doc_id: u64, text: []const u8) !void { + var tokenizer = Tokenizer.init(self.allocator); + defer tokenizer.deinit(); + try tokenizer.tokenize(text); + + var uniques = tokenizer.uniqueTokens(); + defer uniques.deinit(self.allocator); + + for (uniques.items) |u| { + const gop = try self.terms.getOrPut(self.allocator, u.token); + if (!gop.found_existing) { + gop.key_ptr.* = try self.allocator.dupe(u8, u.token); + gop.value_ptr.* = .{}; + } + gop.value_ptr.df += 1; + try gop.value_ptr.posts.append(self.allocator, .{ .doc_id = doc_id, .tf = u.count }); + } + + try self.doc_len.put(self.allocator, doc_id, @intCast(tokenizer.tokens.items.len)); + self.num_docs += 1; + self.total_len += tokenizer.tokens.items.len; + } + + /// Removes a document from the index by scanning postings for its id. + pub fn remove(self: *InvertedIndex, doc_id: u64, text: []const u8) void { + var tokenizer = Tokenizer.init(self.allocator); + defer tokenizer.deinit(); + tokenizer.tokenize(text) catch return; + + var uniques = tokenizer.uniqueTokens(); + defer uniques.deinit(self.allocator); + + for (uniques.items) |u| { + if (self.terms.getPtr(u.token)) |info| { + for (info.posts.items, 0..) |p, idx| { + if (p.doc_id == doc_id) { + _ = info.posts.swapRemove(idx); + info.df -= 1; + break; + } + } + if (info.df == 0) { + if (self.terms.fetchRemove(u.token)) |kv| { + self.allocator.free(kv.key); + var term_value = kv.value; + term_value.posts.deinit(self.allocator); + } + } + } + } + + if (self.doc_len.fetchRemove(doc_id)) |kv| { + self.num_docs -= 1; + self.total_len -= kv.value; + } + } + + pub fn avgDocLen(self: *const InvertedIndex) f64 { + if (self.num_docs == 0) return 0; + return @as(f64, @floatFromInt(self.total_len)) / @as(f64, @floatFromInt(self.num_docs)); + } + + pub fn idf(self: *const InvertedIndex, df: u32) f64 { + const n: f64 = @floatFromInt(self.num_docs); + const d: f64 = @floatFromInt(df); + return @log(1 + (n - d + 0.5) / (d + 0.5)); + } + + pub fn bm25(self: *const InvertedIndex, doc_id: u64, df: u32, freq: u32) f64 { + const avgdl = self.avgDocLen(); + const dl: f64 = @floatFromInt(self.doc_len.get(doc_id) orelse 0); + const t: f64 = @floatFromInt(freq); + const denom = t + k1 * (1 - b + b * dl / (if (avgdl == 0) 1.0 else avgdl)); + return self.idf(df) * (t * (k1 + 1)) / denom; + } + + /// Term frequency for a (doc, term) pair, if present. + pub fn tf(self: *const InvertedIndex, doc_id: u64, term: []const u8) u32 { + const info = self.terms.get(term) orelse return 0; + for (info.posts.items) |p| { + if (p.doc_id == doc_id) return p.tf; + } + return 0; + } + + pub fn postings(self: *const InvertedIndex, term: []const u8) ?[]const TermPosting { + const info = self.terms.get(term) orelse return null; + return info.posts.items; + } + + pub fn docFreq(self: *const InvertedIndex, term: []const u8) u32 { + const info = self.terms.get(term) orelse return 0; + return info.df; + } + + pub fn termCount(self: *const InvertedIndex) usize { + return self.terms.count(); + } + + pub fn docLength(self: *const InvertedIndex, doc_id: u64) u32 { + return self.doc_len.get(doc_id) orelse 0; + } +}; + +const testing = std.testing; + +test "inverted add and search postings" { + var ii = InvertedIndex.init(testing.allocator); + defer ii.deinit(); + + try ii.add(1, "the quick brown fox"); + try ii.add(2, "the lazy dog"); + try ii.add(3, "quick quick fox"); + + const posts = ii.postings("quick").?; + try testing.expectEqual(@as(usize, 2), posts.len); + // doc 1 has tf 1, doc 3 has tf 2 + for (posts) |p| { + if (p.doc_id == 1) try testing.expectEqual(@as(u32, 1), p.tf); + if (p.doc_id == 3) try testing.expectEqual(@as(u32, 2), p.tf); + } + + try testing.expectEqual(@as(u32, 3), ii.num_docs); + try testing.expect(ii.postings("nonexistent") == null); +} + +test "inverted remove document" { + var ii = InvertedIndex.init(testing.allocator); + defer ii.deinit(); + + try ii.add(1, "alpha beta"); + try ii.add(2, "alpha gamma"); + try testing.expectEqual(@as(u32, 2), ii.docFreq("alpha")); + + ii.remove(1, "alpha beta"); + + try testing.expectEqual(@as(u32, 1), ii.docFreq("alpha")); + try testing.expect(ii.docFreq("beta") == 0); + try testing.expect(ii.terms.get("beta") == null); + try testing.expectEqual(@as(u64, 1), ii.num_docs); +} + +test "inverted bm25 prefers rare term in short doc" { + var ii = InvertedIndex.init(testing.allocator); + defer ii.deinit(); + + try ii.add(1, "cat"); + try ii.add(2, "cat dog bird fish bear wolf deer elk"); + + const s1 = ii.bm25(1, ii.docFreq("cat"), ii.tf(1, "cat")); + const s2 = ii.bm25(2, ii.docFreq("cat"), ii.tf(2, "cat")); + // Equal tf, but shorter doc should score higher (bm25 length normalization) + try testing.expect(s1 > s2); +} diff --git a/src/search/numeric.zig b/src/search/numeric.zig new file mode 100644 index 0000000..202b797 --- /dev/null +++ b/src/search/numeric.zig @@ -0,0 +1,84 @@ +const std = @import("std"); + +pub const Allocator = std.mem.Allocator; + +pub const NumEntry = struct { + doc_id: u64, + value: f64, +}; + +/// Per-field numeric index. V1 keeps a flat list and linear-scans for range +/// queries (exact, simple). Can be replaced with a sorted structure later. +pub const NumericIndex = struct { + allocator: Allocator, + entries: std.ArrayListUnmanaged(NumEntry) = .empty, + + pub fn init(allocator: Allocator) NumericIndex { + return .{ .allocator = allocator }; + } + + pub fn deinit(self: *NumericIndex) void { + self.entries.deinit(self.allocator); + } + + pub fn add(self: *NumericIndex, doc_id: u64, value: f64) !void { + try self.entries.append(self.allocator, .{ .doc_id = doc_id, .value = value }); + } + + pub fn remove(self: *NumericIndex, doc_id: u64) void { + for (self.entries.items, 0..) |e, i| { + if (e.doc_id == doc_id) { + _ = self.entries.swapRemove(i); + return; + } + } + } + + /// Collects doc_ids with min <= value <= max into `out`. + pub fn range(self: *const NumericIndex, min: ?f64, max: ?f64, out: *std.ArrayListUnmanaged(u64)) !void { + for (self.entries.items) |e| { + if (min) |mn| { + if (e.value < mn) continue; + } + if (max) |mx| { + if (e.value > mx) continue; + } + try out.append(self.allocator, e.doc_id); + } + } +}; + +const testing = std.testing; + +test "numeric range filtering" { + var ni = NumericIndex.init(testing.allocator); + defer ni.deinit(); + + try ni.add(1, 4.5); + try ni.add(2, 10.0); + try ni.add(3, 25.0); + try ni.add(4, 80.0); + + var out: std.ArrayListUnmanaged(u64) = .empty; + defer out.deinit(testing.allocator); + try ni.range(10.0, 50.0, &out); + + try testing.expectEqual(@as(usize, 2), out.items.len); + try testing.expect(out.items[0] == 2 or out.items[1] == 2); + try testing.expect(out.items[0] == 3 or out.items[1] == 3); +} + +test "numeric remove" { + var ni = NumericIndex.init(testing.allocator); + defer ni.deinit(); + + try ni.add(1, 1.0); + try ni.add(2, 2.0); + ni.remove(1); + + var out: std.ArrayListUnmanaged(u64) = .empty; + defer out.deinit(testing.allocator); + try ni.range(null, null, &out); + try testing.expectEqual(@as(usize, 1), out.items.len); + try testing.expectEqual(@as(u64, 2), out.items[0]); +} diff --git a/src/search/query.zig b/src/search/query.zig new file mode 100644 index 0000000..e4989c4 --- /dev/null +++ b/src/search/query.zig @@ -0,0 +1,379 @@ +const std = @import("std"); + +pub const Allocator = std.mem.Allocator; + +pub const ParseError = error{ QuerySyntax, OutOfMemory }; + +/// Nodes borrow slices from the query string passed to `parse`. +pub const Node = union(enum) { + all, + term: Term, + phrase: Phrase, + field_term: FieldTerm, + numeric: NumericRange, + tag: TagSet, + disjunction: []Node, + conjunction: []Node, + not: *Node, +}; + +pub const Term = struct { + text: []const u8, +}; + +pub const Phrase = struct { + text: []const u8, +}; + +pub const FieldTerm = struct { + field: []const u8, + term: []const u8, +}; + +pub const NumericRange = struct { + field: []const u8, + min: ?f64, + max: ?f64, +}; + +pub const TagSet = struct { + field: []const u8, + tags: []const []const u8, +}; + +pub const Knn = struct { + field: []const u8, + k: usize, + vec_param: []const u8, +}; + +pub const Query = struct { + allocator: Allocator, + filter: Node, + knn: ?Knn = null, + + pub fn deinit(self: *Query) void { + freeNode(self.allocator, &self.filter); + } +}; + +const Scanner = struct { + s: []const u8, + pos: usize = 0, + + fn eof(self: *const Scanner) bool { + return self.pos >= self.s.len; + } + + fn peek(self: *const Scanner) u8 { + return if (self.pos < self.s.len) self.s[self.pos] else 0; + } + + fn skipWs(self: *Scanner) void { + while (self.pos < self.s.len and std.ascii.isWhitespace(self.s[self.pos])) self.pos += 1; + } +}; + +/// Parses a RediSearch query string into a Query AST. +/// Supports: `*`, bare terms, `"phrases"`, `@field:term`, +/// `@tag:{a|b}`, `@num:[min max]`, implicit AND, `|` OR, `-` NOT, +/// and the `=>[KNN k @vec $param]` suffix. +pub fn parse(allocator: Allocator, query: []const u8) ParseError!Query { + if (std.mem.indexOf(u8, query, "=>")) |idx| { + const filter_s = query[0..idx]; + const knn_s = query[idx + 2 ..]; + var sc = Scanner{ .s = filter_s }; + const filter = try parseExpr(allocator, &sc); + const knn = try parseKnn(knn_s); + return .{ .allocator = allocator, .filter = filter, .knn = knn }; + } + + var sc = Scanner{ .s = query }; + const filter = try parseExpr(allocator, &sc); + return .{ .allocator = allocator, .filter = filter }; +} + +fn parseExpr(allocator: Allocator, sc: *Scanner) ParseError!Node { + var conjuncts: std.ArrayListUnmanaged(Node) = .empty; + defer conjuncts.deinit(allocator); + + while (true) { + const conj = try parseConjunct(allocator, sc); + try conjuncts.append(allocator, conj); + sc.skipWs(); + if (sc.peek() == '|') { + sc.pos += 1; + continue; + } + break; + } + + if (conjuncts.items.len == 1) return conjuncts.items[0]; + return .{ .disjunction = try conjuncts.toOwnedSlice(allocator) }; +} + +fn parseConjunct(allocator: Allocator, sc: *Scanner) ParseError!Node { + var factors: std.ArrayListUnmanaged(Node) = .empty; + defer factors.deinit(allocator); + + while (true) { + sc.skipWs(); + if (sc.eof() or sc.peek() == '|' or sc.peek() == ')') break; + const f = try parseFactor(allocator, sc); + try factors.append(allocator, f); + } + + if (factors.items.len == 1) return factors.items[0]; + return .{ .conjunction = try factors.toOwnedSlice(allocator) }; +} + +fn parseFactor(allocator: Allocator, sc: *Scanner) ParseError!Node { + sc.skipWs(); + var negate = false; + if (sc.peek() == '-') { + negate = true; + sc.pos += 1; + } + const atom = try parseAtom(allocator, sc); + if (!negate) return atom; + const boxed = try allocator.create(Node); + boxed.* = atom; + return .{ .not = boxed }; +} + +fn parseAtom(allocator: Allocator, sc: *Scanner) ParseError!Node { + sc.skipWs(); + switch (sc.peek()) { + '(' => { + sc.pos += 1; + const node = try parseExpr(allocator, sc); + sc.skipWs(); + if (sc.peek() == ')') sc.pos += 1; + return node; + }, + '@' => return parseFieldPredicate(allocator, sc), + '"' => { + sc.pos += 1; + const start = sc.pos; + while (sc.pos < sc.s.len and sc.s[sc.pos] != '"') sc.pos += 1; + const text = sc.s[start..sc.pos]; + if (sc.pos < sc.s.len) sc.pos += 1; + return .{ .phrase = .{ .text = text } }; + }, + '*' => { + sc.pos += 1; + return .all; + }, + 0 => return error.QuerySyntax, + else => { + const start = sc.pos; + while (sc.pos < sc.s.len and !isBareStop(sc.s[sc.pos])) sc.pos += 1; + if (sc.pos == start) return error.QuerySyntax; + return .{ .term = .{ .text = sc.s[start..sc.pos] } }; + }, + } +} + +fn parseFieldPredicate(allocator: Allocator, sc: *Scanner) ParseError!Node { + sc.pos += 1; // '@' + const fstart = sc.pos; + while (sc.pos < sc.s.len and sc.s[sc.pos] != ':' and !std.ascii.isWhitespace(sc.s[sc.pos])) sc.pos += 1; + const field = sc.s[fstart..sc.pos]; + if (field.len == 0 or sc.pos >= sc.s.len or sc.s[sc.pos] != ':') return error.QuerySyntax; + sc.pos += 1; + + if (sc.peek() == '{') { + sc.pos += 1; + var tags: std.ArrayListUnmanaged([]const u8) = .empty; + defer tags.deinit(allocator); + while (sc.pos < sc.s.len and sc.s[sc.pos] != '}') { + const tstart = sc.pos; + while (sc.pos < sc.s.len and sc.s[sc.pos] != '|' and sc.s[sc.pos] != '}') sc.pos += 1; + if (sc.pos > tstart) try tags.append(allocator, sc.s[tstart..sc.pos]); + if (sc.pos < sc.s.len and sc.s[sc.pos] == '|') sc.pos += 1; + } + if (sc.pos < sc.s.len) sc.pos += 1; // '}' + return .{ .tag = .{ .field = field, .tags = try tags.toOwnedSlice(allocator) } }; + } + + if (sc.peek() == '[') { + sc.pos += 1; + const min = try parseBound(sc); + const max = try parseBound(sc); + sc.skipWs(); + if (sc.peek() == ']') sc.pos += 1; + return .{ .numeric = .{ .field = field, .min = min, .max = max } }; + } + + const tstart = sc.pos; + while (sc.pos < sc.s.len and !isBareStop(sc.s[sc.pos])) sc.pos += 1; + const term = sc.s[tstart..sc.pos]; + if (term.len == 0) return error.QuerySyntax; + return .{ .field_term = .{ .field = field, .term = term } }; +} + +fn parseBound(sc: *Scanner) ParseError!?f64 { + if (sc.peek() == '(') sc.pos += 1; + if (std.mem.startsWith(u8, sc.s[sc.pos..], "-inf") or std.mem.startsWith(u8, sc.s[sc.pos..], "+inf")) { + sc.pos += 4; + return null; + } + sc.skipWs(); + const start = sc.pos; + while (sc.pos < sc.s.len and !std.ascii.isWhitespace(sc.s[sc.pos]) and sc.s[sc.pos] != ']') sc.pos += 1; + const num_s = sc.s[start..sc.pos]; + return std.fmt.parseFloat(f64, num_s) catch return error.QuerySyntax; +} + +fn parseKnn(s: []const u8) ParseError!Knn { + var sc = Scanner{ .s = s }; + if (sc.peek() == '[') sc.pos += 1; + sc.skipWs(); + if (!std.mem.startsWith(u8, sc.s[sc.pos..], "KNN")) return error.QuerySyntax; + sc.pos += 3; + sc.skipWs(); + const k = try parseUint(&sc); + sc.skipWs(); + if (sc.peek() != '@') return error.QuerySyntax; + sc.pos += 1; + const fstart = sc.pos; + while (sc.pos < sc.s.len and !std.ascii.isWhitespace(sc.s[sc.pos]) and sc.s[sc.pos] != '$') sc.pos += 1; + const field = sc.s[fstart..sc.pos]; + sc.skipWs(); + if (sc.peek() != '$') return error.QuerySyntax; + sc.pos += 1; + const pstart = sc.pos; + while (sc.pos < sc.s.len and !std.ascii.isWhitespace(sc.s[sc.pos]) and sc.s[sc.pos] != ']') sc.pos += 1; + const param = sc.s[pstart..sc.pos]; + + if (k == 0 or field.len == 0 or param.len == 0) return error.QuerySyntax; + return .{ .field = field, .k = k, .vec_param = param }; +} + +fn parseUint(sc: *Scanner) ParseError!usize { + const start = sc.pos; + while (sc.pos < sc.s.len and std.ascii.isDigit(sc.s[sc.pos])) sc.pos += 1; + if (sc.pos == start) return error.QuerySyntax; + return std.fmt.parseInt(usize, sc.s[start..sc.pos], 10) catch return error.QuerySyntax; +} + +inline fn isBareStop(c: u8) bool { + return std.ascii.isWhitespace(c) or c == '|' or c == ')' or c == '('; +} + +fn freeNode(allocator: Allocator, node: *Node) void { + switch (node.*) { + .disjunction => |children| { + for (children) |*child| freeNode(allocator, child); + allocator.free(children); + }, + .conjunction => |children| { + for (children) |*child| freeNode(allocator, child); + allocator.free(children); + }, + .not => |child| { + freeNode(allocator, child); + allocator.destroy(child); + }, + .tag => |t| allocator.free(t.tags), + else => {}, + } +} + +const testing = std.testing; + +fn parseTest(allocator: Allocator, q: []const u8) !Query { + return parse(allocator, q); +} + +test "parse bare term" { + var q = try parseTest(testing.allocator, "hello"); + defer q.deinit(); + try testing.expectEqual(std.meta.activeTag(q.filter), .term); + try testing.expectEqualStrings("hello", q.filter.term.text); +} + +test "parse implicit AND of terms" { + var q = try parseTest(testing.allocator, "hello world"); + defer q.deinit(); + try testing.expectEqual(std.meta.activeTag(q.filter), .conjunction); + try testing.expectEqual(@as(usize, 2), q.filter.conjunction.len); +} + +test "parse OR" { + var q = try parseTest(testing.allocator, "foo|bar"); + defer q.deinit(); + try testing.expectEqual(std.meta.activeTag(q.filter), .disjunction); + try testing.expectEqual(@as(usize, 2), q.filter.disjunction.len); +} + +test "parse fielded term" { + var q = try parseTest(testing.allocator, "@title:zig"); + defer q.deinit(); + try testing.expectEqual(std.meta.activeTag(q.filter), .field_term); + try testing.expectEqualStrings("title", q.filter.field_term.field); + try testing.expectEqualStrings("zig", q.filter.field_term.term); +} + +test "parse phrase" { + var q = try parseTest(testing.allocator, "\"quick fox\""); + defer q.deinit(); + try testing.expectEqual(std.meta.activeTag(q.filter), .phrase); + try testing.expectEqualStrings("quick fox", q.filter.phrase.text); +} + +test "parse numeric range" { + var q = try parseTest(testing.allocator, "@price:[10 50]"); + defer q.deinit(); + try testing.expectEqual(std.meta.activeTag(q.filter), .numeric); + try testing.expectApproxEqAbs(@as(f64, 10), q.filter.numeric.min.?, 0.001); + try testing.expectApproxEqAbs(@as(f64, 50), q.filter.numeric.max.?, 0.001); +} + +test "parse numeric open range" { + var q = try parseTest(testing.allocator, "@price:[-inf 20]"); + defer q.deinit(); + try testing.expectEqual(std.meta.activeTag(q.filter), .numeric); + try testing.expect(q.filter.numeric.min == null); + try testing.expectApproxEqAbs(@as(f64, 20), q.filter.numeric.max.?, 0.001); +} + +test "parse tag set" { + var q = try parseTest(testing.allocator, "@color:{red|blue}"); + defer q.deinit(); + try testing.expectEqual(std.meta.activeTag(q.filter), .tag); + try testing.expectEqual(@as(usize, 2), q.filter.tag.tags.len); + try testing.expectEqualStrings("red", q.filter.tag.tags[0]); + try testing.expectEqualStrings("blue", q.filter.tag.tags[1]); +} + +test "parse negation" { + var q = try parseTest(testing.allocator, "-foo"); + defer q.deinit(); + try testing.expectEqual(std.meta.activeTag(q.filter), .not); + try testing.expectEqual(std.meta.activeTag(q.filter.not.*), .term); +} + +test "parse all" { + var q = try parseTest(testing.allocator, "*"); + defer q.deinit(); + try testing.expectEqual(std.meta.activeTag(q.filter), .all); +} + +test "parse knn suffix" { + var q = try parseTest(testing.allocator, "*=>[KNN 5 @v $B]"); + defer q.deinit(); + try testing.expectEqual(std.meta.activeTag(q.filter), .all); + try testing.expect(q.knn != null); + try testing.expectEqual(@as(usize, 5), q.knn.?.k); + try testing.expectEqualStrings("v", q.knn.?.field); + try testing.expectEqualStrings("B", q.knn.?.vec_param); +} + +test "parse knn with filter" { + var q = try parseTest(testing.allocator, "@title:foo=>[KNN 3 @v $Q]"); + defer q.deinit(); + try testing.expectEqual(std.meta.activeTag(q.filter), .field_term); + try testing.expect(q.knn != null); + try testing.expectEqual(@as(usize, 3), q.knn.?.k); +} diff --git a/src/search/search.zig b/src/search/search.zig new file mode 100644 index 0000000..25182da --- /dev/null +++ b/src/search/search.zig @@ -0,0 +1,306 @@ +const std = @import("std"); +const SearchIndex = @import("index.zig").SearchIndex; +const query_mod = @import("query.zig"); +const Node = query_mod.Node; +const Tokenizer = @import("tokenize.zig").Tokenizer; + +pub const Allocator = std.mem.Allocator; + +pub const SearchOptions = struct { + no_content: bool = false, + with_scores: bool = false, + sort_by: ?[]const u8 = null, + sort_asc: bool = false, + limit_offset: usize = 0, + limit_count: usize = 10, + return_fields: ?[]const []const u8 = null, +}; + +pub const Result = struct { + doc_id: []const u8, + internal_id: u64, + score: f64, + sort_key: ?f64 = null, +}; + +const DocSet = std.AutoHashMapUnmanaged(u64, void); + +const QueryTerm = struct { + field: ?u16, + text: []const u8, +}; + +/// Runs a search query against `index`. `params` supplies values for `$name` +/// references (e.g. the KNN query vector blob). +pub fn search( + allocator: Allocator, + index: *const SearchIndex, + query_str: []const u8, + params: ?*const std.StringHashMapUnmanaged([]const u8), + opts: *const SearchOptions, +) !std.ArrayListUnmanaged(Result) { + var q = try query_mod.parse(allocator, query_str); + defer q.deinit(); + + var candidates: DocSet = .empty; + defer candidates.deinit(allocator); + try collectCandidates(allocator, index, &q.filter, &candidates); + + var results: std.ArrayListUnmanaged(Result) = .empty; + errdefer results.deinit(allocator); + + if (q.knn) |knn| { + const field_idx = index.fieldIndex(knn.field) orelse return error.UnknownField; + const vi = index.getVector(field_idx) orelse return error.UnknownField; + const blob = params.?.get(knn.vec_param) orelse return error.MissingParam; + + var knn_res = try vi.knnFiltered(blob, knn.k, &candidates); + defer knn_res.deinit(allocator); + for (knn_res.items) |kr| { + const doc_id = findDocId(index, kr.doc_id) orelse continue; + try results.append(allocator, .{ + .doc_id = doc_id, + .internal_id = kr.doc_id, + .score = kr.score, + }); + } + } else { + var terms: std.ArrayListUnmanaged(QueryTerm) = .empty; + defer terms.deinit(allocator); + try collectTerms(allocator, index, &q.filter, &terms); + + var it = candidates.iterator(); + while (it.next()) |e| { + const internal_id = e.key_ptr.*; + const doc_id = findDocId(index, internal_id) orelse continue; + const score = scoreDoc(index, terms.items, internal_id); + try results.append(allocator, .{ + .doc_id = doc_id, + .internal_id = internal_id, + .score = score, + }); + } + + try applySort(index, &results, opts); + } + + // LIMIT offset count + if (opts.limit_offset >= results.items.len) { + results.clearRetainingCapacity(); + } else { + const end = @min(results.items.len, opts.limit_offset + opts.limit_count); + const keep = results.items[opts.limit_offset..end]; + // Move the kept window to the front (items overlap within the slice). + std.mem.copyForwards(Result, results.items[0 .. end - opts.limit_offset], keep); + results.shrinkRetainingCapacity(end - opts.limit_offset); + } + + return results; +} + +fn collectCandidates(allocator: Allocator, index: *const SearchIndex, node: *const Node, set: *DocSet) !void { + switch (node.*) { + .all => { + var it = index.docs.iterator(); + while (it.next()) |e| try set.put(allocator, e.value_ptr.internal_id, {}); + }, + .term => |t| { + for (index.fields, 0..) |f, i| { + if (f.field_type != .text) continue; + if (index.getInverted(@intCast(i))) |ii| { + if (ii.postings(t.text)) |posts| { + for (posts) |p| try set.put(allocator, p.doc_id, {}); + } + } + } + }, + .phrase => |p| { + // AND of the phrase's words across all text fields. + var tokenizer = Tokenizer.init(allocator); + defer tokenizer.deinit(); + try tokenizer.tokenize(p.text); + for (tokenizer.tokens.items) |word| { + var word_set: DocSet = .empty; + defer word_set.deinit(allocator); + for (index.fields, 0..) |f, i| { + if (f.field_type != .text) continue; + if (index.getInverted(@intCast(i))) |ii| { + if (ii.postings(word)) |posts| { + for (posts) |post| try word_set.put(allocator, post.doc_id, {}); + } + } + } + if (set.count() == 0) { + // First word: adopt the word set. + var it = word_set.iterator(); + while (it.next()) |e| try set.put(allocator, e.key_ptr.*, {}); + } else { + try intersectWith(allocator, set, &word_set); + } + } + }, + .field_term => |ft| { + const field_idx = index.fieldIndex(ft.field) orelse return error.UnknownField; + const ii = index.getInverted(field_idx) orelse return error.UnknownField; + if (ii.postings(ft.term)) |posts| { + for (posts) |p| try set.put(allocator, p.doc_id, {}); + } + }, + .numeric => |n| { + const field_idx = index.fieldIndex(n.field) orelse return error.UnknownField; + const ni = index.getNumeric(field_idx) orelse return error.UnknownField; + var out: std.ArrayListUnmanaged(u64) = .empty; + defer out.deinit(allocator); + try ni.range(n.min, n.max, &out); + for (out.items) |id| try set.put(allocator, id, {}); + }, + .tag => |t| { + const field_idx = index.fieldIndex(t.field) orelse return error.UnknownField; + const ti = index.getTag(field_idx) orelse return error.UnknownField; + for (t.tags) |tag| { + if (ti.docIds(tag)) |ids| { + var it = ids.iterator(); + while (it.next()) |e| try set.put(allocator, e.key_ptr.*, {}); + } + } + }, + .conjunction => |children| { + // Start from all docs, intersect each child set. + var all: DocSet = .empty; + defer all.deinit(allocator); + var it = index.docs.iterator(); + while (it.next()) |e| try all.put(allocator, e.value_ptr.internal_id, {}); + try intersectWith(allocator, &all, set); + var out = all; + + for (children) |*child| { + var child_set: DocSet = .empty; + defer child_set.deinit(allocator); + try collectCandidates(allocator, index, child, &child_set); + try intersectWith(allocator, &out, &child_set); + } + // Copy `out` back into `set`. + set.clearRetainingCapacity(); + var oit = out.iterator(); + while (oit.next()) |e| try set.put(allocator, e.key_ptr.*, {}); + }, + .disjunction => |children| { + for (children) |*child| try collectCandidates(allocator, index, child, set); + }, + .not => |child| { + var child_set: DocSet = .empty; + defer child_set.deinit(allocator); + try collectCandidates(allocator, index, child, &child_set); + var it = index.docs.iterator(); + while (it.next()) |e| { + if (!child_set.contains(e.value_ptr.internal_id)) { + try set.put(allocator, e.value_ptr.internal_id, {}); + } + } + }, + } +} + +fn intersectWith(allocator: Allocator, a: *DocSet, b: *const DocSet) !void { + var it = a.iterator(); + var to_remove: std.ArrayListUnmanaged(u64) = .empty; + defer to_remove.deinit(allocator); + while (it.next()) |e| { + if (!b.contains(e.key_ptr.*)) try to_remove.append(allocator, e.key_ptr.*); + } + for (to_remove.items) |id| _ = a.remove(id); +} + +fn collectTerms(allocator: Allocator, index: *const SearchIndex, node: *const Node, out: *std.ArrayListUnmanaged(QueryTerm)) !void { + switch (node.*) { + .term => |t| try out.append(allocator, .{ .field = null, .text = t.text }), + .field_term => |ft| { + const field_idx = index.fieldIndex(ft.field) orelse return error.UnknownField; + try out.append(allocator, .{ .field = field_idx, .text = ft.term }); + }, + .phrase => |p| { + var tokenizer = Tokenizer.init(allocator); + defer tokenizer.deinit(); + try tokenizer.tokenize(p.text); + for (tokenizer.tokens.items) |word| try out.append(allocator, .{ .field = null, .text = word }); + }, + .disjunction => |children| { + for (children) |*child| try collectTerms(allocator, index, child, out); + }, + .conjunction => |children| { + for (children) |*child| try collectTerms(allocator, index, child, out); + }, + else => {}, + } +} + +fn scoreDoc(index: *const SearchIndex, terms: []const QueryTerm, internal_id: u64) f64 { + var score: f64 = 0; + for (terms) |t| { + if (t.field) |field_idx| { + if (index.getInverted(field_idx)) |ii| { + score += ii.bm25(internal_id, ii.docFreq(t.text), ii.tf(internal_id, t.text)); + } + } else { + for (index.fields, 0..) |f, i| { + if (f.field_type != .text) continue; + if (index.getInverted(@intCast(i))) |ii| { + score += ii.bm25(internal_id, ii.docFreq(t.text), ii.tf(internal_id, t.text)); + } + } + } + } + return score; +} + +fn findDocId(index: *const SearchIndex, internal_id: u64) ?[]const u8 { + var it = index.docs.iterator(); + while (it.next()) |e| { + if (e.value_ptr.internal_id == internal_id) return e.key_ptr.*; + } + return null; +} + +const SortCtx = struct { + ascending: bool, +}; + +fn scoreLessThan(ctx: SortCtx, a: Result, b: Result) bool { + return if (ctx.ascending) a.score < b.score else a.score > b.score; +} + +fn sortKeyLessThan(ctx: SortCtx, a: Result, b: Result) bool { + const av = a.sort_key orelse -std.math.inf(f64); + const bv = b.sort_key orelse -std.math.inf(f64); + return if (ctx.ascending) av < bv else av > bv; +} + +fn applySort(index: *const SearchIndex, results: *std.ArrayListUnmanaged(Result), opts: *const SearchOptions) !void { + if (opts.sort_by) |sb| { + if (std.ascii.eqlIgnoreCase(sb, "score")) { + const ctx: SortCtx = .{ .ascending = opts.sort_asc }; + std.mem.sort(Result, results.items, ctx, scoreLessThan); + return; + } + const field_idx = index.fieldIndex(sb) orelse return error.UnknownField; + for (results.items) |*r| { + const doc = index.getDocument(r.doc_id) orelse continue; + r.sort_key = docFieldValue(doc, field_idx); + } + const ctx: SortCtx = .{ .ascending = opts.sort_asc }; + std.mem.sort(Result, results.items, ctx, sortKeyLessThan); + return; + } + const ctx: SortCtx = .{ .ascending = false }; + std.mem.sort(Result, results.items, ctx, scoreLessThan); +} + +fn docFieldValue(doc: *const @import("index.zig").Document, field_idx: u16) ?f64 { + const fv = doc.fields[field_idx] orelse return null; + return switch (fv) { + .numeric => |n| n, + .string => |s| std.fmt.parseFloat(f64, s) catch null, + }; +} + +const testing = std.testing; diff --git a/src/search/tag.zig b/src/search/tag.zig new file mode 100644 index 0000000..53cfd8f --- /dev/null +++ b/src/search/tag.zig @@ -0,0 +1,129 @@ +const std = @import("std"); + +pub const Allocator = std.mem.Allocator; + +const TagSet = std.AutoHashMapUnmanaged(u64, void); +const separator: u8 = ','; + +/// Per-field tag index. Tag values are exact-match, case-insensitive. +/// Values are comma-separated (RediSearch default SEPARATOR). +pub const TagIndex = struct { + allocator: Allocator, + map: std.StringHashMapUnmanaged(TagSet) = .empty, + + pub fn init(allocator: Allocator) TagIndex { + return .{ .allocator = allocator }; + } + + pub fn deinit(self: *TagIndex) void { + var it = self.map.iterator(); + while (it.next()) |e| { + self.allocator.free(e.key_ptr.*); + e.value_ptr.deinit(self.allocator); + } + self.map.deinit(self.allocator); + } + + pub fn add(self: *TagIndex, doc_id: u64, value: []const u8) !void { + var tags: std.ArrayListUnmanaged([]const u8) = .empty; + defer tags.deinit(self.allocator); + try splitTags(self.allocator, value, &tags); + + for (tags.items) |tag| { + const gop = try self.map.getOrPut(self.allocator, tag); + if (!gop.found_existing) { + gop.key_ptr.* = try self.allocator.dupe(u8, tag); + gop.value_ptr.* = .empty; + } + try gop.value_ptr.put(self.allocator, doc_id, {}); + } + } + + pub fn remove(self: *TagIndex, doc_id: u64, value: []const u8) void { + var tags: std.ArrayListUnmanaged([]const u8) = .empty; + defer tags.deinit(self.allocator); + splitTags(self.allocator, value, &tags) catch return; + + for (tags.items) |tag| { + if (self.map.getPtr(tag)) |set| { + _ = set.remove(doc_id); + if (set.count() == 0) { + if (self.map.fetchRemove(tag)) |kv| { + self.allocator.free(kv.key); + var tag_value = kv.value; + tag_value.deinit(self.allocator); + } + } + } + } + } + + /// Doc ids for an exact tag (case-insensitive). Null if tag absent. + pub fn docIds(self: *const TagIndex, tag: []const u8) ?*const TagSet { + return self.map.getPtr(tag); + } + + pub fn tagCount(self: *const TagIndex) usize { + return self.map.count(); + } +}; + +fn splitTags(allocator: Allocator, value: []const u8, out: *std.ArrayListUnmanaged([]const u8)) !void { + var start: usize = 0; + var i: usize = 0; + while (i <= value.len) : (i += 1) { + const is_sep = i == value.len or value[i] == separator; + if (is_sep) { + if (i > start) { + try out.append(allocator, value[start..i]); + } + start = i + 1; + } + } +} + +/// Lowercases a tag into a stack buffer for map lookups. +pub fn normalized(allocator: Allocator, tag: []const u8) ![]u8 { + const out = try allocator.dupe(u8, tag); + for (out) |*c| c.* = std.ascii.toLower(c.*); + return out; +} + +const testing = std.testing; + +test "tag add and lookup" { + var ti = TagIndex.init(testing.allocator); + defer ti.deinit(); + + try ti.add(1, "red,blue"); + try ti.add(2, "green"); + try ti.add(3, "red"); + + try testing.expectEqual(@as(usize, 2), ti.docIds("red").?.count()); + try testing.expectEqual(@as(usize, 1), ti.docIds("blue").?.count()); + try testing.expect(ti.docIds("red").?.contains(1)); + try testing.expect(ti.docIds("red").?.contains(3)); +} + +test "tag remove" { + var ti = TagIndex.init(testing.allocator); + defer ti.deinit(); + + try ti.add(1, "red"); + try ti.add(2, "red"); + ti.remove(1, "red"); + + try testing.expectEqual(@as(usize, 1), ti.docIds("red").?.count()); + try testing.expect(ti.docIds("red").?.contains(2)); + try testing.expect(!ti.docIds("red").?.contains(1)); +} + +test "tag empty after removing last doc" { + var ti = TagIndex.init(testing.allocator); + defer ti.deinit(); + + try ti.add(1, "solo"); + ti.remove(1, "solo"); + try testing.expect(ti.docIds("solo") == null); + try testing.expectEqual(@as(usize, 0), ti.tagCount()); +} diff --git a/src/search/tokenize.zig b/src/search/tokenize.zig new file mode 100644 index 0000000..1fd9a84 --- /dev/null +++ b/src/search/tokenize.zig @@ -0,0 +1,190 @@ +const std = @import("std"); + +pub const Allocator = std.mem.Allocator; + +/// Bytes processed per SIMD iteration, tuned to the compile-target CPU +/// (16 on aarch64/SSE2, 32 with AVX2, 64 with AVX-512 for u8). +const VEC_LEN = std.simd.suggestVectorLength(u8) orelse 16; + +/// Widest type that can hold a VEC_LEN-bit separator mask. +const MaskInt = std.meta.Int(.unsigned, VEC_LEN); + +/// Splits text into lowercased tokens, splitting on non-alphanumeric bytes. +/// UTF-8 bytes (>= 0x80) are treated as token characters so unicode words +/// still tokenize as whole units (ASCII-only lowercasing for now). +/// +/// The separator classification is SIMD: VEC_LEN bytes are loaded at once into +/// a `@Vector(VEC_LEN, u8)`, classified with vector compares, and the +/// resulting flag vector is reduced to a bitmask via `@select` + +/// `@reduce(.Add)`. +pub const Tokenizer = struct { + allocator: Allocator, + tokens: std.ArrayListUnmanaged([]const u8) = .empty, + + pub fn init(allocator: Allocator) Tokenizer { + return .{ .allocator = allocator }; + } + + pub fn deinit(self: *Tokenizer) void { + self.reset(); + self.tokens.deinit(self.allocator); + } + + pub fn reset(self: *Tokenizer) void { + for (self.tokens.items) |t| self.allocator.free(t); + self.tokens.clearRetainingCapacity(); + } + + /// Tokenizes `input`. Resulting tokens are owned by this Tokenizer and + /// freed on `reset`/`deinit`. + pub fn tokenize(self: *Tokenizer, input: []const u8) !void { + self.reset(); + var start: usize = 0; + var i: usize = 0; + while (i < input.len) { + const chunk_len = @min(input.len - i, VEC_LEN); + + var chunk: [VEC_LEN]u8 = undefined; + @memset(&chunk, 0); + @memcpy(chunk[0..chunk_len], input[i..][0..chunk_len]); + + const flags: @Vector(VEC_LEN, u8) = sepMask(chunk); + const mask: MaskInt = toMask(flags); + + var b: usize = 0; + while (b < VEC_LEN) : (b += 1) { + const abs_pos = i + b; + // Zero-padded positions beyond the input are treated as separators + // so the final token is flushed at the end of the input. + const sep = if (abs_pos < input.len) (mask & (@as(MaskInt, 1) << @intCast(b))) != 0 else true; + if (sep) { + if (abs_pos > start) { + const token = try self.allocator.dupe(u8, input[start..abs_pos]); + for (token) |*c| c.* = std.ascii.toLower(c.*); + try self.tokens.append(self.allocator, token); + } + start = abs_pos + 1; + } + } + i += chunk_len; + } + } + + /// Counts unique tokens and their term frequencies via linear scan. + /// Returned slices borrow from `self.tokens`. + pub fn uniqueTokens(self: *const Tokenizer) std.ArrayListUnmanaged(TokenCount) { + var uniques: std.ArrayListUnmanaged(TokenCount) = .empty; + for (self.tokens.items) |tok| { + var found = false; + for (uniques.items) |*u| { + if (std.mem.eql(u8, u.token, tok)) { + u.count += 1; + found = true; + break; + } + } + if (!found) { + uniques.append(self.allocator, .{ .token = tok, .count = 1 }) catch break; + } + } + return uniques; + } +}; + +pub const TokenCount = struct { + token: []const u8, + count: u32, +}; + +/// Classifies each of the VEC_LEN bytes as a separator using vector compares. +/// Returns a `@Vector(VEC_LEN, u8)` of 0/1 flags: 1 at separator positions. +/// (Uses `@select` + bitwise ops because Zig's `and`/`or` keywords do not +/// operate element-wise on boolean vectors.) +inline fn sepMask(chunk: [VEC_LEN]u8) @Vector(VEC_LEN, u8) { + const vec: @Vector(VEC_LEN, u8) = chunk; + const one = @as(@Vector(VEC_LEN, u8), @splat(1)); + const zero = @as(@Vector(VEC_LEN, u8), @splat(0)); + const v_zero = @as(@Vector(VEC_LEN, u8), @splat('0')); + const v_nine = @as(@Vector(VEC_LEN, u8), @splat('9')); + const v_cap_a = @as(@Vector(VEC_LEN, u8), @splat('A')); + const v_cap_z = @as(@Vector(VEC_LEN, u8), @splat('Z')); + const v_low_a = @as(@Vector(VEC_LEN, u8), @splat('a')); + const v_low_z = @as(@Vector(VEC_LEN, u8), @splat('z')); + const v_high = @as(@Vector(VEC_LEN, u8), @splat(0x80)); + + const is_digit = @select(u8, vec >= v_zero, one, zero) & @select(u8, vec <= v_nine, one, zero); + const is_upper = @select(u8, vec >= v_cap_a, one, zero) & @select(u8, vec <= v_cap_z, one, zero); + const is_lower = @select(u8, vec >= v_low_a, one, zero) & @select(u8, vec <= v_low_z, one, zero); + const is_high = @select(u8, vec >= v_high, one, zero); + + const is_token = is_digit | is_upper | is_lower | is_high; + return (is_token ^ one) & one; +} + +/// Reduces a 0/1 flag vector to a bitmask: bit k set means position k is a +/// separator. Uses `std.simd.iota` to build the {1<<0, 1<<1, ...} weights. +inline fn toMask(flags: @Vector(VEC_LEN, u8)) MaskInt { + const weights: @Vector(VEC_LEN, MaskInt) = @as(@Vector(VEC_LEN, MaskInt), @splat(1)) << std.simd.iota(MaskInt, VEC_LEN); + const w: @Vector(VEC_LEN, MaskInt) = @intCast(flags); + return @reduce(.Add, w * weights); +} + +const testing = std.testing; + +test "tokenize splits on non-alphanumeric and lowercases" { + var t = Tokenizer.init(testing.allocator); + defer t.deinit(); + try t.tokenize("Hello, World! Redis-Rocks"); + try testing.expectEqual(@as(usize, 4), t.tokens.items.len); + try testing.expectEqualStrings("hello", t.tokens.items[0]); + try testing.expectEqualStrings("world", t.tokens.items[1]); + try testing.expectEqualStrings("redis", t.tokens.items[2]); + try testing.expectEqualStrings("rocks", t.tokens.items[3]); +} + +test "tokenize counts term frequencies" { + var t = Tokenizer.init(testing.allocator); + defer t.deinit(); + try t.tokenize("foo bar foo baz foo"); + var uniques = t.uniqueTokens(); + defer uniques.deinit(testing.allocator); + try testing.expectEqual(@as(usize, 3), uniques.items.len); + var foo_count: u32 = 0; + for (uniques.items) |u| { + if (std.mem.eql(u8, u.token, "foo")) foo_count = u.count; + } + try testing.expectEqual(@as(u32, 3), foo_count); +} + +test "tokenize keeps unicode words whole" { + var t = Tokenizer.init(testing.allocator); + defer t.deinit(); + try t.tokenize("héllo wörld"); + try testing.expectEqual(@as(usize, 2), t.tokens.items.len); + try testing.expectEqualStrings("h\u{e9}llo", t.tokens.items[0]); +} + +test "tokenize handles input longer than one vector" { + var t = Tokenizer.init(testing.allocator); + defer t.deinit(); + try t.tokenize("alpha beta gamma delta epsilon"); + try testing.expectEqual(@as(usize, 5), t.tokens.items.len); + try testing.expectEqualStrings("alpha", t.tokens.items[0]); + try testing.expectEqualStrings("epsilon", t.tokens.items[4]); +} + +test "tokenize handles trailing and leading separators" { + var t = Tokenizer.init(testing.allocator); + defer t.deinit(); + try t.tokenize(" spaced out "); + try testing.expectEqual(@as(usize, 2), t.tokens.items.len); + try testing.expectEqualStrings("spaced", t.tokens.items[0]); + try testing.expectEqualStrings("out", t.tokens.items[1]); +} + +test "tokenize empty input" { + var t = Tokenizer.init(testing.allocator); + defer t.deinit(); + try t.tokenize(""); + try testing.expectEqual(@as(usize, 0), t.tokens.items.len); +} diff --git a/src/search/vector.zig b/src/search/vector.zig new file mode 100644 index 0000000..0cbeb8d --- /dev/null +++ b/src/search/vector.zig @@ -0,0 +1,300 @@ +const std = @import("std"); + +pub const Allocator = std.mem.Allocator; + +pub const VectorType = enum { + float32, + float64, + + pub fn fromSlice(s: []const u8) ?VectorType { + if (std.ascii.eqlIgnoreCase(s, "FLOAT32")) return .float32; + if (std.ascii.eqlIgnoreCase(s, "FLOAT64")) return .float64; + return null; + } + + pub fn elemSize(self: VectorType) usize { + return switch (self) { + .float32 => 4, + .float64 => 8, + }; + } +}; + +pub const DistanceMetric = enum { + l2, + ip, + cosine, + + pub fn fromSlice(s: []const u8) ?DistanceMetric { + if (std.ascii.eqlIgnoreCase(s, "L2")) return .l2; + if (std.ascii.eqlIgnoreCase(s, "IP")) return .ip; + if (std.ascii.eqlIgnoreCase(s, "COSINE")) return .cosine; + return null; + } +}; + +pub const VectorParams = struct { + dim: u16, + typ: VectorType, + metric: DistanceMetric, +}; + +pub const VecEntry = struct { + doc_id: u64, + data: []u8, // owned raw little-endian blob +}; + +/// FLAT (exact) vector index: stores vectors in insertion order and computes +/// distances to every entry on search. O(N*D) exact KNN. +pub const VectorIndex = struct { + allocator: Allocator, + params: VectorParams, + entries: std.ArrayListUnmanaged(VecEntry) = .empty, + + pub fn init(allocator: Allocator, params: VectorParams) VectorIndex { + return .{ .allocator = allocator, .params = params }; + } + + pub fn deinit(self: *VectorIndex) void { + for (self.entries.items) |e| self.allocator.free(e.data); + self.entries.deinit(self.allocator); + } + + pub fn count(self: *const VectorIndex) usize { + return self.entries.items.len; + } + + pub fn add(self: *VectorIndex, doc_id: u64, blob: []const u8) !void { + if (blob.len != @as(usize, self.params.dim) * self.params.typ.elemSize()) { + return error.VectorDimensionMismatch; + } + const owned = try self.allocator.dupe(u8, blob); + errdefer self.allocator.free(owned); + try self.entries.append(self.allocator, .{ .doc_id = doc_id, .data = owned }); + } + + pub fn remove(self: *VectorIndex, doc_id: u64) void { + for (self.entries.items, 0..) |e, i| { + if (e.doc_id == doc_id) { + self.allocator.free(e.data); + _ = self.entries.swapRemove(i); + return; + } + } + } + + /// Higher = more similar. L2 is negated so all metrics share "maximize". + pub fn similarity(self: *const VectorIndex, a: []const u8, b: []const u8) f64 { + const dim: usize = self.params.dim; + switch (self.params.typ) { + .float32 => return similarityFor(f32, self.params.metric, a, b, dim), + .float64 => return similarityFor(f64, self.params.metric, a, b, dim), + } + } + + /// Returns the k nearest entries (sorted best-first) with their scores. + pub fn knn(self: *const VectorIndex, query: []const u8, k: usize) !std.ArrayListUnmanaged(KnnResult) { + var results: std.ArrayListUnmanaged(KnnResult) = .empty; + errdefer results.deinit(self.allocator); + + for (self.entries.items) |e| { + const sim = self.similarity(query, e.data); + try results.append(self.allocator, .{ .doc_id = e.doc_id, .score = sim }); + } + + std.mem.sort(KnnResult, results.items, {}, struct { + fn lt(_: void, a: KnnResult, b: KnnResult) bool { + return a.score > b.score; + } + }.lt); + + if (results.items.len > k) results.shrinkRetainingCapacity(k); + return results; + } + + /// KNN restricted to doc ids present in `candidates`. Returns top-k sorted + /// best-first with their scores. + pub fn knnFiltered( + self: *const VectorIndex, + query: []const u8, + k: usize, + candidates: *const std.AutoHashMapUnmanaged(u64, void), + ) !std.ArrayListUnmanaged(KnnResult) { + var results: std.ArrayListUnmanaged(KnnResult) = .empty; + errdefer results.deinit(self.allocator); + + for (self.entries.items) |e| { + if (!candidates.contains(e.doc_id)) continue; + const sim = self.similarity(query, e.data); + try results.append(self.allocator, .{ .doc_id = e.doc_id, .score = sim }); + } + + std.mem.sort(KnnResult, results.items, {}, struct { + fn lt(_: void, a: KnnResult, b: KnnResult) bool { + return a.score > b.score; + } + }.lt); + + if (results.items.len > k) results.shrinkRetainingCapacity(k); + return results; + } +}; + +pub const KnnResult = struct { + doc_id: u64, + score: f64, +}; + +inline fn readFloat(comptime T: type, blob: []const u8, i: usize) T { + var tmp: [@sizeOf(T)]u8 align(@alignOf(T)) = undefined; + @memcpy(tmp[0..], blob[i * @sizeOf(T) ..][0..@sizeOf(T)]); + return std.mem.bytesToValue(T, &tmp); +} + +/// Native SIMD block size, in elements of `T`, for the compile-target CPU +/// (e.g. 4 on aarch64/SSE2 for f32, 8 with AVX2, 16 with AVX-512). +fn vecLen(comptime T: type) comptime_int { + return std.simd.suggestVectorLength(T) orelse 4; +} + +/// Loads vecLen(T) floats (unaligned byte blob) into a vector. +inline fn loadVec(comptime T: type, blob: []const u8, offset: usize) @Vector(vecLen(T), T) { + const len = comptime vecLen(T); + const bytes = blob[offset * @sizeOf(T) ..][0 .. len * @sizeOf(T)]; + var buf: [len * @sizeOf(T)]u8 align(@alignOf(T)) = undefined; + @memcpy(buf[0..], bytes); + return @bitCast(buf); +} + +fn dot(comptime T: type, a: []const u8, b: []const u8, dim: usize) f64 { + const len = comptime vecLen(T); + var total: f64 = 0; + var i: usize = 0; + while (i + len <= dim) : (i += len) { + const av = loadVec(T, a, i); + const bv = loadVec(T, b, i); + total += @as(f64, @floatCast(@reduce(.Add, av * bv))); + } + while (i < dim) : (i += 1) { + total += @as(f64, @floatCast(readFloat(T, a, i))) * @as(f64, @floatCast(readFloat(T, b, i))); + } + return total; +} + +fn norm(comptime T: type, a: []const u8, dim: usize) f64 { + const len = comptime vecLen(T); + var total: f64 = 0; + var i: usize = 0; + while (i + len <= dim) : (i += len) { + const av = loadVec(T, a, i); + total += @as(f64, @floatCast(@reduce(.Add, av * av))); + } + while (i < dim) : (i += 1) { + const av = readFloat(T, a, i); + total += @as(f64, @floatCast(av)) * @as(f64, @floatCast(av)); + } + return @sqrt(total); +} + +fn l2Sq(comptime T: type, a: []const u8, b: []const u8, dim: usize) f64 { + const len = comptime vecLen(T); + var total: f64 = 0; + var i: usize = 0; + while (i + len <= dim) : (i += len) { + const av = loadVec(T, a, i); + const bv = loadVec(T, b, i); + const diff = av - bv; + total += @as(f64, @floatCast(@reduce(.Add, diff * diff))); + } + while (i < dim) : (i += 1) { + const av = readFloat(T, a, i); + const bv = readFloat(T, b, i); + const d = @as(f64, @floatCast(av - bv)); + total += d * d; + } + return total; +} + +fn similarityFor(comptime T: type, metric: DistanceMetric, a: []const u8, b: []const u8, dim: usize) f64 { + switch (metric) { + .l2 => return -l2Sq(T, a, b, dim), + .ip => return dot(T, a, b, dim), + .cosine => { + const d = dot(T, a, b, dim); + const na = norm(T, a, dim); + const nb = norm(T, b, dim); + if (na == 0 or nb == 0) return 0; + return d / (na * nb); + }, + } +} + +const testing = std.testing; + +fn vecF32(values: []const f32) []const u8 { + return std.mem.sliceAsBytes(values); +} + +test "vector l2 similarity negates distance" { + var vi = VectorIndex.init(testing.allocator, .{ + .dim = 2, + .typ = .float32, + .metric = .l2, + }); + defer vi.deinit(); + + try vi.add(1, vecF32(&[_]f32{ 1, 1 })); + try vi.add(2, vecF32(&[_]f32{ 0, 0 })); + + const query = vecF32(&[_]f32{ 1, 1 }); + const s1 = vi.similarity(query, vi.entries.items[0].data); + const s2 = vi.similarity(query, vi.entries.items[1].data); + try testing.expect(s1 > s2); // identical vector scores highest +} + +test "vector knn returns nearest first" { + var vi = VectorIndex.init(testing.allocator, .{ + .dim = 2, + .typ = .float32, + .metric = .l2, + }); + defer vi.deinit(); + + try vi.add(1, vecF32(&[_]f32{ 0, 0 })); + try vi.add(2, vecF32(&[_]f32{ 10, 10 })); + try vi.add(3, vecF32(&[_]f32{ 1, 1 })); + + const query = vecF32(&[_]f32{ 1, 1 }); + var results = try vi.knn(query, 2); + defer results.deinit(testing.allocator); + + try testing.expectEqual(@as(usize, 2), results.items.len); + try testing.expectEqual(@as(u64, 3), results.items[0].doc_id); // closest + try testing.expectEqual(@as(u64, 1), results.items[1].doc_id); +} + +test "vector rejects wrong blob length" { + var vi = VectorIndex.init(testing.allocator, .{ + .dim = 2, + .typ = .float32, + .metric = .l2, + }); + defer vi.deinit(); + + try testing.expectError(error.VectorDimensionMismatch, vi.add(1, vecF32(&[_]f32{1}))); +} + +test "vector remove" { + var vi = VectorIndex.init(testing.allocator, .{ + .dim = 1, + .typ = .float32, + .metric = .l2, + }); + defer vi.deinit(); + + try vi.add(1, vecF32(&[_]f32{1})); + try vi.add(2, vecF32(&[_]f32{2})); + vi.remove(1); + try testing.expectEqual(@as(usize, 1), vi.count()); + try testing.expectEqual(@as(u64, 2), vi.entries.items[0].doc_id); +} diff --git a/src/server.zig b/src/server.zig index 1c103e9..2b528a4 100644 --- a/src/server.zig +++ b/src/server.zig @@ -5,7 +5,7 @@ const ClientMailbox = @import("client_mailbox.zig").ClientMailbox; const MessageNode = @import("client_mailbox.zig").MessageNode; const freeMessageList = @import("client_mailbox.zig").freeMessageList; const CommandRegistry = @import("./commands/registry.zig").CommandRegistry; -const command_init = @import("./commands/init.zig"); +const command_init = @import("./commands/init_registry.zig"); const Reader = @import("./rdb/zdb.zig").Reader; const Store = @import("store.zig").Store; const pubsub = @import("./commands/pubsub.zig"); diff --git a/src/store.zig b/src/store.zig index f5b9c0a..6a22170 100644 --- a/src/store.zig +++ b/src/store.zig @@ -7,6 +7,7 @@ const Clock = @import("clock.zig"); const Config = @import("config.zig"); const string_match = @import("./util/string_match.zig").string_match; const ScalableBloomFilter = @import("./bloom/bloom.zig").BloomFilter; +const SearchIndex = @import("./search/index.zig").SearchIndex; const assert = std.debug.assert; @@ -24,6 +25,7 @@ pub const ValueType = enum(u8) { short_string, time_series, bloom_filter, + search_index, pub fn toRdbOpcode(self: ValueType) u8 { return @intFromEnum(self); @@ -44,6 +46,7 @@ pub const ZedisValue = union(ValueType) { short_string: ShortString, time_series: *TimeSeries, bloom_filter: *ScalableBloomFilter, + search_index: *SearchIndex, }; pub const ZedisObject = struct { @@ -180,6 +183,7 @@ pub const Store = struct { .short_string => {}, .time_series => {}, .bloom_filter => {}, + .search_index => {}, } return cloned; } @@ -203,6 +207,10 @@ pub const Store = struct { bf_ptr.deinit(); self.allocator.destroy(bf_ptr); }, + .search_index => |si_ptr| { + si_ptr.deinit(); + self.allocator.destroy(si_ptr); + }, } } @@ -459,6 +467,27 @@ pub const Store = struct { } } + pub fn createSearchIndex(self: *Store, key: []const u8, si: SearchIndex) !void { + assert(key.len > 0); + if (self.exists(key)) return error.AlreadyExists; + + const si_ptr = try self.allocator.create(SearchIndex); + errdefer self.allocator.destroy(si_ptr); + + si_ptr.* = si; + try self.putObject(key, .{ .value = .{ .search_index = si_ptr } }); + } + + pub fn getSearchIndex(self: *Store, key: []const u8) !?*SearchIndex { + assert(key.len > 0); + const entry = self.resolveEntry(key, true) orelse return null; + + switch (entry.object.value) { + .search_index => |si_ptr| return si_ptr, + else => return error.WrongType, + } + } + pub fn getSetList(self: *Store, key: []const u8) !*ZedisList { const list = try self.getList(key); if (list == null) return try self.createList(key);