diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 194bf357..60c96adc 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -17,6 +17,11 @@ concurrency: env: CARGO_TERM_COLOR: always + # The size of the shipped code, rounded up to 500 lines, and the number of tests, rounded up + # to 10. A change that crosses either number edits it here, in the same change, where its + # reader sees that the code grew or shrank. Neither ever passes 150,000 lines or 1,000 tests. + SHIPPED_LINES: 11000 + TESTS: 180 jobs: lint: @@ -26,8 +31,21 @@ jobs: - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4 - run: cargo fmt --manifest-path v2/Cargo.toml --all --check - run: cargo clippy --manifest-path v2/Cargo.toml --workspace --all-targets --locked -- -D warnings - - name: Shipped lines stay under 150,000 - run: test "$(find v2/contextdb/src v2/contextdb-cli/src \( -name '*.rs' -o -name '*.pest' \) ! -name tests.rs -exec cat {} + | wc -l)" -le 150000 + - name: Shipped lines are at most the stated number and within 500 of it + run: | + lines=$(find v2/contextdb/src v2/contextdb-cli/src \( -name '*.rs' -o -name '*.pest' \) ! -name tests.rs -exec cat {} + | wc -l) + echo "$lines shipped lines; the stated number is $SHIPPED_LINES" + test "$SHIPPED_LINES" -le 150000 + test "$lines" -le "$SHIPPED_LINES" + test "$lines" -gt "$((SHIPPED_LINES - 500))" + - name: No shipped file is longer than 1,500 lines + run: test -z "$(find v2/contextdb/src v2/contextdb-cli/src -name '*.rs' ! -name tests.rs -exec wc -l {} + | awk '$2 != "total" && $1 > 1500')" + - name: No test-only name in shipped code + run: test -z "$(grep -rnE --include='*.rs' --exclude=tests.rs '_for_test|test[-_]seam|feature *= *"test' v2/contextdb/src v2/contextdb-cli/src)" + - name: No ignored test + run: test -z "$(grep -rn --include='*.rs' '#\[ignore' v2/contextdb v2/contextdb-cli)" + - name: Only the four allow attributes that v2/AGENTS.md names, and no expect attribute + run: test "$(grep -rnE --include='*.rs' '#!?\[ *(allow|expect) *\(' v2/contextdb v2/contextdb-cli | wc -l)" -eq 4 test: strategy: @@ -45,9 +63,15 @@ jobs: - run: cargo test --manifest-path v2/Cargo.toml --workspace --locked --doc - name: The storage library's own tests run: cargo test --manifest-path v2/redb/Cargo.toml --locked --lib --tests - # A test is one the runner lists, or one script file. - - name: Tests stay under 1,000 - run: test "$(( $(cargo nextest list --manifest-path v2/Cargo.toml --workspace --locked -T oneline | wc -l) + $(find v2/contextdb/tests/sql -name '*.sql' | wc -l) ))" -le 1000 + # A test is one the runner lists, or one script file. Counted on Linux, where every test runs. + - name: Tests are at most the stated number and within 10 of it + if: matrix.os == 'ubuntu-latest' + run: | + tests=$(( $(cargo nextest list --manifest-path v2/Cargo.toml --workspace --locked -T oneline | wc -l) + $(find v2/contextdb/tests/sql -name '*.sql' | wc -l) )) + echo "$tests tests; the stated number is $TESTS" + test "$TESTS" -le 1000 + test "$tests" -le "$TESTS" + test "$tests" -gt "$((TESTS - 10))" minimum-rust: runs-on: ubuntu-latest diff --git a/v2/AGENTS.md b/v2/AGENTS.md index b76f75fd..ed6ac801 100644 --- a/v2/AGENTS.md +++ b/v2/AGENTS.md @@ -13,7 +13,7 @@ repeated here; every command is in [`README.md`](README.md). | `contextdb-cli/` | The `contextdb` program, a SQL shell. Every statement and command goes through the library's public calls; it never opens a store's files itself. | | `redb/` | redb 4.1.0 with the changes listed at the top of `redb/README.md`. It keeps its own manifest and lock file and is not a workspace member. | -The library is ten parts, one folder each under `contextdb/src/`: +The library is twelve parts, one folder each under `contextdb/src/`: | Part | Does | Never | |---|---|---| @@ -21,17 +21,29 @@ The library is ten parts, one folder each under `contextdb/src/`: | `system` | the time of day, through `Clock`; the one place in the library that reads a clock | decides anything | | `memory` | the budget and the reservations taken from it | allocates for another part: it counts what they are about to take and says no first | | `sql` | grammar, parser, planner: text to a statement, a statement to a plan | touches storage or knows what a table holds | -| `store` | the folder and data file, how each record is keyed, views of one moment, the reader files (one in `readers/` per reader that is inside a statement, naming the commit it reads) and the locks | interprets SQL or a table's rules | +| `store` | the folder and data file, how each record is keyed, a row's large text, JSON and vectors kept beside the row in slices (text and JSON compressed there when that saves an eighth, `store/record.rs`), views of one moment, the reader files (one in `readers/` per reader that is inside a statement, naming the commit it reads) and the locks | interprets SQL or a table's rules | | `schema` | table declarations: checking, storing, printing them back | changes a row | | `query` | answers a SELECT from one table; holds the one expression evaluator; counts the rows it reads | changes anything | | `write` | applies INSERT, UPDATE and DELETE to a change set; checks types, NOT NULL, the primary key and UNIQUE rules; keeps index entries and row counts | returns a query's result to a caller | | `database` | the handle a store is opened for writing with, its transactions, the one commit function | has a second way to commit | | `read` | a session that reads a store without holding it for writing, from any process | changes anything a writer stored | +| `vector` | the blocks of a VECTOR column (a one-byte copy of every vector), the lists of vectors not yet in a block, the column's partitions and their count, exact search over them, and sealing | decides which rows a reader may see: `query` tells it, row by row, and a row is in an answer only with the distance `query` worked out from the row's stored vector | +| `upkeep` | one cycle of background work when it is called: sealing at most 4,096 waiting vectors into blocks, in one commit; the 4,096 is `SEAL_ROWS` in `upkeep/mod.rs` | starts a thread, reads a clock, or runs unless called | + +The files of `vector/`, each as its opening comment says: + +- `slot.rs`: one vector as a block keeps it (a scale, a rounding error and a signed byte for each number), the range a search gets from it, and the distance every answer is ordered by. +- `block.rs`: the slots of as many vectors of one partition as fit in the block's bytes, each with the key of its row. +- `directory.rs`: the partitions of a VECTOR column, one record for each key tuple under which a row has a vector. +- `search.rs`: the one search: scores the rows in no block exactly, screens the slots of the selected partitions, asks for the true distance of the rows that can reach the answer, and proves the answer. +- `seal.rs`: moving the vectors on the waiting lists into blocks. +- `stage.rs`: what a commit's row changes mean for the vector part, worked out once inside the commit function. +- `tally.rs`: the partitions a transaction's statements have made and emptied, so that the statement that would give a column one partition too many is refused when it runs. ## Where the parts meet A part reaches another only through what that part makes visible: the `pub(crate)` items of its -`mod.rs`, and for `sql` its `ast` and `plan` modules. The types handed across: +`mod.rs`, and for `sql` its `ast` and `plan` modules. A statement form this version does not have is refused in `sql` (`not_supported` in `sql/mod.rs`, called by the parser and the planner), before any table is looked up. The types handed across: | Type | Defined in | Who takes it | |---|---|---| @@ -43,6 +55,14 @@ A part reaches another only through what that part makes visible: the `pub(crate | `Catalog`, `Table` | `schema` | `query`, `write`, `database` and `read` look tables up in it | | `eval`, `matching` | `query` | `write` finds the rows of an UPDATE or DELETE with `query::matching` and works out values with `query::eval` | | `QueryResult` | `query` | `database` and `read` build it and hand it to the caller | +| `PhysicalPlan::VectorSearch`; `Expr::Distance`, `Expr::Row`, `Expr::InList` | `sql` | `query` runs the search, puts its distance wherever the select list writes `<=>`, and evaluates `IN` lists and row values | +| `VectorColumn` (a VECTOR column's number, partition key, limit and mode), from `Table::vectors`, `Table::vector` and `Catalog::vector` | `schema` | `query`, `write` and `vector` | +| `staged`, `vectors` and `vector_number` on `Read`; `put_vector`, `delete_vector` and `put_vector_number` on `ChangeSet`; `VectorRecord`, `VectorNumber`, `VectorId` | `store` | `vector` reads and stages its blocks, waiting entries, partition entries and counts through them; `store` keeps them and never interprets one | +| `search`, `Nearest`, `distance` | `vector` | `query` calls `search` with the partitions the filter selects and answers it, row by row, with `distance` of the row's stored vector, or with nothing for a row that is gone, filtered out or has no vector | +| `Tally`, the partitions a transaction has started and emptied | `vector` | `database` holds one beside each change set and forgets a refused statement's part of it; `write` notes every row it stages and settles the tally as a statement's last act, which is where a statement over `MAX_PARTITIONS` is refused | +| `stage`, `seal` | `vector` | `database` calls `stage` inside the one commit function, after the commit's other rules have passed; `upkeep` calls `seal` | +| `cycle`, `MaintenanceReport` | `upkeep` | `database` runs one cycle for `run_maintenance_cycle` and commits it through the one commit function | +| `Store::count_cache_under` | `store` | `database` has the store check its page cache against a new limit, on the one path that sets the memory limit | ## Invariants and the tests that guard them @@ -59,17 +79,27 @@ Every path in this guide is written from `v2/`. | The memory limit holds | `contextdb/tests/memory_ceiling.rs`: `a_million_rows_load_and_scan_under_64_mib` | | A refused statement changes nothing | `contextdb/tests/memory.rs`: `a_statement_that_would_cross_the_limit_is_refused_by_name_and_committed_rows_stay`; `contextdb/tests/transactions.rs`: `a_refused_statement_changes_nothing_and_leaves_the_transaction_usable` | | A damaged page is reported as damage | `contextdb/tests/crash/main.rs`: `a_damaged_page_is_refused_as_damage_by_name`; `redb/tests/contextdb_changes.rs`: `a_page_with_a_wrong_checksum_is_refused_when_read` | +| The stored bytes of a large value, a vector record or a table's vector declaration that disagree with themselves are reported as damage, never a panic or another value | `contextdb/src/store/tests.rs`: `a_plain_text_whose_slices_hold_fewer_bytes_than_its_row_records_is_damage`, `a_compressed_text_that_decodes_to_bytes_that_are_not_text_is_damage`; `contextdb/src/vector/tests.rs`: `a_pending_record_of_a_row_with_no_vector_is_reported_as_damage`, `a_write_to_a_vector_column_that_is_no_column_of_its_table_is_reported_as_damage` | | Counted work: a count reads no rows, a key lookup reads one | `contextdb/tests/counted_work.rs`: `count_of_a_whole_table_reads_no_rows`, `a_lookup_by_primary_key_reads_one_row` | +| Every row with a vector is found by a search, whether its vector waits on a list or is in a block, and after a kill inside a commit or a seal | `contextdb/tests/vectors.rs`: `exact_search_returns_the_true_nearest_vectors`; `contextdb/tests/vectors_upkeep.rs`: `sealing_advances_only_when_asked_and_in_bounded_cycles`; `contextdb/tests/crash/main.rs`: `a_kill_inside_a_commit_of_vectors_keeps_rows_and_vectors_together`, `a_kill_at_any_storage_call_of_a_seal_leaves_every_vector_found` | +| A row whose vector is NULL, and a deleted row, is never in an answer | `contextdb/tests/vectors.rs`: `exact_search_returns_the_true_nearest_vectors`, `exact_search_reads_its_own_snapshot` | +| A partition exists exactly while a row with a vector is in it, and `MAX_PARTITIONS` counts those and nothing else | `contextdb/tests/vectors_partitions.rs`: `vector_partition_limit_counts_and_refuses`, `the_partition_cap_counts_live_partitions_while_a_reader_holds_an_earlier_state` | +| No answer holds a row twice, whatever the row's vector has been | `contextdb/tests/vectors.rs`: `a_row_is_in_an_answer_once_after_its_vector_changes` | +| A store's data file is no longer than a recorded number of bytes for a fixed content | `contextdb/tests/vectors_upkeep.rs`: `a_small_fixed_store_stays_within_its_recorded_size` | +| A text or JSON value reads back byte for byte whether or not it is stored compressed; one that does not shrink by an eighth takes the room it takes uncompressed, and text that repeats takes less than its own bytes | `contextdb/tests/large_values.rs`: `large_text_and_json_come_back_byte_for_byte`, `a_value_that_does_not_shrink_is_kept_as_it_is`, `a_store_of_repeating_text_is_smaller_than_the_text` | +| No value is produced past the memory limit: a compressed value whose real size does not fit is refused before its bytes exist, and compression never refuses a commit | `contextdb/tests/large_values.rs`: `a_compressed_value_too_large_for_the_limit_is_refused`, `a_commit_with_no_room_to_compress_is_written_plain`; `contextdb/tests/memory_ceiling.rs`: `a_refused_compressed_value_is_never_produced` | +| A compressed value whose stored bytes do not give the length it records is reported as store damage | `contextdb/tests/large_values.rs`: `a_damaged_length_of_a_compressed_value_is_reported_as_damage` | ## Where a change lives | Change | Folder | Test it needs | |---|---|---| | A new SQL statement or clause | `contextdb/src/sql/` (`grammar.pest`, `parser/`, the `Statement` type in `ast.rs`, `plan.rs`, `planner.rs`; take its first words out of the grammar's list of what a later version reads), then the part that runs it: `schema`, `write` or `query` | a script under `contextdb/tests/sql/` | +| A new vector clause | `contextdb/src/sql/` reads it, as any clause; `contextdb/src/schema/` checks, stores and prints back one that is declared with the column; `contextdb/src/query/search.rs` settles what a search is asked (mode, vector, partitions, filter); `contextdb/src/vector/` only when blocks, partitions or the search itself change | `contextdb/tests/vectors_statements.rs` for what is read and refused, `contextdb/tests/vectors.rs` for the rows a search returns | | A new refusal (error kind) | `Error` in `contextdb/src/types/mod.rs`; how the program prints it in `contextdb-cli/src/output.rs` | `contextdb/tests/refusals.rs` for its fields and that nothing changed; `contextdb/tests/sql/refusals.sql` for its name | | What a statement returns | `contextdb/src/query/` for a SELECT, `contextdb/src/write/` for a row count | the expected output of the scripts under `contextdb/tests/sql/`, rewritten and the diff read | | How a store is laid out on disk | `contextdb/src/store/` (keys, records, the data file, the folder's files) | `contextdb/tests/values.rs` and `contextdb/tests/store.rs` for reopen, `contextdb/tests/crash/` for a kill | -| The storage library | `redb/src/`; `redb/README.md` lists every change against upstream, so add yours there. The storage library decides which pages may be reused; the `store` part only tells it the oldest commit a reader still holds, and keeps the reader files and the lock | `redb/tests/contextdb_changes.rs` | +| The storage library | `redb/src/`; `redb/README.md` lists every change against upstream, so add yours there. The storage library decides which pages may be reused; the `store` part only tells it the oldest commit a reader still holds, and keeps the reader files and the lock. Two of its rules keep a store close to the size of what it holds, and a change there keeps both: a leaf written in key order is left full, and past 1 MiB the file grows by a sixteenth of its usable bytes (`a_table_written_in_key_order_has_full_leaves`, `a_file_is_at_most_a_sixteenth_longer_than_its_pages`) | `redb/tests/contextdb_changes.rs` | | A new shell option or dot command | `contextdb-cli/src/main.rs` for an option, `contextdb-cli/src/shell.rs` for a command | `contextdb-cli/tests/shell.rs` | | A new public call | the part that owns it, re-exported in `contextdb/src/lib.rs`, described in `contextdb/README.md` | a Rust test under `contextdb/tests/` | @@ -87,12 +117,47 @@ Two lint settings hold this up: - `clippy.toml` refuses `std::thread::sleep`, starting a thread, `Instant::now`, `SystemTime::now` and `uuid::Uuid::new_v4`: the engine does not sleep, start a thread or draw a random number, and reads the time only through `Clock`. - `Cargo.toml` denies `clippy::iter_over_hash_type` for both members: a hash map is walked in a different order in every process, so nothing is built by walking one. +## Making a change + +A defect is fixed by changing or removing code where its promise is owned, not by adding +around it. Before writing a line: + +1. **Name what broke.** Find the row of "Invariants and the tests that guard them" that the + defect breaks, or write the one sentence that should be a row. A defect that cannot be + said in one sentence is not understood yet. +2. **Find the part that owns it** in the parts table, and fix it there, once. A check added + at each caller, a flag, an option, a retry, a second path or a special case for one input + means the fix is in the wrong part: stop and look again. +3. **Fix the kind, not the one.** Look for the same mistake beside it in that part; a defect + is usually one of several of its kind. +4. **Remove what the fix makes unnecessary** in the same change: a guard an earlier fix + added, a helper nothing calls, a comment that is no longer true. Prefer making a wrong + state impossible where it enters (a type, a limit, a refusal when it is written) to + handling it wherever it could arrive. +5. **Extend a test that guards the invariant** rather than add one beside it. A new test is + for an invariant, or a case of one, that no test covers, and the table's row names it. A + test changes only when the promise it states has changed, never to let a fix pass. +6. **Keep it small.** A fix that adds more than about fifty shipped lines, or changes more + than two parts, is a redesign: stop, and write in the change's description what the + invariant is and why its owner cannot hold the fix. + +The check holds the line where no judgement is needed (`.github/workflows/check.yml`): + +- It states the size of the shipped code, rounded up to 500 lines, and the number of tests, + rounded up to 10. A change that crosses either edits the number in the same change, and + says why in its description; a change that removes enough lowers it. Neither number ever + passes 150,000 lines or 1,000 tests. +- It refuses a shipped file longer than 1,500 lines, a test-only name in shipped code, an + ignored test, and an `allow` attribute beyond the four named below. +- The lint settings above refuse a sleep, a thread, a clock read, a random id and a walk + over a hash map. + ## Do not add - A second way to do one thing: there is one commit function, one expression evaluator, one `Error`. - A test-only item in shipped code: no test-only name, feature or `pub` door; tests use what a user has. - A check another part already makes: `write` checks a table's rules, `memory` says no, `store` only keeps. -- A lint switched off. Outside `redb/` there are four `allow` attributes and they stay the only ones: the machine clock in `contextdb/src/system/mod.rs` (the one place the time of day is read); the timer in `contextdb/examples/measure.rs` (measuring elapsed time is what that program is for); and two in `contextdb/tests/support/mod.rs`, one on its one thread (it moves a child's output lines to a channel so that a wait can give up) and `dead_code` for the whole file (each test program uses only some of its helpers). +- A lint switched off. Outside `redb/` there are four `allow` attributes and they stay the only ones: the machine clock in `contextdb/src/system/mod.rs` (the one place the time of day is read); the timer in `contextdb/examples/support/mod.rs` (measuring elapsed time is what the programs under `contextdb/examples/` are for); and two in `contextdb/tests/support/mod.rs`, one on its one thread (it moves a child's output lines to a channel so that a wait can give up) and `dead_code` for the whole file (each test program uses only some of its helpers). ## Fast commands diff --git a/v2/Cargo.lock b/v2/Cargo.lock index 94b65371..e6bb3b9d 100644 --- a/v2/Cargo.lock +++ b/v2/Cargo.lock @@ -155,6 +155,7 @@ name = "contextdb" version = "2.0.0" dependencies = [ "contextdb-redb", + "lz4_flex", "pest", "pest_derive", "serde_json", @@ -266,6 +267,12 @@ version = "0.4.35" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d15f86d12da44cd8d4ec3c73984537eb7f4cadcb9776a575a0a8a05ca11005c" +[[package]] +name = "lz4_flex" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226" + [[package]] name = "memchr" version = "2.8.3" diff --git a/v2/README.md b/v2/README.md index d0029a4f..d4022b40 100644 --- a/v2/README.md +++ b/v2/README.md @@ -27,8 +27,12 @@ cargo +1.89 check --manifest-path v2/Cargo.toml --workspace --locked ``` The last line checks the oldest Rust the packages declare (`rustup toolchain install 1.89` -first). The workflow also fails when the library and the program together pass 150,000 lines or -the tests pass 1,000; the two lines that count them are in the workflow file. +first). The workflow also states the size of the library and the program, rounded up to 500 +lines, and the number of tests, rounded up to 10, and fails when either is passed or is no +longer within that distance: a change that grows or shrinks the code past a boundary edits +the number in the workflow file, in the same change. It also refuses a shipped file longer +than 1,500 lines, a test-only name in shipped code, an ignored test and an `allow` attribute +beyond the four `AGENTS.md` names. `AGENTS.md` says how a change is made so that these hold. ## One test diff --git a/v2/contextdb-cli/src/main.rs b/v2/contextdb-cli/src/main.rs index 82603005..40a0c493 100644 --- a/v2/contextdb-cli/src/main.rs +++ b/v2/contextdb-cli/src/main.rs @@ -51,16 +51,26 @@ EXIT STATUS: 2 the command line was refused; nothing was opened MEMORY: - --memory-limit is this session's ceiling, in bytes or with K, M or G. Without it a - --write session has the limit the store remembers and a read session has none. - The store keeps 8M for its page cache and 1M is set aside before that, so the - smallest limit that opens a store is 9M. + --memory-limit is this session's ceiling, in bytes or with K, M or G; the store does + not remember it. Without it a --write session has the limit the store remembers and + a read session has none. + A session sets aside 1M and its page cache, sized when it starts: a quarter of its + limit and never less than 8M, which is also its size with no limit. So the smallest + limit that opens a store is 9M. `SET MEMORY_LIMIT 512M;` (or NONE) declares the limit the store remembers; it needs --write, and inside a transaction it still takes effect at once and ROLLBACK does not - undo it. `SHOW MEMORY_LIMIT;` reports this session's limit and what is in use. + undo it. The smallest limit a session takes is 1M plus its page cache plus what it + holds, which is 9M in a session started with no limit; a limit below it is refused. + To go below the page cache of a session started under a larger limit, start with the + smaller one, `--memory-limit 64M`, and then run `SET MEMORY_LIMIT 64M;`. + `SHOW MEMORY_LIMIT;` reports this session's limit and what is in use: with no limit, + `none` and 9437184 bytes used, the 1M and the 8M page cache. JSON (--json), one document a line: a SELECT or SHOW MEMORY_LIMIT {"result":{"columns":[...],"rows":[{...}]}} + a vector is an array of numbers, each at the full + precision of the stored 32-bit value, so a number + may print with more digits than were typed any other statement {"rows_affected":N} .tables {"tables":{"items":[...]}} .schema {"schema":{"table":...,"ddl":...}} diff --git a/v2/contextdb-cli/tests/shell.rs b/v2/contextdb-cli/tests/shell.rs index 6e9b49d1..69001312 100644 --- a/v2/contextdb-cli/tests/shell.rs +++ b/v2/contextdb-cli/tests/shell.rs @@ -94,6 +94,49 @@ fn listing(folder: &Path) -> Vec { names } +/// The tables a run printed for a person, in order: each is its header and then its rows, cell +/// by cell. A result with no rows prints no table and is an empty list here. +fn tables(stdout: &str) -> Vec>> { + let mut tables = Vec::new(); + let mut table = Vec::new(); + for line in stdout.lines() { + let cells = line + .strip_prefix('|') + .and_then(|rest| rest.strip_suffix('|')); + if let Some(cells) = cells { + let cells = cells.split('|').map(|cell| cell.trim().to_owned()); + table.push(cells.collect()); + } else if line.starts_with('(') { + tables.push(std::mem::take(&mut table)); + } + } + tables +} + +/// The rows of a printed search whose columns are `id` and `distance`: each row's key and its +/// distance as a number, however the program writes a whole REAL. +fn nearest(table: &[Vec]) -> Vec<(i64, f64)> { + assert_eq!(table[0], ["id", "distance"], "{table:?}"); + table[1..] + .iter() + .map(|row| (row[0].parse().unwrap(), row[1].parse().unwrap())) + .collect() +} + +/// A vector in a JSON document: an array of numbers. +fn numbers(vector: &Json) -> Vec { + vector + .as_array() + .unwrap_or_else(|| panic!("not a JSON array: {vector}")) + .iter() + .map(|number| { + number + .as_f64() + .unwrap_or_else(|| panic!("not a number: {number}")) as f32 + }) + .collect() +} + const CREATE_NOTES: &str = "CREATE TABLE notes (id INTEGER PRIMARY KEY, body TEXT);\n"; const SELECT_NOTES: &str = "SELECT id, body FROM notes ORDER BY id;\n"; @@ -636,6 +679,20 @@ fn the_memory_limit_option_is_the_sessions_ceiling() { documents(&run.stdout)[0]["result"]["rows"][0]["limit"].clone() }; assert_eq!(limit(&[arg(&store), "--write", "--json"]), json!("none")); + // With no limit a session counts 9 MiB as used before any statement: 1 MiB and a page + // cache of 8 MiB. A read shell with no limit counts the same. + for args in [ + &[arg(&store), "--write", "--json"][..], + &[arg(&store), "--json"], + ] { + let run = shell(&dir, args, "SHOW MEMORY_LIMIT;\n"); + assert_eq!(run.code, Some(0), "{}", run.stderr); + assert_eq!( + documents(&run.stdout)[0]["result"]["rows"][0]["used"], + json!(9_437_184_u64), + "used with no limit" + ); + } for (size, bytes) in [ ("65536K", 67_108_864_u64), ("64M", 67_108_864), @@ -730,3 +787,227 @@ fn the_memory_limit_option_is_the_sessions_ceiling() { ); } } + +#[test] +fn a_similarity_search_runs_through_the_program() { + let dir = folder("a-similarity-search-runs"); + let store = dir.join("store"); + // From [1, 0, 0] row 1 is at distance 0, row 5 between 0 and 1, row 2 at 1 and row 3 at 2. + // Row 4 has no vector. + let declare = "\ +CREATE TABLE t ( + id INTEGER PRIMARY KEY, + kind TEXT NOT NULL, + v VECTOR(3) PARTITION_KEY (kind) SEARCH_MODE EXACT +); +INSERT INTO t (id, kind, v) +VALUES (1, 'a', [1.0, 0.0, 0.0]), (2, 'a', [0.0, 1.0, 0.0]), (3, 'b', [-1.0, 0.0, 0.0]), + (4, 'b', NULL), (5, 'b', [0.1, 0.2, 0.3]); +"; + let search = + "SELECT id, v <=> [1.0, 0.0, 0.0] AS distance FROM t ORDER BY distance LIMIT 10;\n"; + let rows = "SELECT id, v FROM t ORDER BY id;\n"; + + // Declared, written and searched in one piped write session. + let run = shell( + &dir, + &[arg(&store), "--write"], + &format!("{declare}{search}{rows}"), + ); + assert_eq!(run.code, Some(0), "{}", run.stderr); + assert_eq!(run.errors(), Vec::<&str>::new()); + let printed = tables(&run.stdout); + assert_eq!(printed.len(), 2, "{}", run.stdout); + + // The rows that have a vector, nearest first, each with its distance under the name the + // statement gave it. + let found = nearest(&printed[0]); + let ids: Vec = found.iter().map(|(id, _)| *id).collect(); + assert_eq!(ids, [1, 5, 2, 3], "{}", run.stdout); + assert_eq!( + (found[0].1, found[2].1, found[3].1), + (0.0, 1.0, 2.0), + "{}", + run.stdout + ); + assert!(0.0 < found[1].1 && found[1].1 < 1.0, "{}", run.stdout); + + // A vector prints as its numbers in brackets, and no vector as NULL. + assert_eq!(printed[1][0], ["id", "v"], "{}", run.stdout); + assert_eq!(printed[1][1], ["1", "[1.0, 0.0, 0.0]"], "{}", run.stdout); + assert_eq!(printed[1][4], ["4", "NULL"], "{}", run.stdout); + assert_eq!(printed[1][5], ["5", "[0.1, 0.2, 0.3]"], "{}", run.stdout); + + // For a program the distance is a number, a vector an array of numbers and no vector null. + let run = shell(&dir, &[arg(&store), "--json"], &format!("{search}{rows}")); + assert_eq!(run.code, Some(0), "{}", run.stderr); + let results = documents(&run.stdout); + assert_eq!(results.len(), 2, "{}", run.stdout); + assert_eq!(results[0]["result"]["columns"], json!(["id", "distance"])); + let found = results[0]["result"]["rows"].as_array().unwrap(); + let ids: Vec = found + .iter() + .map(|row| row["id"].as_i64().unwrap()) + .collect(); + assert_eq!(ids, [1, 5, 2, 3], "{}", run.stdout); + for row in found { + assert!(row["distance"].is_number(), "{row}"); + } + let distance = |row: usize| found[row]["distance"].as_f64().unwrap(); + assert_eq!( + (distance(0), distance(2), distance(3)), + (0.0, 1.0, 2.0), + "{}", + run.stdout + ); + assert!(0.0 < distance(1) && distance(1) < 1.0, "{}", run.stdout); + assert_eq!(results[1]["result"]["columns"], json!(["id", "v"])); + let stored = results[1]["result"]["rows"].as_array().unwrap(); + assert_eq!(stored.len(), 5, "{}", run.stdout); + assert_eq!(numbers(&stored[0]["v"]), [1.0, 0.0, 0.0]); + assert!(stored[3]["v"].is_null(), "{}", stored[3]); + assert_eq!(numbers(&stored[4]["v"]), [0.1, 0.2, 0.3]); + + // `.schema` prints the VECTOR column with its clauses in the order they are declared in, + // and with the limit of 256 partitions that a key declared without one has. + let schema = shell(&dir, &[arg(&store)], ".schema t\n"); + assert_eq!(schema.code, Some(0), "{}", schema.stderr); + let words = schema + .stdout + .split_whitespace() + .collect::>() + .join(" "); + assert!( + words.contains("v VECTOR(3) PARTITION_KEY (kind) MAX_PARTITIONS 256 SEARCH_MODE EXACT"), + "{}", + schema.stdout + ); + + // What it printed declares the same table again in another store. + let copy = dir.join("copy"); + let made = shell(&dir, &[arg(©), "--write"], &schema.stdout); + assert_eq!(made.code, Some(0), "{}", made.stderr); + assert_eq!( + shell(&dir, &[arg(©)], ".schema t\n").stdout, + schema.stdout + ); +} + +#[test] +fn indexed_search_is_refused_by_name_through_the_program() { + let dir = folder("indexed-search-is-refused"); + let store = dir.join("store"); + let made = shell( + &dir, + &[arg(&store), "--write"], + "CREATE TABLE passages (id INTEGER PRIMARY KEY, embedding VECTOR(3));\n\ + INSERT INTO passages (id, embedding) VALUES (1, [0.0, 1.0, 0.0]), (2, [1.0, 0.0, 0.0]);\n", + ); + assert_eq!(made.code, Some(0), "{}", made.stderr); + + // The same search asked for INDEXED and then for EXACT. + let search = |mode: &str| { + format!( + "SELECT id FROM passages ORDER BY embedding <=> [1.0, 0.0, 0.0] USE VECTOR {mode} \ + LIMIT 2;\n" + ) + }; + let input = format!("{}{}", search("INDEXED"), search("EXACT")); + let exact = "+----+\n| id |\n+----+\n| 2 |\n| 1 |\n+----+\n(2 rows)\n"; + + // Reading or writing, this version keeps no index and says so. + for mode in [&[][..], &["--write"][..]] { + let mut args = vec![arg(&store)]; + args.extend_from_slice(mode); + + // For a person: one line on the error stream, the refusal's sentence naming the column, + // and a piped run that fails. + let run = shell(&dir, &args, &input); + assert_eq!(run.code, Some(1), "{mode:?}: {}", run.stderr); + let lines: Vec<&str> = run.stderr.lines().collect(); + assert_eq!(lines.len(), 1, "{mode:?}: {}", run.stderr); + assert!( + lines[0].starts_with("Error: ") + && lines[0].contains("passages.embedding") + && lines[0].contains("INDEXED"), + "{mode:?}: {}", + run.stderr + ); + // INDEXED gave no rows under any name, and the statement after it still ran. + assert_eq!(run.stdout, exact, "{mode:?}"); + + // For a program: the refusal's kind. + args.push("--json"); + let run = shell(&dir, &args, &input); + assert_eq!(run.code, Some(1), "{mode:?}: {}", run.stderr); + let refusals = documents(&run.stderr); + assert_eq!(refusals.len(), 1, "{mode:?}: {}", run.stderr); + let error = &refusals[0]["error"]; + assert_eq!(error["class"], "sql", "{mode:?}: {error}"); + assert_eq!( + error["detail"]["kind"], "vector_indexed_route_unavailable", + "{mode:?}: {error}" + ); + assert!( + error["message"] + .as_str() + .is_some_and(|message| message.contains("passages.embedding")), + "{mode:?}: {error}" + ); + assert_eq!( + documents(&run.stdout), + [json!({"result": {"columns": ["id"], "rows": [{"id": 2}, {"id": 1}]}})], + "{mode:?}" + ); + } +} + +#[test] +fn a_read_shell_searches_while_a_writer_holds_the_store() { + let dir = folder("a-read-shell-searches-while-a-writer-holds"); + let store = dir.join("store"); + + // This test is the writer, and holds the store until the end. From [1, 0, 0] row 1 is at + // distance 1 and row 2 at 2. + let writer = Database::open(&store).unwrap(); + for sql in [ + "CREATE TABLE passages (id INTEGER PRIMARY KEY, embedding VECTOR(3))", + "INSERT INTO passages (id, embedding) VALUES (1, [0.0, 1.0, 0.0]), (2, [-1.0, 0.0, 0.0])", + ] { + writer.execute(sql, ¶ms! {}).unwrap(); + } + let search = "SELECT id, embedding <=> [1.0, 0.0, 0.0] AS distance FROM passages \ + ORDER BY distance LIMIT 2;\n"; + + // The store really is held: a second writer is refused, naming this process. + let second = shell(&dir, &[arg(&store), "--write"], search); + assert_eq!(second.code, Some(1), "{}", second.stderr); + assert!( + second.stderr.contains(&std::process::id().to_string()), + "{}", + second.stderr + ); + + // A read shell needs nothing from the writer: it searches and gets the rows. + let run = shell(&dir, &[arg(&store)], search); + assert_eq!(run.code, Some(0), "{}", run.stderr); + let printed = tables(&run.stdout); + assert_eq!(printed.len(), 1, "{}", run.stdout); + assert_eq!(nearest(&printed[0]), [(1, 1.0), (2, 2.0)], "{}", run.stdout); + + // The writer was not held up by it, and the next read shell searches its next commit, in + // which row 3 is the nearest. + writer + .execute( + "INSERT INTO passages (id, embedding) VALUES (3, [1.0, 0.0, 0.0])", + ¶ms! {}, + ) + .unwrap(); + let run = shell(&dir, &[arg(&store)], search); + assert_eq!(run.code, Some(0), "{}", run.stderr); + let printed = tables(&run.stdout); + assert_eq!(printed.len(), 1, "{}", run.stdout); + assert_eq!(nearest(&printed[0]), [(3, 0.0), (1, 1.0)], "{}", run.stdout); + + writer.close().unwrap(); +} diff --git a/v2/contextdb/Cargo.toml b/v2/contextdb/Cargo.toml index 1002a1d0..111a7b5b 100644 --- a/v2/contextdb/Cargo.toml +++ b/v2/contextdb/Cargo.toml @@ -9,6 +9,7 @@ repository = "https://github.com/context-graph-ai/contextdb" readme = "README.md" [dependencies] +lz4_flex = { version = "0.14", default-features = false, features = ["std", "safe-encode", "safe-decode"] } pest = "2.8.6" pest_derive = "2.8.6" redb = { package = "contextdb-redb", version = "4.1.0", path = "../redb" } diff --git a/v2/contextdb/README.md b/v2/contextdb/README.md index 0d75e9cb..42032c29 100644 --- a/v2/contextdb/README.md +++ b/v2/contextdb/README.md @@ -2,16 +2,23 @@ An embedded SQL database: tables in one folder, inside your own process. +An application that keeps its records in one store and its embeddings in a separate vector +library has to keep the two consistent by hand; here a vector is a column, a row and its vector +commit together, and one `SELECT` filters rows and orders them by distance. + This is the first version of contextdb 2. What it has: - tables with typed columns, a primary key, `NOT NULL`, `UNIQUE` and `DEFAULT`, and indexes; - `INSERT`, `UPDATE`, `DELETE` and `SELECT` over one table; +- vector search: the rows nearest a vector by cosine distance, each with its distance, exact, + over vectors that a column may keep in partitions; - transactions; - one process writing a store while any number of other processes read it; - a memory limit that holds. -What it does not have yet: vector search, relationships between rows, joins, grouping, and sync -between machines. A statement that uses one of them is refused and changes nothing. +What it does not have yet: an index for vector search (every search is exact), relationships +between rows, joins, grouping, and sync between machines. A statement that uses one of them is +refused and changes nothing. It is not on crates.io yet. Clone and use it from that checkout: @@ -26,30 +33,42 @@ use contextdb::{params, Database, Uuid, Value}; fn main() -> contextdb::Result<()> { let db = Database::open("./notes")?; // a folder, created if it is missing db.execute( - "CREATE TABLE notes (id UUID PRIMARY KEY, body TEXT NOT NULL, made TIMESTAMP DEFAULT NOW())", + "CREATE TABLE notes (id UUID PRIMARY KEY, body TEXT NOT NULL, embedding VECTOR(3), + made TIMESTAMP DEFAULT NOW())", ¶ms! {}, )?; - db.execute( - "INSERT INTO notes (id, body) VALUES ($id, $body)", - ¶ms! { "id" => Uuid::new_v4(), "body" => "buy milk" }, - )?; + // A real embedding has hundreds of numbers and comes from a model; three show the shape. + let notes = [("buy milk", vec![0.9, 0.1, 0.0]), ("call the plumber", vec![0.0, 0.2, 0.9])]; + for (body, embedding) in notes { + db.execute( + "INSERT INTO notes (id, body, embedding) VALUES ($id, $body, $embedding)", + ¶ms! { "id" => Uuid::new_v4(), "body" => body, "embedding" => embedding }, + )?; + } let tx = db.transaction()?; // rolls back if dropped tx.execute("UPDATE notes SET body = 'buy oat milk' WHERE body = 'buy milk'", ¶ms! {})?; tx.commit()?; - let found = db.execute("SELECT body, made FROM notes ORDER BY made LIMIT 10", ¶ms! {})?; + // The notes nearest a vector, nearest first, each with its distance. + let found = db.execute( + "SELECT body, embedding <=> $q AS distance FROM notes ORDER BY embedding <=> $q LIMIT 10", + ¶ms! { "q" => vec![1.0, 0.0, 0.0] }, + )?; for row in found.rows() { - if let Value::Text(body) = &row[0] { - println!("{body}"); + if let (Value::Text(body), Value::Float64(distance)) = (&row[0], &row[1]) { + println!("{body}: {distance:.3}"); } } Ok(()) } ``` -It prints `buy oat milk`. This page is also the crate's documentation, and the documentation -tests compile this program as it is written here. Rust 1.89 or newer; Linux and macOS. Licence: +It prints `buy oat milk: 0.006` and then `call the plumber: 1.000`. After writing vectors, a host +calls `db.run_maintenance_cycle()?` until its report's `vector.maintenance_vectors_remaining` is +0, because this version does not do it by itself; a search before that returns the same rows, +slower ([Sealing](#vector-search) below). This page is also the crate's +documentation, and the documentation tests compile this program as it is written here. Rust 1.89 or newer; Linux and macOS. Licence: Apache-2.0. It is developed in the open at by the context-graph-ai organisation. ## Statements @@ -57,6 +76,8 @@ Apache-2.0. It is developed in the open at vector [USE VECTOR AUTO | EXACT | INDEXED] LIMIT k; SET MEMORY_LIMIT 512M; SET MEMORY_LIMIT NONE; SHOW MEMORY_LIMIT; ``` @@ -72,15 +95,19 @@ SET MEMORY_LIMIT 512M; SET MEMORY_LIMIT NONE; SHOW MEMORY_LIMIT; - A condition or a value is built from column names, literals, `$name` values, `= != <> < <= > >=`, `AND`, `OR`, `NOT`, `IS NULL`, `IS NOT NULL`, `+ - * /`, parentheses and `NOW()`. A comparison with NULL is not true, as in standard SQL. +- `x IN (a, b, ...)` and `x NOT IN (a, b, ...)` test a value against a list, and + `(x, y) IN ((a, b), (c, d), ...)` tests several columns at once, both as standard SQL has + them. A list is written out: `IN (SELECT ...)` is refused. - `SELECT` also takes `COUNT(*)`, `COUNT(column)`, `SUM`, `MIN` and `MAX` over the rows that match, and `AS` to name a column of the result. - `ORDER BY` puts NULL last ascending and first descending. +- The second `SELECT` is a vector search; [Vector search](#vector-search) below says what it + returns. - `--` and `/* */` comments are read past. -Not in this version, and refused: `JOIN`, `GROUP BY`, `DISTINCT`, `IN`, `LIKE`, `BETWEEN`, -`OFFSET`, `UNION`, subqueries, table aliases, `ALTER`, `DROP`, a primary key over several -columns, and similarity search (`<=>`). Most are refused by name; this is the sentence the -refusal prints: +Not in this version, and refused: `JOIN`, `GROUP BY`, `DISTINCT`, `LIKE`, `BETWEEN`, +`OFFSET`, `UNION`, subqueries, table aliases, `ALTER`, `DROP`, and a primary key over several +columns. Most are refused by name; this is the sentence the refusal prints: ```text plan error: LIKE is not supported in this version of the library @@ -97,12 +124,77 @@ plan error: LIKE is not supported in this version of the library | `UUID` | a UUID, written as quoted text | `Value::Uuid` | | `TIMESTAMP` | milliseconds since the first moment of 1970, UTC; written as ISO 8601 text (`'2026-01-01T00:00:01Z'`) or taken from `NOW()` | `Value::Timestamp` | | `JSON` | a JSON value, written as quoted text; the value is kept, not the text's spacing | `Value::Json` | -| `VECTOR(n)` | a vector of exactly `n` numbers, written `[0.1, 0.2, 0.3]`; stored and returned, not searchable in this version | `Value::Vector` | +| `VECTOR(n)` | a vector of exactly `n` finite numbers, each kept as an `f32`, written `[0.1, 0.2, 0.3]`; searched with `<=>` | `Value::Vector` | A value must be of its column's type: text into an `INTEGER` column is refused, and so is NULL into a `NOT NULL` column, on `INSERT` and on `UPDATE` alike. A whole number is accepted into a `REAL` column. +## Vector search + +A `VECTOR(n)` column is searched by ordering its rows by their distance from a vector and +taking the first `k`: + +```sql +SELECT id, embedding <=> $q AS distance FROM notes ORDER BY distance LIMIT 10; +``` + +- `embedding <=> $q` is the cosine distance between the row's vector and `$q`, a `REAL`. Lower + is closer: 0 for vectors that point the same way, 1 for vectors at a right angle, 2 for + opposite ones, as Postgres with pgvector reads the operator. Rows come nearest first, and + rows at the same distance in primary-key order. +- A search returns the columns the statement wrote and no other. To get a row's distance, write + the expression in the select list, alone or inside another expression, and name it as you + like: `1 - (embedding <=> $q) AS similarity` is a similarity. `ORDER BY` takes the expression + itself or the name the select list gives it. +- The vector on the right is a `$name` bound to a `Vec`, a literal `[1.0, 0.0, 0.0]`, or + the same as quoted text, `'[1, 0, 0]'`. It has the column's dimension and finite numbers. +- `WHERE` is applied first: the answer is the `k` nearest of the rows that pass it. A row whose + vector is NULL is in no answer. The distance to a vector of all zeros is 1. +- Every search is exact: it returns the rows that comparing `$q` with every stored vector + returns. + +Only `<=>`, cosine, is here: pgvector's `<->` and `<#>` are a parse error. `<=>` has a value +only for the rows of a search, so it is refused everywhere else: in `WHERE`, +in the select list of a statement that does not order by it, in an `INSERT`, `UPDATE` or +`DELETE`, with a vector other than the one the statement orders by, and in an `ORDER BY` that +is `DESC`, has a second item or has no `LIMIT`; a search takes no `OFFSET` either. + +**Sealing.** Beside every stored vector the engine keeps a one-byte copy of each of its numbers +and packs the copies of many vectors into a block, which a search screens quickly to rule most +rows out without reading them. A vector that was just written is on the waiting list, the +vectors not yet in a block, which a search compares one by one by reading their rows. +`Database::run_maintenance_cycle` seals them, which is to move them from the list into blocks: +one call seals at most 4,096 in one commit and returns how many it sealed and how many still +wait; nothing in this version calls it on its own. A search returns the same rows with the same distances whether or not vectors are sealed; sealed, +it reads far fewer rows, and the result's `QueryResult::rows_examined` says how many it read. + +**Partitions.** `PARTITION_KEY (column, ...)` on a `VECTOR` column lays its vectors out by the +values of the named columns, one partition for each combination that has a row with a vector, +so that a search which filters on the key reads only that part. A partition is not a wall: a +search with no filter on the key reads every partition, and which rows a caller may see is the +`WHERE`'s business. A search whose `WHERE` gives the key columns values with `=` or `IN` reads +only the partitions named, up to 4,096 of them (a filter that names more reads every +partition); the whole condition is still applied to every row, so partitions change the work +and never the answer. A key column is `NOT NULL` and a `UUID`, `TEXT`, +`INTEGER`, `BOOLEAN` or `TIMESTAMP`. + +`MAX_PARTITIONS m`, written after the key when the table is declared, is the most partitions +the column may have at one time, 256 when it is not written. An `INSERT` or `UPDATE` that would +take the column past it is refused with `Error::VectorPartitionLimitExceeded`, which names the +column and its limit; the statement changes nothing, and inside a transaction the transaction +goes on. A column with no `PARTITION_KEY` is one partition and has no limit. + +**Modes.** `SEARCH_MODE` on the column says how its searches run when a statement does not, and +`USE VECTOR`, written after the `ORDER BY` item, says it for one statement. `AUTO`, which is +the default, and `EXACT` both run the exact search. `INDEXED` asks for an index, and this +version has none: a column may be declared with it, and a search under it is refused before +anything is read, unless the statement itself says `USE VECTOR EXACT` or `AUTO`. + +```text +no vector index ships in this version, so notes.embedding cannot be searched INDEXED; use EXACT or AUTO +``` + ## Binding values `$name` in a statement is the entry `"name"` of the map passed with it. `params!` writes the @@ -114,7 +206,9 @@ any other value is written as a `Value`, for example `Value::Timestamp(1_767_225 ## Refusals -Every call that can fail returns `contextdb::Result` with one `contextdb::Error`. A program +Every call that can fail returns `contextdb::Result` with one `contextdb::Error`, except +`Database::memory_limit`, `memory_used`, `table_names` and `committed_watermark`, which return +a plain value and panic on a handle that was closed. A program matches on the kind (`Error::UniqueViolation { table, column }`, `Error::ColumnNotNullable`, `Error::TableNotFound`, `Error::DatabaseLocked`, `Error::MemoryBudgetExceeded`, and the rest listed on the `Error` type); a person reads the sentence it prints: @@ -155,23 +249,61 @@ answering when the writer stops, is killed or starts again. ## The memory limit `db.set_memory_limit(Some(bytes))` or `SET MEMORY_LIMIT 512M` declares how much memory the -engine may use for data in the writing process, and the store remembers it until it is changed -or set to `NONE`. A statement that would pass the limit is refused and committed rows stay as -they were: +engine may use for data in the writing process, and the store remembers it: every later open +for writing that names no limit of its own is under it, until it is changed or set to `NONE`. +A statement that would pass the limit is refused and committed rows stay as they were: ```text memory limit reached during transaction: 30000210 bytes were needed and 7340032 bytes of the 16777216-byte limit were free. select or write less at once, or raise the limit with SET MEMORY_LIMIT or the memory_limit open option ``` -With nothing declared there is no ceiling. A reader has a limit of its own, given when it opens, -and none unless one is given. Every open sets 9 MiB of the limit aside, 8 MiB for the page cache -and 1 MiB for small fixed things, so 9 MiB is the smallest limit that opens a store and a -statement needs room above it. +With nothing declared there is no ceiling. A limit given when a store is opened, as +`DatabaseOpenOptions::memory_limit` or the shell's `--memory-limit`, is that session's limit and +no more: it holds until the session closes, and the store does not remember it. The options +type is `#[non_exhaustive]`, so start from its default and assign: + +```rust,no_run +use contextdb::{Database, DatabaseOpenOptions}; + +fn main() -> contextdb::Result<()> { + let mut options = DatabaseOpenOptions::default(); + options.memory_limit = Some(64 << 20); // 64 MiB, for this session only + let db = Database::open_with_options("./notes", options)?; + db.close() +} +``` + +A reader has a +limit of its own, given when it opens, and none unless one is given: it never takes the limit +the store remembers. + +Every open sets aside 1 MiB for small fixed things and its page cache, and counts both as used +from the start. The cache is sized once, when the store is opened, and keeps that size until +the session closes: a quarter of the limit in force and never less than 8 MiB, which is also +its size when there is no limit. So 9 MiB is the smallest limit that opens a store, and a +statement needs room above what is set aside. + +A limit declared in an open session has to hold with the cache that session already has: the +smallest one it takes is 1 MiB, plus its cache, plus what its open transactions and results +hold. In a store opened with no limit in force that is 9 MiB plus what is held, so +`SET MEMORY_LIMIT 64M` is accepted at once. A smaller limit is refused with +`Error::MemoryBudgetExceeded`, changes nothing and is not remembered; the bytes the refusal +says were needed are the smallest limit that session takes. Any other limit holds at once for +what statements take, and the next open sizes its cache from it. + +Only a limit below the cache of a session opened under a larger limit needs two steps: open the +store with the smaller limit, which sizes the cache from it, and declare it in that session, +which makes the store remember it: + +```text +echo "SET MEMORY_LIMIT 64M;" | contextdb ./notes --write --memory-limit 64M +``` `SHOW MEMORY_LIMIT` returns one row about the session that asks, a writer's or a reader's: -`limit` (its limit in bytes, or `none`), `used` (bytes in use now, the 9 MiB included), -`available` (what is left under the limit), `startup_ceiling` (the limit the session was opened -under) and `ceiling_refresh_failures` (always 0 in this version). +`limit` (its limit in bytes, or `none`), `used` (bytes in use now, the 1 MiB and the page cache +included), `available` (what is left under the limit), `startup_ceiling` (the limit the session +was opened under) and `ceiling_refresh_failures` (always 0 in this version). In a session with +no limit it reports `none` as the limit and 9437184 bytes, which is 9 MiB, as used. ## The store on disk @@ -182,7 +314,9 @@ notes/ readers/ one small file per reader, while a statement of its runs ``` -Nothing is written anywhere else. While no process has the folder open for writing it is a plain +Nothing is written anywhere else. Text and JSON values from about 1 KiB up to 1 MiB are stored +compressed when that saves an eighth or more, and read back byte for byte, with nothing to +switch on or off, so a store of long text is usually smaller than the text. While no process has the folder open for writing it is a plain folder: copy it, move it or delete it with ordinary tools. Do not copy it while a writer is running; this version has no command that copies a store beside a live writer. @@ -225,7 +359,9 @@ piped run exits 1 if any statement or command was refused, an unknown command an a table that is not there included, and 0 otherwise. However the input arrives, the exit status is 1 when the store could not be opened or closed, and 2 when the command line itself is refused. `--json` prints one JSON document per statement or command and each refusal as a JSON document on the -error stream. `contextdb --help` has the +error stream. A vector in a result is a JSON array of numbers, each written at the full +precision of the 32-bit value that is stored, so a number may print with more digits than +were typed, where the table a person reads shows `[0.1, ...]`. `contextdb --help` has the options, the exit statuses and the JSON forms. ## Numbers @@ -236,16 +372,18 @@ nothing else. | What | Measured | |---|---| -| 20,000 rows of 1.6 KB, one commit per row: time to insert | 17.2 s | -| That store on disk | 60.3 MiB | -| Open that store | 0.631 ms | -| `SELECT COUNT(*)` over it | 0.026 ms, 0 rows read | -| A million rows under a 64 MiB limit: time to load | 11.9 s | -| Scan the million rows | 246 ms | -| Most memory the engine had allocated during that load and scan | 13.4 MiB | -| Most memory the process held, as the operating system reports it | 18.0 MiB | -| First build of a program that uses the library | 5 s, 46 packages, a 2.7 MiB program | -| From a fresh clone to a passing `cargo test` | 2.1 min | -| Lines of Rust and grammar in the library and the program | 7,369 | -| Tests | 98 | -| Lines that differ from redb 4.1.0 in the storage library under the engine, which is redb 4.1.0 with the changes listed in `v2/redb/README.md` | 577 | +| 20,000 rows of 1.6 KB, one commit per row: time to insert | 11.0 s | +| That store on disk (its rows repeat their text, which compresses far better than real text) | 2.4 MiB | +| Open that store | 0.606 ms | +| `SELECT COUNT(*)` over it | 0.025 ms, 0 rows read | +| A million rows under a 64 MiB limit: time to load | 9.7 s | +| Scan the million rows | 239.167 ms | +| Most memory the engine had allocated during that load and scan | 16.4 MiB | +| Most memory the process held, as the operating system reports it | 21.3 MiB | +| First build of a program that uses the library | 6 s, 47 packages, a 3.0 MiB program | +| From a fresh clone to a passing `cargo test` | 2.7 min | +| Ten nearest of 20,000 vectors, sealed, through SQL: 95th percentile | 1.500 ms | +| That vector store on disk, and how many times the bytes it holds | 44.3 MiB, 1.50 times | +| Lines of Rust and grammar in the library and the program | 10,809 | +| Tests | 171 | +| Lines that differ from redb 4.1.0 in the storage library under the engine, which is redb 4.1.0 with the changes listed in `v2/redb/README.md` | 626 | diff --git a/v2/contextdb/examples/measure.rs b/v2/contextdb/examples/measure.rs index c19a110b..f59e3875 100644 --- a/v2/contextdb/examples/measure.rs +++ b/v2/contextdb/examples/measure.rs @@ -1,83 +1,37 @@ //! Prints the numbers this library is held to, one per line as `name value unit`. //! -//! `measure readme ` runs the README's first program at 20,000 rows of 1.6 KB, one -//! commit per row, and prints the store's size on disk, how long opening it takes and how long -//! `COUNT(*)` takes. `measure million ` loads a million small rows under a 64 MiB limit -//! in commits of 10,000, scans them, and prints the most memory the process had allocated beside -//! the most the operating system saw it hold. The folder must be missing or empty. +//! - `measure readme ` inserts 20,000 notes of 1.6 KB, one +//! commit per row, and prints the store's size on disk, how long opening it takes and how +//! long `COUNT(*)` takes. +//! - `measure million ` loads a million small rows under a 64 MiB limit in commits of +//! 10,000, scans them, and prints the most memory the process had allocated beside the most +//! the operating system saw it hold. +//! - `measure vectors ` loads 20,000 rows with a vector of 384 numbers made from the +//! row number, 500 to a commit, seals them, searches for the 10 nearest of 100 further +//! vectors after 5 searches that are not timed, and prints the search time and the store's +//! size beside the bytes held. +//! - `measure sessions ` loads every line that is +//! not empty of the first `count` files the list names, 500 rows to a commit, and prints the +//! files' bytes, the bytes put in (each line's and 8 for its number) and the store's size. +//! It prints nothing of a file but those counts. +//! +//! The store's folder must be missing or empty. A store's size is the bytes of every file in +//! its folder after it was closed. //! //! This program measures. It never passes or fails, no test reads what it prints, and //! `numbers.sh` is what runs it. -use std::alloc::{GlobalAlloc, Layout, System}; -use std::fmt::Display; +mod support; + +use std::fs::File; +use std::io::{BufRead, BufReader}; use std::path::Path; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::{Duration, Instant}; use contextdb::{Database, DatabaseOpenOptions, Result, Uuid, params}; - -/// The system allocator, counting the bytes handed out and not yet given back. -struct Counting; - -static IN_USE: AtomicUsize = AtomicUsize::new(0); -static PEAK: AtomicUsize = AtomicUsize::new(0); - -// SAFETY: every request goes to the system allocator unchanged; this only counts its bytes. -unsafe impl GlobalAlloc for Counting { - unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - let in_use = IN_USE.fetch_add(layout.size(), Ordering::Relaxed) + layout.size(); - PEAK.fetch_max(in_use, Ordering::Relaxed); - unsafe { System.alloc(layout) } - } - - unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) { - IN_USE.fetch_sub(layout.size(), Ordering::Relaxed); - unsafe { System.dealloc(pointer, layout) } - } -} +use support::{bytes_on_disk, line, mebibytes, milliseconds, percentile, ratio, seconds, timed}; #[global_allocator] -static ALLOCATOR: Counting = Counting; - -/// Runs `work` and says how long it took. Apart from the library's `system` part, this is the -/// only place in the workspace that reads a clock. -#[allow( - clippy::disallowed_methods, - reason = "measuring elapsed time is what this program is for" -)] -fn timed(work: impl FnOnce() -> T) -> (T, Duration) { - let start = Instant::now(); - let result = work(); - (result, start.elapsed()) -} - -fn line(name: &str, value: impl Display, unit: &str) { - println!("{name} {value} {unit}"); -} - -fn milliseconds(duration: Duration) -> String { - format!("{:.3}", duration.as_secs_f64() * 1000.0) -} - -fn mebibytes(bytes: u64) -> String { - format!("{:.1}", bytes as f64 / 1_048_576.0) -} - -/// The bytes of every file under a folder. -fn bytes_on_disk(folder: &Path) -> std::io::Result { - let mut bytes = 0; - for entry in std::fs::read_dir(folder)? { - let entry = entry?; - let meta = entry.metadata()?; - bytes += if meta.is_dir() { - bytes_on_disk(&entry.path())? - } else { - meta.len() - }; - } - Ok(bytes) -} +static ALLOCATOR: support::Counting = support::Counting; /// The most memory the operating system has seen this process hold, where it says so. fn peak_resident_bytes() -> Option { @@ -87,7 +41,7 @@ fn peak_resident_bytes() -> Option { Some(kibibytes * 1024) } -/// The README's first program, at 20,000 rows of 1.6 KB. +/// 20,000 notes of 1.6 KB, one commit per row. fn readme(folder: &Path) -> Result<()> { const ROWS: u128 = 20_000; let db = Database::open(folder)?; @@ -108,7 +62,7 @@ fn readme(folder: &Path) -> Result<()> { inserted?; db.close()?; line("readme_rows", ROWS, "rows"); - line("readme_insert", format!("{:.1}", insert.as_secs_f64()), "s"); + line("readme_insert", seconds(insert), "s"); let on_disk = bytes_on_disk(folder).expect("the store's folder can be listed"); line("readme_store_size", mebibytes(on_disk), "MiB"); @@ -153,11 +107,7 @@ fn million(folder: &Path) -> Result<()> { loaded?; let (sum, scan) = timed(|| db.execute("SELECT SUM(n) FROM small", ¶ms! {})); let sum = sum?; - line( - "million_rows_load", - format!("{:.1}", load.as_secs_f64()), - "s", - ); + line("million_rows_load", seconds(load), "s"); line("million_rows_scan", milliseconds(scan), "ms"); line( "million_rows_scan_rows_examined", @@ -166,7 +116,7 @@ fn million(folder: &Path) -> Result<()> { ); line( "million_rows_peak_allocated", - mebibytes(PEAK.load(Ordering::Relaxed) as u64), + mebibytes(support::take_peak_allocated()), "MiB", ); if let Some(bytes) = peak_resident_bytes() { @@ -175,29 +125,170 @@ fn million(folder: &Path) -> Result<()> { db.close() } -fn main() -> Result<()> { - let arguments: Vec = std::env::args().skip(1).collect(); - let [what, folder] = arguments.as_slice() else { - eprintln!("usage: measure readme|million "); - std::process::exit(2); +/// The numbers a vector is made of: the same for the same `seed` on every machine, each from +/// -1 up to 1. +fn vector_of(seed: u64, dimensions: usize) -> Vec { + let mut state = seed; + let mut next = || { + // SplitMix64. + state = state.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = state; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) }; - let folder = Path::new(folder); - if folder - .read_dir() - .is_ok_and(|mut entries| entries.next().is_some()) - { - eprintln!( - "{} is not empty; give a folder that is missing or empty", - folder.display() - ); - std::process::exit(2); + (0..dimensions) + .map(|_| (next() >> 40) as f32 / (1u64 << 23) as f32 - 1.0) + .collect() +} + +/// 20,000 vectors of 384 numbers, sealed, and the 10 nearest of 100 others. +fn vectors(folder: &Path) -> Result<()> { + const ROWS: i64 = 20_000; + const PER_COMMIT: i64 = 500; + const DIMENSIONS: usize = 384; + let db = Database::open(folder)?; + db.execute( + "CREATE TABLE items (id INTEGER PRIMARY KEY, embedding VECTOR(384))", + ¶ms! {}, + )?; + for first in (1..=ROWS).step_by(PER_COMMIT as usize) { + let tx = db.transaction()?; + for id in first..first + PER_COMMIT { + tx.execute( + "INSERT INTO items (id, embedding) VALUES ($id, $embedding)", + ¶ms! { "id" => id, "embedding" => vector_of(id as u64, DIMENSIONS) }, + )?; + } + tx.commit()?; } - match what.as_str() { - "readme" => readme(folder), - "million" => million(folder), - _ => { - eprintln!("usage: measure readme|million "); - std::process::exit(2); + while db + .run_maintenance_cycle()? + .vector + .maintenance_vectors_remaining + > 0 + {} + let mut times = Vec::new(); + for n in -5..100 { + let q = vector_of((ROWS + 6 + n) as u64, DIMENSIONS); + let (found, time) = timed(|| { + db.execute( + "SELECT id, embedding <=> $q AS distance FROM items ORDER BY distance LIMIT 10", + ¶ms! { "q" => q }, + ) + }); + found?; + if n >= 0 { + times.push(time); } } + db.close()?; + line( + "vectors_20000_find_p95", + milliseconds(percentile(×, 0.95)), + "ms", + ); + let on_disk = bytes_on_disk(folder).expect("the store's folder can be listed"); + line("vectors_20000_store_size", mebibytes(on_disk), "MiB"); + let held = ROWS as u64 * (8 + 4 * DIMENSIONS as u64); + line("vectors_20000_size_over_held", ratio(on_disk, held), "x"); + Ok(()) +} + +/// Stops the program with a sentence that holds nothing of a session file. +fn stop(sentence: &str, position: usize) -> ! { + eprintln!("file {position} of the list {sentence}"); + std::process::exit(1); +} + +/// The lines of the first `count` files the list names, one row each. +fn sessions(corpus: &Path, list: &Path, count: usize, folder: &Path) -> Result<()> { + const PER_COMMIT: u32 = 500; + let names = std::fs::read_to_string(list).expect("the list file can be read"); + let names = names.lines().filter(|name| !name.is_empty()).take(count); + let db = Database::open(folder)?; + db.execute( + "CREATE TABLE session_lines (id INTEGER PRIMARY KEY, body TEXT NOT NULL)", + ¶ms! {}, + )?; + let (mut files, mut source, mut put_in, mut id) = (0usize, 0u64, 0u64, 0i64); + let (loaded, load) = timed(|| -> Result<()> { + let mut tx = db.transaction()?; + let mut in_commit = 0; + for name in names { + files += 1; + let Ok(file) = File::open(corpus.join(name)) else { + stop("cannot be opened", files); + }; + let mut file = BufReader::new(file); + loop { + let mut bytes = Vec::new(); + match file.read_until(b'\n', &mut bytes) { + Ok(0) => break, + Ok(read) => source += read as u64, + Err(_) => stop("cannot be read", files), + } + if bytes.last() == Some(&b'\n') { + bytes.pop(); + } + if bytes.is_empty() { + continue; + } + let Ok(body) = String::from_utf8(bytes) else { + stop("has a line that is not UTF-8", files); + }; + id += 1; + put_in += body.len() as u64 + 8; + tx.execute( + "INSERT INTO session_lines (id, body) VALUES ($id, $body)", + ¶ms! { "id" => id, "body" => body }, + )?; + in_commit += 1; + if in_commit == PER_COMMIT { + std::mem::replace(&mut tx, db.transaction()?).commit()?; + in_commit = 0; + } + } + } + tx.commit() + }); + loaded?; + db.close()?; + let on_disk = bytes_on_disk(folder).expect("the store's folder can be listed"); + let name = |what: &str| format!("sessions_{files}_{what}"); + line(&name("source_size"), mebibytes(source), "MiB"); + line(&name("rows"), id, "rows"); + line(&name("load"), seconds(load), "s"); + line(&name("put_in_size"), mebibytes(put_in), "MiB"); + line(&name("store_size"), mebibytes(on_disk), "MiB"); + line(&name("size_over_source"), ratio(on_disk, source), "x"); + line(&name("size_over_put_in"), ratio(on_disk, put_in), "x"); + Ok(()) +} + +fn usage() -> ! { + eprintln!( + "usage: measure readme|million|vectors \n \ + measure sessions " + ); + std::process::exit(2); +} + +fn main() -> Result<()> { + let arguments: Vec = std::env::args().skip(1).collect(); + let arguments: Vec<&str> = arguments.iter().map(String::as_str).collect(); + let Some(folder) = arguments.last().map(Path::new) else { + usage(); + }; + support::must_be_empty(folder); + match arguments.as_slice() { + ["readme", _] => readme(folder), + ["million", _] => million(folder), + ["vectors", _] => vectors(folder), + ["sessions", corpus, list, count, _] => match count.parse() { + Ok(count) => sessions(Path::new(corpus), Path::new(list), count, folder), + Err(_) => usage(), + }, + _ => usage(), + } } diff --git a/v2/contextdb/examples/passages.rs b/v2/contextdb/examples/passages.rs new file mode 100644 index 00000000..2b27bf7b --- /dev/null +++ b/v2/contextdb/examples/passages.rs @@ -0,0 +1,367 @@ +//! Loads a dump of search passages into a store and measures searches on them, printing one +//! number per line as `name value unit`. +//! +//! `passages ` +//! +//! The dump folder holds `rows.bin` (36 bytes for each passage, numbers low byte first: a row +//! number `u64`, the passage's id, the index `u32` of its Context, the index `u32` of its +//! (Context, generation) pair, and flags `u32`: 1 current, 2 searchable, 4 has a vector), +//! `vectors.f32` (768 `f32` for each passage, in the same order), `contexts.txt` (`index uuid +//! name` on each line), `pairs.txt` (`index context-index generation-uuid active`) and +//! `queries.f32` (768 `f32` for each question). The answers folder holds `engine-one.jsonl`, +//! `engine-several.jsonl` and `engine-all.jsonl`: on each line the ten answers recorded for one +//! question, `{"q": n, "ids": [...], "scores": [...]}`, a score being one minus the distance, +//! rounded to 32 bits. The store folder must be missing or empty. +//! +//! What it does, in order: +//! +//! 1. opens the store with no memory limit and inserts the passages, 500 to a commit; +//! 2. searches all Contexts for every question once through a read session, with every vector +//! still waiting to be sealed; +//! 3. calls `run_maintenance_cycle` until nothing waits, closes the store and takes its size; +//! 4. through a read session with no memory limit, then one under 1 GiB, then one under +//! 64 MiB, searches one Context (the first), several (the first three) and all of them: 5 +//! searches that are not timed, then every question, each `execute` call timed whole; +//! 5. compares the answers of the session with no limit with the recorded ones. +//! +//! A time named `p50` or `p95` is the one at position `round((count - 1) * share)` of the +//! times in order. "Held" is the bytes of the values put in: 50 for a passage's five other +//! columns and 3,072 for its vector. +//! +//! This program measures. It never passes or fails, no test reads what it prints, and +//! `numbers.sh` is what runs it. It prints nothing of a passage but counts, sizes and times. + +mod support; + +use std::collections::{BTreeSet, HashMap}; +use std::fs::File; +use std::io::{BufReader, Read}; +use std::path::Path; +use std::time::Duration; + +use contextdb::serde_json::Value as Json; +use contextdb::{Database, ReadSession, ReadSessionOptions, Uuid, Value, params}; +use support::{bytes_on_disk, line, mebibytes, milliseconds, percentile, ratio, seconds, timed}; + +#[global_allocator] +static ALLOCATOR: support::Counting = support::Counting; + +type Fallible = Result>; + +const DIMENSIONS: usize = 768; +const PER_COMMIT: usize = 500; +const UNTIMED: usize = 5; +const SCOPES: [&str; 3] = ["one", "several", "all"]; + +/// What the dump says beside its passages. +struct Dump { + /// Each Context's id, by index. + contexts: Vec, + /// Each pair's Context index, generation, and whether it is the Context's active one. + pairs: Vec<(usize, Uuid, bool)>, + /// The question vectors. + questions: Vec>, +} + +/// What the searches of one scope gave: for each question the time, the rows read, and the +/// ids and distances returned. +#[derive(Default)] +struct Finds { + first: Duration, + times: Vec, + examined: Vec, + answers: Vec>, +} + +fn floats(bytes: &[u8]) -> Vec { + let four = bytes.chunks_exact(4); + four.map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])) + .collect() +} + +fn read_dump(folder: &Path) -> Fallible { + let mut contexts = Vec::new(); + for text in std::fs::read_to_string(folder.join("contexts.txt"))?.lines() { + contexts.push(text.split(' ').nth(1).ok_or("contexts.txt")?.parse()?); + } + let mut pairs = Vec::new(); + for text in std::fs::read_to_string(folder.join("pairs.txt"))?.lines() { + let fields: Vec<&str> = text.split(' ').collect(); + let [_, context, generation, active] = fields.as_slice() else { + return Err("pairs.txt".into()); + }; + pairs.push((context.parse()?, generation.parse()?, *active == "1")); + } + let questions = floats(&std::fs::read(folder.join("queries.f32"))?); + let questions = questions.chunks_exact(DIMENSIONS).map(<[f32]>::to_vec); + Ok(Dump { + contexts, + pairs, + questions: questions.collect(), + }) +} + +/// Inserts every passage of the dump. Returns how many there were and the bytes held. +fn load(db: &Database, folder: &Path, dump: &Dump) -> Fallible<(u64, u64)> { + let rows = File::open(folder.join("rows.bin"))?; + let passages = rows.metadata()?.len() / 36; + let mut rows = BufReader::new(rows); + let mut vectors = BufReader::new(File::open(folder.join("vectors.f32"))?); + let (mut record, mut vector) = ([0u8; 36], vec![0u8; 4 * DIMENSIONS]); + let mut held = 0; + let mut tx = db.transaction()?; + for passage in 1..=passages as usize { + rows.read_exact(&mut record)?; + vectors.read_exact(&mut vector)?; + let number = |at: usize| { + u32::from_le_bytes([record[at], record[at + 1], record[at + 2], record[at + 3]]) + }; + let (context, generation, _) = dump.pairs[number(28) as usize]; + let flags = number(32); + let embedding = if flags & 4 == 0 { + Value::Null + } else { + Value::Vector(floats(&vector)) + }; + held += 3 * 16 + + 2 + + if flags & 4 == 0 { + 0 + } else { + 4 * DIMENSIONS as u64 + }; + tx.execute( + "INSERT INTO event_search_passages \ + (id, context_id, generation_id, vector_text, is_current, is_searchable) \ + VALUES ($id, $context, $generation, $vector, $current, $searchable)", + ¶ms! { + "id" => Uuid::from_slice(&record[8..24])?, + "context" => dump.contexts[context], + "generation" => generation, + "vector" => embedding, + "current" => flags & 1 != 0, + "searchable" => flags & 2 != 0, + }, + )?; + if passage % PER_COMMIT == 0 { + std::mem::replace(&mut tx, db.transaction()?).commit()?; + } + } + tx.commit()?; + Ok((passages, held)) +} + +/// The search statement for the active pairs of the scope's Contexts, with its values. +fn statement(dump: &Dump, scope: &[usize], question: &[f32]) -> (String, HashMap) { + let mut values = params! { + "query" => question.to_vec(), "is_current" => true, "is_searchable" => true, + }; + let active = dump.pairs.iter().filter(|pair| pair.2); + let wanted = active.filter(|pair| scope.contains(&pair.0)); + let pairs: BTreeSet<(Uuid, Uuid)> = + wanted.map(|pair| (dump.contexts[pair.0], pair.1)).collect(); + let mut names = Vec::new(); + for (i, (context, generation)) in pairs.into_iter().enumerate() { + values.insert(format!("search_context_{i}"), Value::Uuid(context)); + values.insert(format!("search_generation_{i}"), Value::Uuid(generation)); + names.push(format!("($search_context_{i}, $search_generation_{i})")); + } + let filter = if names.len() == 1 { + "context_id = $search_context_0 AND generation_id = $search_generation_0".to_owned() + } else { + format!("(context_id, generation_id) IN ({})", names.join(", ")) + }; + let sql = format!( + "SELECT id, context_id, generation_id, vector_text <=> $query AS distance \ + FROM event_search_passages WHERE {filter} \ + AND is_current = $is_current AND is_searchable = $is_searchable \ + ORDER BY distance LIMIT 10" + ); + (sql, values) +} + +/// The Contexts of a scope, by index. +fn contexts_of(scope: &str, dump: &Dump) -> Vec { + match scope { + "one" => vec![0], + "several" => vec![0, 1, 2], + _ => (0..dump.contexts.len()).collect(), + } +} + +/// Searches a scope for every question, after `untimed` searches that are not recorded. +fn find(session: &ReadSession, dump: &Dump, scope: &str, untimed: usize) -> Fallible { + let scope = contexts_of(scope, dump); + let mut finds = Finds::default(); + for run in 0..untimed + dump.questions.len() { + let question = run.checked_sub(untimed); + let asked = question.unwrap_or(run % dump.questions.len()); + let (sql, values) = statement(dump, &scope, &dump.questions[asked]); + let (found, time) = timed(|| session.execute(&sql, &values)); + let found = found?; + if run == 0 { + finds.first = time; + } + if question.is_none() { + continue; + } + let answer = found.rows().iter().map(|row| match (&row[0], &row[3]) { + (Value::Uuid(id), Value::Float64(distance)) => Ok((*id, *distance)), + _ => Err("a search row is not an id and a distance"), + }); + finds.answers.push(answer.collect::>()?); + finds.times.push(time); + finds.examined.push(found.rows_examined()); + } + Ok(finds) +} + +/// The three scopes through a read session opened under `limit`. +fn find_under(store: &Path, limit: Option, dump: &Dump) -> Fallible> { + let mut options = ReadSessionOptions::default(); + options.memory_limit = limit; + let session = ReadSession::open_with_options(store, options)?; + let each = SCOPES.iter(); + each.map(|scope| find(&session, dump, scope, UNTIMED)) + .collect() +} + +/// Compares a scope's answers with the recorded ones: how many questions have the same ten +/// ids, how many have them in the same order, and the largest difference between a recorded +/// score and one minus the distance, rounded to 32 bits, of the same id. +fn compare(answers: &[Vec<(Uuid, f64)>], recorded: &Path) -> Fallible<(usize, usize, f64)> { + let (mut same_ten, mut same_order, mut gap) = (0, 0, 0f64); + for text in std::fs::read_to_string(recorded)?.lines() { + let recorded: Json = contextdb::serde_json::from_str(text)?; + let (Some(question), Some(ids), Some(scores)) = ( + recorded["q"].as_u64(), + recorded["ids"].as_array(), + recorded["scores"].as_array(), + ) else { + continue; // the summary line + }; + let ids = ids.iter().map(|id| id.as_str().unwrap_or_default().parse()); + let ids: Vec = ids.collect::>()?; + let answer = answers + .get(question as usize) + .ok_or("a question out of range")?; + let ours: Vec = answer.iter().map(|(id, _)| *id).collect(); + same_order += usize::from(ours == ids); + let as_set = |ids: &[Uuid]| ids.iter().copied().collect::>(); + same_ten += usize::from(as_set(&ours) == as_set(&ids)); + for (id, distance) in answer { + if let Some(at) = ids.iter().position(|recorded| recorded == id) { + let score = scores[at].as_f64().ok_or("a score that is not a number")?; + gap = gap.max((f64::from((1.0 - distance) as f32) - score).abs()); + } + } + } + Ok((same_ten, same_order, gap)) +} + +fn main() -> Fallible<()> { + let arguments: Vec = std::env::args().skip(1).collect(); + let [dump_folder, answers, store] = arguments.as_slice() else { + eprintln!("usage: passages "); + std::process::exit(2); + }; + let (dump_folder, answers, store) = + (Path::new(dump_folder), Path::new(answers), Path::new(store)); + support::must_be_empty(store); + let dump = read_dump(dump_folder)?; + + let db = Database::open(store)?; + db.execute( + "CREATE TABLE event_search_passages ( + id UUID PRIMARY KEY, + context_id UUID NOT NULL, + generation_id UUID NOT NULL, + vector_text VECTOR(768) PARTITION_KEY (context_id, generation_id) + MAX_PARTITIONS 4096 SEARCH_MODE AUTO, + is_current BOOLEAN NOT NULL, + is_searchable BOOLEAN NOT NULL DEFAULT true + )", + ¶ms! {}, + )?; + support::take_peak_allocated(); + let (loaded, load_time) = timed(|| load(&db, dump_folder, &dump)); + let (passages, held) = loaded?; + let load_peak = support::take_peak_allocated(); + let name = |what: &str| format!("passages_{passages}_{what}"); + line(&name("load"), seconds(load_time), "s"); + + let unsealed = find(&ReadSession::open(store)?, &dump, "all", 0)?; + + let mut cycles = 0; + let (sealed, seal_time) = timed(|| -> contextdb::Result<()> { + loop { + cycles += 1; + let report = db.run_maintenance_cycle()?; + if report.vector.maintenance_vectors_remaining == 0 { + return Ok(()); + } + } + }); + sealed?; + db.close()?; + let on_disk = bytes_on_disk(store)?; + line(&name("seal"), seconds(seal_time), "s"); + line(&name("seal_cycles"), cycles, "cycles"); + line(&name("store_size"), mebibytes(on_disk), "MiB"); + line( + &name("bytes_per_passage"), + on_disk / passages.max(1), + "bytes", + ); + line(&name("size_over_held"), ratio(on_disk, held), "x"); + let p95 = |times: &[Duration]| milliseconds(percentile(times, 0.95)); + line(&name("find_unsealed_all_p95"), p95(&unsealed.times), "ms"); + let rows = percentile(&unsealed.examined, 0.5); + line(&name("find_unsealed_all_rows_examined_p50"), rows, "rows"); + + support::take_peak_allocated(); + let unlimited = find_under(store, None, &dump)?; + for (scope, finds) in SCOPES.iter().zip(&unlimited) { + let p50 = milliseconds(percentile(&finds.times, 0.5)); + line(&name(&format!("find_{scope}_p50")), p50, "ms"); + line(&name(&format!("find_{scope}_p95")), p95(&finds.times), "ms"); + } + let all = &unlimited[2]; + line(&name("find_all_first"), milliseconds(all.first), "ms"); + line( + &name("find_all_max"), + milliseconds(percentile(&all.times, 1.0)), + "ms", + ); + for (limit, bytes) in [("1g", 1 << 30), ("64m", 64 << 20)] { + let limited = find_under(store, Some(bytes), &dump)?; + let name = name(&format!("find_all_p95_limit_{limit}")); + line(&name, p95(&limited[2].times), "ms"); + } + for (what, share) in [("p50", 0.5), ("p95", 0.95), ("max", 1.0)] { + let rows = percentile(&all.examined, share); + line( + &name(&format!("find_all_rows_examined_{what}")), + rows, + "rows", + ); + } + + let mut compared = Vec::new(); + for (scope, finds) in SCOPES.iter().zip(&unlimited) { + let recorded = answers.join(format!("engine-{scope}.jsonl")); + compared.push(compare(&finds.answers, &recorded)?); + } + for (scope, (same_ten, _, _)) in SCOPES.iter().zip(&compared) { + line(&name(&format!("same_ten_{scope}")), same_ten, "of_100"); + } + for (scope, (_, same_order, _)) in SCOPES.iter().zip(&compared) { + line(&name(&format!("same_order_{scope}")), same_order, "of_100"); + } + let gap = compared.iter().map(|(_, _, gap)| *gap).fold(0.0, f64::max); + line(&name("distance_gap_max"), gap, "-"); + line(&name("load_peak_allocated"), mebibytes(load_peak), "MiB"); + let find_peak = support::take_peak_allocated(); + line(&name("find_peak_allocated"), mebibytes(find_peak), "MiB"); + Ok(()) +} diff --git a/v2/contextdb/examples/support/mod.rs b/v2/contextdb/examples/support/mod.rs new file mode 100644 index 00000000..19e57b98 --- /dev/null +++ b/v2/contextdb/examples/support/mod.rs @@ -0,0 +1,105 @@ +//! What the measuring programs of this folder share: an allocator that counts, the one timer, +//! the size of a folder, and the `name value unit` lines they print. + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::fmt::Display; +use std::path::Path; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; + +/// The system allocator, counting the bytes handed out and not yet given back. A program makes +/// it its global allocator. +pub struct Counting; + +static IN_USE: AtomicUsize = AtomicUsize::new(0); +static PEAK: AtomicUsize = AtomicUsize::new(0); + +// SAFETY: every request goes to the system allocator unchanged; this only counts its bytes. +unsafe impl GlobalAlloc for Counting { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let in_use = IN_USE.fetch_add(layout.size(), Ordering::Relaxed) + layout.size(); + PEAK.fetch_max(in_use, Ordering::Relaxed); + unsafe { System.alloc(layout) } + } + + unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) { + IN_USE.fetch_sub(layout.size(), Ordering::Relaxed); + unsafe { System.dealloc(pointer, layout) } + } +} + +/// The most bytes the process had allocated at one moment since the last call, or since it +/// started. The next count starts from what is allocated now. +pub fn take_peak_allocated() -> u64 { + PEAK.swap(IN_USE.load(Ordering::Relaxed), Ordering::Relaxed) as u64 +} + +/// Runs `work` and says how long it took. Apart from the library's `system` part, this is the +/// only place in the workspace that reads a clock. +#[allow( + clippy::disallowed_methods, + reason = "measuring elapsed time is what these programs are for" +)] +pub fn timed(work: impl FnOnce() -> T) -> (T, Duration) { + let start = Instant::now(); + let result = work(); + (result, start.elapsed()) +} + +pub fn line(name: &str, value: impl Display, unit: &str) { + println!("{name} {value} {unit}"); +} + +pub fn milliseconds(duration: Duration) -> String { + format!("{:.3}", duration.as_secs_f64() * 1000.0) +} + +pub fn seconds(duration: Duration) -> String { + format!("{:.1}", duration.as_secs_f64()) +} + +pub fn mebibytes(bytes: u64) -> String { + format!("{:.1}", bytes as f64 / 1_048_576.0) +} + +/// One number over another, to two places. +pub fn ratio(over: u64, under: u64) -> String { + format!("{:.2}", over as f64 / under as f64) +} + +/// The value a `share` of the values lie at or under: the one at position +/// `round((count - 1) * share)` when they are put in order. There must be at least one. +pub fn percentile(values: &[T], share: f64) -> T { + let mut sorted = values.to_vec(); + sorted.sort_unstable(); + sorted[((sorted.len() - 1) as f64 * share).round() as usize] +} + +/// The bytes of every file under a folder, by length. +pub fn bytes_on_disk(folder: &Path) -> std::io::Result { + let mut bytes = 0; + for entry in std::fs::read_dir(folder)? { + let entry = entry?; + let meta = entry.metadata()?; + bytes += if meta.is_dir() { + bytes_on_disk(&entry.path())? + } else { + meta.len() + }; + } + Ok(bytes) +} + +/// Stops the program with a sentence unless the folder is missing or empty. +pub fn must_be_empty(folder: &Path) { + if folder + .read_dir() + .is_ok_and(|mut entries| entries.next().is_some()) + { + eprintln!( + "{} is not empty; give a folder that is missing or empty", + folder.display() + ); + std::process::exit(2); + } +} diff --git a/v2/contextdb/src/database/mod.rs b/v2/contextdb/src/database/mod.rs index f9d1e0e1..939ed370 100644 --- a/v2/contextdb/src/database/mod.rs +++ b/v2/contextdb/src/database/mod.rs @@ -19,6 +19,8 @@ use crate::sql::{self, Effect, Plan}; use crate::store::{ChangeSet, Read, Store, View, WrapStorage}; use crate::system::{self, Clock}; use crate::types::{Error, Result, TxId, Value}; +use crate::upkeep::{self, MaintenanceReport}; +use crate::vector::{self, Tally}; use crate::write; /// What a host may say when it opens a store for writing. @@ -42,8 +44,11 @@ pub struct DatabaseOpenOptions { /// ever declared. A limit given here holds until the store is closed and leaves what the /// store remembers as it is. /// - /// Every open sets 9 MiB aside out of the limit, 8 MiB of it for the store's page cache. A - /// smaller limit is refused with [`crate::Error::MemoryBudgetExceeded`] and opens no store. + /// Every open sets aside, out of the limit, 1 MiB and the store's page cache. The cache is + /// a quarter of the limit in force and never less than 8 MiB, which is also its size when + /// no limit is in force; it is chosen when the store is opened and keeps its size until the + /// store is closed. A limit below 9 MiB is refused with + /// [`crate::Error::MemoryBudgetExceeded`] and opens no store. /// /// A [`crate::ReadSession`] never takes the limit the store remembers: `None` in its own /// options means no limit at all, so a reader that needs a ceiling is given one there. @@ -108,6 +113,9 @@ struct Staged { catalog: Arc, /// The declarations it ran, in order, for a commit that has to run them again. declared: Vec, + /// The vector partitions its statements have made and emptied, for the limit a statement + /// is refused by. The commit function does not read it: it counts again. + partitions: Tally, } fn planned(sql: &str) -> Result { @@ -146,6 +154,7 @@ impl Open { changes: ChangeSet::new(view, &self.budget)?, catalog, declared: Vec::new(), + partitions: Tally::new(&self.budget)?, }) } @@ -174,8 +183,17 @@ impl Open { return Err(Error::MixedSchemaAndRowWrite); } let (catalog, changes) = (&staged.catalog, &mut staged.changes); - write::apply(plan, params, catalog, changes, &self.budget, now) - .map(QueryResult::affected) + let partitions = &mut staged.partitions; + write::apply( + plan, + params, + catalog, + changes, + partitions, + &self.budget, + now, + ) + .map(QueryResult::affected) } Effect::Reads => { if let Plan::ShowMemoryLimit = plan { @@ -211,12 +229,14 @@ impl Open { /// its declarations again over the newest moment, because a table's and an index's number /// are counted on from the catalog and a new index is filled from the rows that are there; /// a refusal of that second run is the commit's refusal. One that wrote rows is laid over - /// the newest moment and `write` checks its rows against it. + /// the newest moment and `write` checks its rows against it. Then the vector part stages + /// what the rows mean for it, so a row and its place in a search commit together. fn commit(&self, _held: &MutexGuard<'_, ()>, staged: Staged) -> Result<()> { let Staged { mut changes, mut catalog, declared, + partitions: _, } = staged; let (view, tables) = self.newest()?; if changes.commit() != view.commit() { @@ -243,6 +263,7 @@ impl Open { if changes.is_empty() { return Ok(()); } + vector::stage(&mut changes, &catalog, &self.budget)?; let commit = self.store.commit(changes)?; // The commit is on disk whatever the read after it says. When that read fails no view // is kept and the next call reads again, so none answers from before this commit. @@ -259,6 +280,9 @@ impl Open { /// hold in this open and is committed for the store to remember, or neither happens. fn set_memory_limit(&self, limit: Option) -> Result<()> { let held = locked(&self.committing); + // Refused here, before anything changes, when the page cache of this open and what is + // in use do not fit under the limit. + self.store.count_cache_under(limit)?; let mut staged = self.begin()?; staged.changes.put_memory_limit(limit)?; let before = self.budget.limit(); @@ -389,13 +413,18 @@ impl Database { /// Declares the memory limit in bytes, or no limit for `None`. /// /// The limit takes hold at once and is committed, so the store remembers it for every - /// later open whose options name no limit, exactly as `SET MEMORY_LIMIT` does. + /// later open whose options name no limit, exactly as `SET MEMORY_LIMIT` does. The page + /// cache keeps the size this open gave it; the next open sizes it from the limit. /// /// # Errors /// - /// [`crate::Error::MemoryBudgetExceeded`] when the limit is below what is in use now, which - /// is never less than the 9 MiB every open sets aside. [`crate::Error::StoreIo`] when the - /// commit cannot be written. In both cases the limit stays what it was. + /// [`crate::Error::MemoryBudgetExceeded`] when the limit is below the smallest this open + /// can take: 1 MiB, its page cache, and what open transactions and results not yet dropped + /// hold. That is never less than 9 MiB, and is 9 MiB plus what is held for a store opened + /// with no limit in force. The refusal names it; a limit below the cache of a store opened + /// under a larger limit is given as [`DatabaseOpenOptions::memory_limit`] when the store is + /// opened. [`crate::Error::StoreIo`] when the commit cannot be written. In both cases the + /// limit stays what it was. /// [`crate::Error::DatabaseClosed`] on a handle that was closed. pub fn set_memory_limit(&self, limit: Option) -> Result<()> { self.opened()? @@ -467,6 +496,41 @@ impl Database { locked(&self.must_be_open().newest).commit } + /// Runs one cycle of upkeep and returns what it did. + /// + /// A cycle moves at most 4,096 vectors from the pending lists of VECTOR columns into blocks, + /// in one commit, and reports how many it moved and how many still wait. Searches give the + /// same rows before and after a cycle; they read fewer rows after it, because a vector in a + /// block is ruled in or out from its one-byte copy and a pending vector is read from its + /// row. A store with nothing pending is left as it is, and no commit is made. + /// + /// Under a memory limit with no room for all of that, a cycle seals the partitions it has + /// room for, at least one, and the next cycle goes on from there. + /// + /// Nothing in this version calls a cycle on its own: a host that searches a VECTOR column + /// calls this after it has written, and again while + /// `vector.maintenance_vectors_remaining` is not zero. Other calls on the handle wait for + /// the cycle's commit as they wait for any other commit; readers do not wait. + /// + /// # Errors + /// + /// [`crate::Error::MemoryBudgetExceeded`] when the memory limit has no room to seal the + /// pending vectors of one partition: a call holds one block's bytes (256 KiB less 64) for its + /// whole run, the keys of the pending vectors it collected (at most 4,096, each key's bytes + /// plus 24), the partition's record, and one row with its vector at a time; + /// [`crate::Error::StoreIo`] or + /// [`crate::Error::StoreCorrupted`] when the store cannot be read or written; + /// [`crate::Error::DatabaseClosed`] on a handle that was closed. In each case nothing was + /// sealed and every vector is still found by a search. + pub fn run_maintenance_cycle(&self) -> Result { + let open = self.opened()?; + let held = locked(&open.committing); + let mut staged = open.begin()?; + let report = upkeep::cycle(&mut staged.changes, &staged.catalog, &open.budget)?; + open.commit(&held, staged)?; + Ok(report) + } + /// Lets go of the store, so that the next writer can open it. Every commit that returned is /// already on disk, and closing writes nothing. While a call is still running on another /// thread, the store is let go when that call returns. Closing a handle that was closed @@ -514,6 +578,7 @@ impl Transaction<'_> { let result = open.run(&plan, params, staged); if result.is_err() { staged.changes.rewind(mark); + staged.partitions.forget(); } result } diff --git a/v2/contextdb/src/lib.rs b/v2/contextdb/src/lib.rs index 77df6f5a..0492c21a 100644 --- a/v2/contextdb/src/lib.rs +++ b/v2/contextdb/src/lib.rs @@ -9,6 +9,8 @@ mod sql; mod store; mod system; mod types; +mod upkeep; +mod vector; mod write; pub use database::{Database, DatabaseOpenOptions, Transaction}; @@ -16,7 +18,11 @@ pub use query::QueryResult; pub use read::{ReadSession, ReadSessionOptions}; pub use store::{StorageBackend, WrapStorage}; pub use system::Clock; -pub use types::{ColumnType, Error, Result, TxId, Uuid, Value}; +pub use types::{ + ColumnType, Error, Result, TxId, Uuid, Value, VectorIndexRef, VectorPartitionDeclarationIssue, + VectorPolicyDeclarationIssue, +}; +pub use upkeep::{MaintenanceReport, VectorMaintenanceReport}; /// The JSON package whose `Value` a [`Value::Json`] holds, so a host needs no copy of its own. pub use serde_json; diff --git a/v2/contextdb/src/memory/mod.rs b/v2/contextdb/src/memory/mod.rs index 719d1c54..d8e87225 100644 --- a/v2/contextdb/src/memory/mod.rs +++ b/v2/contextdb/src/memory/mod.rs @@ -28,9 +28,10 @@ const NO_CEILING: u64 = u64::MAX; /// names one. With no limit there is no ceiling, and what is reserved is still counted. /// /// Three kinds of memory are counted against it: small fixed things, covered by an allowance -/// of 1 MiB set aside in [`Budget::new`]; the page cache of 8 MiB, reserved once when the store -/// is opened; and everything that grows with data, which must hold a [`Reservation`]. A limit -/// below those 9 MiB therefore cannot open any store: the open is refused with +/// of 1 MiB set aside in [`Budget::new`]; the page cache, reserved once when the store is +/// opened; and everything that grows with data, which must hold a [`Reservation`]. The page +/// cache is a quarter of the limit and never less than 8 MiB, which is also its size with no +/// limit. A limit below 9 MiB therefore cannot open any store: the open is refused with /// `MemoryBudgetExceeded`. #[derive(Clone)] pub(crate) struct Budget { diff --git a/v2/contextdb/src/query/access.rs b/v2/contextdb/src/query/access.rs index bfe6428c..9410e380 100644 --- a/v2/contextdb/src/query/access.rs +++ b/v2/contextdb/src/query/access.rs @@ -1,15 +1,16 @@ -//! Index choice: which rows of a table a statement has to read, and reading them. +//! Index choice: which rows of a table a statement has to read, reading them, and working out +//! the select list for a row that was read. use std::collections::HashMap; use std::ops::Bound; -use super::expression::{column_index, eval}; +use super::expression::{column_index, eval, eval_found, most_bytes, reads}; use super::{heap_bytes, row_overhead}; use crate::memory::Reservation; use crate::schema::Table; use crate::sql::Expr; use crate::sql::ast::{BinOp, SortDirection as Asked}; -use crate::sql::plan::SortKey; +use crate::sql::plan::{ProjectColumn, SortKey}; use crate::store::{Cell, IndexId, Key, Read, Row}; use crate::types::{ColumnType, Result, SortDirection, Value}; @@ -65,7 +66,7 @@ fn starting_after(key: Key) -> Bound { /// The value as a key of a column of this type holds it, when there is one such value: only /// then is looking the value up the same as comparing every row with it. -fn as_key(column: ColumnType, value: Value) -> Option { +pub(super) fn as_key(column: ColumnType, value: Value) -> Option { match (column, value) { (ColumnType::Integer, value @ Value::Int64(_)) | (ColumnType::Real, value @ Value::Float64(_)) @@ -300,16 +301,9 @@ impl<'a> Source<'a> { let mut examined = 0; let mut offer = |key: Key, row: Row| { examined += 1; - let before = reservation.bytes(); match self.values(row, reservation)? { - Some(values) => { - let bytes = reservation.bytes() - before; - each(key, values, bytes, reservation) - } - None => { - reservation.shrink(reservation.bytes() - before); - Ok(true) - } + Some((values, bytes)) => each(key, values, bytes, reservation), + None => Ok(true), } }; let table = self.table.id; @@ -341,9 +335,83 @@ impl<'a> Source<'a> { Ok(examined) } - /// The values of a stored row if it passes the filter, charged to `reservation`. A column - /// the statement does not read is left as NULL. - fn values(&self, row: Row, reservation: &mut Reservation) -> Result>> { + /// The row with this key if it is there and passes the filter, as [`Source::walk`] hands a + /// row over: its values, and the bytes of `reservation` it holds. + pub(super) fn row( + &self, + key: &Key, + reservation: &mut Reservation, + ) -> Result, u64)>> { + match self.read.row(self.table.id, key)? { + Some(row) => self.values(row, reservation), + None => Ok(None), + } + } + + /// For each item of a select list, the column whose value it moves out of a row: the column + /// it names, when no other item of the list reads that column. Any other item copies what + /// it reads. Empty for `SELECT *`, which hands the row over as it is. + pub(super) fn moved(&self, columns: Option<&[ProjectColumn]>) -> Result>> { + let columns = columns.unwrap_or_default(); + let mut readers = vec![0; self.table.meta.columns.len()]; + for column in columns { + let mut read = vec![false; readers.len()]; + reads(&column.expr, self.table, self.params, &mut read)?; + for (readers, read) in readers.iter_mut().zip(read) { + *readers += usize::from(read); + } + } + let alone = |column: &ProjectColumn| match &column.expr { + Expr::Column(named) => column_index(self.table, named).ok(), + _ => None, + }; + let moved = columns.iter().map(alone); + Ok(moved.map(|at| at.filter(|at| readers[*at] == 1)).collect()) + } + + /// The select list worked out for one row this source handed over, or the row itself for + /// `SELECT *`. `moved` is [`Source::moved`] of the list, and `distance` is the row's distance + /// when a similarity search found it. The memory of the values is asked for before they + /// are made, and the row's bytes are given back once they are. A value that is moved is + /// held once: what it takes passes from the row to the result and is not asked for again. + pub(super) fn output( + &self, + columns: Option<&[ProjectColumn]>, + moved: &[Option], + mut row: Vec, + bytes: u64, + distance: Option, + reservation: &mut Reservation, + ) -> Result> { + let Some(columns) = columns else { + return Ok(row); + }; + let most = columns + .iter() + .map(|column| most_bytes(&column.expr, self.table, &row, self.params)); + let passed = moved.iter().flatten().map(|at| heap_bytes(&row[*at])); + let passed = passed.sum::().min(bytes); + reservation.grow(row_overhead(columns.len()) + most.sum::() - passed)?; + let values = columns + .iter() + .zip(moved) + .map(|(column, moved)| match moved { + Some(at) => Ok(std::mem::replace(&mut row[*at], Value::Null)), + None => { + let (params, now) = (self.params, self.now); + eval_found(&column.expr, self.table, &row, distance, params, now) + } + }) + .collect::>>()?; + reservation.shrink(bytes - passed); + Ok(values) + } + + /// The values of a stored row if it passes the filter, and the bytes of `reservation` they + /// hold; a row that does not pass holds none. A column the statement does not read is left + /// as NULL. + fn values(&self, row: Row, reservation: &mut Reservation) -> Result, u64)>> { + let before = reservation.bytes(); let mut cells = row.cells; reservation.grow(row_overhead(cells.len()))?; let mut values = vec![Value::Null; cells.len()]; @@ -351,10 +419,11 @@ impl<'a> Source<'a> { if let Some(filter) = self.filter && eval(filter, self.table, &values, self.params, self.now)? != Value::Bool(true) { + reservation.shrink(reservation.bytes() - before); return Ok(None); } self.take(&mut cells, &mut values, &self.rest, reservation)?; - Ok(Some(values)) + Ok(Some((values, reservation.bytes() - before))) } /// Moves the marked cells into `values`, loading a value kept beside its row. The memory diff --git a/v2/contextdb/src/query/expression.rs b/v2/contextdb/src/query/expression.rs index c0a95719..a22d7dff 100644 --- a/v2/contextdb/src/query/expression.rs +++ b/v2/contextdb/src/query/expression.rs @@ -15,6 +15,9 @@ use crate::types::{Error, Result, Uuid, Value}; /// one time in milliseconds, which `NOW()` answers with. A comparison with NULL, and one between /// values of types that do not compare, is unknown: it gives NULL, and a filter keeps a row /// only for TRUE. +/// +/// A distance, `column <=> vector`, has no value here: only a similarity search gives a row one +/// ([`eval_found`]), so it is refused by name. pub(crate) fn eval( expr: &Expr, table: &Table, @@ -22,7 +25,20 @@ pub(crate) fn eval( params: &HashMap, now: i64, ) -> Result { - let value = |expr: &Expr| eval(expr, table, row, params, now); + eval_found(expr, table, row, None, params, now) +} + +/// [`eval`] for a row a similarity search found. `distance` is the distance the search worked +/// out for the row, and every `<=>` in the expression is that number as it is. +pub(super) fn eval_found( + expr: &Expr, + table: &Table, + row: &[Value], + distance: Option, + params: &HashMap, + now: i64, +) -> Result { + let value = |expr: &Expr| eval_found(expr, table, row, distance, params, now); Ok(match expr { Expr::Column(column) => { let at = column_index(table, column)?; @@ -92,9 +108,70 @@ pub(crate) fn eval( } _ => return Err(refused(name)), }, + Expr::InList { + expr, + list, + negated, + } => { + // TRUE when an item equals the left side, else unknown when any equality was. + let left = side(expr, value)?; + let mut found = Some(false); + for item in list { + let item = side(item, value)?; + if item.len() != left.len() { + return Err(Error::PlanError( + "a row value beside IN needs as many values as its partner".to_owned(), + )); + } + match rows_equal(&left, &item) { + Some(true) => { + found = Some(true); + break; + } + Some(false) => {} + None => found = None, + } + } + found.map_or(Value::Null, |found| Value::Bool(found != *negated)) + } + Expr::Row(_) => { + return Err(Error::PlanError( + "a row value, (a, b), stands only on either side of IN".to_owned(), + )); + } + Expr::Distance { .. } => Value::Float64(distance.ok_or_else(|| { + Error::PlanError( + "<=> has a value only for the rows of a similarity search, SELECT ... ORDER BY \ + column <=> vector LIMIT k, and in that statement's select list" + .to_owned(), + ) + })?), }) } +/// The values of one side of IN: those of a row value, or the one value of anything else. +fn side(expr: &Expr, value: impl Fn(&Expr) -> Result) -> Result> { + match expr { + Expr::Row(items) => items.iter().map(value).collect(), + one => Ok(vec![value(one)?]), + } +} + +/// Whether two row values of one length are equal, as standard SQL has it: FALSE as soon as one +/// pair of columns is definitely different, whatever a NULL beside it holds; TRUE when every +/// pair is equal; otherwise unknown. One value is a row of one column. +fn rows_equal(left: &[Value], right: &[Value]) -> Option { + let mut equal = Some(true); + for (left, right) in left.iter().zip(right) { + match compare(left, right) { + Some(Ordering::Equal) => {} + Some(_) => return Some(false), + None => equal = None, + } + } + equal +} + /// The most memory the value of an expression can take outside its own cell. It is known /// before the expression is worked out, so the memory is asked for before the value is made. pub(super) fn most_bytes( @@ -123,8 +200,10 @@ pub(super) fn most_bytes( /// Marks in `read` each column the expression reads. /// /// Refuses, before any row is read, what [`eval`] would refuse at the first row: a column the -/// table does not have, a `$name` nothing is bound to, and a form this version does not work -/// out. A statement is therefore refused or not whatever rows its table holds. +/// table does not have, a `$name` nothing is bound to, and a function this version does not +/// work out. A statement is therefore refused or not whatever rows its table holds. A distance +/// and a row value are walked and not judged: the planner refuses the ones a SELECT cannot +/// hold, and [`eval`] any other when it meets it. pub(super) fn reads( expr: &Expr, table: &Table, @@ -154,6 +233,22 @@ pub(super) fn reads( reads(arg, table, params, read)?; } } + // The search gives a row's distance: the column on the left is looked up, not read. + Expr::Distance { left, right } => { + reads(left, table, params, &mut read.to_vec())?; + reads(right, table, params, read)?; + } + Expr::Row(items) => { + for item in items { + reads(item, table, params, read)?; + } + } + Expr::InList { expr, list, .. } => { + reads(expr, table, params, read)?; + for item in list { + reads(item, table, params, read)?; + } + } } Ok(()) } diff --git a/v2/contextdb/src/query/mod.rs b/v2/contextdb/src/query/mod.rs index 414a7587..ad87e195 100644 --- a/v2/contextdb/src/query/mod.rs +++ b/v2/contextdb/src/query/mod.rs @@ -1,11 +1,13 @@ //! query: answers a SELECT from one table: which rows to read and through which index, the -//! filter, the select list, COUNT, SUM, MIN and MAX, ORDER BY and LIMIT. It holds the one -//! expression evaluator, and counts exactly the rows it reads. +//! filter, the select list, COUNT, SUM, MIN and MAX, ORDER BY and LIMIT, and a similarity +//! search, whose rows the vector part finds. It holds the one expression evaluator, and counts +//! exactly the rows it reads. //! //! It never changes anything. mod access; mod expression; +mod search; #[cfg(test)] mod tests; @@ -14,6 +16,7 @@ use std::collections::HashMap; use std::fmt; pub(crate) use expression::eval; +pub(crate) use search::vector_from_text; use crate::memory::{Budget, Reservation}; use crate::schema::{Catalog, Table}; @@ -129,6 +132,9 @@ pub(crate) fn memory_limit_shown(budget: &Budget, opened_under: Option) -> /// table is answered from the kept count. Loads a value kept beside its row only when a selected /// or filtered column needs it. The result, and any sort or grouping on the way to it, grows /// through reservations on `budget` and refuses with `MemoryBudgetExceeded` before it grows. +/// +/// A similarity search, `ORDER BY column <=> vector LIMIT k`, is answered by the vector part +/// from the rows this part lets it see, and returns the nearest rows first. pub(crate) fn run( plan: &Plan, params: &HashMap, @@ -137,6 +143,9 @@ pub(crate) fn run( budget: &Budget, now: i64, ) -> Result { + if let Some(found) = search::run(plan, params, catalog, read, budget, now)? { + return Ok(found); + } let select = Select::of(plan)?; let table = catalog .table(select.table) @@ -152,12 +161,8 @@ pub(crate) fn run( reads(&key.expr, table, params, &mut wanted)?; } let mut aggregates = Vec::new(); - let columns = match select.columns { - None => { - wanted.fill(true); - let declared = table.meta.columns.iter(); - declared.map(|column| column.name.clone()).collect() - } + match select.columns { + None => wanted.fill(true), Some(columns) => { for column in columns { let Some(aggregate) = Aggregate::of(&column.expr)? else { @@ -173,19 +178,9 @@ pub(crate) fn run( if !aggregates.is_empty() && aggregates.len() < columns.len() { return Err(unsupported("a column beside an aggregate (GROUP BY)")); } - columns - .iter() - .map(|column| match (&column.alias, &column.expr) { - (Some(alias), _) => alias.clone(), - (None, Expr::Column(column)) => column.column.clone(), - (None, Expr::FunctionCall { name, .. }) if !aggregates.is_empty() => { - name.clone() - } - (None, _) => "expr".to_owned(), - }) - .collect() } - }; + } + let columns = names(table, select.columns, !aggregates.is_empty()); let source = Source::new(table, select.filter, params, now, read, filtered, &wanted); let limit = select.limit.map_or(usize::MAX, |limit| { @@ -217,24 +212,11 @@ pub(crate) fn run( rows.push(aggregates.into_iter().map(Aggregate::answer).collect()); } } else if limit > 0 { - // The select list worked out for one row. The memory of its values is asked for before - // they are made, and the row's bytes are given back once they are. - let output = - |row: Vec, bytes: u64, reservation: &mut Reservation| -> Result> { - let Some(columns) = select.columns else { - return Ok(row); - }; - let most = columns - .iter() - .map(|column| most_bytes(&column.expr, table, &row, params)); - reservation.grow(row_overhead(columns.len()) + most.sum::())?; - let values = columns - .iter() - .map(|column| eval(&column.expr, table, &row, params, now)) - .collect::>>()?; - reservation.shrink(bytes); - Ok(values) - }; + // The select list worked out for one row. + let moved = source.moved(select.columns)?; + let output = |row: Vec, bytes: u64, reservation: &mut Reservation| { + source.output(select.columns, &moved, row, bytes, None, reservation) + }; let (path, in_order) = source.choose(select.order)?; if in_order { examined = source.walk(&path, &mut reservation, |_, row, bytes, reservation| { @@ -330,6 +312,25 @@ pub(crate) fn matching( Ok((found, reservation)) } +/// What a result calls its columns: the table's own names for `SELECT *` (`None`), otherwise +/// each item's `AS` name, a bare column's name, an aggregate's name, and `expr` for any other +/// expression. +fn names(table: &Table, columns: Option<&[ProjectColumn]>, aggregates: bool) -> Vec { + let Some(columns) = columns else { + let declared = table.meta.columns.iter(); + return declared.map(|column| column.name.clone()).collect(); + }; + columns + .iter() + .map(|column| match (&column.alias, &column.expr) { + (Some(alias), _) => alias.clone(), + (None, Expr::Column(column)) => column.column.clone(), + (None, Expr::FunctionCall { name, .. }) if aggregates => name.clone(), + (None, _) => "expr".to_owned(), + }) + .collect() +} + /// One planned SELECT, as the planner nests it: LIMIT over the select list over ORDER BY over /// one table with its filter. struct Select<'a> { diff --git a/v2/contextdb/src/query/search.rs b/v2/contextdb/src/query/search.rs new file mode 100644 index 00000000..757c304c --- /dev/null +++ b/v2/contextdb/src/query/search.rs @@ -0,0 +1,300 @@ +//! A similarity search: the rows of a table that pass the statement's filter and lie nearest a +//! vector. This part settles what is searched for, which partitions the filter selects and +//! which rows may be seen, and works out each row's distance from its stored vector; the vector +//! part finds the rows. + +use std::collections::HashMap; + +use super::access::{Source, as_key}; +use super::expression::{column_index, eval, reads}; +use super::{QueryResult, names}; +use crate::memory::Budget; +use crate::schema::{Catalog, Table}; +use crate::sql::ast::{BinOp, ColumnRef}; +use crate::sql::{Expr, Plan}; +use crate::store::{Key, Read}; +use crate::types::{ColumnType, Error, Result, Value, VectorIndexRef, VectorSearchMode}; +use crate::vector::{self, Nearest}; + +/// The most key tuples a filter may name; past that the search reads every partition. +const MOST_TUPLES: usize = 4096; + +/// Answers a planned similarity search, a `VectorSearch` under the select list unless the +/// statement is `SELECT *`: the `k` rows that pass the filter and have a vector nearest the one +/// searched for, nearest first, rows at one distance in key order. `None` for any other plan. +/// +/// The column, the mode, the vector and every name of the statement are settled before a row +/// is read, so a refusal does not depend on what the table holds. Then the vector part searches +/// the partitions the filter selects and asks, row by row, for the distance: the row is read, +/// the whole filter applied, and the distance worked out from the row's stored vector. Each row +/// of the answer is read once more for the select list, in which every `<=>` is the distance +/// the row was ordered by. Every one of those reads is counted. +pub(super) fn run( + plan: &Plan, + params: &HashMap, + catalog: &Catalog, + read: &dyn Read, + budget: &Budget, + now: i64, +) -> Result> { + let (plan, list) = match plan { + Plan::Project { input, columns } => (&**input, Some(columns.as_slice())), + plan => (plan, None), + }; + let Plan::VectorSearch { + table, + column, + query, + k, + filter, + mode, + } = plan + else { + return Ok(None); + }; + let filter = filter.as_ref(); + let table = catalog + .table(table) + .ok_or_else(|| Error::TableNotFound(table.clone()))?; + let named = ColumnRef { + table: None, + column: column.clone(), + }; + let at = column_index(table, &named)?; + let index = VectorIndexRef { + table: table.name.clone(), + column: named.column, + }; + let Some(vector) = table.vector(column) else { + return Err(Error::UnknownVectorIndex { index }); + }; + if mode.unwrap_or_else(|| vector.mode()) == VectorSearchMode::Indexed { + return Err(Error::VectorIndexedRouteUnavailable { index }); + } + let near = match eval(query, table, &[], params, now)? { + Value::Vector(vector) => Some(vector), + Value::Text(text) => vector_from_text(&text), + _ => None, + }; + let Some(near) = near else { + return Err(Error::PlanError( + "the right side of <=> is not a vector: write [1.0, 0.0], its text '[1, 0]', or a \ + $name bound to a vector" + .to_owned(), + )); + }; + let actual = near.len(); + if let ColumnType::Vector(expected) = table.meta.columns[at].column_type + && expected != actual + { + return Err(Error::VectorIndexDimensionMismatch { + index, + expected, + actual, + }); + } + if let Some(position) = near.iter().position(|number| !number.is_finite()) { + return Err(Error::InvalidVectorValue { index, position }); + } + + let width = table.meta.columns.len(); + let (mut filtered, mut wanted) = (vec![false; width], vec![false; width]); + if let Some(filter) = filter { + reads(filter, table, params, &mut filtered)?; + } + match list { + None => wanted.fill(true), + Some(list) => { + for column in list { + reads(&column.expr, table, params, &mut wanted)?; + } + } + } + let columns = names(table, list, false); + let k = usize::try_from(*k).unwrap_or(usize::MAX); + if k == 0 { + return Ok(Some(QueryResult::from_rows(columns, Vec::new()))); + } + + let key = vector.partition_key.iter(); + let key = key.filter_map(|name| table.meta.columns.iter().position(|c| c.name == *name)); + let key = key.collect(); + let scope = Scope { + table, + key, + params, + now, + } + .keys(filter)?; + // A row the vector part asks about is read for the filter and its vector alone; a row of + // the answer for what the select list reads, so a large value is loaded only then. + let mut its_vector = vec![false; width]; + its_vector[at] = true; + let asked = Source::new(table, filter, params, now, read, filtered, &its_vector); + let returned = Source::new(table, None, params, now, read, vec![false; width], &wanted); + // What the vector part holds is given back once the rows are built; the rows stay counted + // with the result. + let mut searching = budget.reserve(0, "query")?; + let mut reservation = budget.reserve(0, "query")?; + let mut examined = 0; + + let nearest = Nearest { + table: table.id, + vector: vector.id, + scope: &scope, + query: &near, + k, + }; + let mut distance_of = |key: &Key| -> Result> { + examined += 1; + let Some((row, bytes)) = asked.row(key, &mut reservation)? else { + return Ok(None); + }; + let distance = match &row[at] { + Value::Vector(stored) => Some(vector::distance(stored, &near)), + _ => None, + }; + reservation.shrink(bytes); + Ok(distance) + }; + let hits = vector::search(&nearest, read, &mut searching, &mut distance_of)?; + + let moved = returned.moved(list)?; + let mut rows = Vec::new(); + for hit in &hits { + examined += 1; + if let Some((row, bytes)) = returned.row(&hit.key, &mut reservation)? { + let distance = Some(hit.distance); + rows.push(returned.output(list, &moved, row, bytes, distance, &mut reservation)?); + } + } + Ok(Some(QueryResult { + columns, + rows, + rows_affected: 0, + rows_examined: examined, + _reservation: Some(reservation), + })) +} + +/// The values a filter fixes for the PARTITION_KEY columns, in their order; `None` for a +/// column it leaves open. +type Tuple = Vec>; + +/// Reads a filter for the partitions it selects. +struct Scope<'a> { + table: &'a Table, + /// Where each PARTITION_KEY column is in a row, in declared order. + key: Vec, + params: &'a HashMap, + now: i64, +} + +impl Scope<'_> { + /// The partitions to search: for each key tuple the filter names, the key of the leading + /// key columns it fixes, which every partition with those values begins with. Empty for + /// every partition: no filter, a tuple that fixes no leading column, or more than 4,096 + /// tuples. It may select too much and never too little: the whole filter is still applied + /// to every row. + fn keys(&self, filter: Option<&Expr>) -> Result> { + let (Some(filter), false) = (filter, self.key.is_empty()) else { + return Ok(Vec::new()); + }; + let tuples = self.said(filter)?; + if tuples.len() > MOST_TUPLES { + return Ok(Vec::new()); + } + let mut keys = Vec::with_capacity(tuples.len()); + for tuple in tuples { + let leading: Vec = tuple.into_iter().map_while(|value| value).collect(); + if leading.is_empty() { + return Ok(Vec::new()); + } + keys.push(Key::from_values(&leading, &[])?); + } + // A partition is named once, and not again under a key that begins as another does. + keys.sort(); + keys.dedup_by(|later, kept| later.starts_with(kept)); + Ok(keys) + } + + /// What a filter says about the key columns: a row passes it only under one of these + /// tuples. `AND` joins what its two sides say, `OR` gives the tuples of both, and + /// `column = value`, `column IN (values)` and `(columns) IN (row values)` say which values; + /// anything else leaves every column open. + fn said(&self, filter: &Expr) -> Result> { + let open = || vec![vec![None; self.key.len()]]; + Ok(match filter { + Expr::BinaryOp { left, op, right } if *op == BinOp::And => { + let (left, right) = (self.said(left)?, self.said(right)?); + if left.len().saturating_mul(right.len()) > MOST_TUPLES { + return Ok(open()); + } + let join = |left: &Tuple, right: &Tuple| -> Tuple { + let both = left.iter().zip(right); + both.map(|(left, right)| left.clone().or_else(|| right.clone())) + .collect() + }; + let pairs = left.iter().flat_map(|l| right.iter().map(move |r| (l, r))); + pairs.map(|(left, right)| join(left, right)).collect() + } + Expr::BinaryOp { left, op, right } if *op == BinOp::Or => { + [self.said(left)?, self.said(right)?].concat() + } + Expr::BinaryOp { left, op, right } if *op == BinOp::Eq => { + let (column, value) = match **left { + Expr::Column(_) => (left, right), + _ => (right, left), + }; + vec![self.tuple(sides(column), sides(value))?] + } + Expr::InList { + expr, + list, + negated: false, + } => { + let tuples = list.iter(); + let tuples = tuples.map(|item| self.tuple(sides(expr), sides(item))); + tuples.collect::>()? + } + _ => open(), + }) + } + + /// The tuple `columns = values` fixes, pair by pair. A pair counts when it is a key column + /// beside a literal or a `$name` whose value a key of the column's type holds; any other + /// pair leaves its column open. + fn tuple(&self, columns: &[Expr], values: &[Expr]) -> Result { + let mut tuple = vec![None; self.key.len()]; + for pair in columns.iter().zip(values) { + let (Expr::Column(column), value @ (Expr::Literal(_) | Expr::Parameter(_))) = pair + else { + continue; + }; + let at = column_index(self.table, column)?; + if let Some(place) = self.key.iter().position(|key| *key == at) { + let value = eval(value, self.table, &[], self.params, self.now)?; + tuple[place] = as_key(self.table.meta.columns[at].column_type, value); + } + } + Ok(tuple) + } +} + +/// The expressions of one side of `=` or IN: those of a row value, or the one expression. +fn sides(expr: &Expr) -> &[Expr] { + match expr { + Expr::Row(items) => items, + one => std::slice::from_ref(one), + } +} + +/// Reads the text spelling of a vector, `[0.1, 0.2, 0.3]`: the one reader of it, for a value +/// written into a VECTOR column and for the vector of a search. +pub(crate) fn vector_from_text(text: &str) -> Option> { + let inside = text.trim().strip_prefix('[')?.strip_suffix(']')?; + inside + .split(',') + .map(|number| number.trim().parse().ok()) + .collect() +} diff --git a/v2/contextdb/src/query/tests.rs b/v2/contextdb/src/query/tests.rs index 19290c4c..41ff41fd 100644 --- a/v2/contextdb/src/query/tests.rs +++ b/v2/contextdb/src/query/tests.rs @@ -19,6 +19,7 @@ fn worked_out(expr: &Expr) -> Value { indexes: Vec::new(), }, index_ids: Vec::new(), + vectors: Vec::new(), }; eval(expr, &table, &[], &HashMap::new(), 0).expect("an expression of literals") } diff --git a/v2/contextdb/src/read/mod.rs b/v2/contextdb/src/read/mod.rs index 0cc5d7f3..6b751556 100644 --- a/v2/contextdb/src/read/mod.rs +++ b/v2/contextdb/src/read/mod.rs @@ -26,9 +26,10 @@ pub struct ReadSessionOptions { /// The most memory, in bytes, the session may use for its page cache and for the results /// it has returned and that are still alive. `None`, the default, is no ceiling: a reader /// does not take the limit the store remembers for its writer, so a reader that needs a - /// cap sets it here. Every open sets aside 8 MiB for the page cache and 1 MiB for the - /// small fixed needs of a statement, so a limit under 9 MiB is refused at open with - /// [`crate::Error::MemoryBudgetExceeded`]. + /// cap sets it here. Every open sets aside its page cache and 1 MiB for the small fixed + /// needs of a statement. The page cache is a quarter of the limit and never less than + /// 8 MiB, so a limit under 9 MiB is refused at open with + /// [`crate::Error::MemoryBudgetExceeded`]; with no limit the page cache is 8 MiB. pub memory_limit: Option, /// Where the session gets the time that `NOW()` in a statement answers with. `None`, the /// default, is the machine's clock. @@ -85,7 +86,7 @@ impl ReadSession { /// # Errors /// /// Those of [`ReadSession::open`], and [`crate::Error::MemoryBudgetExceeded`] when - /// [`ReadSessionOptions::memory_limit`] is under the 9 MiB every open sets aside. + /// [`ReadSessionOptions::memory_limit`] is under 9 MiB, the least an open sets aside. pub fn open_with_options( path: impl AsRef, options: ReadSessionOptions, diff --git a/v2/contextdb/src/schema/check.rs b/v2/contextdb/src/schema/check.rs index 4fff4efd..b48abbac 100644 --- a/v2/contextdb/src/schema/check.rs +++ b/v2/contextdb/src/schema/check.rs @@ -1,20 +1,16 @@ //! Checking a declaration: what CREATE TABLE and CREATE INDEX may say, and the table each gives. -use super::{Catalog, RESERVED_PREFIXES, Table}; -use crate::sql::ast::{self, DataType, Expr, Literal, UnaryOp}; +use super::{Catalog, RESERVED_PREFIXES, Table, VectorColumn}; +use crate::sql::ast::{self, DataType, Expr, Literal}; +use crate::sql::not_supported; use crate::sql::plan::{CreateIndexPlan, CreateTablePlan}; -use crate::store::{IndexId, TableId}; +use crate::store::{IndexId, TableId, VectorId}; use crate::types::{ - ColumnDef, ColumnDefault, ColumnType, Error, IndexDecl, Result, SortDirection, TableMeta, Value, + ColumnDef, ColumnDefault, ColumnType, Error, IndexDecl, Result, SortDirection, TableMeta, + Value, VectorIndexRef, VectorPartitionDeclarationIssue as PartitionIssue, + VectorPolicyDeclarationIssue, }; -/// The refusal of a statement or clause that arrives with a later version. -pub(super) fn not_supported(clause: &str) -> Error { - Error::PlanError(format!( - "{clause} is not supported in this version of the library" - )) -} - fn invalid(reason: String) -> Error { Error::SchemaInvalid { reason } } @@ -76,6 +72,17 @@ pub(super) fn new_table(create: &CreateTablePlan, catalog: &Catalog) -> Result Result Result { + let index = VectorIndexRef { + table: table.to_owned(), + column: declared.name.clone(), + }; + if declared.data_type == DataType::Vector(0) { + let issue = VectorPolicyDeclarationIssue::DimensionsNotPositive; + return Err(Error::InvalidVectorPolicyDeclaration { index, issue }); + } + let max_partitions = partitions(declared, columns) + .map_err(|issue| Error::InvalidVectorPartitionDeclaration { index, issue })?; + Ok(VectorColumn { + id, + column: declared.name.clone(), + partition_key: declared.vector.partition_key.clone(), + max_partitions, + search_mode: declared.vector.search_mode, }) } +/// Checks the PARTITION_KEY and the MAX_PARTITIONS of a VECTOR column and gives the limit it +/// declares, `None` when it declares none. +fn partitions( + declared: &ast::ColumnDef, + columns: &[ColumnDef], +) -> std::result::Result, PartitionIssue> { + let key = &declared.vector.partition_key; + for (at, name) in key.iter().enumerate() { + let column = columns.iter().find(|column| column.name == *name); + let issue = match column { + _ if key[..at].contains(name) => PartitionIssue::DuplicatePartitionKeyColumn, + _ if *name == declared.name => PartitionIssue::VectorColumnInPartitionKey, + None => PartitionIssue::UnknownPartitionKeyColumn, + // Two values of a key column are one partition or they are not: a REAL has two + // zeros and a NaN, and JSON and VECTOR values have no order to be keyed by. + Some(ColumnDef { + column_type: ColumnType::Real | ColumnType::Json | ColumnType::Vector(_), + .. + }) => PartitionIssue::UnsupportedPartitionKeyColumnType, + // The primary key column holds no NULL, written NOT NULL or not. + Some(column) if column.nullable && !column.primary_key => { + PartitionIssue::NullablePartitionKeyColumn + } + Some(_) => continue, + }; + return Err(issue); + } + let Some(written) = &declared.vector.max_partitions else { + return Ok(None); + }; + if key.is_empty() { + return Err(PartitionIssue::MaxPartitionsWithoutPartitionKey); + } + match written.parse::() { + Ok(limit) if limit > 0 => Ok(Some(limit)), + Ok(_) => Err(PartitionIssue::MaxPartitionsNotPositive), + Err(_) if written.starts_with('-') => Err(PartitionIssue::MaxPartitionsNotPositive), + Err(_) => Err(PartitionIssue::MaxPartitionsOutOfRange), + } +} + /// The table a CREATE INDEX names, with the new index as its last. pub(super) fn with_index(create: &CreateIndexPlan, catalog: &Catalog) -> Result
{ let reserved = RESERVED_PREFIXES @@ -206,14 +282,6 @@ pub(super) fn default(expr: &Expr) -> Option { Expr::Literal(Literal::Integer(whole)) => Value::Int64(*whole), Expr::Literal(Literal::Real(real)) => Value::Float64(*real), Expr::Literal(Literal::Text(words)) => Value::Text(words.clone()), - Expr::UnaryOp { - op: UnaryOp::Neg, - operand, - } => match **operand { - Expr::Literal(Literal::Integer(whole)) => Value::Int64(-whole), - Expr::Literal(Literal::Real(real)) => Value::Float64(-real), - _ => return None, - }, _ => return None, }; Some(ColumnDefault::Literal(value)) diff --git a/v2/contextdb/src/schema/mod.rs b/v2/contextdb/src/schema/mod.rs index 0faace21..34260dda 100644 --- a/v2/contextdb/src/schema/mod.rs +++ b/v2/contextdb/src/schema/mod.rs @@ -3,7 +3,8 @@ //! It never changes a row. //! //! A table's declaration is kept as one schema record under the table's name: its number, its -//! columns and rules as they were declared, and its indexes with their numbers. +//! columns and rules as they were declared, its indexes with their numbers, and each VECTOR +//! column's number and clauses. //! //! A table's indexes come in one order: the index behind the primary key (`__pk_`), //! the index behind each UNIQUE column in column order (`__unique_`; the key column has @@ -20,8 +21,8 @@ mod tests; use std::collections::BTreeMap; use crate::sql::Plan; -use crate::store::{ChangeSet, IndexId, TableId, View}; -use crate::types::{Error, Result, TableMeta}; +use crate::store::{ChangeSet, IndexId, Read, TableId, VectorId, View}; +use crate::types::{ColumnType, Error, Result, TableMeta, VectorIndexRef, VectorSearchMode}; /// What the name of an index the engine makes for itself begins with. An index of CREATE INDEX /// may not take one of these. @@ -35,6 +36,58 @@ pub(crate) struct Table { pub meta: TableMeta, /// One id per entry of `meta.indexes`, in the same order. pub index_ids: Vec, + /// One entry per VECTOR column, in column order. + pub vectors: Vec, +} + +/// How many partitions a VECTOR column with a PARTITION_KEY may have when its declaration +/// names no MAX_PARTITIONS. +const DEFAULT_MAX_PARTITIONS: u32 = 256; + +/// What a table keeps about one of its VECTOR columns: its number in the store and the clauses +/// it was declared with. The column's dimension is in its type. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct VectorColumn { + pub id: VectorId, + /// The column's name. + pub column: String, + /// The columns of PARTITION_KEY, in declared order; empty when none was declared. + pub partition_key: Vec, + /// MAX_PARTITIONS as declared; `None` when the declaration names none. + pub max_partitions: Option, + /// SEARCH_MODE as declared; `None` when the declaration names none. + pub search_mode: Option, +} + +impl VectorColumn { + /// The most partitions the column may have: what was declared, 256 when a PARTITION_KEY + /// was declared with no MAX_PARTITIONS, and `None` for a column with no PARTITION_KEY, + /// which has its one partition and no limit to meet. + pub(crate) fn limit(&self) -> Option { + let partitioned = !self.partition_key.is_empty(); + partitioned.then(|| self.max_partitions.unwrap_or(DEFAULT_MAX_PARTITIONS)) + } + + /// The mode a search of the column runs in when its statement has no USE VECTOR. + pub(crate) fn mode(&self) -> VectorSearchMode { + self.search_mode.unwrap_or_default() + } +} + +impl Table { + /// What the table keeps about the VECTOR column of this name, `None` when it has no such + /// column or the column is of another type. + pub(crate) fn vector(&self, column: &str) -> Option<&VectorColumn> { + self.vectors.iter().find(|vector| vector.column == column) + } + + /// The column as a refusal names it. + pub(crate) fn vector_ref(&self, vector: &VectorColumn) -> VectorIndexRef { + VectorIndexRef { + table: self.name.clone(), + column: vector.column.clone(), + } + } } /// Every table declared in a store at one moment. @@ -49,6 +102,14 @@ impl Catalog { let mut tables = BTreeMap::new(); for bytes in view.schema_records()? { let table = record::decode(&bytes); + // No declaration names, as a VECTOR column, a column that is not one of its table. + let is_vector = |vector: &VectorColumn| { + let column = table.meta.columns.iter().find(|c| c.name == vector.column); + column.is_some_and(|c| matches!(c.column_type, ColumnType::Vector(_))) + }; + if !table.vectors.iter().all(is_vector) { + return Err(view.damaged("a schema record's VECTOR column is not one")); + } tables.insert(table.name.clone(), table); } Ok(Catalog { tables }) @@ -63,6 +124,14 @@ impl Catalog { pub(crate) fn tables(&self) -> Vec<&Table> { self.tables.values().collect() } + + /// The table that has the VECTOR column with this number, and what it keeps about it. + pub(crate) fn vector(&self, id: VectorId) -> Option<(&Table, &VectorColumn)> { + self.tables.values().find_map(|table| { + let vector = table.vectors.iter().find(|vector| vector.id == id)?; + Some((table, vector)) + }) + } } /// Applies one declaration (CREATE TABLE, CREATE INDEX, ALTER TABLE, DROP) to a change set. @@ -88,7 +157,7 @@ pub(crate) fn declare(plan: &Plan, catalog: &Catalog, changes: &mut ChangeSet) - }); } Plan::CreateIndex(create) => check::with_index(create, catalog)?, - _ => return Err(check::not_supported("this declaration")), + _ => return Err(crate::sql::not_supported("this declaration")), }; changes.put_schema(&table.name, record::encode(&table))?; tables.insert(table.name.clone(), table); diff --git a/v2/contextdb/src/schema/record.rs b/v2/contextdb/src/schema/record.rs index 0650e5cd..a74cf567 100644 --- a/v2/contextdb/src/schema/record.rs +++ b/v2/contextdb/src/schema/record.rs @@ -8,10 +8,11 @@ //! every page it reads. Where a byte says which of several kinds follows, the last kind is read //! for any byte the earlier ones do not claim. -use super::Table; -use crate::store::{IndexId, TableId}; +use super::{Table, VectorColumn}; +use crate::store::{IndexId, TableId, VectorId}; use crate::types::{ ColumnDef, ColumnDefault, ColumnType, IndexDecl, SortDirection, TableMeta, Value, + VectorSearchMode, }; /// The schema record of a table. @@ -77,6 +78,23 @@ pub(super) fn encode(table: &Table) -> Vec { out.push(u8::from(*direction == SortDirection::Desc)); } } + number(&mut out, table.vectors.len()); + for vector in &table.vectors { + number(&mut out, vector.id.0 as usize); + text(&mut out, &vector.column); + number(&mut out, vector.partition_key.len()); + for column in &vector.partition_key { + text(&mut out, column); + } + // A declared MAX_PARTITIONS is never zero, so zero says that none was declared. + number(&mut out, vector.max_partitions.unwrap_or(0) as usize); + out.push(match vector.search_mode { + None => 0, + Some(VectorSearchMode::Auto) => 1, + Some(VectorSearchMode::Exact) => 2, + Some(VectorSearchMode::Indexed) => 3, + }); + } out } @@ -133,6 +151,18 @@ pub(super) fn decode(record: &[u8]) -> Table { }); IndexDecl { name, columns } }); + let vectors = read.list(|read| VectorColumn { + id: VectorId(read.number() as u32), + column: read.text(), + partition_key: read.list(Reader::text), + max_partitions: Some(read.number() as u32).filter(|limit| *limit != 0), + search_mode: match read.byte() { + 0 => None, + 1 => Some(VectorSearchMode::Auto), + 2 => Some(VectorSearchMode::Exact), + _ => Some(VectorSearchMode::Indexed), + }, + }); Table { id, name, @@ -142,6 +172,7 @@ pub(super) fn decode(record: &[u8]) -> Table { indexes, }, index_ids, + vectors, } } diff --git a/v2/contextdb/src/schema/render.rs b/v2/contextdb/src/schema/render.rs index 0e0ed6bd..5920090f 100644 --- a/v2/contextdb/src/schema/render.rs +++ b/v2/contextdb/src/schema/render.rs @@ -1,14 +1,18 @@ //! Printing a declaration back as SQL. use super::{RESERVED_PREFIXES, Table}; -use crate::types::{ColumnDef, ColumnDefault, ColumnType, SortDirection, Value}; +use crate::types::{ColumnDef, ColumnDefault, ColumnType, SortDirection, Value, VectorSearchMode}; /// A table's declaration as SQL: the CREATE TABLE, then one CREATE INDEX for each index the /// author made, every statement ended by `;` and a new line. A clause left at its default is /// not printed, and neither are the indexes the engine keeps behind a key or a UNIQUE rule. pub(super) fn declaration(table: &Table) -> String { let meta = &table.meta; - let mut lines: Vec = meta.columns.iter().map(column).collect(); + let mut lines: Vec = meta + .columns + .iter() + .map(|declared| column(table, declared)) + .collect(); for rule in &meta.unique_constraints { let columns: Vec = rule.iter().map(|column| name(column)).collect(); lines.push(format!("UNIQUE ({})", columns.join(", "))); @@ -40,7 +44,7 @@ pub(super) fn declaration(table: &Table) -> String { } /// One column's line: its name, its type, and the clauses it was given. -fn column(column: &ColumnDef) -> String { +fn column(table: &Table, column: &ColumnDef) -> String { let mut text = name(&column.name); text.push(' '); match column.column_type { @@ -53,6 +57,23 @@ fn column(column: &ColumnDef) -> String { ColumnType::Timestamp => text.push_str("TIMESTAMP"), ColumnType::Vector(size) => text.push_str(&format!("VECTOR({size})")), } + if let Some(vector) = table.vector(&column.name) { + // The limit on partitions refuses writes, so it is printed for every column with a + // PARTITION_KEY, whether or not it was written. SEARCH_MODE is printed only when it was. + if let Some(limit) = vector.limit() { + let key: Vec = vector.partition_key.iter().map(|key| name(key)).collect(); + text.push_str(&format!( + " PARTITION_KEY ({}) MAX_PARTITIONS {limit}", + key.join(", ") + )); + } + text.push_str(match vector.search_mode { + None => "", + Some(VectorSearchMode::Auto) => " SEARCH_MODE AUTO", + Some(VectorSearchMode::Exact) => " SEARCH_MODE EXACT", + Some(VectorSearchMode::Indexed) => " SEARCH_MODE INDEXED", + }); + } if column.primary_key { text.push_str(" PRIMARY KEY"); } diff --git a/v2/contextdb/src/schema/tests.rs b/v2/contextdb/src/schema/tests.rs index b648333b..2166feb4 100644 --- a/v2/contextdb/src/schema/tests.rs +++ b/v2/contextdb/src/schema/tests.rs @@ -1,7 +1,7 @@ use std::collections::BTreeMap; use super::{Catalog, Table, check, record, render}; -use crate::sql::ast::{self, DataType, Expr, Literal, UnaryOp}; +use crate::sql::ast::{self, DataType, Expr, Literal}; use crate::sql::plan::{CreateIndexPlan, CreateTablePlan}; use crate::store::{IndexId, TableId}; use crate::types::{ @@ -85,6 +85,7 @@ fn notes() -> Table { ], }, index_ids: [7, 8, 9, 12, 15].map(IndexId).to_vec(), + vectors: Vec::new(), } } @@ -121,32 +122,49 @@ CREATE INDEX by_body ON notes (body); ); } -#[test] -fn a_default_is_kept_in_the_type_the_author_wrote() { - let literal = Expr::Literal; - let minus = |literal| Expr::UnaryOp { - op: UnaryOp::Neg, - operand: Box::new(Expr::Literal(literal)), +/// The DEFAULT of a column as the parser hands it over, for this text after the word. +fn written_default(text: &str) -> Expr { + let statement = format!("CREATE TABLE t (c INTEGER DEFAULT {text})"); + let parsed = match crate::sql::parse(&statement) { + Ok(ast::Statement::CreateTable(create)) => create, + other => panic!("{statement}\nwas read as {other:?}"), }; - let now = Expr::FunctionCall { - name: "now".to_owned(), - args: Vec::new(), + let [column] = parsed.columns.as_slice() else { + panic!( + "{statement}\nwas read with {} columns", + parsed.columns.len() + ); }; - for (expr, value) in [ - (literal(Literal::Null), Value::Null), - (literal(Literal::Bool(true)), Value::Bool(true)), - (literal(Literal::Integer(3)), Value::Int64(3)), - (literal(Literal::Real(0.0)), Value::Float64(0.0)), - (literal(Literal::Text("it's".to_owned())), "it's".into()), - (minus(Literal::Integer(1)), Value::Int64(-1)), - (minus(Literal::Real(2.5)), Value::Float64(-2.5)), + column + .default + .clone() + .unwrap_or_else(|| panic!("{statement}\nwas read with no default")) +} + +#[test] +fn a_default_is_kept_in_the_type_the_author_wrote() { + for (text, value) in [ + ("NULL", Value::Null), + ("TRUE", Value::Bool(true)), + ("3", Value::Int64(3)), + ("0.0", Value::Float64(0.0)), + ("'it''s'", "it's".into()), + ("-1", Value::Int64(-1)), + ("-2.5", Value::Float64(-2.5)), ] { - assert_eq!(check::default(&expr), Some(ColumnDefault::Literal(value))); + assert_eq!( + check::default(&written_default(text)), + Some(ColumnDefault::Literal(value)), + "DEFAULT {text}" + ); } - assert_eq!(check::default(&now), Some(ColumnDefault::Now)); + assert_eq!( + check::default(&written_default("NOW()")), + Some(ColumnDefault::Now) + ); // Anything that is not a literal or NOW() is no default this version keeps. - assert_eq!(check::default(&Expr::Parameter("p".to_owned())), None); - assert_eq!(check::default(&minus(Literal::Text("x".to_owned()))), None); + assert_eq!(check::default(&written_default("$p")), None); + assert_eq!(check::default(&written_default("-'x'")), None); } /// A column as the parser hands over one that wrote no clause. @@ -158,6 +176,7 @@ fn declared(name: &str, data_type: DataType) -> ast::ColumnDef { primary_key: false, unique: false, default: None, + vector: ast::VectorClauses::default(), } } diff --git a/v2/contextdb/src/sql/ast.rs b/v2/contextdb/src/sql/ast.rs index fbc96b75..074ef877 100644 --- a/v2/contextdb/src/sql/ast.rs +++ b/v2/contextdb/src/sql/ast.rs @@ -1,6 +1,6 @@ //! A statement as it was written: what `parse` returns and `plan` reads. -pub use crate::types::SortDirection; +pub use crate::types::{SortDirection, VectorSearchMode}; #[derive(Debug, Clone)] pub enum Statement { @@ -26,7 +26,11 @@ pub struct Select { pub columns: Vec, pub table: String, pub where_clause: Option, + /// The ORDER BY items. A similarity search has one, a distance ([`Expr::Distance`]) or the + /// name the select list gives one. pub order_by: Vec, + /// `USE VECTOR AUTO | EXACT | INDEXED`, when it was written. + pub use_vector: Option, pub limit: Option, } @@ -66,6 +70,26 @@ pub enum Expr { expr: Box, negated: bool, }, + /// `left <=> right`: the cosine distance between a VECTOR column and a vector, a REAL that + /// is lower for a closer row. It binds tighter than a comparison and looser than `+` and + /// `-`, as in Postgres. This version works it out for the rows of a similarity search: the + /// planner reads it as the one ORDER BY key of a SELECT with LIMIT and, written the same + /// way, inside that statement's select list; anywhere else it is refused by name. + Distance { + left: Box, + right: Box, + }, + /// A row value, `(a, b)`: two or more expressions in parentheses. It is read wherever an + /// expression is, and has a meaning only on either side of `IN`: the planner refuses it + /// anywhere else. + Row(Vec), + /// `expr [NOT] IN (item, ...)`. Beside a row value every item is a row value of the same + /// length: the planner refuses any other list. + InList { + expr: Box, + list: Vec, + negated: bool, + }, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -140,6 +164,21 @@ pub struct ColumnDef { pub primary_key: bool, pub unique: bool, pub default: Option, + /// The clauses after `VECTOR(n)`. Left as it starts on a column of any other type. + pub vector: VectorClauses, +} + +/// The clauses a VECTOR column may carry, as they were written. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct VectorClauses { + /// The columns of `PARTITION_KEY (...)`, in the order written; empty when it was not. + pub partition_key: Vec, + /// The number after `MAX_PARTITIONS`, kept as written: `schema` judges it, so that zero, a + /// negative number and one too large are refused by their own kind and not as a syntax + /// error. + pub max_partitions: Option, + /// `SEARCH_MODE AUTO | EXACT | INDEXED`, when it was written. + pub search_mode: Option, } #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/v2/contextdb/src/sql/grammar.pest b/v2/contextdb/src/sql/grammar.pest index 909af2ae..827cc715 100644 --- a/v2/contextdb/src/sql/grammar.pest +++ b/v2/contextdb/src/sql/grammar.pest @@ -20,28 +20,32 @@ stmt = _{ // first words alone and the rest of the text is left unread, so that the statement is refused by // that name and not as a syntax error. `later` stands only where nothing this version reads can // begin with one of these words: after a whole statement or clause, and after an expression, -// where IN, LIKE, BETWEEN and the distance operator `<=>` would follow. DISTINCT and GRAPH_TABLE -// stand where a name could, so they are looked for first. +// where LIKE and BETWEEN would follow. DISTINCT and GRAPH_TABLE stand where a name could, so +// they are looked for first. later = { later_name ~ ANY* } later_name = @{ - "<=>" - | ( - ^"ALTER" | ^"DROP" | ^"WITH" | ^"BEGIN" | ^"COMMIT" | ^"ROLLBACK" | ^"PURGE" | ^"DISCARD" + ( + ^"SHOW" ~ SPACE+ ~ (^"VECTOR_INDEXES" | ^"VECTOR_PARTITIONS") + | ^"ALTER" | ^"DROP" | ^"WITH" | ^"BEGIN" | ^"COMMIT" | ^"ROLLBACK" | ^"PURGE" | ^"DISCARD" | ^"DECLARE" | ^"DISTINCT" | ^"GRAPH_TABLE" | ((^"LEFT" | ^"INNER" | ^"CROSS") ~ SPACE+)? ~ ^"JOIN" | ^"GROUP" ~ SPACE+ ~ ^"BY" | ^"HAVING" | ^"UNION" | ^"OFFSET" - | ^"USE" ~ SPACE+ ~ (^"VECTOR" | ^"RANK") + | ^"USE" ~ SPACE+ ~ ^"RANK" | ^"ON" ~ SPACE+ ~ ^"CONFLICT" - | (^"NOT" ~ SPACE+)? ~ (^"IN" | ^"LIKE" | ^"BETWEEN") + | (^"NOT" ~ SPACE+)? ~ (^"LIKE" | ^"BETWEEN") | ^"FOREIGN" ~ SPACE+ ~ ^"KEY" | ^"REFERENCES" | ^"EXPIRES" | ^"IMMUTABLE" | ^"RANK_POLICY" | ^"CONTEXT_ID" | ^"SCOPE_LABEL_READ" | ^"SCOPE_LABEL" | ^"ACL" - | ^"PARTITION_KEY" | ^"MAX_PARTITIONS" | ^"SEARCH_MODE" | ^"AUTO_INDEX_AT" | ^"SEARCH" - | ^"IVF" | ^"CONSOLIDATION" + | ^"AUTO_INDEX_AT" | ^"SEARCH" | ^"IVF" | ^"CONSOLIDATION" | ^"RETAIN" | ^"HISTORY" | ^"SYNC" | ^"STATE_MACHINE" | ^"STATE" ~ SPACE+ ~ ^"MACHINE" | ^"DAG" | ^"PROPAGATE" | ^"DELIVERY" ~ SPACE+ ~ ^"MANIFEST" | ^"EDGE" ~ SPACE+ ~ ^"DISCARD" | ^"MAINTAIN" ~ SPACE+ ~ ^"COUNT" ) ~ !ident_char } +// Two forms of a later version that begin with words this version reads, so each is known by +// where it stands: a SELECT as the list of IN, and the options after the mode of USE VECTOR. +// Like `later`, each leaves the rest of the text unread. +subquery = @{ ^"SELECT" ~ !ident_char } +vector_options = { "(" } set_memory_limit = { ^"SET" ~ ^"MEMORY_LIMIT" ~ "="? ~ ("'" ~ memory_limit_value ~ "'" | memory_limit_value) } show_memory_limit = { ^"SHOW" ~ ^"MEMORY_LIMIT" } @@ -55,7 +59,8 @@ select_stmt = { | select_list ~ ^"FROM" ~ ( &(^"GRAPH_TABLE" ~ "(") ~ later - | identifier ~ where_clause? ~ order_by_clause? ~ limit_clause? ~ later? + | identifier ~ where_clause? ~ order_by_clause? ~ use_vector_clause? ~ limit_clause? + ~ later? ) ) } @@ -66,6 +71,8 @@ where_clause = { ^"WHERE" ~ expr } order_by_clause = { ^"ORDER" ~ ^"BY" ~ order_item ~ ("," ~ order_item)* } order_item = { expr ~ sort_dir? } sort_dir = { ^"ASC" | ^"DESC" } +use_vector_clause = { ^"USE" ~ ^"VECTOR" ~ search_mode ~ (vector_options ~ ANY*)? } +search_mode = @{ (^"AUTO" | ^"EXACT" | ^"INDEXED") ~ !ident_char } limit_clause = { ^"LIMIT" ~ integer } create_table_stmt = { @@ -83,7 +90,13 @@ data_type = { | ^"UUID" | ^"TEXT" | ^"INTEGER" | ^"INT" | ^"REAL" | ^"FLOAT" | ^"BOOLEAN" | ^"BOOL" | ^"TIMESTAMP" | ^"JSON" } -vector_type = { ^"VECTOR" ~ "(" ~ integer ~ ")" } +vector_type = { + ^"VECTOR" ~ "(" ~ integer ~ ")" + ~ partition_key? ~ (^"MAX_PARTITIONS" ~ partition_count)? ~ (^"SEARCH_MODE" ~ search_mode)? +} +partition_key = { ^"PARTITION_KEY" ~ "(" ~ identifier ~ ("," ~ identifier)* ~ ")" } +// Read with its sign, so that `schema` refuses zero and a negative number by their kind. +partition_count = @{ "-"? ~ ASCII_DIGIT+ } column_constraint = _{ not_null | nullable_marker | primary_key | unique | default_clause } not_null = { ^"NOT" ~ ^"NULL" } nullable_marker = { ^"NULL" } @@ -112,16 +125,20 @@ update_stmt = { ^"UPDATE" ~ identifier ~ ^"SET" ~ assignment ~ ("," ~ assignment // One level of rule here is one level of the parser's recursion for every parenthesis, so the // boolean operators share one flat rule and the arithmetic operators another, and the builder // applies their precedence: NOT, then AND, then OR; `*` and `/` before `+` and `-`. -expr = { not_op* ~ additive_expr ~ comparison_suffix? ~ (bool_op ~ not_op* ~ additive_expr ~ comparison_suffix?)* } +expr = { not_op* ~ distance_expr ~ comparison_suffix? ~ (bool_op ~ not_op* ~ distance_expr ~ comparison_suffix?)* } bool_op = _{ and_op | or_op } or_op = @{ ^"OR" ~ !ident_char } and_op = @{ ^"AND" ~ !ident_char } not_op = @{ ^"NOT" ~ !ident_char } -comparison_suffix = _{ is_null_suffix | cmp_suffix | later } -cmp_suffix = { cmp_op ~ additive_expr } +comparison_suffix = _{ is_null_suffix | in_suffix | cmp_suffix | later } +cmp_suffix = { cmp_op ~ distance_expr } cmp_op = { "=" | "!=" | "<>" | "<=" | ">=" | "<" | ">" } is_null_suffix = { ^"IS" ~ not_op? ~ ^"NULL" } +in_suffix = { not_op? ~ ^"IN" ~ "(" ~ (subquery ~ ANY* | expr ~ ("," ~ expr)* ~ ")") } + +// The distance operator binds tighter than a comparison and looser than `+` and `-`. +distance_expr = { additive_expr ~ ("<=>" ~ additive_expr)? } additive_expr = { unary_minus* ~ primary_expr ~ (math_op ~ unary_minus* ~ primary_expr)* } math_op = { "+" | "-" | "*" | "/" } @@ -137,8 +154,12 @@ primary_expr = _{ | string | vector_lit | column_ref - | "(" ~ expr ~ ")" + | parenthesized } +// One expression in parentheses is that expression; two or more are a row value. One rule reads +// both: a rule that read the first expression, failed and read it again would double the work +// at every level of nesting. +parenthesized = { "(" ~ expr ~ ("," ~ expr)* ~ ")" } vector_lit = { "[" ~ vector_component ~ ("," ~ vector_component)* ~ "]" } vector_component = @{ "-"? ~ ASCII_DIGIT+ ~ ("." ~ ASCII_DIGIT+)? } function_call = { identifier ~ "(" ~ (star | (expr ~ ("," ~ expr)*))? ~ ")" } diff --git a/v2/contextdb/src/sql/mod.rs b/v2/contextdb/src/sql/mod.rs index 91a5f3c7..efd9c3cc 100644 --- a/v2/contextdb/src/sql/mod.rs +++ b/v2/contextdb/src/sql/mod.rs @@ -14,7 +14,7 @@ mod planner; pub(crate) use ast::{Expr, Statement}; pub(crate) use plan::{Effect, PhysicalPlan as Plan}; -use crate::types::Result; +use crate::types::{Error, Result}; /// Reads one SQL statement. Refuses with `ParseError`, or with `PlanError` when the text begins /// a statement or a clause that a later version of the library reads. Costs time in proportion @@ -27,3 +27,11 @@ pub(crate) fn parse(text: &str) -> Result { pub(crate) fn plan(statement: &Statement) -> Result { planner::plan(statement) } + +/// The refusal of a statement, a clause or a form that a later version of the library reads, +/// by its name. +pub(crate) fn not_supported(what: &str) -> Error { + Error::PlanError(format!( + "{what} is not supported in this version of the library" + )) +} diff --git a/v2/contextdb/src/sql/parser/expr.rs b/v2/contextdb/src/sql/parser/expr.rs index 15249888..be63b271 100644 --- a/v2/contextdb/src/sql/parser/expr.rs +++ b/v2/contextdb/src/sql/parser/expr.rs @@ -1,9 +1,9 @@ //! Expressions. The grammar reads the boolean operators as one flat run and the arithmetic //! operators as another; the functions here apply their precedence. //! -//! `expr`, `comparison_expr`, `additive_expr` and `primary_expr` call each other once for every -//! level of nesting, to the limit, on the caller's stack. So each is kept small: what is not on -//! that path is in a function of its own, and a list is gathered with a loop. +//! `expr`, `comparison_expr`, `distance_expr`, `additive_expr` and `primary_expr` call each other +//! once for every level of nesting, to the limit, on the caller's stack. So each is kept small: +//! what is not on that path is in a function of its own, and a list is gathered with a loop. use super::{MAX_EXPR_NESTING, Pair, Rule, depth_error, has, ident, number, string}; use crate::sql::ast::*; @@ -58,6 +58,12 @@ fn depth_of(expr: &Expr) -> usize { Expr::UnaryOp { operand, .. } => depth_of(operand), Expr::FunctionCall { args, .. } => args.iter().map(depth_of).max().unwrap_or(0), Expr::IsNull { expr, .. } => depth_of(expr), + Expr::Distance { left, right } => depth_of(left).max(depth_of(right)), + Expr::Row(items) => items.iter().map(depth_of).max().unwrap_or(0), + Expr::InList { expr, list, .. } => { + let items = list.iter().map(depth_of).max().unwrap_or(0); + depth_of(expr).max(items) + } Expr::Column(_) | Expr::Literal(_) | Expr::Parameter(_) => return 0, }; below + 1 @@ -91,7 +97,7 @@ fn not(operand: Expr) -> Expr { /// One comparison: its left side and what the grammar read after it, if anything. fn comparison_expr(left: Pair<'_>, suffix: Option>) -> Result { - let left = Box::new(additive_expr(left)?); + let left = Box::new(distance_expr(left)?); let Some(suffix) = suffix else { return Ok(*left); }; @@ -101,6 +107,13 @@ fn comparison_expr(left: Pair<'_>, suffix: Option>) -> Result { negated: has(&suffix, Rule::not_op), }); } + if suffix.as_rule() == Rule::in_suffix { + return Ok(Expr::InList { + expr: left, + negated: has(&suffix, Rule::not_op), + list: exprs(suffix)?, + }); + } let mut parts = suffix.into_inner(); let op = match parts.next().expect("an operator").as_str() { "=" => BinOp::Eq, @@ -113,7 +126,20 @@ fn comparison_expr(left: Pair<'_>, suffix: Option>) -> Result { Ok(Expr::BinaryOp { left, op, - right: Box::new(additive_expr(parts.next().expect("an operand"))?), + right: Box::new(distance_expr(parts.next().expect("an operand"))?), + }) +} + +/// `sum [<=> sum]`: a sum alone, or the distance between two. +fn distance_expr(pair: Pair<'_>) -> Result { + let mut sides = pair.into_inner(); + let left = additive_expr(sides.next().expect("an operand"))?; + let Some(right) = sides.next() else { + return Ok(left); + }; + Ok(Expr::Distance { + left: Box::new(left), + right: Box::new(additive_expr(right)?), }) } @@ -163,10 +189,7 @@ fn additive_expr(pair: Pair<'_>) -> Result { _ => { let mut operand = primary_expr(part)?; for _ in 0..std::mem::take(&mut minuses) { - operand = Expr::UnaryOp { - op: UnaryOp::Neg, - operand: Box::new(operand), - }; + operand = negated(operand); } product = Some(match product.take() { Some((left, times_or_over)) => (arithmetic(times_or_over, left, operand)?, ""), @@ -182,6 +205,19 @@ fn additive_expr(pair: Pair<'_>) -> Result { } } +/// `-operand`. A number written with a minus sign is that number: a literal, as a number +/// written without a sign is, so whatever reads a literal reads it. +fn negated(operand: Expr) -> Expr { + match operand { + Expr::Literal(Literal::Integer(whole)) => Expr::Literal(Literal::Integer(-whole)), + Expr::Literal(Literal::Real(real)) => Expr::Literal(Literal::Real(-real)), + operand => Expr::UnaryOp { + op: UnaryOp::Neg, + operand: Box::new(operand), + }, + } +} + /// The `*` of `SELECT *` and of `COUNT(*)`: a column of that name. pub(super) fn star() -> Expr { Expr::Column(ColumnRef { @@ -192,8 +228,14 @@ pub(super) fn star() -> Expr { fn primary_expr(primary: Pair<'_>) -> Result { match primary.as_rule() { - // A parenthesised expression. - Rule::expr => expr(primary), + // One expression in parentheses is that expression; two or more are a row value. + Rule::parenthesized => { + let mut items = exprs(primary)?; + Ok(match items.len() { + 1 => items.swap_remove(0), + _ => Expr::Row(items), + }) + } Rule::function_call => function_call(primary), Rule::column_ref => Ok(column_ref(primary)), _ => literal(&primary), diff --git a/v2/contextdb/src/sql/parser/mod.rs b/v2/contextdb/src/sql/parser/mod.rs index 3a06ff0a..287970cd 100644 --- a/v2/contextdb/src/sql/parser/mod.rs +++ b/v2/contextdb/src/sql/parser/mod.rs @@ -13,6 +13,7 @@ use pest::Parser; use pest::error::{InputLocation, LineColLocation}; use super::ast::*; +use super::not_supported; use crate::types::{Error, Result}; #[derive(pest_derive::Parser)] @@ -29,16 +30,17 @@ pub(super) fn parse(text: &str) -> Result { check_source_nesting(text)?; let read = SqlParser::parse(Rule::statement, text).map_err(|error| syntax_error(text, &error))?; - if let Some(later) = read - .clone() - .flatten() - .find(|part| part.as_rule() == Rule::later_name) - { - let words: Vec<&str> = later.as_str().split_whitespace().collect(); - return Err(Error::PlanError(format!( - "{} is not supported in this version of the library", - words.join(" ").to_ascii_uppercase() - ))); + for part in read.clone().flatten() { + let later = match part.as_rule() { + Rule::later_name => { + let words: Vec<&str> = part.as_str().split_whitespace().collect(); + words.join(" ").to_ascii_uppercase() + } + Rule::subquery => "IN (SELECT ...)".to_owned(), + Rule::vector_options => "USE VECTOR with options".to_owned(), + _ => continue, + }; + return Err(not_supported(&later)); } let pair = read .flatten() @@ -288,6 +290,15 @@ fn where_clause(pair: Pair<'_>) -> Result> { .transpose() } +/// AUTO, EXACT or INDEXED, after SEARCH_MODE and after USE VECTOR. +fn search_mode(pair: &Pair<'_>) -> VectorSearchMode { + match pair.as_str().to_ascii_uppercase().as_str() { + "AUTO" => VectorSearchMode::Auto, + "EXACT" => VectorSearchMode::Exact, + _ => VectorSearchMode::Indexed, + } +} + /// ASC when nothing is written. fn direction(written: Option>) -> SortDirection { match written { @@ -302,6 +313,7 @@ fn select(pair: Pair<'_>) -> Result { }); } } + Rule::use_vector_clause => { + let mode = part.into_inner().next().expect("a mode"); + select.use_vector = Some(search_mode(&mode)); + } Rule::limit_clause => { let count = part.into_inner().next().expect("a count"); select.limit = Some(number(&count, "LIMIT")?); @@ -407,11 +423,23 @@ fn column_def(pair: Pair<'_>) -> Result { let name = ident(&parts.next().expect("a column name")); let type_pair = parts.next().expect("a column type"); let type_name = type_pair.as_str().to_ascii_uppercase(); + let mut vector = VectorClauses::default(); let data_type = match (type_pair.into_inner().next(), type_name.as_str()) { - (Some(vector), _) => DataType::Vector(number( - &vector.into_inner().next().expect("a dimension"), - "the VECTOR dimension", - )?), + (Some(declared), _) => { + let mut parts = declared.into_inner(); + let dimension = parts.next().expect("a dimension"); + for clause in parts { + match clause.as_rule() { + Rule::partition_key => vector.partition_key = idents(clause), + Rule::partition_count => { + vector.max_partitions = Some(clause.as_str().to_owned()); + } + // The one clause left: SEARCH_MODE. + _ => vector.search_mode = Some(search_mode(&clause)), + } + } + DataType::Vector(number(&dimension, "the VECTOR dimension")?) + } (None, "UUID") => DataType::Uuid, (None, "TEXT") => DataType::Text, (None, "INTEGER" | "INT") => DataType::Integer, @@ -427,6 +455,7 @@ fn column_def(pair: Pair<'_>) -> Result { primary_key: false, unique: false, default: None, + vector, }; // A constraint written twice is refused. let once = |written: bool, clause: &str| match written { diff --git a/v2/contextdb/src/sql/parser/tests.rs b/v2/contextdb/src/sql/parser/tests.rs index ddc52424..82aaa818 100644 --- a/v2/contextdb/src/sql/parser/tests.rs +++ b/v2/contextdb/src/sql/parser/tests.rs @@ -1,8 +1,10 @@ use super::parse; -use crate::sql::ast::{DataType, Expr, Literal, Statement}; +use crate::sql::ast::{ + BinOp, ColumnRef, DataType, Expr, Literal, Select, Statement, VectorClauses, +}; use crate::sql::plan; use crate::sql::plan::{Effect, PhysicalPlan}; -use crate::types::Error; +use crate::types::{Error, SortDirection, VectorSearchMode}; fn refusal(sql: &str) -> String { match parse(sql) { @@ -217,8 +219,6 @@ fn what_a_later_version_reads_is_refused_by_its_name() { "SELECT a FROM GRAPH_TABLE(g MATCH (a)-[e]->(b))", "GRAPH_TABLE", ), - ("SELECT id FROM t ORDER BY v <=> $q LIMIT 3", "<=>"), - ("SELECT id FROM t WHERE a NOT IN (1, 2)", "NOT IN"), ("SELECT id FROM t WHERE a LIKE 'x%'", "LIKE"), ("UPDATE t SET a = 1 WHERE b BETWEEN 1 AND 2", "BETWEEN"), ( @@ -260,5 +260,314 @@ fn what_a_later_version_reads_is_refused_by_its_name() { planned("SELECT state, history FROM t WHERE immutable ORDER BY state"); // What is neither read nor named is a syntax error. assert!(refusal("SELECT a FROM t x").contains("unexpected 'x'")); - assert!(refusal("SELECT (a, b) FROM t").contains("unexpected ','")); + // A row value is read wherever an expression is. In a select list it is refused by name + // when the statement is planned. + let read = parse("SELECT (a, b) FROM t") + .unwrap_or_else(|error| panic!("a row value in a select list was not read: {error:?}")); + match plan(&read) { + Err(Error::PlanError(text)) => { + assert!(text.to_lowercase().contains("row value"), "{text}"); + } + other => panic!("a row value in a select list was planned as {other:?}"), + } +} + +fn select(sql: &str) -> Select { + match parse(sql) { + Ok(Statement::Select(select)) => select, + other => panic!("{sql}\nwas not read as a SELECT: {other:?}"), + } +} + +/// True for ` <=> $`. +fn is_distance(expr: &Expr, column: &str, parameter: &str) -> bool { + let Expr::Distance { left, right } = expr else { + return false; + }; + let of_column = matches!( + &**left, + Expr::Column(ColumnRef { table: None, column: name }) if name == column + ); + of_column && matches!(&**right, Expr::Parameter(name) if name == parameter) +} + +#[test] +fn a_distance_is_read_as_an_expression() { + // In a select item under a name, with the search ordered by that name. + let read = select("SELECT id, v <=> $q AS distance FROM t ORDER BY distance LIMIT 3"); + assert!(is_distance(&read.columns[1].expr, "v", "q")); + assert_eq!(read.columns[1].alias.as_deref(), Some("distance")); + assert!(matches!( + &read.order_by[..], + [item] if item.direction == SortDirection::Asc + && matches!( + &item.expr, + Expr::Column(ColumnRef { table: None, column }) if column == "distance" + ) + )); + assert_eq!((read.limit, read.use_vector), (Some(3), None)); + + // Under no name, and as the ORDER BY item itself, with ASC written or not. + for order in ["v <=> $q", "v <=> $q ASC", "v<=>$q asc"] { + let read = select(&format!( + "SELECT id, v <=> $q FROM t ORDER BY {order} LIMIT 3" + )); + assert!(is_distance(&read.columns[1].expr, "v", "q"), "{order}"); + assert_eq!(read.columns[1].alias, None, "{order}"); + assert!( + matches!( + &read.order_by[..], + [item] if item.direction == SortDirection::Asc && is_distance(&item.expr, "v", "q") + ), + "{order}" + ); + } + + // Inside another expression: one minus the distance. + let read = select("SELECT id, 1 - (v <=> $q) AS similarity FROM t ORDER BY v <=> $q LIMIT 3"); + assert!(matches!( + &read.columns[1].expr, + Expr::FunctionCall { name, args } + if name == "__sub" + && matches!(&args[..], [Expr::Literal(Literal::Integer(1)), distance] + if is_distance(distance, "v", "q")) + )); + assert_eq!(read.columns[1].alias.as_deref(), Some("similarity")); + + // The vector on the right is a literal, its quoted text or a bound name. + let read = select("SELECT id FROM t ORDER BY v <=> [1.0, 0.0] LIMIT 3"); + assert!(matches!( + &read.order_by[0].expr, + Expr::Distance { right, .. } + if matches!(&**right, Expr::Literal(Literal::Vector(vector)) if vector == &[1.0, 0.0]) + )); + let read = select("SELECT id FROM t ORDER BY v <=> '[1,0]' LIMIT 3"); + assert!(matches!( + &read.order_by[0].expr, + Expr::Distance { right, .. } + if matches!(&**right, Expr::Literal(Literal::Text(text)) if text == "[1,0]") + )); +} + +#[test] +fn a_distance_binds_tighter_than_a_comparison_and_looser_than_a_sum() { + let item = |sql: &str| { + let read = select(&format!("SELECT {sql} FROM t")); + format!("{:?}", read.columns[0].expr) + }; + assert_eq!(item("v <=> $q < 0.5"), item("(v <=> $q) < 0.5")); + assert_eq!(item("0.5 >= v <=> $q"), item("0.5 >= (v <=> $q)")); + assert_eq!(item("v <=> $q + $r"), item("v <=> ($q + $r)")); + assert_eq!(item("a + b <=> $q - $r"), item("(a + b) <=> ($q - $r)")); + assert_eq!(item("1 - v <=> $q"), item("(1 - v) <=> $q")); + assert_eq!( + item("v <=> $q < w <=> $q AND NOT v <=> $q = 1"), + item("((v <=> $q) < (w <=> $q)) AND (NOT ((v <=> $q) = 1))") + ); + // The first of them is a comparison of a distance, and the third a distance to a sum. + let read = select("SELECT v <=> $q < 0.5, v <=> $q + $r FROM t"); + assert!(matches!( + &read.columns[0].expr, + Expr::BinaryOp { left, op: BinOp::Lt, .. } if is_distance(left, "v", "q") + )); + assert!(matches!( + &read.columns[1].expr, + Expr::Distance { right, .. } + if matches!(&**right, Expr::FunctionCall { name, .. } if name == "__add") + )); +} + +#[test] +fn a_search_is_planned_as_one_step_with_its_limit() { + // Under a projection when the statement names its columns: the search holds the column, + // the vector, the limit, the filter and the mode, and no limit wraps it. + let planned_search = planned( + "SELECT id, v <=> $q AS distance FROM t WHERE a = 1 ORDER BY distance USE VECTOR EXACT \ + LIMIT 3", + ); + assert_eq!(planned_search.effect(), Effect::Reads); + let (input, columns) = match planned_search { + PhysicalPlan::Project { input, columns } => (input, columns), + other => panic!("a search with a select list is not under a projection: {other:?}"), + }; + assert!(is_distance(&columns[1].expr, "v", "q")); + assert_eq!(columns[1].alias.as_deref(), Some("distance")); + assert!(matches!( + &*input, + PhysicalPlan::VectorSearch { + table, + column, + query: Expr::Parameter(query), + k: 3, + filter: Some(_), + mode: Some(VectorSearchMode::Exact), + } if table == "t" && column == "v" && query == "q" + )); + // `SELECT *` is the search itself, and with no USE VECTOR the mode is left to the column. + for order in ["v <=> $q", "v <=> $q ASC"] { + let sql = format!("SELECT * FROM t ORDER BY {order} LIMIT 10"); + assert!( + matches!( + planned(&sql), + PhysicalPlan::VectorSearch { + k: 10, + filter: None, + mode: None, + .. + } + ), + "{sql}" + ); + } + for (mode, read) in [ + ("AUTO", VectorSearchMode::Auto), + ("exact", VectorSearchMode::Exact), + ("Indexed", VectorSearchMode::Indexed), + ] { + let sql = format!("SELECT * FROM t ORDER BY v <=> $q USE VECTOR {mode} LIMIT 1"); + assert_eq!(select(&sql).use_vector, Some(read), "{sql}"); + assert!( + matches!( + planned(&sql), + PhysicalPlan::VectorSearch { mode: Some(found), .. } if found == read + ), + "{sql}" + ); + } + // What is refused without looking at any table, when the statement is read or planned. + let refused = |sql: &str| match parse(sql) { + Ok(statement) => plan(&statement).expect_err(sql), + Err(error) => error, + }; + // A search with no LIMIT has no kind of its own: it is refused as every statement of a + // later version is. + match refused("SELECT id FROM t ORDER BY v <=> $q") { + Error::PlanError(text) => assert_eq!( + text, + "a vector search without LIMIT is not supported in this version of the library" + ), + other => panic!("a search with no LIMIT was refused as {other:?}"), + } + assert!(matches!( + refused("SELECT id FROM t ORDER BY id USE VECTOR EXACT LIMIT 3"), + Error::UseVectorRequiresVectorOrder + )); + for sql in [ + "SELECT id FROM t ORDER BY v <=> $q DESC LIMIT 3", + "SELECT id FROM t ORDER BY v <=> $q, id LIMIT 3", + "SELECT id FROM t ORDER BY $q <=> v LIMIT 3", + "SELECT id FROM t WHERE v <=> $q < 0.5", + "SELECT id, v <=> $q AS distance FROM t", + "SELECT id, v <=> $q AS distance FROM t ORDER BY id LIMIT 3", + ] { + match refused(sql) { + Error::PlanError(text) => assert!(text.contains("<=>"), "{sql}: {text}"), + other => panic!("{sql}\nwas refused as {other:?}"), + } + } +} + +#[test] +fn a_row_value_and_any_right_side_of_a_distance_are_read_and_refused_when_planned() { + // Each is read by the parser, and refused by the planner in a sentence that names it. + // The refusal of ROW_VECTOR(...) names the function and the `<=>` it stands beside. + let names: [(&str, &[&str]); 4] = [ + ("SELECT id FROM t WHERE (a, b) = (1, 2)", &["row value"]), + ( + "SELECT id FROM t WHERE (a, b) IN ((1, 2, 3))", + &["row value"], + ), + ("SELECT id FROM t ORDER BY v <=> w LIMIT 3", &["<=>"]), + ( + "SELECT id FROM t ORDER BY v <=> ROW_VECTOR('t', 'v', 1) LIMIT 3", + &["<=>", "row_vector"], + ), + ]; + for (sql, names) in names { + let read = parse(sql).unwrap_or_else(|error| panic!("{sql}\nwas not read: {error:?}")); + match plan(&read) { + Err(Error::PlanError(text)) => { + for name in names { + assert!(text.to_lowercase().contains(name), "{sql}: {text}"); + } + } + other => panic!("{sql}\nwas planned as {other:?}"), + } + } + // The two statements that describe an index and its partitions are a later version's. + for (sql, name) in [ + ("SHOW VECTOR_INDEXES", "VECTOR_INDEXES"), + ("show vector_partitions docs", "VECTOR_PARTITIONS"), + ] { + match parse(sql) { + Err(Error::PlanError(text)) => assert!( + text.to_uppercase().contains(name) + && text.ends_with("is not supported in this version of the library"), + "{sql}: {text}" + ), + other => panic!("{sql}\nwas not refused by name: {other:?}"), + } + } +} + +#[test] +fn the_clauses_of_a_vector_column_and_in_lists_are_read() { + let Statement::CreateTable(table) = parse( + "CREATE TABLE t (id INTEGER PRIMARY KEY, ctx TEXT NOT NULL, kind INT NOT NULL, \ + plain VECTOR(3), v VECTOR(4) PARTITION_KEY (ctx, kind) MAX_PARTITIONS 8 \ + SEARCH_MODE EXACT NOT NULL, w VECTOR(2) partition_key (ctx) max_partitions -1, \ + x VECTOR(2) SEARCH_MODE indexed)", + ) + .unwrap() else { + panic!("not a CREATE TABLE"); + }; + let clauses: Vec<&VectorClauses> = table.columns.iter().map(|column| &column.vector).collect(); + let keyed = |key: &[&str], max: Option<&str>, mode: Option| VectorClauses { + partition_key: key.iter().map(|column| (*column).to_owned()).collect(), + max_partitions: max.map(str::to_owned), + search_mode: mode, + }; + let none = VectorClauses::default(); + assert_eq!( + clauses, + [ + &none, + &none, + &none, + &none, + &keyed(&["ctx", "kind"], Some("8"), Some(VectorSearchMode::Exact)), + // The number is kept as it was written: it is `schema` that judges it. + &keyed(&["ctx"], Some("-1"), None), + &keyed(&[], None, Some(VectorSearchMode::Indexed)), + ] + ); + assert_eq!(table.columns[4].data_type, DataType::Vector(4)); + assert!(!table.columns[4].nullable && table.columns[3].nullable); + + // IN and NOT IN over a list, and over a list of row values. + let filter = |sql: &str| { + let read = select(&format!("SELECT id FROM t WHERE {sql}")); + read.where_clause.expect("a WHERE") + }; + assert!(matches!( + filter("a IN (1, 2, $p)"), + Expr::InList { expr, list, negated: false } + if matches!(&*expr, Expr::Column(_)) && list.len() == 3 + )); + assert!(matches!( + filter("a not in ('x')"), + Expr::InList { list, negated: true, .. } if list.len() == 1 + )); + assert!(matches!( + filter("(a, b) IN ((1, 'x'), (2, 'y'))"), + Expr::InList { expr, list, negated: false } + if matches!(&*expr, Expr::Row(row) if row.len() == 2) + && list.len() == 2 + && list.iter().all(|item| matches!(item, Expr::Row(row) if row.len() == 2)) + )); + // A list is an operand of AND and OR like any comparison. + assert!(matches!( + filter("a IN (1, 2) AND NOT b IN (3) OR c = 4"), + Expr::BinaryOp { op: BinOp::Or, .. } + )); } diff --git a/v2/contextdb/src/sql/plan.rs b/v2/contextdb/src/sql/plan.rs index 14aded19..41b1a68f 100644 --- a/v2/contextdb/src/sql/plan.rs +++ b/v2/contextdb/src/sql/plan.rs @@ -1,6 +1,6 @@ //! The planned statement: what `schema`, `write` and `query` each take. -use super::ast::{ColumnDef, Expr, SetMemoryLimitValue, SortDirection}; +use super::ast::{ColumnDef, Expr, SetMemoryLimitValue, SortDirection, VectorSearchMode}; #[derive(Debug, Clone)] pub enum PhysicalPlan { @@ -26,6 +26,23 @@ pub enum PhysicalPlan { input: Box, count: u64, }, + /// The `k` rows of a table that pass the filter and lie nearest `query` by cosine distance + /// in a VECTOR column, nearest first. `query` decides which partitions the filter selects + /// and the vector part finds the rows. It stands where `Sort` over `Scan` stands in other + /// plans; its LIMIT is `k`, so no `Limit` wraps it. A `Project` over it may hold the + /// statement's distance expression, `column <=> query`, any number of times: each stands + /// for the distance the search found for the row. + VectorSearch { + table: String, + /// The VECTOR column on the left of `<=>`. + column: String, + /// The right of `<=>`. + query: Expr, + k: u64, + filter: Option, + /// What `USE VECTOR` asked for; `None` leaves it to the column's declaration. + mode: Option, + }, SetMemoryLimit(SetMemoryLimitValue), ShowMemoryLimit, } @@ -61,6 +78,7 @@ impl PhysicalPlan { | PhysicalPlan::Sort { .. } | PhysicalPlan::Project { .. } | PhysicalPlan::Limit { .. } + | PhysicalPlan::VectorSearch { .. } | PhysicalPlan::ShowMemoryLimit => Effect::Reads, } } diff --git a/v2/contextdb/src/sql/planner.rs b/v2/contextdb/src/sql/planner.rs index 2e7d52b7..cdbbed2e 100644 --- a/v2/contextdb/src/sql/planner.rs +++ b/v2/contextdb/src/sql/planner.rs @@ -4,7 +4,8 @@ //! chooses no index: a SELECT's plan names the table to read and the filter to apply, and //! `query` decides how to reach the rows. -use super::ast::{ColumnRef, Expr, Select, SelectColumn, Statement}; +use super::ast::{ColumnRef, Expr, Literal, Select, SelectColumn, SortDirection, Statement}; +use super::not_supported; use super::plan::*; use crate::types::{Error, Result}; @@ -25,20 +26,31 @@ pub(super) fn plan(statement: &Statement) -> Result { table: index.table.clone(), columns: index.columns.clone(), }), - Statement::Insert(insert) => PhysicalPlan::Insert(InsertPlan { - table: insert.table.clone(), - columns: insert.columns.clone(), - values: insert.values.clone(), - }), - Statement::Delete(delete) => PhysicalPlan::Delete(DeletePlan { - table: delete.table.clone(), - where_clause: delete.where_clause.clone(), - }), - Statement::Update(update) => PhysicalPlan::Update(UpdatePlan { - table: update.table.clone(), - assignments: update.assignments.clone(), - where_clause: update.where_clause.clone(), - }), + Statement::Insert(insert) => { + insert.values.iter().flatten().try_for_each(rows_stand)?; + PhysicalPlan::Insert(InsertPlan { + table: insert.table.clone(), + columns: insert.columns.clone(), + values: insert.values.clone(), + }) + } + Statement::Delete(delete) => { + delete.where_clause.iter().try_for_each(in_a_write)?; + PhysicalPlan::Delete(DeletePlan { + table: delete.table.clone(), + where_clause: delete.where_clause.clone(), + }) + } + Statement::Update(update) => { + let values = update.assignments.iter().map(|(_, value)| value); + let mut written = values.chain(&update.where_clause); + written.try_for_each(in_a_write)?; + PhysicalPlan::Update(UpdatePlan { + table: update.table.clone(), + assignments: update.assignments.clone(), + where_clause: update.where_clause.clone(), + }) + } Statement::Select(select) => plan_select(select)?, Statement::SetMemoryLimit(limit) => PhysicalPlan::SetMemoryLimit(limit.clone()), Statement::ShowMemoryLimit => PhysicalPlan::ShowMemoryLimit, @@ -46,38 +58,74 @@ pub(super) fn plan(statement: &Statement) -> Result { } /// A scan with its filter, sorted, then projected unless the statement is `SELECT *`, then cut -/// to the limit. +/// to the limit. A similarity search stands where the scan, the sort and the limit would. fn plan_select(select: &Select) -> Result { - let mut current = PhysicalPlan::Scan { - table: select.table.clone(), - filter: select.where_clause.clone(), - }; - if !select.order_by.is_empty() { - let mut keys = Vec::with_capacity(select.order_by.len()); - for item in &select.order_by { - let expr = resolve_output_alias(&item.expr, &select.columns); + let written = select.columns.iter().map(|column| &column.expr); + let written = written.chain(&select.where_clause); + let written = written.chain(select.order_by.iter().map(|item| &item.expr)); + let written: Vec<&Expr> = written.collect(); + written.iter().copied().try_for_each(rows_stand)?; + let keys: Vec = select + .order_by + .iter() + .map(|item| SortKey { + expr: resolve_output_alias(&item.expr, &select.columns), + direction: item.direction, + }) + .collect(); + // A distance among several ORDER BY items is a search of a later version. + let distance = |key: &SortKey| matches!(key.expr, Expr::Distance { .. }); + if keys.len() > 1 && keys.iter().any(distance) { + return Err(not_supported("ORDER BY <=> with a second sort key")); + } + // A similarity search: the one ORDER BY item is a distance. + let mut search = None; + if let [key] = keys.as_slice() + && let Expr::Distance { left, right } = &key.expr + { + search = Some(plan_search(select, left, right, key)?); + } + let searching = search.is_some(); + let mut current = match search { + Some(search) => search, + None => { + if select.use_vector.is_some() { + return Err(Error::UseVectorRequiresVectorOrder); + } + for expr in written { + distances(expr, &mut |_, _| { + Err(not_supported( + "<=> anywhere but as the one ORDER BY item of a SELECT with LIMIT, and in \ + that SELECT's select list,", + )) + })?; + } // A key `query` cannot work out for every row is refused here: sorting by it would // treat every row as equal and hand the rows back unsorted. - let sortable = match &expr { + let sortable = |key: &SortKey| match &key.expr { Expr::Column(_) => true, Expr::FunctionCall { name, .. } => SORTABLE_FUNCTIONS .iter() .any(|known| name.eq_ignore_ascii_case(known)), _ => false, }; - if !sortable { + if !keys.iter().all(sortable) { return Err(Error::OrderByExpressionNotSupported); } - keys.push(SortKey { - expr, - direction: item.direction, - }); + let scan = PhysicalPlan::Scan { + table: select.table.clone(), + filter: select.where_clause.clone(), + }; + if keys.is_empty() { + scan + } else { + PhysicalPlan::Sort { + input: Box::new(scan), + keys, + } + } } - current = PhysicalPlan::Sort { - input: Box::new(current), - keys, - }; - } + }; let star = matches!( select.columns.as_slice(), [SelectColumn { expr: Expr::Column(ColumnRef { table: None, column }), alias: None }] @@ -96,7 +144,7 @@ fn plan_select(select: &Select) -> Result { .collect(), }; } - if let Some(count) = select.limit { + if let (Some(count), false) = (select.limit, searching) { current = PhysicalPlan::Limit { input: Box::new(current), count, @@ -120,3 +168,149 @@ fn resolve_output_alias(expr: &Expr, columns: &[SelectColumn]) -> Expr { }; aliased.map_or_else(|| expr.clone(), |select| select.expr.clone()) } + +/// The search a SELECT asks for when its one ORDER BY item is `left <=> right`: the rows nearest +/// `right` in the column `left`. Refuses the forms of a search, and of a distance beside it, +/// that a later version reads. +fn plan_search(select: &Select, left: &Expr, right: &Expr, key: &SortKey) -> Result { + let Some(k) = select.limit else { + return Err(not_supported("a vector search without LIMIT")); + }; + if key.direction == SortDirection::Desc { + return Err(not_supported("ORDER BY <=> DESC")); + } + let Some(column) = searched_column(select, left)? else { + // No search: the statement is the sort it is written as, which `query` refuses by the + // column's name once it has found the table. + let (table, filter, keys) = (select.table.clone(), None, vec![key.clone()]); + let input = Box::new(PhysicalPlan::Scan { table, filter }); + return Ok(PhysicalPlan::Sort { input, keys }); + }; + match right { + Expr::Parameter(_) | Expr::Literal(Literal::Text(_) | Literal::Vector(_)) => {} + // `+` and its like are calls too, under names no statement writes. + Expr::FunctionCall { name, .. } if !name.starts_with("__") => { + let what = format!("<=> with the function {name} on its right"); + return Err(not_supported(&what)); + } + _ => { + return Err(not_supported( + "<=> with anything but a vector literal, its quoted text or a $name on its right", + )); + } + } + if let Some(filter) = &select.where_clause { + distances(filter, &mut |_, _| Err(not_supported("<=> in WHERE")))?; + } + for item in &select.columns { + distances(&item.expr, &mut |other: &Expr, vector: &Expr| { + let Some(other) = searched_column(select, other)? else { + return Ok(()); + }; + if other == column && same_vector(vector, right) { + return Ok(()); + } + Err(not_supported(&format!( + "<=> on the column {other} in the select list, where it is not the distance the \ + search is ordered by," + ))) + })?; + } + Ok(PhysicalPlan::VectorSearch { + table: select.table.clone(), + column: column.to_owned(), + query: right.clone(), + k, + filter: select.where_clause.clone(), + mode: select.use_vector, + }) +} + +/// The column on the left of a distance: a column of the statement's table, by its name alone or +/// with the table's before it. `None` under the name of any other table: `query` says that it +/// is not found, as it does anywhere else in the statement, after it has found the table. +fn searched_column<'a>(select: &Select, left: &'a Expr) -> Result> { + match left { + Expr::Column(ColumnRef { table, column }) + if table.as_ref().is_none_or(|table| *table == select.table) => + { + Ok(Some(column.as_str())) + } + Expr::Column(_) => Ok(None), + _ => Err(not_supported( + "<=> with anything but a column of the table on its left", + )), + } +} + +/// Whether two right sides of `<=>` are written the same: the same `$name`, the same quoted text +/// or the same vector literal. +fn same_vector(one: &Expr, other: &Expr) -> bool { + match (one, other) { + (Expr::Parameter(one), Expr::Parameter(other)) => one == other, + (Expr::Literal(Literal::Text(one)), Expr::Literal(Literal::Text(other))) => one == other, + (Expr::Literal(Literal::Vector(one)), Expr::Literal(Literal::Vector(other))) => { + one == other + } + _ => false, + } +} + +/// The expressions directly inside `expr`. +fn inside(expr: &Expr) -> Vec<&Expr> { + match expr { + Expr::Column(_) | Expr::Literal(_) | Expr::Parameter(_) => Vec::new(), + Expr::BinaryOp { left, right, .. } | Expr::Distance { left, right } => { + vec![&**left, &**right] + } + Expr::UnaryOp { operand: inner, .. } | Expr::IsNull { expr: inner, .. } => vec![&**inner], + Expr::FunctionCall { args: items, .. } | Expr::Row(items) => items.iter().collect(), + Expr::InList { expr, list, .. } => std::iter::once(&**expr).chain(list).collect(), + } +} + +/// Hands `each` the left and right side of every distance written in `expr`, at any depth. +fn distances(expr: &Expr, each: &mut dyn FnMut(&Expr, &Expr) -> Result<()>) -> Result<()> { + if let Expr::Distance { left, right } = expr { + each(left, right)?; + } + inside(expr) + .into_iter() + .try_for_each(|inner| distances(inner, each)) +} + +/// What an expression of an UPDATE or DELETE must pass: [`rows_stand`], and no distance, which +/// only a search gives a row. Refused here, the statement is refused whatever its table holds. +fn in_a_write(expr: &Expr) -> Result<()> { + rows_stand(expr)?; + distances(expr, &mut |_, _| { + Err(not_supported("<=> in an UPDATE or DELETE")) + }) +} + +/// Refuses a row value that stands anywhere but on either side of `IN`, and one beside `IN` +/// that has another number of values than its partner. +fn rows_stand(expr: &Expr) -> Result<()> { + let (left, list) = match expr { + Expr::InList { expr, list, .. } => (&**expr, list), + Expr::Row(_) => return Err(not_supported("a row value anywhere but beside IN")), + _ => return inside(expr).into_iter().try_for_each(rows_stand), + }; + let width = |side: &Expr| match side { + Expr::Row(values) => values.len(), + _ => 1, + }; + for side in std::iter::once(left).chain(list) { + if width(side) != width(left) { + return Err(Error::PlanError( + "a row value beside IN must have as many values as its partner".to_owned(), + )); + } + // What a side holds is checked as any expression is: a row value holds no row value. + match side { + Expr::Row(values) => values.iter().try_for_each(rows_stand)?, + _ => rows_stand(side)?, + } + } + Ok(()) +} diff --git a/v2/contextdb/src/store/disk.rs b/v2/contextdb/src/store/disk.rs index a1550b81..b1e684f7 100644 --- a/v2/contextdb/src/store/disk.rs +++ b/v2/contextdb/src/store/disk.rs @@ -1,14 +1,18 @@ //! What the data file holds, how one committed moment of it is read, and how a commit is //! written into it. //! -//! Five tables of the storage library hold everything. Rows and index entries are keyed by the -//! number of their table or index and then their key, so one table's rows are one run of keys. -//! Large values sit under a place number of their own. The last table holds the few numbers the -//! store keeps: the newest commit, the commit that last changed the schema, the next free place, -//! the declared memory limit, and each table's row count. +//! Seven tables of the storage library hold everything. Rows and index entries are keyed by +//! the number of their table or index and then their key, so one table's rows are one run of +//! keys. A large value is cut into slices of 243 bytes, each under a place number of its own: +//! the slices of one value are under consecutive places, and its row says how many bytes there +//! are and which place is the first. The vector part's blocks have a table to themselves, and +//! its partition entries and pending entries share another. The last table holds the few +//! numbers the store keeps: the newest commit, the commit that last changed the schema, the +//! next free place, the declared memory limit, each table's row count, and the vector part's +//! three numbers. use std::io::{self, ErrorKind}; -use std::ops::Bound; +use std::ops::{Bound, Range}; use std::path::Path; use std::sync::Arc; @@ -17,6 +21,7 @@ use redb::{ReadOnlyTable, ReadTransaction, ReadableDatabase, ReadableTable, Tabl use super::folder::Hold; use super::keys::{self, Span, stored}; use super::{Beside, Cell, ChangeSet, EntryIter, IndexId, Key, Read, Row, RowIter, TableId, View}; +use super::{VectorNumber, VectorRecord}; use super::{record, split}; use crate::memory::Reservation; use crate::types::{Error, Result, TxId, Value}; @@ -28,12 +33,25 @@ pub(super) type Stored = std::result::Result; const ROWS: TableDefinition<&[u8], &[u8]> = TableDefinition::new("rows"); /// Index entries: index number, entry and row key, to the length of the entry. const INDEX: TableDefinition<&[u8], &[u8]> = TableDefinition::new("index"); -/// Large values: place, to the value as `record` writes it. +/// Large values: place, to one slice of the record `record` stores for the value. const BESIDE: TableDefinition = TableDefinition::new("beside"); /// Schema records by name, as `schema` wrote them. const SCHEMA: TableDefinition<&str, &[u8]> = TableDefinition::new("schema"); /// The numbers the store keeps. const NUMBERS: TableDefinition = TableDefinition::new("numbers"); +/// The vector part's blocks, under their keys as `keys::vector` writes them, and nothing else. +/// The storage library rewrites a whole leaf when any record in it changes, and it divides a +/// leaf by bytes, so a small record that shared a leaf with a block of 256 KiB would have the +/// block rewritten every time it changed. +const BLOCKS: TableDefinition<&[u8], &[u8]> = TableDefinition::new("vector_blocks"); +/// The vector part's partition entries and pending entries, under their keys as `keys::vector` +/// writes them. +const VECTORS: TableDefinition<&[u8], &[u8]> = TableDefinition::new("vectors"); + +/// The bytes of one slice of a large value; the last slice of a value may be shorter. A leaf +/// of the storage library is 4 bytes and, for each slice, 8 of key, 4 of length and the slice, +/// so sixteen slices fill a page of 4 KiB: 4,084 of its 4,096 bytes. +const SLICE: usize = 243; const COMMIT: u64 = 0; const SCHEMA_COMMIT: u64 = 1; @@ -46,6 +64,21 @@ fn count_of(table: u32) -> u64 { (1 << 32) | u64::from(table) } +/// Where a number of the vector part is kept among the numbers. +pub(super) fn vector_number_at(number: VectorNumber) -> u64 { + match number { + VectorNumber::Partitions(column) => (2 << 32) | u64::from(column.0), + VectorNumber::NextPartition(column) => (3 << 32) | u64::from(column.0), + VectorNumber::Pending => 4 << 32, + } +} + +/// The places of the slices of a large value: its cell gives them, so nothing is read to +/// find them. +fn places(kept: &Beside) -> Range { + kept.place..kept.place + kept.bytes.div_ceil(SLICE as u64) +} + /// Tells a failure of the storage library or of a file as the engine tells it: as damage when /// stored data failed its check, and otherwise as a failure of `operation`, which is what the /// store was doing. @@ -78,10 +111,25 @@ pub(super) fn failed( } } +/// What a record that is not one the store wrote is told as: a failed check. +fn unsound() -> redb::Error { + redb::Error::Corrupted("a stored record is not one the store wrote".to_owned()) +} + fn number(numbers: &impl ReadableTable, key: u64) -> Stored { Ok(numbers.get(key)?.map_or(0, |number| number.value())) } +/// The memory limit the store remembers, in bytes; `None` when none is declared. +fn memory_limit(numbers: &impl ReadableTable) -> Stored> { + Ok(numbers.get(MEMORY_LIMIT)?.map(|limit| limit.value())) +} + +/// The memory limit the newest commit of a data file remembers. +pub(super) fn remembered_limit(data: &redb::Database) -> Stored> { + memory_limit(&data.begin_read()?.open_table(NUMBERS)?) +} + /// Makes the tables in a data file that does not have them yet. pub(super) fn prepare(data: &redb::Database) -> Stored<()> { match data.begin_read()?.open_table(NUMBERS) { @@ -95,17 +143,21 @@ pub(super) fn prepare(data: &redb::Database) -> Stored<()> { write.open_table(BESIDE)?; write.open_table(SCHEMA)?; write.open_table(NUMBERS)?; + write.open_table(BLOCKS)?; + write.open_table(VECTORS)?; write.commit()?; Ok(()) } -/// One committed moment: the five tables as one read of the storage library shows them. +/// One committed moment: the seven tables as one read of the storage library shows them. pub(super) struct Moment { rows: ReadOnlyTable<&'static [u8], &'static [u8]>, index: ReadOnlyTable<&'static [u8], &'static [u8]>, beside: ReadOnlyTable, schema: ReadOnlyTable<&'static str, &'static [u8]>, numbers: ReadOnlyTable, + blocks: ReadOnlyTable<&'static [u8], &'static [u8]>, + vectors: ReadOnlyTable<&'static [u8], &'static [u8]>, commit: TxId, folder: Arc, /// In a reading process, the file that keeps this moment's pages from being reused. @@ -124,6 +176,8 @@ impl Moment { index: read.open_table(INDEX)?, beside: read.open_table(BESIDE)?, schema: read.open_table(SCHEMA)?, + blocks: read.open_table(BLOCKS)?, + vectors: read.open_table(VECTORS)?, commit: TxId(number(&numbers, COMMIT)?), numbers, folder, @@ -144,7 +198,7 @@ impl Moment { read: F, ) -> Result> + use> where - F: Fn(&[u8]) -> T, + F: Fn(&[u8]) -> Option, { let folder = self.folder.clone(); let range = table @@ -152,7 +206,8 @@ impl Moment { .map_err(|error| failed(&folder, "read", error))?; Ok(range.map(move |entry| { let (key, record) = entry.map_err(|error| failed(&folder, "read", error))?; - Ok((Key(key.value()[4..].to_vec()), read(record.value()))) + let read = read(record.value()).ok_or_else(|| failed(&folder, "read", unsound()))?; + Ok((Key(key.value()[4..].to_vec()), read)) })) } @@ -162,8 +217,10 @@ impl Moment { table: u32, span: &Span, ) -> Result> + use<>> { - self.scan(&self.rows, table, span, |row| Row { - cells: record::decode_row(row), + self.scan(&self.rows, table, span, |row| { + Some(Row { + cells: record::decode_row(row), + }) }) } @@ -175,7 +232,7 @@ impl Moment { span: &Span, ) -> Result> + use<>> { self.scan(&self.index, index, span, |length| { - u32::from_le_bytes(std::array::from_fn(|at| length[at])) as usize + Some(u32::from_le_bytes(*length.first_chunk()?) as usize) }) } @@ -195,10 +252,29 @@ impl Moment { })) } - fn kept_beside(&self, place: u64) -> Stored { - let kept = self.beside.get(place)?; - let kept = kept.expect("the commit that keeps a row keeps its large values"); - Ok(record::decode_value(kept.value())) + /// Tells of a record that is not what was stored: damage, told as a failed check of the + /// storage library is. + fn damaged(&self, reason: &str) -> Error { + let reason = redb::Error::Corrupted(reason.to_owned()); + failed(&self.folder, "read", reason) + } + + /// Loads a large value: grows `reservation` by the bytes stored for it, gathers its slices + /// with one range read into one buffer of that length, and has `record` make the value. + fn kept_beside(&self, kept: &Beside, reservation: &mut Reservation) -> Result { + reservation.grow(kept.bytes)?; + let mut stored = Vec::with_capacity(kept.bytes as usize); + let slices = self.beside.range::(places(kept)); + for slice in slices.map_err(|error| failed(&self.folder, "read", error))? { + let (_, slice) = slice.map_err(|error| failed(&self.folder, "read", error))?; + stored.extend_from_slice(slice.value()); + } + if stored.len() as u64 != kept.bytes { + return Err(self.damaged("a large value's slices are not the bytes its row records")); + } + // A record its own bytes do not support is damage. + let loaded = record::load(stored, reservation)?; + loaded.map_err(|reason| self.damaged(reason)) } } @@ -226,8 +302,52 @@ impl Read for View { } fn load(&self, beside: &Beside, reservation: &mut Reservation) -> Result { - reservation.grow(beside.bytes)?; - self.moment.told(self.moment.kept_beside(beside.place)) + let held = reservation.bytes(); + let loaded = self.moment.kept_beside(beside, reservation); + if loaded.is_err() { + // A load that made no value holds nothing: what it grew is given back. + reservation.shrink(reservation.bytes() - held); + } + loaded + } + + fn staged(&self, _table: TableId, _each: &mut dyn FnMut(&Key) -> Result) -> Result<()> { + Ok(()) + } + + fn vectors( + &self, + from: VectorRecord<'_>, + each: &mut dyn FnMut(VectorRecord<'_>, &[u8]) -> Result, + ) -> Result<()> { + let moment = &*self.moment; + // Every key of a kind begins with the kind's byte, so the kind ends before the next. + let start = keys::vector(from); + let end = [start[0] + 1]; + let table = match start[0] { + keys::BLOCK => &moment.blocks, + _ => &moment.vectors, + }; + let records = table.range::<&[u8]>(start.as_slice()..end.as_slice()); + let mut key = Key(Vec::new()); + for record in records.map_err(|error| failed(&moment.folder, "read", error))? { + let (stored, value) = record.map_err(|error| failed(&moment.folder, "read", error))?; + // The value is the bytes of the storage library's own page, lent for the call. + let record = keys::vector_of(stored.value(), &mut key).ok_or_else(unsound); + if !each(moment.told(record)?, value.value())? { + break; + } + } + Ok(()) + } + + fn vector_number(&self, number: VectorNumber) -> Result { + let at = vector_number_at(number); + self.moment.told(self::number(&self.moment.numbers, at)) + } + + fn damaged(&self, reason: &str) -> Error { + self.moment.damaged(reason) } } @@ -247,15 +367,13 @@ impl View { /// The memory limit the store remembers, in bytes; `None` when none is declared. pub(crate) fn memory_limit(&self) -> Result> { - let kept = self.moment.numbers.get(MEMORY_LIMIT); - let kept = kept.map(|kept| kept.map(|limit| limit.value())); - self.moment.told(kept.map_err(redb::Error::from)) + self.moment.told(memory_limit(&self.moment.numbers)) } } /// Writes a change set as one transaction of the storage library and waits for the disk. /// Returns the number of the commit it made. -pub(super) fn write(data: &redb::Database, changes: &ChangeSet) -> Stored { +pub(super) fn write(data: &redb::Database, changes: &mut ChangeSet) -> Stored { // The storage library has a commit option, `set_quick_repair`, that lets reopening after a // kill skip its repair. It is not used: with it every commit rewrites the library's whole // free-page table, about 133 KB for each GiB of file. So an open after a kill reads every @@ -287,9 +405,15 @@ pub(super) fn write(data: &redb::Database, changes: &ChangeSet) -> Stored { let before = match &row.values { Some(values) => { let row = record::encode_row(values, |large| -> Stored { - beside.insert(free, large)?; - free += 1; - Ok(free - 1) + // The compressor's buffer is counted until the value is written. + let room = record::compress(large, &mut changes.reserved); + let first = free; + for slice in large.chunks(SLICE) { + beside.insert(free, slice)?; + free += 1; + } + changes.reserved.shrink(room); + Ok(first) })?; let before = rows.insert(key.as_slice(), row.as_slice())?; before.map(|row| record::decode_row(row.value())) @@ -301,7 +425,9 @@ pub(super) fn write(data: &redb::Database, changes: &ChangeSet) -> Stored { }; for cell in before.into_iter().flatten() { if let Cell::Beside(kept) = cell { - beside.remove(kept.place)?; + for place in places(&kept) { + beside.remove(place)?; + } } } } @@ -342,6 +468,22 @@ pub(super) fn write(data: &redb::Database, changes: &ChangeSet) -> Stored { Some(None) => drop(numbers.remove(MEMORY_LIMIT)?), None => {} } + + let mut blocks = write.open_table(BLOCKS)?; + let mut vectors = write.open_table(VECTORS)?; + for (key, value) in &changes.vectors { + let table = match key.0[0] { + keys::BLOCK => &mut blocks, + _ => &mut vectors, + }; + match value { + Some(value) => drop(table.insert(key.0.as_slice(), value.as_slice())?), + None => drop(table.remove(key.0.as_slice())?), + } + } + for (at, number) in &changes.vector_numbers { + numbers.insert(*at, *number)?; + } } write.commit()?; Ok(commit) diff --git a/v2/contextdb/src/store/folder.rs b/v2/contextdb/src/store/folder.rs index d0a96b17..081ef588 100644 --- a/v2/contextdb/src/store/folder.rs +++ b/v2/contextdb/src/store/folder.rs @@ -6,6 +6,11 @@ //! holds it, shared, for the moment it takes to find the newest commit and leave a file naming //! it; the writer holds it, alone, while it looks through those files. So the writer never //! sees a reader between its choice of a commit and the file that says so. +//! +//! Such a lock belongs to the open file, and a child process has a copy of every file its +//! parent had open when it was started, until it replaces itself with its program. So every +//! lock is let go by unlocking the file before it is closed: a copy a child still has then +//! holds nothing. use std::fs::{self, File, OpenOptions, TryLockError}; use std::io::{self, Read as _, Write as _}; @@ -83,9 +88,18 @@ impl StorageBackend for Plain { } } +/// A file this process holds locked. Dropping it unlocks the file and then closes it. +pub(super) struct Locked(File); + +impl Drop for Locked { + fn drop(&mut self) { + let _ = self.0.unlock(); + } +} + /// Takes the folder's `lock` for this process and writes the process id into it. When another /// process holds it, answers at once with the id that process wrote. -pub(super) fn take_lock(folder: &Path) -> io::Result> { +pub(super) fn take_lock(folder: &Path) -> io::Result> { let mut lock = OpenOptions::new() .read(true) .write(true) @@ -94,8 +108,10 @@ pub(super) fn take_lock(folder: &Path) -> io::Result> { .open(folder.join("lock"))?; match lock.try_lock() { Ok(()) => { - lock.set_len(0)?; - lock.write_all(std::process::id().to_string().as_bytes())?; + let mut lock = Locked(lock); + lock.0.set_len(0)?; + lock.0 + .write_all(std::process::id().to_string().as_bytes())?; Ok(Ok(lock)) } Err(TryLockError::WouldBlock) => { @@ -112,7 +128,7 @@ pub(super) fn take_lock(folder: &Path) -> io::Result> { pub(super) struct Hold { pub(super) commit: u64, path: PathBuf, - _locked: File, + _locked: Locked, } impl Hold { @@ -126,9 +142,10 @@ impl Hold { NEXT.fetch_add(1, Ordering::Relaxed) ); let path = readers.join(name); - let mut file = File::create(&path)?; + let file = File::create(&path)?; file.lock()?; - file.write_all(commit.to_string().as_bytes())?; + let mut file = Locked(file); + file.0.write_all(commit.to_string().as_bytes())?; Ok(Hold { commit, path, @@ -164,7 +181,10 @@ fn look_through(readers: &Path) -> io::Result> { }; match file.try_lock() { // Nobody holds it: its reader is gone, and may have removed it already. - Ok(()) => drop(fs::remove_file(&path)), + Ok(()) => { + let _ = fs::remove_file(&path); + file.unlock()?; + } Err(TryLockError::WouldBlock) => { let mut named = String::new(); file.read_to_string(&mut named)?; diff --git a/v2/contextdb/src/store/keys.rs b/v2/contextdb/src/store/keys.rs index aec5d20f..f382547f 100644 --- a/v2/contextdb/src/store/keys.rs +++ b/v2/contextdb/src/store/keys.rs @@ -11,7 +11,7 @@ use std::ops::Bound; -use super::Key; +use super::{Key, VectorId, VectorRecord}; use crate::types::{Error, Result, SortDirection, Value}; /// No value's form begins with this byte, turned or not, so a key followed by it sorts after @@ -111,3 +111,50 @@ pub(super) fn span(from: Bound<&Key>, to: Bound<&Key>) -> Span { pub(super) fn stored(id: u32, key: &[u8]) -> Vec { [&id.to_be_bytes(), key].concat() } + +/// The first byte of the stored key of a block of the vector part. Blocks are kept in a table +/// of their own. +pub(super) const BLOCK: u8 = 0; +const PARTITION: u8 = 1; +const PENDING: u8 = 2; + +/// A record of the vector part as the data file keys it: one byte for its kind, the column, +/// and then the rest of its fields in order, every number most significant byte first. So the +/// records of one kind are one run of keys in the order of their fields: the columns follow +/// each other by number, and within a column a partition's pending rows or blocks are one run. +pub(super) fn vector(record: VectorRecord<'_>) -> Vec { + let number = u32::to_be_bytes; + match record { + VectorRecord::Block(column, partition, block) => { + let numbers = [number(column.0), number(partition), number(block)]; + [&[BLOCK][..], numbers.as_flattened()].concat() + } + VectorRecord::Partition(column, tuple) => { + [&[PARTITION][..], &number(column.0), &tuple.0].concat() + } + VectorRecord::Pending(column, partition, row) => { + let numbers = [number(column.0), number(partition)]; + [&[PENDING][..], numbers.as_flattened(), &row.0].concat() + } + } +} + +/// The record a stored key of the vector part stands for. A key tuple or a row key is read +/// into `key`, which the record then borrows. +pub(super) fn vector_of<'a>(stored: &[u8], key: &'a mut Key) -> Option> { + let number = |at: usize| Some(u32::from_be_bytes(*stored.get(at..)?.first_chunk()?)); + let column = VectorId(number(1)?); + key.0.clear(); + Some(match stored[0] { + BLOCK => VectorRecord::Block(column, number(5)?, number(9)?), + PARTITION => { + key.0.extend_from_slice(&stored[5..]); + VectorRecord::Partition(column, key) + } + _ => { + let partition = number(5)?; + key.0.extend_from_slice(&stored[9..]); + VectorRecord::Pending(column, partition, key) + } + }) +} diff --git a/v2/contextdb/src/store/mod.rs b/v2/contextdb/src/store/mod.rs index d4a6b62c..7f9d6b6e 100644 --- a/v2/contextdb/src/store/mod.rs +++ b/v2/contextdb/src/store/mod.rs @@ -1,8 +1,9 @@ //! store: the folder and the data file; how each kind of record is keyed; views of one moment; //! reader files and locks. //! -//! It never interprets SQL or a table's rules: it keeps rows, index entries, counts and schema -//! records under the keys it is given and hands them back exactly as they were written. +//! It never interprets SQL or a table's rules: it keeps rows, index entries, counts, schema +//! records and the vector part's records under the keys it is given and hands them back exactly +//! as they were written. //! //! Two of the things the parts meet at are defined here: [`Read`], the view to read from, and //! [`ChangeSet`], the change set to write with. @@ -28,12 +29,16 @@ use redb::{ReadTransaction, ReadableDatabase}; use crate::memory::{Budget, Reservation}; use crate::types::{Error, Result, SortDirection, TxId, Value}; use disk::{Moment, Stored}; -use folder::{DataFile, Hold}; +use folder::{DataFile, Hold, Locked}; use keys::Span; -/// The size of the page cache of one open store, in bytes. It is reserved from the budget when -/// the store is opened and does not change. -const CACHE_BYTES: u64 = 8 << 20; +/// The size of the page cache of a store opened under this memory limit, in bytes: a quarter +/// of the limit and never less than 8 MiB, which is also its size when there is no limit. It is +/// reserved from the budget when the store is opened and does not change until the store is +/// closed, because the storage library fixes its cache when the data file is opened. +fn cache_under(limit: Option) -> u64 { + (limit.unwrap_or(0) / 4).max(8 << 20) +} /// What a refusal names when the memory limit cannot hold the page cache. const CACHE_NEEDED: &str = "page cache needed to open the store"; @@ -43,11 +48,14 @@ const CACHE_NEEDED: &str = "page cache needed to open the store"; /// /// It is given in [`crate::DatabaseOpenOptions::storage`] or /// [`crate::ReadSessionOptions::storage`]; with none given the file is read and written as it -/// is. The library opens and locks the data file itself, calls the function once for each open -/// with the open file, and from then on reads and writes the file only through the -/// [`StorageBackend`] the function returned. A session opened for reading hands over a file -/// opened for reading only. A host wraps the file to keep the bytes somewhere else or to change -/// them on the way; a test wraps it to count the calls, to fail one or to stop at one. +/// is. The library opens and locks the data file itself, calls the function with the open file +/// each time it has opened it, and from then on reads and writes the file only through the +/// [`StorageBackend`] the function returned. That is once for each open, and twice for a writer +/// that names no memory limit over a store that remembers one: the data file is opened to +/// read the limit, and again with the page cache the limit calls for when that is another +/// size. A session opened for reading hands over a file opened for reading only. A host wraps +/// the file to keep the bytes somewhere else or to change them on the way; a test wraps it to +/// count the calls, to fail one or to stop at one. /// /// An error returned by the function, or by a call of what it returned, is reported as /// [`crate::Error::StoreIo`]. The two exceptions are the kinds `UnexpectedEof` and @@ -64,6 +72,40 @@ pub(crate) struct TableId(pub u32); #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] pub(crate) struct IndexId(pub u32); +/// A VECTOR column's number inside one store, given by `schema` like a [`TableId`]. Every record +/// the vector part keeps for the column begins with it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub(crate) struct VectorId(pub u32); + +/// A record the vector part keeps. `store` decides the bytes of its key and where it is held; +/// the value is the vector part's own and comes back exactly as it was put. +/// +/// Records of one kind are in the order of their fields, left to right: every record of one +/// column is one run of keys, and inside it every pending row or block of one partition is. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum VectorRecord<'a> { + /// What the column keeps about one partition, under the partition's key tuple: the values + /// of the PARTITION_KEY columns encoded as [`Key::from_values`] encodes them, ascending. A + /// column with no PARTITION_KEY has one partition, under the key of no values. + Partition(VectorId, &'a Key), + /// A row whose vector is in no block yet: the partition's number, then the row's key. Its + /// value is empty. + Pending(VectorId, u32, &'a Key), + /// One block of a partition: the partition's number, then the block's. + Block(VectorId, u32, u32), +} + +/// A number the store keeps for the vector part. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum VectorNumber { + /// How many partitions a column has. + Partitions(VectorId), + /// The number the column's next new partition takes. It only ever rises. + NextPartition(VectorId), + /// How many pending records there are, over every column. + Pending, +} + /// A primary key or an index entry, encoded so that comparing the bytes compares the values. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] pub(crate) struct Key(Vec); @@ -83,6 +125,22 @@ impl Key { keys::encode(values, directions).map(Key) } + /// The bytes of the key, for a part that keeps keys inside a record of its own. + pub(crate) fn as_bytes(&self) -> &[u8] { + &self.0 + } + + /// The key that [`Key::as_bytes`] gave these bytes for. + pub(crate) fn from_bytes(bytes: &[u8]) -> Key { + Key(bytes.to_vec()) + } + + /// True when this key begins with every value of `prefix`, in order. Every key begins with + /// the key of no values. + pub(crate) fn starts_with(&self, prefix: &Key) -> bool { + self.0.starts_with(&prefix.0) + } + /// The first key after every key that begins with this one, `None` when there is none. pub(crate) fn prefix_end(&self) -> Option { let last = self.0.iter().rposition(|byte| *byte != 0xFF)?; @@ -107,10 +165,13 @@ pub(crate) enum Cell { Beside(Beside), } -/// Where a large value is kept and how many bytes loading it takes. +/// Where a large value is kept and how many bytes are stored for it. #[derive(Debug, Clone, PartialEq)] pub(crate) struct Beside { + /// The bytes of the value's stored record. A large text or JSON value may be stored + /// compressed, so this can be less than the bytes the loaded value takes. bytes: u64, + /// The place of the first slice of the record; the others are under the places after it. place: u64, } @@ -153,10 +214,46 @@ pub(crate) trait Read { /// How many rows the table holds at this moment. One small record, never a scan. fn count(&self, table: TableId) -> Result; - /// Loads a value kept beside its row. Grows `reservation` by the bytes the value takes - /// first and refuses with `MemoryBudgetExceeded` before reading anything if that does not - /// fit. + /// Loads a value kept beside its row. Grows `reservation` by the bytes stored for the + /// value first and refuses with `MemoryBudgetExceeded` before reading anything if that does + /// not fit. A value that is stored compressed can be refused a second time, after its + /// stored bytes are read and before any byte of the value exists: the reservation is then + /// grown by the bytes the value really takes. + /// + /// After a load the reservation holds the bytes the value takes, however it was stored; + /// while a JSON or vector value is made it holds them twice. A load that fails gives back + /// all it grew. A stored length the stored bytes do not support is `StoreCorrupted`. fn load(&self, beside: &Beside, reservation: &mut Reservation) -> Result; + + /// Calls `each` with the key of every row of the table that this reader has staged over + /// its committed moment, in key order, until `each` answers false: the rows a transaction + /// has written or removed and not yet committed. A [`View`] has none. + fn staged(&self, table: TableId, each: &mut dyn FnMut(&Key) -> Result) -> Result<()>; + + /// Calls `each` with the vector part's records of the kind of `from`, in key order, from + /// `from` itself or the first record after it, until `each` answers false or the kind has + /// no more records. Records of other columns follow those of `from`'s column, so the + /// caller stops where its own run ends. + /// + /// The value is lent, not copied: a block is handed over in place from the page cache, and + /// costs one read of the file when it is not there. A partition's key tuple compares with + /// `from` as a whole key, so the tuple of no values starts at a column's first partition. + /// + /// Vector records are staged only inside the commit function and by upkeep, after the last + /// read of their change set. A [`ChangeSet`] therefore answers from its committed moment. + fn vectors( + &self, + from: VectorRecord<'_>, + each: &mut dyn FnMut(VectorRecord<'_>, &[u8]) -> Result, + ) -> Result<()>; + + /// One of the numbers kept for the vector part; zero when it was never put. + fn vector_number(&self, number: VectorNumber) -> Result; + + /// The refusal of a record that its own bytes, or the records it goes with, do not + /// support: `StoreCorrupted`, naming the store's folder, with this reason. A part that + /// keeps records of its own tells the damage it finds in them with this. + fn damaged(&self, reason: &str) -> Error; } /// One committed moment of a store. @@ -189,6 +286,11 @@ pub(crate) struct ChangeSet { schema: BTreeMap>, /// The memory limit the store is to remember, when one was staged: `Some(None)` for none. memory_limit: Option>, + /// Staged records of the vector part by their stored key: the value, or `None` for a + /// removal. + vectors: BTreeMap>>, + /// Staged numbers of the vector part by their place among the store's numbers. + vector_numbers: BTreeMap, /// Every staging call of a statement in order, each with what it put aside and the bytes it /// was counted as, so that the calls after a mark can be undone. calls: Vec<(Undo, u64)>, @@ -212,6 +314,8 @@ enum Undo { Entry(u32, Key, Option>), Count(u32, i64), Schema(String, Option>), + Vector(Key, Option>>), + VectorNumber(u64, Option), } /// What one staging call is counted as beside its key and its values: its place in the map @@ -228,6 +332,8 @@ impl ChangeSet { counts: BTreeMap::new(), schema: BTreeMap::new(), memory_limit: None, + vectors: BTreeMap::new(), + vector_numbers: BTreeMap::new(), calls: Vec::new(), reserved: budget.reserve(0, "transaction")?, }) @@ -255,13 +361,16 @@ impl ChangeSet { pub(crate) fn rewind(&mut self, mark: usize) { for (undo, bytes) in self.calls.drain(mark..).rev() { match undo { - Undo::Row(table, key, before) => put_back(&mut self.rows, table, key, before), + Undo::Row(table, key, before) => { + put_back(self.rows.entry(table).or_default(), key, before); + } Undo::Entry(index, joined, before) => { - put_back(&mut self.index, index, joined, before); + put_back(self.index.entry(index).or_default(), joined, before); } Undo::Count(table, delta) => *self.counts.entry(table).or_default() -= delta, - Undo::Schema(name, Some(before)) => drop(self.schema.insert(name, before)), - Undo::Schema(name, None) => drop(self.schema.remove(&name)), + Undo::Schema(name, before) => put_back(&mut self.schema, name, before), + Undo::Vector(key, before) => put_back(&mut self.vectors, key, before), + Undo::VectorNumber(at, before) => put_back(&mut self.vector_numbers, at, before), } self.reserved.shrink(bytes); } @@ -349,6 +458,38 @@ impl ChangeSet { Ok(()) } + /// Stages a record of the vector part, new or in place of the record with that key. The + /// value's bytes are charged to the change set's reservation until it is committed or + /// dropped. + pub(crate) fn put_vector(&mut self, record: VectorRecord<'_>, value: Vec) -> Result<()> { + self.stage_vector(record, Some(value)) + } + + /// Stages the removal of a record of the vector part. Removing a record that is not there + /// changes nothing. + pub(crate) fn delete_vector(&mut self, record: VectorRecord<'_>) -> Result<()> { + self.stage_vector(record, None) + } + + fn stage_vector(&mut self, record: VectorRecord<'_>, value: Option>) -> Result<()> { + let key = Key(keys::vector(record)); + let held = value.as_ref().map_or(0, Vec::len); + let bytes = self.counted(&key, held as u64)?; + let before = self.vectors.insert(key.clone(), value); + self.calls.push((Undo::Vector(key, before), bytes)); + Ok(()) + } + + /// Stages one of the numbers kept for the vector part. + pub(crate) fn put_vector_number(&mut self, number: VectorNumber, value: u64) -> Result<()> { + self.reserved.grow(STAGED_BYTES)?; + let at = disk::vector_number_at(number); + let before = self.vector_numbers.insert(at, value); + self.calls + .push((Undo::VectorNumber(at, before), STAGED_BYTES)); + Ok(()) + } + /// Counts one staging call against the budget before anything is kept: its key twice, once /// in the map and once in the list of calls, and what its values hold. Returns the bytes. fn counted(&mut self, key: &Key, held: u64) -> Result { @@ -374,8 +515,7 @@ impl ChangeSet { } /// Puts a staged record back as it was before one staging call. -fn put_back(staged: &mut Staged, id: u32, key: Key, before: Option) { - let records = staged.entry(id).or_default(); +fn put_back(records: &mut BTreeMap, key: K, before: Option) { match before { Some(before) => drop(records.insert(key, before)), None => drop(records.remove(&key)), @@ -496,12 +636,39 @@ impl Read for ChangeSet { fn load(&self, beside: &Beside, reservation: &mut Reservation) -> Result { self.base.load(beside, reservation) } + + fn staged(&self, table: TableId, each: &mut dyn FnMut(&Key) -> Result) -> Result<()> { + for key in self.rows.get(&table.0).into_iter().flat_map(BTreeMap::keys) { + if !each(key)? { + break; + } + } + Ok(()) + } + + fn vectors( + &self, + from: VectorRecord<'_>, + each: &mut dyn FnMut(VectorRecord<'_>, &[u8]) -> Result, + ) -> Result<()> { + self.base.vectors(from, each) + } + + fn vector_number(&self, number: VectorNumber) -> Result { + self.base.vector_number(number) + } + + fn damaged(&self, reason: &str) -> Error { + self.base.damaged(reason) + } } /// One open store folder, held either for writing or for reading. pub(crate) struct Store { held: HeldFor, folder: Arc, + /// The budget the page cache is reserved from. + budget: Budget, _cache: Reservation, } @@ -510,7 +677,7 @@ pub(crate) struct Store { enum HeldFor { Writing { data: redb::Database, - _lock: File, + _lock: Locked, }, Reading { data: redb::ReadOnlyDatabase, @@ -556,10 +723,17 @@ impl Fixing { impl HeldFor { /// Takes the folder for writing, or answers with the id of the process that holds it. + /// With the folder comes the size of the page cache its data file was opened with. + /// + /// `limit` is the memory limit of the open's budget, and the cache is the one it calls + /// for. With none, the limit in force is the one the store remembers, which only the open + /// data file tells: the file is opened with the cache of no limit and, when the remembered + /// limit calls for another size, let go and opened again with that size. fn writing( folder: &Path, storage: Option<&WrapStorage>, - ) -> Stored> { + limit: Option, + ) -> Stored> { fs::create_dir_all(folder)?; let lock = match folder::take_lock(folder)? { Ok(lock) => lock, @@ -567,9 +741,26 @@ impl HeldFor { }; let readers = folder.join("readers"); fs::create_dir_all(&readers)?; - let gate = File::open(&readers)?; // Removes what readers that are gone left behind. - folder::oldest_held(&readers, &gate)?; + folder::oldest_held(&readers, &File::open(&readers)?)?; + let mut cache = cache_under(limit); + let mut data = HeldFor::data(folder, storage, cache)?; + disk::prepare(&data)?; + if limit.is_none() { + let remembered = cache_under(disk::remembered_limit(&data)?); + if remembered != cache { + drop(data); + cache = remembered; + data = HeldFor::data(folder, storage, cache)?; + } + } + Ok(Ok((HeldFor::Writing { data, _lock: lock }, cache))) + } + + /// Opens the data file of a folder held for writing, with a page cache of `cache` bytes. + fn data(folder: &Path, storage: Option<&WrapStorage>, cache: u64) -> Stored { + let readers = folder.join("readers"); + let gate = File::open(&readers)?; let file = OpenOptions::new() .read(true) .write(true) @@ -577,25 +768,29 @@ impl HeldFor { .truncate(false) .open(folder.join("data"))?; let data = redb::Builder::new() - .set_cache_size(CACHE_BYTES as usize) + .set_cache_size(cache as usize) // When the files cannot be looked through, every commit counts as held. .set_held_commits_callback(move || { folder::oldest_held(&readers, &gate).unwrap_or(Some(0)) }) .create_with_backend(DataFile::over(file, storage)?)?; - disk::prepare(&data)?; - Ok(Ok(HeldFor::Writing { data, _lock: lock })) + Ok(data) } - /// Opens the folder for reading, or answers `None` when it has no data file. - fn reading(folder: &Path, storage: Option<&WrapStorage>) -> Stored> { + /// Opens the folder for reading with a page cache of `cache` bytes, or answers `None` when + /// it has no data file. + fn reading( + folder: &Path, + storage: Option<&WrapStorage>, + cache: u64, + ) -> Stored> { let file = match File::open(folder.join("data")) { Ok(file) => file, Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(None), Err(error) => return Err(error.into()), }; let data = redb::Builder::new() - .set_cache_size(CACHE_BYTES as usize) + .set_cache_size(cache as usize) .open_read_only_with_backend(DataFile::over(file, storage)?)?; let readers = folder.join("readers"); let fixing = Mutex::new(Fixing { @@ -612,21 +807,33 @@ impl Store { /// /// Takes the folder's `lock` for the life of the value and refuses at once with /// `DatabaseLocked`, naming the holder, when another process has it. Clears what a killed - /// reader or writer left behind. Reserves the page cache from `budget`. Costs a few - /// milliseconds whatever the store's size: nothing is loaded. + /// reader or writer left behind. Costs a few milliseconds whatever the store's size: + /// nothing is loaded. + /// + /// Reserves the page cache from `budget`: the cache the budget's limit calls for, and with + /// no limit there the cache that the limit the store remembers calls for. The caller sets + /// the remembered limit on the budget afterwards, and the cache fits under it. pub(crate) fn open_for_writing( folder: &Path, budget: &Budget, storage: Option<&WrapStorage>, ) -> Result { - let cache = budget.reserve(CACHE_BYTES, CACHE_NEEDED)?; - let held = HeldFor::writing(folder, storage); + let limit = budget.limit(); + let mut cache = budget.reserve(cache_under(limit), CACHE_NEEDED)?; + let held = HeldFor::writing(folder, storage, limit); match held.map_err(|error| disk::failed(folder, "open", error))? { - Ok(held) => Ok(Store { - held, - folder: folder.into(), - _cache: cache, - }), + Ok((held, opened_with)) => { + // Another size is opened only with no limit in force, so nothing refuses this. + if opened_with != cache.bytes() { + cache = budget.reserve(opened_with, CACHE_NEEDED)?; + } + Ok(Store { + held, + folder: folder.into(), + budget: budget.clone(), + _cache: cache, + }) + } Err(holder_pid) => Err(Error::DatabaseLocked { holder_pid, path: folder.to_path_buf(), @@ -636,13 +843,16 @@ impl Store { /// Opens the folder for reading. Never creates anything: refuses with `StoreMissing` when /// there is no store. Works whether or not a writer is running, and never makes a writer - /// fail; a view holds one up only for the instant it has the short lock. + /// fail; a view holds one up only for the instant it has the short lock. Reserves the page + /// cache that the budget's limit calls for: a reader does not take the limit the store + /// remembers. pub(crate) fn open_for_reading( folder: &Path, budget: &Budget, storage: Option<&WrapStorage>, ) -> Result { - let held = HeldFor::reading(folder, storage) + let cache = cache_under(budget.limit()); + let held = HeldFor::reading(folder, storage, cache) .map_err(|error| disk::failed(folder, "open", error))?; let Some(held) = held else { return Err(Error::StoreMissing { @@ -652,10 +862,41 @@ impl Store { Ok(Store { held, folder: folder.into(), - _cache: budget.reserve(CACHE_BYTES, CACHE_NEEDED)?, + budget: budget.clone(), + _cache: budget.reserve(cache, CACHE_NEEDED)?, }) } + /// Checks a memory limit that is about to take hold against the page cache of this open. + /// + /// The cache keeps the size it was opened with until the store is closed, and the budget + /// counts all of it. A declared limit must hold, so the smallest limit this open can take + /// is what the budget counts now: 1 MiB, the cache, and what open transactions and live + /// results hold. A limit below that is refused with `MemoryBudgetExceeded`, operation + /// `set_limit`, and nothing changes: the refusal names the smallest limit as the bytes + /// needed and says to open the store with the limit as an open option. A limit at or above + /// it, or none, is accepted: once `database` has set it on the budget it holds at once, + /// the store remembers it, and it sizes the cache at the next open. `database` calls this + /// on the one path of `SET MEMORY_LIMIT`, before anything changes. + pub(crate) fn count_cache_under(&self, limit: Option) -> Result<()> { + let smallest = self.budget.used(); + match limit { + Some(limit) if limit < smallest => Err(Error::MemoryBudgetExceeded { + subsystem: "store", + operation: "set_limit", + requested_bytes: smallest, + available_bytes: limit, + budget_limit_bytes: limit, + hint: "the bytes needed are the smallest limit this session can take, because \ + its page cache keeps the size it was opened with; for a lower limit, \ + open the store with the limit as an open option, which sizes the \ + page cache from the limit, and declare the limit in that session for \ + the store to remember it", + }), + _ => Ok(()), + } + } + /// The newest committed moment. /// /// On a reading store this takes the short shared lock, reads which commit is newest and @@ -682,16 +923,18 @@ impl Store { /// Writes a change set as one commit and waits for the disk. /// - /// Rows, the values kept beside them, index entries, counts, schema records and the list of - /// changed keys go in one storage transaction: after a kill either all of it is there or - /// none. The caller holds the commit lock and has laid the change set over the newest + /// Rows, the values kept beside them, index entries, counts, schema records, the vector + /// part's records and numbers and the list of changed keys go in one storage transaction: + /// after a kill either all of it is there or none. A large text or JSON value is compressed + /// on its way in, in a buffer counted against the change set's reservation while it is + /// used. The caller holds the commit lock and has laid the change set over the newest /// moment; one over an older moment is refused with `StoreIo` and writes nothing. Returns /// the commit's number. Only a store opened for writing is given a change set to commit. - pub(crate) fn commit(&self, changes: ChangeSet) -> Result { + pub(crate) fn commit(&self, mut changes: ChangeSet) -> Result { let HeldFor::Writing { data, .. } = &self.held else { unreachable!("a store opened for reading is given nothing to commit") }; - match disk::write(data, &changes) { + match disk::write(data, &mut changes) { Ok(commit) => Ok(TxId(commit)), Err(error) => Err(disk::failed(&self.folder, "commit", error)), } diff --git a/v2/contextdb/src/store/record.rs b/v2/contextdb/src/store/record.rs index a63b41fd..f3bc4f70 100644 --- a/v2/contextdb/src/store/record.rs +++ b/v2/contextdb/src/store/record.rs @@ -1,15 +1,37 @@ //! How a row and its values are written as bytes. //! //! A row is its cells one after another. A cell is one tag byte and then the value; text, JSON -//! and vectors carry their length in bytes first. A large value is written the same way, but -//! beside the row, and the row keeps a cell that says where and how many bytes. +//! and vectors carry their length in bytes first, as a `u64`. That is a value's record. A +//! text, JSON or vector value whose record is more than 1,024 bytes is kept beside the row, and +//! the row keeps a cell that says where and how many bytes are stored for it. +//! +//! What is stored for a large value is its record in one of two forms, and the first byte says +//! which. The plain form is the record itself. The compressed form is the tag with its top bit +//! set, the same eight bytes of length, and one LZ4 block that decodes to exactly that many +//! bytes and runs to the end. A text or JSON value of at most 1 MiB is stored compressed when +//! that saves an eighth of its bytes; a longer one and a vector are always stored plain, and a +//! value that stays in its row is never touched. Either form reads back as the same value. + +use lz4_flex::block::{compress_into, decompress_into, get_maximum_output_size}; use super::{Beside, Cell}; -use crate::types::{TxId, Uuid, Value}; +use crate::memory::Reservation; +use crate::types::{self, TxId, Uuid, Value}; /// A text, JSON or vector value that takes more bytes than this is kept beside its row. const BESIDE_FROM: usize = 1024; +/// The tag and the length in front of a text, JSON or vector value. +const FRONT: usize = 9; + +/// Set in the first byte of a stored record whose value is one LZ4 block. +const COMPRESSED: u8 = 0x80; + +/// The longest text or JSON value that is compressed. Reading a compressed value holds its +/// stored bytes and its real bytes together for a moment, so this keeps that extra under 1 MiB +/// and a longer value reads holding no more than its own bytes. +const COMPRESS_TO: usize = 1 << 20; + const NULL: u8 = 0; const BOOL: u8 = 1; const INT64: u8 = 2; @@ -22,29 +44,65 @@ const VECTOR: u8 = 8; const TXID: u8 = 9; const BESIDE: u8 = 10; -/// Writes a row. Each large value is handed, already written, to `beside`, which keeps it and -/// answers with its place. +/// Writes a row. The record of each large value is handed to `beside`, which may put the +/// record's other form in its place, keeps it, and answers with its place. The row's cell +/// holds the length of what was kept. pub(super) fn encode_row( values: &[Value], - mut beside: impl FnMut(&[u8]) -> Result, + mut beside: impl FnMut(&mut Vec) -> Result, ) -> Result, E> { let mut row = Vec::new(); + // The record of one value that can be large, so that a large one never passes through + // the row. + let mut record = Vec::new(); for value in values { - let at = row.len(); - put_value(&mut row, value); - let grows = matches!(value, Value::Text(_) | Value::Json(_) | Value::Vector(_)); - if grows && row.len() - at > BESIDE_FROM { - let bytes = (row.len() - at) as u64; - let place = beside(&row[at..])?; - row.truncate(at); + if !matches!(value, Value::Text(_) | Value::Json(_) | Value::Vector(_)) { + put_value(&mut row, value); + continue; + } + record.clear(); + put_value(&mut record, value); + if record.len() > BESIDE_FROM { + let place = beside(&mut record)?; row.push(BESIDE); - row.extend(bytes.to_le_bytes()); + row.extend((record.len() as u64).to_le_bytes()); row.extend(place.to_le_bytes()); + } else { + row.extend(&record); } } Ok(row) } +/// Puts the compressed form of a large value's plain record in its place, when the value is +/// one that is compressed and that saves an eighth of its bytes. +/// +/// The compressor works in a buffer of 20 bytes more than 1.1 times the value. `reservation` +/// is grown by that first, and the bytes it was grown by are returned for the caller to give +/// back once the record is written. When the budget has no room for the buffer the record +/// stays plain: compression never refuses a commit. +pub(super) fn compress(record: &mut Vec, reservation: &mut Reservation) -> u64 { + let len = record.len() - FRONT; + let room = get_maximum_output_size(len); + if record[0] == VECTOR || len > COMPRESS_TO || reservation.grow(room as u64).is_err() { + return 0; + } + let mut block = vec![0; room]; + // A buffer of that size holds whatever the value compresses to, so nothing fails here. + if let Ok(size) = compress_into(&record[FRONT..], &mut block) + && size <= len - len / 8 + { + // The block moves up to let the tag and the length in before it. They fit: the block + // is no longer than the value, and the buffer is at least 20 bytes longer. + block.copy_within(..size, FRONT); + block[..FRONT].copy_from_slice(&record[..FRONT]); + block[0] |= COMPRESSED; + block.truncate(FRONT + size); + *record = block; + } + room as u64 +} + fn put_value(out: &mut Vec, value: &Value) { match value { Value::Null => out.push(NULL), @@ -96,7 +154,65 @@ pub(super) fn decode_row(mut row: &[u8]) -> Vec { cells } -/// Reads back one value that was kept beside its row. +/// Makes a large value from its stored record, in either form. `reservation` already holds +/// the record's bytes, and holds the bytes of the value's plain record when the value is made. +/// +/// A compressed record is decoded first. The reservation is grown by the bytes of the plain +/// record before one of them exists, so a value that does not fit under the memory limit is +/// refused and never produced; the stored bytes are given back once it is decoded. A text +/// keeps the plain record's buffer as its own. While a JSON or vector value is made, the +/// record and the value both exist, and the reservation holds the bytes twice for that moment. +/// +/// The inner error is the reason a record's own bytes do not support the length it gives: +/// the store is damaged. Nothing is allocated from such a length. +pub(super) fn load( + mut stored: Vec, + reservation: &mut Reservation, +) -> types::Result> { + let Some(front) = stored.first_chunk::() else { + return Ok(Err("a large value is shorter than its tag and length")); + }; + if front[0] & COMPRESSED != 0 { + let len = u64::from_le_bytes(take_array(&mut &front[1..])); + // No LZ4 block decodes to more than 255 times its own bytes. + if len > 255 * (stored.len() - FRONT) as u64 { + return Ok(Err( + "a compressed large value gives a length its bytes cannot decode to", + )); + } + reservation.grow(FRONT as u64 + len)?; + let mut plain = vec![0; FRONT + len as usize]; + plain[..FRONT].copy_from_slice(front); + plain[0] &= !COMPRESSED; + let decoded = decompress_into(&stored[FRONT..], &mut plain[FRONT..]); + if decoded.ok() != Some(len as usize) { + return Ok(Err("a compressed large value decodes to another length")); + } + reservation.shrink(stored.len() as u64); + stored = plain; + } + let len = u64::from_le_bytes(take_array(&mut &stored[1..])); + if len != (stored.len() - FRONT) as u64 { + return Ok(Err("a large value gives a length its bytes are not")); + } + if stored[0] == TEXT { + stored.drain(..FRONT); + let text = String::from_utf8(stored).map(Value::Text); + return Ok(text.map_err(|_| "a large text is not text")); + } + let twice = stored.len() as u64; + reservation.grow(twice)?; + let json = &stored[FRONT..]; + let value = match stored[0] { + JSON => serde_json::from_slice(json).ok().map(Value::Json), + VECTOR => Some(decode_value(&stored)), + _ => None, + }; + reservation.shrink(twice); + Ok(value.ok_or("a large value is not text, JSON or a vector")) +} + +/// Reads one value back from its plain record. pub(super) fn decode_value(bytes: &[u8]) -> Value { take_value(bytes[0], &mut &bytes[1..]) } diff --git a/v2/contextdb/src/store/tests.rs b/v2/contextdb/src/store/tests.rs index ef3fbf48..a76a9f84 100644 --- a/v2/contextdb/src/store/tests.rs +++ b/v2/contextdb/src/store/tests.rs @@ -1,10 +1,20 @@ //! Checks of the pure parts of the store: the order keys take as bytes, the span a pair of -//! bounds stands for, and a row written and read back. +//! bounds stands for, and a row written and read back. And the store's own check of a +//! compressed value's recorded length: the storage library checks every page it reads, so a +//! length changed on its way from the file never reaches that check, and the tests of damage +//! here write the wrong bytes into the data file as a commit of the storage library: a wrong +//! length, a large value with a slice missing, and a block that decodes to the recorded length +//! of bytes that are no value of its kind. use std::ops::Bound; +use std::path::{Path, PathBuf}; -use super::{Cell, Key, keys, record}; -use crate::types::{SortDirection, TxId, Uuid, Value}; +use lz4_flex::block::{compress_into, get_maximum_output_size}; +use redb::ReadableTable; + +use super::{Cell, ChangeSet, Key, Read, Store, TableId, keys, record}; +use crate::memory::Budget; +use crate::types::{Error, SortDirection, TxId, Uuid, Value}; fn key(values: &[Value]) -> Key { Key::from_values(values, &[]).expect("these values can be part of a key") @@ -161,3 +171,329 @@ fn a_row_reads_back_as_written_with_its_large_values_beside_it() { .collect(); assert_eq!(read, values); } + +/// The storage library's table of large values, as `disk` declares it: place to slice. +const BESIDE: redb::TableDefinition = redb::TableDefinition::new("beside"); + +#[test] +fn a_compressed_value_whose_recorded_length_is_wrong_is_damage_and_takes_no_memory() { + const MIB: u64 = 1 << 20; + // 16 KiB of one sentence, which the store keeps compressed in a block of about a hundredth + // of that. + const LEN: u64 = 16 << 10; + let folder = + std::env::temp_dir().join(format!("contextdb-wrong-length-{}", std::process::id())); + if folder.exists() { + std::fs::remove_dir_all(&folder).expect("clear the folder an earlier run left"); + } + let sentence = "A store keeps what it is given and hands it back as it was. "; + let body: String = sentence.chars().cycle().take(LEN as usize).collect(); + let (table, row_key) = (TableId(1), key(&[Value::Int64(1)])); + + // One row whose text goes beside it, written under no limit so that it is compressed. + let budget = Budget::new(None).expect("a budget"); + let store = Store::open_for_writing(&folder, &budget, None).expect("open a fresh store"); + let mut changes = ChangeSet::new(store.view().expect("a view"), &budget).expect("begin"); + let row = vec![Value::Int64(1), text(&body)]; + changes + .put_row(table, row_key.clone(), row, true) + .expect("stage"); + store.commit(changes).expect("commit"); + drop(store); + + // Opens the store under a limit that leaves `room` bytes once it is open, and loads the + // text. Returns what the load gave, where the value is kept, the bytes the statement + // holds after the load, and whether what the budget counts is back where it was once + // the statement has let go. + let load = |room: u64| { + let budget = Budget::new(Some(9 * MIB + room)).expect("a budget"); + let store = Store::open_for_writing(&folder, &budget, None).expect("open the store"); + let view = store.view().expect("a view"); + let row = view.row(table, &row_key).expect("read the row"); + let Some(Cell::Beside(beside)) = row.expect("the row is there").cells.pop() else { + panic!("a text of {LEN} bytes was kept in its row"); + }; + let before = budget.used(); + let mut reservation = budget.reserve(0, "query").expect("a reservation"); + let loaded = view.load(&beside, &mut reservation); + let held = reservation.bytes(); + drop(reservation); + (loaded, beside, held, budget.used() == before) + }; + let (loaded, beside, held, given_back) = load(MIB); + assert!(loaded.expect("load the value") == text(&body)); + assert_eq!(held, 9 + LEN, "held for the loaded value"); + assert!(given_back, "counted once the value is let go"); + // The compressed record: nine bytes, then one block. + let block = beside.bytes - 9; + assert!( + beside.bytes < LEN / 50, + "{} bytes are stored for a repeating text of {LEN}: it is not compressed", + beside.bytes + ); + // The room every load below is given: what reading the value takes at its most, the + // stored record and the value's own record together, one byte and a little more. No + // block decodes to more than 255 times its own bytes, and that much does not fit. + let room = beside.bytes + 9 + LEN + 1 + 64; + assert!(LEN <= 255 * block && 9 + 255 * block > room); + + // Writes this length where the record keeps its own, as a commit of the storage library, + // which then hands the slice over as a page that passes its check. + let record_length = |wrong: u64| { + let data = redb::Database::open(folder.join("data")).expect("open the data file"); + let write = data.begin_write().expect("begin"); + { + let mut slices = write.open_table(BESIDE).expect("the large values"); + let first = slices.get(beside.place).expect("read the first slice"); + let mut first = first.expect("the first slice is there").value().to_vec(); + assert_eq!(first[0], 0x84, "the first byte of a compressed text"); + first[1..9].copy_from_slice(&wrong.to_le_bytes()); + slices + .insert(beside.place, first.as_slice()) + .expect("write the first slice"); + } + write.commit().expect("commit"); + }; + // The outcome of a load that must make no value: the reason of the store's damage, which + // names the folder and the large value, or `None` for the memory limit's refusal. Either + // way the statement holds nothing afterwards and nothing stays counted. + let damage = |wrong: u64| -> Option { + record_length(wrong); + let (loaded, _, held, given_back) = load(room); + assert_eq!(held, 0, "held after a recorded length of {wrong}"); + assert!(given_back, "counted after a recorded length of {wrong}"); + match loaded { + Err(Error::StoreCorrupted { path, reason }) => { + assert_eq!(path, folder, "a recorded length of {wrong}"); + assert!(reason.contains("large value"), "{wrong}: {reason}"); + Some(reason) + } + Err(Error::MemoryBudgetExceeded { operation, .. }) => { + assert_eq!(operation, "query", "a recorded length of {wrong}"); + None + } + other => panic!("a recorded length of {wrong} gave {other:?}"), + } + }; + + // The value reads under that room as it is. + let (loaded, _, held, _) = load(room); + assert!(loaded.expect("load the value under the room") == text(&body)); + assert_eq!(held, 9 + LEN); + + // A length the block cannot decode to, from one more than 255 times its bytes to the + // largest number there is, is damage before any memory is asked for. The limit has room + // for none of them: a length that was asked for would be the limit's refusal, and one + // that was taken without asking would end the test. + for wrong in [255 * block + 1, 1 << 40, 1 << 62, u64::MAX - 8, u64::MAX] { + assert!( + damage(wrong).is_some(), + "memory was asked for a recorded length of {wrong}" + ); + } + // The most a wrong length can ask for is 255 times the block, and it is asked for before + // it is taken: with no room for it, it is the limit's refusal, as for any value that + // does not fit. + assert!( + damage(255 * block).is_none(), + "no memory was asked for a recorded length of 255 times the block" + ); + // A length there is room for, and that the block does not decode to, is damage once the + // block is decoded: one byte more than the value has, one byte less, half of it, none. + for wrong in [LEN + 1, LEN - 1, LEN / 2, 0] { + assert!( + damage(wrong).is_some(), + "a recorded length of {wrong} was refused for want of room" + ); + } + + // With its own length back, the value is as it was written. + record_length(LEN); + let (loaded, _, held, _) = load(room); + assert!(loaded.expect("load the value") == text(&body)); + assert_eq!(held, 9 + LEN); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +/// The bytes of one slice of a large value, as `disk` cuts them. +const SLICE: u64 = 243; + +/// The one row of the stores the tests below write: table 1, key 1. +fn the_row() -> (TableId, Key) { + (TableId(1), key(&[Value::Int64(1)])) +} + +/// A folder of this name that holds nothing, for one test. +fn fresh_folder(name: &str) -> PathBuf { + let folder = std::env::temp_dir().join(format!("contextdb-{name}-{}", std::process::id())); + if folder.exists() { + std::fs::remove_dir_all(&folder).expect("clear the folder an earlier run left"); + } + folder +} + +/// Writes a store of one row, the number 1 and then `value`, under no memory limit, so that a +/// value the store compresses is compressed, and closes it. +fn store_one_row(folder: &Path, value: Value) { + let (table, row_key) = the_row(); + let budget = Budget::new(None).expect("a budget"); + let store = Store::open_for_writing(folder, &budget, None).expect("open a fresh store"); + let mut changes = ChangeSet::new(store.view().expect("a view"), &budget).expect("begin"); + let row = vec![Value::Int64(1), value]; + changes.put_row(table, row_key, row, true).expect("stage"); + store.commit(changes).expect("commit"); +} + +/// Opens the store and loads the large value of its one row. Returns what the load gave, where +/// the value is kept, the bytes the statement holds after the load, and whether what the budget +/// counts is back where it was once the statement has let go. +fn load_the_value(folder: &Path) -> (crate::types::Result, super::Beside, u64, bool) { + let (table, row_key) = the_row(); + let budget = Budget::new(None).expect("a budget"); + let store = Store::open_for_writing(folder, &budget, None).expect("open the store"); + let view = store.view().expect("a view"); + let row = view.row(table, &row_key).expect("read the row"); + let Some(Cell::Beside(beside)) = row.expect("the row is there").cells.pop() else { + panic!("the value was kept in its row"); + }; + let before = budget.used(); + let mut reservation = budget.reserve(0, "query").expect("a reservation"); + let loaded = view.load(&beside, &mut reservation); + let held = reservation.bytes(); + drop(reservation); + (loaded, beside, held, budget.used() == before) +} + +/// Asserts that loading the large value of the store's one row is the store's damage, naming +/// the folder, and that the statement then holds nothing and nothing stays counted. `what` +/// says what was done to the value. +fn loading_is_damage(folder: &Path, what: &str) { + let (loaded, _, held, given_back) = load_the_value(folder); + match loaded { + Err(Error::StoreCorrupted { path, .. }) => assert_eq!(path, folder, "{what}"), + Ok(Value::Text(text)) => panic!("{what} was read as a text of {} bytes", text.len()), + Ok(_) => panic!("{what} was read as a value"), + Err(other) => panic!("{what} gave {other:?}"), + } + assert_eq!(held, 0, "held after the refusal: {what}"); + assert!(given_back, "counted after the refusal: {what}"); +} + +/// Letters in a fixed scatter, which no compressor makes shorter, so that a text or JSON value +/// of them is stored plain. +fn scattered(count: usize) -> String { + const LETTERS: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"; + let mut state = 0x9E37_79B9_7F4A_7C15_u64; + let letters = (0..count).map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + char::from(LETTERS[(state >> 58) as usize]) + }); + letters.collect() +} + +/// Stores `value`, whose plain record is nine bytes and then `plain` more and which is stored +/// as that, removes the last of its slices as a commit of the storage library, and asserts that +/// reading the value is then damage: the row still says how many bytes are stored for the +/// value, and the slices that are left hold fewer. +fn a_value_without_its_last_slice_is_damage(name: &str, value: &Value, plain: usize) { + let folder = fresh_folder(name); + store_one_row(&folder, value.clone()); + let (loaded, beside, _, _) = load_the_value(&folder); + assert!(loaded.expect("load the value as it was written") == *value); + assert_eq!(beside.bytes, 9 + plain as u64, "the value is stored plain"); + let slices = beside.bytes.div_ceil(SLICE); + assert!(slices > 1 && beside.bytes % SLICE != 0, "{slices} slices"); + + let data = redb::Database::open(folder.join("data")).expect("open the data file"); + let write = data.begin_write().expect("begin"); + { + let mut kept = write.open_table(BESIDE).expect("the large values"); + let last = kept.remove(beside.place + slices - 1); + assert!(last.expect("remove the last slice").is_some()); + } + write.commit().expect("commit"); + drop(data); + + loading_is_damage(&folder, "a large value whose last slice is gone"); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +#[test] +fn a_plain_text_whose_slices_hold_fewer_bytes_than_its_row_records_is_damage() { + let letters = scattered(2000); + a_value_without_its_last_slice_is_damage("short-text", &text(&letters), 2000); +} + +#[test] +fn a_plain_json_whose_slices_hold_fewer_bytes_than_its_row_records_is_damage() { + // A JSON text is its letters between two quotation marks. + let json = Value::Json(serde_json::Value::String(scattered(2000))); + a_value_without_its_last_slice_is_damage("short-json", &json, 2002); +} + +#[test] +fn a_vector_whose_slices_hold_fewer_bytes_than_its_row_records_is_damage() { + let numbers: Vec = (0..384_u16).map(|n| f32::from(n) / 8.0).collect(); + a_value_without_its_last_slice_is_damage("short-vector", &Value::Vector(numbers), 4 * 384); +} + +/// Stores `value`, which the store compresses, and puts in the place of its block, as a commit +/// of the storage library, the block of `wrong`: as many bytes as the value has, which are no +/// value of its kind. The stored record keeps its tag, its recorded length and its own number +/// of bytes, so the block decodes to exactly the length the record gives. Asserts that reading +/// the value is then damage. +fn a_block_of_other_bytes_is_damage(name: &str, value: &Value, wrong: &[u8]) { + let folder = fresh_folder(name); + store_one_row(&folder, value.clone()); + let (loaded, beside, _, _) = load_the_value(&folder); + assert!(loaded.expect("load the value as it was written") == *value); + + let data = redb::Database::open(folder.join("data")).expect("open the data file"); + let write = data.begin_write().expect("begin"); + { + let mut kept = write.open_table(BESIDE).expect("the large values"); + let first = kept.get(beside.place).expect("read the first slice"); + let first = first.expect("the first slice is there").value().to_vec(); + assert!(first[0] & 0x80 != 0, "the value is stored compressed"); + let recorded = u64::from_le_bytes(std::array::from_fn(|at| first[1 + at])); + assert_eq!(recorded, wrong.len() as u64, "the length the record gives"); + + // The same tag and length, and then the block of the wrong bytes. + let mut record = vec![0; 9 + get_maximum_output_size(wrong.len())]; + record[..9].copy_from_slice(&first[..9]); + let block = compress_into(wrong, &mut record[9..]).expect("the buffer holds any block"); + record.truncate(9 + block); + assert_eq!( + record.len() as u64, + beside.bytes, + "the wrong bytes make a record as long as the value's" + ); + for (place, slice) in (beside.place..).zip(record.chunks(SLICE as usize)) { + kept.insert(place, slice).expect("write a slice"); + } + } + write.commit().expect("commit"); + drop(data); + + loading_is_damage(&folder, "a block that decodes to bytes of another kind"); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +#[test] +fn a_compressed_text_that_decodes_to_bytes_that_are_not_text_is_damage() { + const LEN: usize = 16 << 10; + // No text holds the byte 0xFF. + a_block_of_other_bytes_is_damage("not-text", &text(&"a".repeat(LEN)), &[0xFF; LEN]); +} + +#[test] +fn a_compressed_json_that_decodes_to_bytes_that_are_not_json_is_damage() { + const LEN: usize = 16 << 10; + // The value is one JSON text; the wrong bytes have a brace where each of its quotation + // marks is, which is no JSON. + let json = Value::Json(serde_json::Value::String("a".repeat(LEN - 2))); + let wrong = format!("{{{}}}", "a".repeat(LEN - 2)); + a_block_of_other_bytes_is_damage("not-json", &json, wrong.as_bytes()); +} diff --git a/v2/contextdb/src/types/mod.rs b/v2/contextdb/src/types/mod.rs index 7ca56c57..2be9229e 100644 --- a/v2/contextdb/src/types/mod.rs +++ b/v2/contextdb/src/types/mod.rs @@ -214,6 +214,49 @@ pub enum Error { budget_limit_bytes: u64, hint: &'static str, }, + /// A vector written into a VECTOR column, or the vector of a search, holds a number that + /// is not finite. `position` counts its numbers from zero. Returned by the INSERT or UPDATE + /// that wrote it, or by the SELECT that searched with it. + InvalidVectorValue { + index: VectorIndexRef, + position: usize, + }, + /// The vector of a search has a different length from the dimension its VECTOR column + /// declares. `expected` is the declared dimension, `actual` the length given. + VectorIndexDimensionMismatch { + index: VectorIndexRef, + expected: usize, + actual: usize, + }, + /// The column on the left of `<=>` is not a VECTOR column. + UnknownVectorIndex { index: VectorIndexRef }, + /// A `PARTITION_KEY` or `MAX_PARTITIONS` clause that cannot stand; `issue` says why. + /// Returned by `CREATE TABLE`, which then declares nothing. + InvalidVectorPartitionDeclaration { + index: VectorIndexRef, + issue: VectorPartitionDeclarationIssue, + }, + /// A VECTOR column whose declaration cannot stand; `issue` says why. Returned by + /// `CREATE TABLE`, which then declares nothing. + InvalidVectorPolicyDeclaration { + index: VectorIndexRef, + issue: VectorPolicyDeclarationIssue, + }, + /// A write would have given the column more partitions than its `MAX_PARTITIONS` allows: + /// a row with a vector arrived under partition key values no row with a vector has yet. + /// Returned by the INSERT or UPDATE that brought the row, which then changes nothing and + /// leaves its transaction usable. Returned by a commit only when another transaction made + /// a partition while this one was open; nothing of that commit is written. + VectorPartitionLimitExceeded { + index: VectorIndexRef, + max_partitions: u32, + }, + /// A search asked for `INDEXED`, with `USE VECTOR INDEXED` or through the column's + /// `SEARCH_MODE INDEXED`, and this version keeps no index. It never answers such a search + /// with an exact scan under the indexed name; `EXACT` and `AUTO` answer. + VectorIndexedRouteUnavailable { index: VectorIndexRef }, + /// `USE VECTOR` was written in a SELECT that does not order by `<=>`. + UseVectorRequiresVectorOrder, } impl fmt::Display for Error { @@ -313,6 +356,44 @@ impl fmt::Display for Error { "memory limit reached during {operation}: {requested_bytes} bytes were needed and \ {available_bytes} bytes of the {budget_limit_bytes}-byte limit were free. {hint}" ), + Error::InvalidVectorValue { index, position } => write!( + f, + "vector {index} contains a non-finite component at position {position}" + ), + Error::VectorIndexDimensionMismatch { + index, + expected, + actual, + } => write!( + f, + "vector index dimension mismatch on {index}: expected {expected}, got {actual}" + ), + Error::UnknownVectorIndex { index } => write!(f, "unknown vector index {index}"), + Error::InvalidVectorPartitionDeclaration { index, issue } => { + write!( + f, + "invalid vector partition declaration on {index}: {issue}" + ) + } + Error::InvalidVectorPolicyDeclaration { index, issue } => { + write!(f, "invalid vector policy declaration on {index}: {issue}") + } + Error::VectorPartitionLimitExceeded { + index, + max_partitions, + } => write!( + f, + "vector partition limit reached on {index}: MAX_PARTITIONS is {max_partitions}; \ + use fewer partition-key values, or declare the table with a larger MAX_PARTITIONS" + ), + Error::VectorIndexedRouteUnavailable { index } => write!( + f, + "no vector index ships in this version, so {index} cannot be searched INDEXED; \ + use EXACT or AUTO" + ), + Error::UseVectorRequiresVectorOrder => { + f.write_str("USE VECTOR requires ORDER BY embedding <=> $param in the same query") + } } } } @@ -343,6 +424,94 @@ pub enum ColumnType { Timestamp, } +/// A VECTOR column as a refusal names it: its table, then its own name. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct VectorIndexRef { + pub table: String, + pub column: String, +} + +impl fmt::Display for VectorIndexRef { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}.{}", self.table, self.column) + } +} + +/// Why a `PARTITION_KEY` or `MAX_PARTITIONS` clause was refused. None of these carries a value +/// a statement bound, so printing one never echoes data. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum VectorPartitionDeclarationIssue { + /// A column of the key is not a column of the table. + UnknownPartitionKeyColumn, + /// A column of the key may hold NULL. + NullablePartitionKeyColumn, + /// A column of the key is REAL, JSON or VECTOR. + UnsupportedPartitionKeyColumnType, + /// The key names a column twice. + DuplicatePartitionKeyColumn, + /// The key names the VECTOR column it belongs to. + VectorColumnInPartitionKey, + /// `MAX_PARTITIONS` was written on a column with no `PARTITION_KEY`. + MaxPartitionsWithoutPartitionKey, + /// `MAX_PARTITIONS` is zero or negative. + MaxPartitionsNotPositive, + /// `MAX_PARTITIONS` is larger than the engine keeps. + MaxPartitionsOutOfRange, +} + +impl fmt::Display for VectorPartitionDeclarationIssue { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::UnknownPartitionKeyColumn => { + "PARTITION_KEY references a column that does not exist in the same table" + } + Self::NullablePartitionKeyColumn => "every PARTITION_KEY column must be NOT NULL", + Self::UnsupportedPartitionKeyColumnType => { + "PARTITION_KEY columns must use UUID, TEXT, INTEGER, BOOLEAN or TIMESTAMP" + } + Self::DuplicatePartitionKeyColumn => "a PARTITION_KEY column may appear only once", + Self::VectorColumnInPartitionKey => { + "PARTITION_KEY may not contain the vector column itself" + } + Self::MaxPartitionsWithoutPartitionKey => "MAX_PARTITIONS requires PARTITION_KEY", + Self::MaxPartitionsNotPositive => "MAX_PARTITIONS must be a positive integer", + Self::MaxPartitionsOutOfRange => { + "MAX_PARTITIONS is outside the persisted numeric range" + } + }) + } +} + +/// Why a VECTOR column's declaration was refused. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum VectorPolicyDeclarationIssue { + /// `VECTOR(0)`. + DimensionsNotPositive, +} + +impl fmt::Display for VectorPolicyDeclarationIssue { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::DimensionsNotPositive => "vector dimensions must be at least 1", + }) + } +} + +/// How a similarity search is answered: what `SEARCH_MODE` declares for a column and what +/// `USE VECTOR` asks for in one statement. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum VectorSearchMode { + /// The engine chooses. This version has no index, so it answers exactly. + #[default] + Auto, + /// Every allowed vector is compared: the true nearest rows. + Exact, + /// Only through an index; refused when there is none. + Indexed, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum SortDirection { #[default] diff --git a/v2/contextdb/src/upkeep/mod.rs b/v2/contextdb/src/upkeep/mod.rs new file mode 100644 index 00000000..a663cf59 --- /dev/null +++ b/v2/contextdb/src/upkeep/mod.rs @@ -0,0 +1,50 @@ +//! upkeep: background work in small steps. This version has one kind of work, sealing: moving +//! the vectors that wait on a column's pending lists into blocks. +//! +//! It never blocks a write or a query: a cycle does a bounded amount of work in one commit, and +//! what it has not done waits for the next cycle. Nothing here starts a thread or reads a +//! clock; a cycle runs when it is called. + +use crate::memory::Budget; +use crate::schema::Catalog; +use crate::store::ChangeSet; +use crate::types::Result; +use crate::vector::{self, Sealed}; + +/// The most vectors one cycle seals. +pub(crate) const SEAL_ROWS: usize = 4096; + +/// What one cycle of upkeep did, as [`crate::Database::run_maintenance_cycle`] returns it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +#[non_exhaustive] +pub struct MaintenanceReport { + /// What the cycle did for vector search. + pub vector: VectorMaintenanceReport, +} + +/// The part of a [`MaintenanceReport`] about vector search. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +#[non_exhaustive] +pub struct VectorMaintenanceReport { + /// How many vectors this cycle moved from pending lists into blocks. + pub maintenance_vectors_done: u64, + /// How many vectors are still on a pending list after this cycle, over every column of the + /// store. Zero means a further cycle has nothing to seal. + pub maintenance_vectors_remaining: u64, +} + +/// Stages one cycle of upkeep into `changes`, for the caller to commit, and says what it +/// staged. The caller holds the commit lock from before `changes` was begun. +pub(crate) fn cycle( + changes: &mut ChangeSet, + catalog: &Catalog, + budget: &Budget, +) -> Result { + let Sealed { done, remaining } = vector::seal(changes, catalog, budget, SEAL_ROWS)?; + Ok(MaintenanceReport { + vector: VectorMaintenanceReport { + maintenance_vectors_done: done, + maintenance_vectors_remaining: remaining, + }, + }) +} diff --git a/v2/contextdb/src/vector/block.rs b/v2/contextdb/src/vector/block.rs new file mode 100644 index 00000000..584d6e31 --- /dev/null +++ b/v2/contextdb/src/vector/block.rs @@ -0,0 +1,166 @@ +//! A block: the slots of as many vectors of one partition as fit in the block's bytes +//! (328 at 768 numbers a vector), each with the key of its row. +//! +//! A block is `count` as four bytes, then `count` slots of one length, then `count` numbers of +//! four bytes, each the end of a key within the bytes that follow, then the keys one after +//! another. Every number is low byte first. The slots come first and are of one length, so a +//! search strides over them without reading a key. + +use super::slot::{slot_len, write_slot}; +use crate::store::Key; + +/// The most bytes a block is written with: 64 less than 256 KiB. The storage library keeps a +/// record in a run of pages whose length is a power of two, so a block of this size fills its +/// run and one byte more would take a run twice as long. +pub(super) const BLOCK_BYTES: usize = (256 << 10) - 64; + +/// Bytes of a block's count, and of each key end. +const NUMBER: usize = 4; + +/// What a store is told when a stored block is not [`BlockReader::is_whole`]. +pub(super) const BROKEN: &str = "a block of the vector part does not hold what it counts"; + +/// How many of `keys`, in their order, a block of `bytes` bytes takes within [`BLOCK_BYTES`], +/// and the bytes of those keys together. A block with no slot takes one key of any length. +pub(super) fn takes(dimension: usize, bytes: usize, keys: &[Key]) -> (usize, usize) { + let (mut count, mut key_bytes) = (0, 0); + for key in keys { + let more = key_bytes + key.as_bytes().len(); + let with = bytes + (count + 1) * (slot_len(dimension) + NUMBER) + more; + if with > BLOCK_BYTES && (bytes > NUMBER || count > 0) { + break; + } + (count, key_bytes) = (count + 1, more); + } + (count, key_bytes) +} + +/// A block being put together. +pub(super) struct BlockWriter { + dimension: usize, + slots: Vec, + ends: Vec, + keys: Vec, +} + +impl BlockWriter { + /// A writer that starts with the slots and keys of the block `from`, or with none for + /// `None`, and how many of `keys` it [`takes`]. Its lists are made here with room for those + /// and never grow. A block that takes none is not copied: its writer has no slot. + pub(super) fn new(dimension: usize, from: Option<&[u8]>, keys: &[Key]) -> (BlockWriter, usize) { + let (count, key_bytes) = takes(dimension, from.map_or(NUMBER, <[u8]>::len), keys); + let block = from.filter(|_| count > 0).unwrap_or(&[0; NUMBER]); + let reader = BlockReader::new(block, dimension); + let mut slots = Vec::with_capacity(reader.ends_at() - NUMBER + count * slot_len(dimension)); + let mut ends = Vec::with_capacity(reader.len() + count); + let mut keys = Vec::with_capacity(block.len() - reader.keys_at() + key_bytes); + slots.extend_from_slice(&block[NUMBER..reader.ends_at()]); + ends.extend((0..reader.len()).map(|at| reader.end(at) as u32)); + keys.extend_from_slice(&block[reader.keys_at()..]); + let writer = BlockWriter { + dimension, + slots, + ends, + keys, + }; + (writer, count) + } + + /// Bytes the block has now. + pub(super) fn bytes(&self) -> usize { + NUMBER + self.slots.len() + NUMBER * self.ends.len() + self.keys.len() + } + + /// Adds the slot of `vector` and the key of its row. + pub(super) fn push(&mut self, key: &Key, vector: &[f32]) { + let at = self.slots.len(); + self.slots.resize(at + slot_len(self.dimension), 0); + write_slot(vector, &mut self.slots[at..]); + self.keys.extend_from_slice(key.as_bytes()); + self.ends.push(self.keys.len() as u32); + } + + /// True when the writer has no slot. + pub(super) fn is_empty(&self) -> bool { + self.ends.is_empty() + } + + /// The block's bytes. + pub(super) fn finish(self) -> Vec { + let mut block = Vec::with_capacity(self.bytes()); + block.extend_from_slice(&(self.ends.len() as u32).to_le_bytes()); + block.extend_from_slice(&self.slots); + block.extend(self.ends.iter().flat_map(|end| end.to_le_bytes())); + block.extend_from_slice(&self.keys); + block + } +} + +/// A block as it was written, read in place. +pub(super) struct BlockReader<'a> { + bytes: &'a [u8], + dimension: usize, +} + +impl<'a> BlockReader<'a> { + /// Reads a block [`BlockWriter::finish`] gave for vectors of this dimension. + pub(super) fn new(bytes: &'a [u8], dimension: usize) -> BlockReader<'a> { + BlockReader { bytes, dimension } + } + + /// True when the bytes are a block's: its count, that many slots and key ends, no end + /// before the one ahead of it, and the last where the block ends. Anything else is damage, + /// and is asked before a slot or a key is read, which would read past such bytes. + pub(super) fn is_whole(&self) -> bool { + let Some(count) = self.bytes.first_chunk::() else { + return false; + }; + let each = slot_len(self.dimension) + NUMBER; + let keys_at = (u32::from_le_bytes(*count) as usize).checked_mul(each); + let keys_at = keys_at.and_then(|slots| slots.checked_add(NUMBER)); + let Some(keys_at) = keys_at.filter(|at| *at <= self.bytes.len()) else { + return false; + }; + let (ends, _) = self.bytes[self.ends_at()..keys_at].as_chunks::(); + let mut ends = ends.iter().map(|end| u32::from_le_bytes(*end) as usize); + ends.clone().is_sorted() && ends.next_back().unwrap_or(0) == self.bytes.len() - keys_at + } + + /// The number of four bytes that begins at this byte of the block. + fn number(&self, at: usize) -> usize { + let bytes = self.bytes; + u32::from_le_bytes([bytes[at], bytes[at + 1], bytes[at + 2], bytes[at + 3]]) as usize + } + + /// Where the key ends begin: after the count and the slots. + fn ends_at(&self) -> usize { + NUMBER + self.len() * slot_len(self.dimension) + } + + /// Where the keys begin: after the key ends. + fn keys_at(&self) -> usize { + self.ends_at() + NUMBER * self.len() + } + + /// Where the key at this place ends, counted from where the keys begin. + fn end(&self, at: usize) -> usize { + self.number(self.ends_at() + NUMBER * at) + } + + /// How many slots the block has. + pub(super) fn len(&self) -> usize { + self.number(0) + } + + /// The slot at this place, counted from zero. + pub(super) fn slot(&self, at: usize) -> &'a [u8] { + let len = slot_len(self.dimension); + &self.bytes[NUMBER + at * len..][..len] + } + + /// The key of the row the slot at this place was written for. + pub(super) fn key(&self, at: usize) -> Key { + let start = if at == 0 { 0 } else { self.end(at - 1) }; + Key::from_bytes(&self.bytes[self.keys_at()..][start..self.end(at)]) + } +} diff --git a/v2/contextdb/src/vector/directory.rs b/v2/contextdb/src/vector/directory.rs new file mode 100644 index 00000000..af2c061b --- /dev/null +++ b/v2/contextdb/src/vector/directory.rs @@ -0,0 +1,102 @@ +//! The partitions of a VECTOR column: one record for each key tuple under which a row has a +//! vector. + +use std::collections::BTreeMap; + +use crate::store::{Key, Read, VectorId, VectorRecord}; +use crate::types::Result; + +/// What a column keeps about one partition. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub(super) struct Partition { + /// The partition's number within its column. Pending rows and blocks are kept under it. + pub number: u32, + /// How many rows have a vector in the partition, pending or in a block. + pub live: u64, + /// How many blocks the partition has. They are numbered from zero, and every one but the + /// last is full. + pub blocks: u32, +} + +impl Partition { + /// The record of the partition: its three numbers, low byte first. + pub(super) fn encode(&self) -> Vec { + let mut record = Vec::with_capacity(16); + record.extend(self.number.to_le_bytes()); + record.extend(self.live.to_le_bytes()); + record.extend(self.blocks.to_le_bytes()); + record + } + + /// The partition a record was written from. + pub(super) fn decode(record: &[u8]) -> Partition { + Partition { + number: u32::from_le_bytes(std::array::from_fn(|at| record[at])), + live: u64::from_le_bytes(std::array::from_fn(|at| record[4 + at])), + blocks: u32::from_le_bytes(std::array::from_fn(|at| record[12 + at])), + } + } +} + +/// What a store is told when a count the vector part keeps is below what it counts. +pub(super) const SHORT: &str = "a count of the vector part is below what it counts"; + +/// The partition a stored record was written from. A record of another length is damage. +fn checked(read: &dyn Read, record: &[u8]) -> Result { + if record.len() != 16 { + return Err(read.damaged("a partition record of the vector part is not 16 bytes")); + } + Ok(Partition::decode(record)) +} + +/// The partition a column has under a key tuple at the committed moment, `None` when it has +/// none. +pub(super) fn partition( + read: &dyn Read, + vector: VectorId, + tuple: &Key, +) -> Result> { + let record = VectorRecord::Partition(vector, tuple); + kept(read, record, |entry| checked(read, entry))?.transpose() +} + +/// What `read` gives for the value of one record of the vector part at the committed moment, +/// `None` when the record is not there. +pub(super) fn kept( + from: &dyn Read, + record: VectorRecord<'_>, + read: impl Fn(&[u8]) -> T, +) -> Result> { + let mut found = None; + from.vectors(record, &mut |first, value| { + found = (first == record).then(|| read(value)); + Ok(false) + })?; + Ok(found) +} + +/// The partitions a search covers: those whose key tuple begins with one of `scope`, or every +/// partition of the column when `scope` is empty. Each comes once, in the order of the +/// numbers, whatever the order of `scope`. A tuple no row has a vector under selects nothing. +pub(super) fn selected(read: &dyn Read, vector: VectorId, scope: &[Key]) -> Result> { + // Every key tuple begins with the key of no values. + let every = [Key::from_values(&[], &[])?]; + let scope = if scope.is_empty() { &every } else { scope }; + let mut found = BTreeMap::new(); + for prefix in scope { + read.vectors( + VectorRecord::Partition(vector, prefix), + &mut |record, entry| match record { + VectorRecord::Partition(column, tuple) + if column == vector && tuple.starts_with(prefix) => + { + let partition = checked(read, entry)?; + found.insert(partition.number, partition); + Ok(true) + } + _ => Ok(false), + }, + )?; + } + Ok(found.into_values().collect()) +} diff --git a/v2/contextdb/src/vector/mod.rs b/v2/contextdb/src/vector/mod.rs new file mode 100644 index 00000000..691fc8b5 --- /dev/null +++ b/v2/contextdb/src/vector/mod.rs @@ -0,0 +1,31 @@ +//! vector: the blocks of a VECTOR column, exact search over them, and its partitions. +//! +//! It never decides which rows a reader may see: `query` tells it, row by row, and a row is in +//! an answer only with the distance `query` worked out from the row's stored vector. +//! +//! Every vector is stored once, at full precision, with its row. For search this part keeps a +//! second, smaller copy: per column and partition, one signed byte for each number of a vector, +//! packed into blocks. A vector that is in no block yet is on its partition's pending list. +//! Blocks and lists are derived from rows: a row that is gone, or whose vector has changed, +//! leaves a slot behind that no answer can come from. +//! +//! A search screens every block of the partitions it was given with the one-byte copies. That +//! gives each vector a range its true similarity is proven to lie in, and the search asks +//! `query` for the true distance of only the rows whose range can still reach the answer. The +//! rows it returns are the ones a comparison of every stored vector returns. + +mod block; +mod directory; +mod seal; +mod search; +mod slot; +mod stage; +mod tally; +#[cfg(test)] +mod tests; + +pub(crate) use seal::{Sealed, seal}; +pub(crate) use search::{Nearest, search}; +pub(crate) use slot::distance; +pub(crate) use stage::stage; +pub(crate) use tally::Tally; diff --git a/v2/contextdb/src/vector/seal.rs b/v2/contextdb/src/vector/seal.rs new file mode 100644 index 00000000..6ad7e03c --- /dev/null +++ b/v2/contextdb/src/vector/seal.rs @@ -0,0 +1,210 @@ +//! Sealing: moving the vectors on pending lists into blocks. + +use super::block::{BLOCK_BYTES, BROKEN, BlockReader, BlockWriter}; +use super::directory; +use super::tally::{position, tuples}; +use crate::memory::{Budget, Reservation}; +use crate::schema::Catalog; +use crate::store::{Cell, ChangeSet, Key, Read, Row, VectorId, VectorNumber, VectorRecord}; +use crate::types::{ColumnType, Error, Result, Value}; + +/// What one call of [`seal`] did. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct Sealed { + /// How many vectors it moved into blocks. + pub done: u64, + /// How many are still pending, over every column. + pub remaining: u64, +} + +/// What a store is told when a pending record is not of a row that has a vector in a +/// partition of a VECTOR column, which [`super::stage`] keeps true of every one. +const WAITING: &str = "a pending record of the vector part is of no row with a vector"; + +/// Stages the sealing of the first `most` pending vectors, in the order of column number, +/// partition number and row key. +/// +/// The vectors of a partition are added to its last block while that block has room and to new +/// blocks after it, so every block of a partition but the last is full. Each is read from its +/// row at full precision, one at a time, and its pending record is removed. Nothing is staged +/// when nothing is pending. +/// +/// Every byte is reserved before it is taken. The blocks, the removals and the partition +/// records it stages are charged to the change set's reservation as they are staged. What it +/// works with is held by one reservation it takes from `budget`, named `maintenance`: the keys +/// of the pending records it collected, one block being put together (`BLOCK_BYTES`, taken once +/// and kept for the call), the partition record it is adding to, and one row with its vector, +/// given back before the next is read. +/// +/// A pending record, a partition record, a block or a count that does not agree with the rows +/// is `StoreCorrupted`, and nothing is staged. +/// +/// It works a partition at a time. When the memory limit has no room for a further partition, +/// that partition is left as it was and the call ends with those before it staged. It is +/// refused only when there is no room for the first partition, and then nothing is staged. +pub(crate) fn seal( + changes: &mut ChangeSet, + catalog: &Catalog, + budget: &Budget, + most: usize, +) -> Result { + let mut sealed = Sealed { + done: 0, + remaining: changes.vector_number(VectorNumber::Pending)?, + }; + if sealed.remaining == 0 { + return Ok(sealed); + } + // One block is put together at a time, so its bytes are taken once. + let mut held = budget.reserve(BLOCK_BYTES as u64, "maintenance")?; + let mut next = Some((VectorId(0), 0)); + while let Some(from) = next + && (sealed.done as usize) < most + { + let mark = changes.mark(); + let room = most - sealed.done as usize; + match seal_partition(changes, catalog, &mut held, from, room, sealed.remaining) { + Ok((done, following)) => { + sealed.done += done; + sealed.remaining -= done; + next = following; + } + Err(refused) => { + changes.rewind(mark); + let no_room = matches!(refused, Error::MemoryBudgetExceeded { .. }); + if sealed.done == 0 || !no_room { + return Err(refused); + } + break; + } + } + held.shrink(held.bytes() - BLOCK_BYTES as u64); + } + Ok(sealed) +} + +/// Stages the sealing of one partition: the first that has a pending record at `from`, a +/// column and a partition number, or after it. Seals at most `most` of its vectors, of the +/// `remaining` that are pending over every column. Answers how many it sealed and the next +/// column and partition number with a pending record, `None` when it met none. +fn seal_partition( + changes: &mut ChangeSet, + catalog: &Catalog, + held: &mut Reservation, + from: (VectorId, u32), + most: usize, + remaining: u64, +) -> Result<(u64, Option<(VectorId, u32)>)> { + let first = Key::from_values(&[], &[])?; + let (mut found, mut following) = (None, None); + let mut keys: Vec = Vec::new(); + changes.vectors( + VectorRecord::Pending(from.0, from.1, &first), + &mut |record, _| { + let VectorRecord::Pending(vector, number, key) = record else { + return Ok(false); + }; + if *found.get_or_insert((vector, number)) != (vector, number) { + following = Some((vector, number)); + return Ok(false); + } + // No place left: the list is made again with as many more, the old held until then. + if keys.len() == keys.capacity() { + let more = keys.capacity().max(4).min(most - keys.len()); + held.grow(((keys.len() + more) * size_of::()) as u64)?; + keys.reserve_exact(more); + held.shrink((keys.len() * size_of::()) as u64); + } + held.grow(key.as_bytes().len() as u64)?; + keys.push(key.clone()); + Ok(keys.len() < most) + }, + )?; + let Some((vector, number)) = found else { + return Ok((0, None)); + }; + let waiting = catalog.vector(vector); + let (table, column) = waiting.ok_or_else(|| changes.damaged(WAITING))?; + let at = position(table, &column.column); + let ColumnType::Vector(dimension) = table.meta.columns[at].column_type else { + unreachable!("a VECTOR column is of the type VECTOR") + }; + + // The partition is found through the key tuple of its first pending row. + let row = changes.row(table.id, &keys[0])?; + let (tuple, _) = tuples(table, column, row.as_ref(), None, &*changes, held)?; + let tuple = tuple.ok_or_else(|| changes.damaged(WAITING))?; + held.grow(tuple.as_bytes().len() as u64 + 64)?; + let partition = directory::partition(&*changes, vector, &tuple)?; + let mut partition = partition.ok_or_else(|| changes.damaged(WAITING))?; + + // Its last block is added to while it has room. A writer is made, after the block before + // it is staged, with room for exactly the slots still to come, which `left` counts. + let mut block = partition.blocks.saturating_sub(1); + let last = VectorRecord::Block(vector, number, block); + let begun = |last: &[u8]| { + let whole = BlockReader::new(last, dimension).is_whole(); + whole.then(|| BlockWriter::new(dimension, Some(last), &keys)) + }; + let (mut writer, mut left) = match directory::kept(&*changes, last, begun)? { + Some(begun) => begun.ok_or_else(|| changes.damaged(BROKEN))?, + None => BlockWriter::new(dimension, None, &keys), + }; + // A row is its cells and, when the row itself holds it, its vector; a vector kept beside + // its row is counted as it is loaded. + let width = table.meta.columns.len(); + let row_bytes = size_of::() + width * size_of::() + 4 * dimension; + for (done, key) in keys.iter().enumerate() { + let taken = held.bytes(); + held.grow(row_bytes as u64)?; + let row = changes.row(table.id, key)?; + let mut row = row.ok_or_else(|| changes.damaged(WAITING))?; + let numbers = match row.cells.swap_remove(at) { + Cell::Here(numbers) => numbers, + Cell::Beside(beside) => changes.load(&beside, held)?, + }; + let Value::Vector(numbers) = numbers else { + return Err(changes.damaged(WAITING)); + }; + if left == 0 { + // A block that was full before this call is left as it is. + if !writer.is_empty() { + let record = VectorRecord::Block(vector, number, block); + put_block(changes, held, record, writer)?; + } + block += 1; + (writer, left) = BlockWriter::new(dimension, None, &keys[done..]); + } + writer.push(key, &numbers); + left -= 1; + drop((row, numbers)); + held.shrink(held.bytes() - taken); + changes.delete_vector(VectorRecord::Pending(vector, number, key))?; + } + let record = VectorRecord::Block(vector, number, block); + put_block(changes, held, record, writer)?; + partition.blocks = block + 1; + changes.put_vector(VectorRecord::Partition(vector, &tuple), partition.encode())?; + let sealed = keys.len() as u64; + let left = remaining.checked_sub(sealed); + let left = left.ok_or_else(|| changes.damaged(directory::SHORT))?; + changes.put_vector_number(VectorNumber::Pending, left)?; + Ok((sealed, following)) +} + +/// Stages a block of `bytes` bytes. `finish` makes those bytes while the writer still holds +/// its own, so they are asked for before it. When it returns the writer's own are gone, and +/// the finished block is the one block the call keeps `BLOCK_BYTES` for: what was asked is +/// given back, and the change set counts the block once, as the record it stages. +fn put_block( + changes: &mut ChangeSet, + held: &mut Reservation, + record: VectorRecord<'_>, + writer: BlockWriter, +) -> Result<()> { + let bytes = writer.bytes(); + held.grow(bytes as u64)?; + let block = writer.finish(); + held.shrink(bytes as u64); + changes.put_vector(record, block) +} diff --git a/v2/contextdb/src/vector/search.rs b/v2/contextdb/src/vector/search.rs new file mode 100644 index 00000000..9aae8ebb --- /dev/null +++ b/v2/contextdb/src/vector/search.rs @@ -0,0 +1,365 @@ +//! The one search: score the rows that are in no block exactly, screen every slot of the +//! selected partitions, ask for the true distance of the rows whose range can reach the +//! answer, and prove the answer before returning it. + +use std::cmp::{Ordering, Reverse}; +use std::collections::{BTreeSet, BinaryHeap}; + +use super::block::{BROKEN, BlockReader}; +use super::directory::{self, Partition}; +use super::slot::Screen; +use crate::memory::Reservation; +use crate::store::{Key, Read, TableId, VectorId, VectorRecord}; +use crate::types::Result; + +/// What a search is asked for. +pub(crate) struct Nearest<'a> { + /// The table of the column, for the rows a transaction has staged. + pub table: TableId, + pub vector: VectorId, + /// The partitions to search, as key tuples or the first values of key tuples; empty for + /// every partition of the column. + pub scope: &'a [Key], + /// The vector to be near: of the column's dimension, every number finite. + pub query: &'a [f32], + /// How many rows to return at most. + pub k: usize, +} + +/// One row of an answer. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct Hit { + pub key: Key, + /// The row's distance, as `distance_of` gave it. + pub distance: f64, +} + +/// Bytes reserved for a row of the answer besides its key. +const PAIR: u64 = 32; + +/// Bytes reserved for a lower bound. A search holds at most as many as the answer has rows. +const LOW: u64 = 8; + +/// The `k` nearest rows, nearest first; rows at one distance come in the order of their keys. +/// +/// `distance_of` is `query`'s: given a row's key it reads the row and answers with its distance +/// from the vector of the search, or with `None` when the row is not there, does not pass the +/// statement's filter or has no vector. This function calls it once for every row the reader +/// has staged, once for every pending row of the selected partitions, and for the slots it +/// could not rule out, best range first. It never returns a row `distance_of` did not give a +/// distance for. +/// +/// A row is in an answer once. Several entries can lead to one row: a row the reader has staged +/// and the pending record or the slot of the row it replaces, or the slots of two vectors the +/// row has had. `distance_of` answers from the row's stored vector, so every one of them gives +/// the same distance and key. The answer holds each key once: a pair whose key it already holds +/// is dropped when it arrives, and a pair whose row was pushed out earlier cannot come back, +/// because it is greater than every pair the answer holds by then. An entry left by an earlier +/// vector therefore costs one row read and changes nothing. +/// +/// The answer is exactly the `k` least of `(distance, key)` over the rows of the selected +/// partitions that `distance_of` gives a distance for. Blocks are read in place from the page +/// cache; the list of slots still in the running and the answer grow through `reservation`. +/// When the answer is returned, `reservation` holds 32 bytes and the key's for each of its +/// rows and nothing else of the search. +pub(crate) fn search( + nearest: &Nearest<'_>, + read: &dyn Read, + reservation: &mut Reservation, + distance_of: &mut dyn FnMut(&Key) -> Result>, +) -> Result> { + if nearest.k == 0 { + return Ok(Vec::new()); + } + let partitions = directory::selected(read, nearest.vector, nearest.scope)?; + // The vector at length one that the screen keeps: four bytes a number. + let unit = 4 * nearest.query.len() as u64; + reservation.grow(unit)?; + let screen = Screen::new(nearest.query); + let mut found = Found { + k: nearest.k, + answer: BTreeSet::new(), + lows: BinaryHeap::new(), + reservation, + distance_of, + }; + + // The rows that are in no block, exactly: those the reader has staged, then the pending + // rows of each partition, which end where the records of the next partition begin. + read.staged(nearest.table, &mut |key| found.score(key).map(|()| true))?; + let first = Key::from_values(&[], &[])?; + for partition in &partitions { + let run = (nearest.vector, partition.number); + let from = VectorRecord::Pending(run.0, run.1, &first); + read.vectors(from, &mut |record, _| match record { + VectorRecord::Pending(vector, number, key) if (vector, number) == run => { + found.score(key).map(|()| true) + } + _ => Ok(false), + })?; + } + + // The slots. Every slot the pass left out had a `high` below the threshold it ended with, + // so when the last row of a full answer is at least that similar, no slot left out could + // have entered, and when there never was a threshold none was left out. Otherwise a row + // that is gone, filtered out or counted twice raised the threshold too far: one more pass + // keeps every slot that reaches the last row found, or every slot when rows are missing. + let mut kept = Kept::default(); + let threshold = found.pass(nearest, &partitions, read, &screen, &mut kept, None)?; + let last = found.last(); + if threshold > f64::NEG_INFINITY && last.is_none_or(|last| last < threshold) { + let floor = last.unwrap_or(f64::NEG_INFINITY); + found.pass(nearest, &partitions, read, &screen, &mut kept, Some(floor))?; + } + + let lows = LOW * found.lows.len() as u64; + found.reservation.shrink(unit + kept.reserved() + lows); + let hits = found.answer.into_iter(); + Ok(hits + .map(|(Order(distance), key)| Hit { key, distance }) + .collect()) +} + +/// A distance or a similarity where numbers are put in order. None is NaN: a distance is +/// between vectors of finite numbers. +#[derive(Clone, Copy)] +struct Order(f64); + +impl Ord for Order { + fn cmp(&self, other: &Order) -> Ordering { + self.0.total_cmp(&other.0) + } +} + +impl PartialOrd for Order { + fn partial_cmp(&self, other: &Order) -> Option { + Some(self.cmp(other)) + } +} + +impl PartialEq for Order { + fn eq(&self, other: &Order) -> bool { + self.cmp(other).is_eq() + } +} + +impl Eq for Order {} + +/// What a search has found so far. +struct Found<'a> { + k: usize, + /// The answer so far: the `k` least pairs of distance and key among the rows scored. + answer: BTreeSet<(Order, Key)>, + /// The `k` greatest similarities proven so far to be reached by a row or by a slot's + /// vector, the least of them on top. With `k` of them, a slot whose `high` is below the + /// least cannot be among the `k` nearest of the vectors seen. + lows: BinaryHeap>, + reservation: &'a mut Reservation, + distance_of: &'a mut dyn FnMut(&Key) -> Result>, +} + +impl Found<'_> { + /// Asks for the distance of one row and offers it to the answer, and its similarity to + /// the lower bounds when the answer took it. + fn score(&mut self, key: &Key) -> Result<()> { + if let Some(distance) = (self.distance_of)(key)? + && self.offer(distance, key)? + { + self.offer_low(1.0 - distance)?; + } + Ok(()) + } + + /// Offers a row to the answer and says whether it went in. This is the one way in, and + /// its one rule: a pair whose key the answer holds is dropped; any other goes in when the + /// answer has room, or when it is less than the greatest pair held, which it replaces. + /// + /// A key always comes with the same distance, so the pair is what the answer is asked for. + fn offer(&mut self, distance: f64, key: &Key) -> Result { + let pair = (Order(distance), key.clone()); + if self.answer.contains(&pair) { + return Ok(false); + } + let full = self.answer.len() == self.k; + if full && self.answer.last().is_some_and(|last| pair > *last) { + return Ok(false); + } + self.reservation.grow(PAIR + key.as_bytes().len() as u64)?; + if full && let Some((_, out)) = self.answer.pop_last() { + self.reservation.shrink(PAIR + out.as_bytes().len() as u64); + } + self.answer.insert(pair); + Ok(true) + } + + /// Offers a similarity that a row or a slot's vector is proven to reach. + fn offer_low(&mut self, low: f64) -> Result<()> { + if self.lows.len() < self.k { + self.reservation.grow(LOW)?; + self.lows.push(Reverse(Order(low))); + } else if let Some(mut least) = self.lows.peek_mut() + && Order(low) > least.0 + { + *least = Reverse(Order(low)); + } + Ok(()) + } + + /// What a slot's `high` must reach to stay in the running: the least of `k` lower bounds, + /// and nothing while there are fewer. + fn threshold(&self) -> f64 { + let least = self.lows.peek().filter(|_| self.lows.len() == self.k); + least.map_or(f64::NEG_INFINITY, |Reverse(Order(low))| *low) + } + + /// The similarity of the last row of the answer once it has `k` rows. + fn last(&self) -> Option { + let last = self.answer.last().filter(|_| self.answer.len() == self.k); + last.map(|(Order(distance), _)| 1.0 - distance) + } + + /// Screens every slot of the partitions and visits the ones kept; returns the threshold + /// the screen ended with. + /// + /// A slot is kept, with its `high` and a copy of its key, unless its `high` is below the + /// threshold. With no `floor` the threshold is [`Found::threshold`] and rises as kept + /// slots offer their `low`; with one it is the floor throughout. Kept slots are visited + /// by `high`, greatest first, and the visit ends at the first whose `high` is below the + /// similarity of the last row of a full answer: neither it nor any after it can enter, + /// while one that equals it may, on its key. + fn pass( + &mut self, + nearest: &Nearest<'_>, + partitions: &[Partition], + read: &dyn Read, + screen: &Screen, + kept: &mut Kept, + floor: Option, + ) -> Result { + kept.clear(); + let mut threshold = floor.unwrap_or_else(|| self.threshold()); + for partition in partitions { + let run = (nearest.vector, partition.number); + let from = VectorRecord::Block(run.0, run.1, 0); + read.vectors(from, &mut |record, bytes| match record { + VectorRecord::Block(vector, number, _) if (vector, number) == run => { + let block = BlockReader::new(bytes, nearest.query.len()); + if !block.is_whole() { + return Err(read.damaged(BROKEN)); + } + for at in 0..block.len() { + let (low, high) = screen.range(block.slot(at)); + if f64::from(high) < threshold { + continue; + } + kept.keep(high, &block.key(at), threshold, self.reservation)?; + if floor.is_none() { + self.offer_low(f64::from(low))?; + threshold = self.threshold(); + } + } + Ok(true) + } + _ => Ok(false), + })?; + } + kept.drop_below(threshold); + let by_high = |a: &Slot, b: &Slot| b.high.total_cmp(&a.high).then(a.start.cmp(&b.start)); + kept.slots.sort_unstable_by(by_high); + for slot in &kept.slots { + if self.last().is_some_and(|last| f64::from(slot.high) < last) { + break; + } + let key = &kept.keys[slot.start..][..slot.len as usize]; + self.score(&Key::from_bytes(key))?; + } + Ok(threshold) + } +} + +/// A slot still in the running: the top of its range, and where its key is in [`Kept::keys`]. +struct Slot { + high: f32, + len: u32, + start: usize, +} + +/// The slots still in the running. +#[derive(Default)] +struct Kept { + slots: Vec, + /// The keys of the slots, one after another. + keys: Vec, + /// How many slots, and how many bytes of keys, the reservation holds room for. + room: usize, + key_room: usize, +} + +impl Kept { + /// Slots there is room for at first. + const FIRST: usize = 1024; + + /// Bytes the reservation holds for the list. + fn reserved(&self) -> u64 { + (self.room * size_of::() + self.key_room) as u64 + } + + /// Empties the list and keeps its room. + fn clear(&mut self) { + self.slots.clear(); + self.keys.clear(); + } + + /// Adds a slot. When the list is full, the slots now below the threshold are dropped + /// first, and the room is doubled only if more than half are left. + fn keep( + &mut self, + high: f32, + key: &Key, + threshold: f64, + reservation: &mut Reservation, + ) -> Result<()> { + if self.slots.len() == self.room { + self.drop_below(threshold); + let more = match self.room { + 0 => Kept::FIRST, + room if self.slots.len() > room / 2 => room, + _ => 0, + }; + reservation.grow((more * size_of::()) as u64)?; + self.slots + .reserve_exact(self.room + more - self.slots.len()); + self.room += more; + } + let key = key.as_bytes(); + let start = self.keys.len(); + if start + key.len() > self.key_room { + let more = (start + key.len() - self.key_room).max(self.key_room); + reservation.grow(more as u64)?; + self.keys.reserve_exact(self.key_room + more - start); + self.key_room += more; + } + self.keys.extend_from_slice(key); + // The key is out of a block, where four bytes say where it ends: its length fits. + let len = key.len() as u32; + self.slots.push(Slot { high, len, start }); + Ok(()) + } + + /// Drops the slots whose `high` is below the threshold and closes the gaps their keys + /// leave. The slots are in the order they were kept in, the order their keys lie in, so + /// a key that stays only ever moves towards the start. + fn drop_below(&mut self, threshold: f64) { + let (keys, mut end) = (&mut self.keys, 0); + self.slots.retain_mut(|slot| { + let stays = f64::from(slot.high) >= threshold; + if stays { + keys.copy_within(slot.start..slot.start + slot.len as usize, end); + slot.start = end; + end += slot.len as usize; + } + stays + }); + keys.truncate(end); + } +} diff --git a/v2/contextdb/src/vector/slot.rs b/v2/contextdb/src/vector/slot.rs new file mode 100644 index 00000000..ef576d84 --- /dev/null +++ b/v2/contextdb/src/vector/slot.rs @@ -0,0 +1,192 @@ +//! One vector as a block keeps it, the range a search gets from it, and the distance that +//! every answer is ordered by. +//! +//! A slot is eight bytes and then one signed byte for each number of the vector: the scale as +//! an `f32`, the rounding error as an `f32`, then the codes. The codes are of the vector +//! brought to length one, `u = x / |x|`: `scale = max |u_i| / 127` and `code_i = round(u_i / +//! scale)`. The rounding error is the length of `u - scale * code`, worked out in 64 bits and +//! rounded up, so it is never smaller than the truth. A vector of all zeros has scale and +//! error zero. + +/// Bytes of a slot before its codes. +pub(super) const SLOT_HEADER: usize = 8; + +/// What a figure worked out in 64 bits is multiplied by so that it is not below the truth: the +/// roundings of that arithmetic come to less than a thousand millionth of the figure for any +/// vector of fewer than ten million numbers. +const UP: f64 = 1.0 + 1e-9; + +/// Bytes the slot of a vector of this dimension takes. +pub(super) fn slot_len(dimension: usize) -> usize { + SLOT_HEADER + dimension +} + +/// The cosine distance between a stored vector and the vector of a search. Every answer is +/// ordered by this number and by nothing else, so a search and a comparison of every stored +/// vector agree to the last bit. +/// +/// Each number is widened to 64 bits before it is multiplied, and the three sums (the dot +/// product and the two squared lengths) are added in the order of the numbers. The distance is +/// `1 - dot / (|stored| * |query|)`, and 1 when either vector is all zeros. +pub(crate) fn distance(stored: &[f32], query: &[f32]) -> f64 { + let (mut dot, mut aa, mut bb) = (0.0_f64, 0.0_f64, 0.0_f64); + for (x, y) in stored.iter().zip(query) { + let (x, y) = (f64::from(*x), f64::from(*y)); + dot += x * y; + aa += x * x; + bb += y * y; + } + if aa == 0.0 || bb == 0.0 { + 1.0 + } else { + 1.0 - dot / (aa.sqrt() * bb.sqrt()) + } +} + +/// The length of a vector, in 64 bits: the square of the largest `f32` does not fit in 32. +fn length(vector: &[f32]) -> f64 { + let squares = vector.iter().map(|x| f64::from(*x) * f64::from(*x)); + squares.sum::().sqrt() +} + +/// Writes the slot of `vector` into `out`, which is `slot_len(vector.len())` bytes long. Every +/// number of the vector is finite: `write` refuses any other before it is stored. +/// +/// The codes are worked out with the scale as it is stored, and so is the error, which is +/// then rounded up twice: by [`UP`] for the arithmetic, and to the next `f32`. +pub(super) fn write_slot(vector: &[f32], out: &mut [u8]) { + let length = length(vector); + if length == 0.0 { + out.fill(0); + return; + } + let unit = |x: &f32| f64::from(*x) / length; + let largest = vector.iter().map(|x| unit(x).abs()).fold(0.0, f64::max); + let scale = (largest / 127.0) as f32; + let mut off = 0.0_f64; + for (x, code) in vector.iter().zip(&mut out[SLOT_HEADER..]) { + let exact = unit(x); + let rounded = (exact / f64::from(scale)).round().clamp(-127.0, 127.0); + *code = (rounded as i8).cast_unsigned(); + let by = exact - rounded * f64::from(scale); + off += by * by; + } + let error = ((off.sqrt() * UP) as f32).next_up(); + out[..4].copy_from_slice(&scale.to_le_bytes()); + out[4..SLOT_HEADER].copy_from_slice(&error.to_le_bytes()); +} + +/// The codes of a slot multiplied by the numbers of `unit` and added up, in 32 bits. Sixteen +/// sums run side by side, so the compiler turns the loop into the widest instructions of the +/// machine the library is built for. +fn product(codes: &[u8], unit: &[f32]) -> f32 { + let mut lanes = [0.0_f32; 16]; + let (code_runs, code_rest) = codes.as_chunks::<16>(); + let (unit_runs, unit_rest) = unit.as_chunks::<16>(); + for (codes, unit) in code_runs.iter().zip(unit_runs) { + for ((lane, code), number) in lanes.iter_mut().zip(codes).zip(unit) { + *lane += f32::from(code.cast_signed()) * number; + } + } + for ((lane, code), number) in lanes.iter_mut().zip(code_rest).zip(unit_rest) { + *lane += f32::from(code.cast_signed()) * number; + } + lanes.iter().sum() +} + +/// What a search works out once from its vector and then uses for every slot it screens. +pub(super) struct Screen { + /// The vector of the search brought to length one, in 32 bits. + unit: Vec, + /// The length of `unit`, rounded up. + length: f64, + /// The most that the 32-bit arithmetic of one slot's product can be off by, for a slot of + /// scale one. + slack: f64, + /// The most a similarity can be off by for reasons that do not depend on the slot: `unit` + /// is not exactly of length one, and [`distance`] rounds too. + drift: f64, +} + +impl Screen { + /// `query` is of the column's dimension and every number of it is finite. + pub(super) fn new(query: &[f32]) -> Screen { + let n = query.len(); + let length = length(query); + let mut unit = Vec::with_capacity(n); + let (mut squares, mut sum, mut off) = (0.0_f64, 0.0_f64, 0.0_f64); + for x in query { + // A vector of zeros has no direction: its `unit` is zeros too. + let exact = if length == 0.0 { + 0.0 + } else { + f64::from(*x) / length + }; + let kept = exact as f32; + unit.push(kept); + let kept = f64::from(kept); + squares += kept * kept; + sum += kept.abs(); + off += (kept - exact) * (kept - exact); + } + // One slot's product is rounded this many times, one after another: a multiplication, + // the additions of its lane, the sixteen lanes added up, the multiplication by the scale. + let roundings = (n.div_ceil(16) + 18) as f64 * f64::from(f32::EPSILON) / 2.0; + let gamma = roundings / (1.0 - roundings); + Screen { + unit, + length: squares.sqrt() * UP, + slack: 2.0 * 128.0 * gamma * sum * UP, + drift: off.sqrt() * UP + 8.0 * (n + 4) as f64 * f64::EPSILON / 2.0, + } + } + + /// The range that the similarity of a slot's vector is proven to lie in: `(low, high)` + /// with `low <= 1 - distance(x, query) <= high` for the vector `x` the slot was written + /// from. + /// + /// Proof. Write `u = x / |x|` and `w = query / |query|`, `s`, `c` and `r` for the slot's + /// scale, codes and error, `t` for `unit`, which is `w` rounded to 32 bits, and `n` for + /// the dimension. In exact arithmetic `u.w - s (c.t) = u.(w - t) + (u - s c).t`, so by + /// Cauchy-Schwarz, with `|u| = 1`, + /// + /// `|u.w - s (c.t)| <= |w - t| + |u - s c| |t|`. + /// + /// What is compared is not `u.w` and `s (c.t)` but numbers worked out from them, and + /// each rounding between has its term: + /// + /// - `|u - s c| |t|`. [`write_slot`] measured `r` against `u` as 64-bit arithmetic gives + /// it, which is within `(n + 3) 2^-54` of `u`, and rounded `r` up; `length` is `|t|` + /// rounded up. The term is `r * length`. + /// - `s (c.t)`. [`product`] and the multiplication by `s` round at most `ceil(n / 16) + 18` + /// times in 32 bits, on terms no larger than `127 s |t_i|`. With `g` for that count + /// times `2^-24` and `gamma = g / (1 - g)`, what is worked out is within + /// `127 gamma s |t|_1` of `s (c.t)`. The term is `s * slack`, twice that and more. + /// - `|w - t|`. [`Screen::new`] measured `t` against `w` as 64-bit arithmetic gives it, + /// rounded up, and that `w` is within `(n + 3) 2^-54` of `w`. + /// - `u.w`. [`distance`] multiplies exactly, two numbers of 24 bits in 53, and rounds + /// `n - 1` times in each of its three sums and four times after them; taking it from + /// one and the result from one again rounds twice more. So `1 - distance` is within + /// `(2n + 6) 2^-53` of `u.w`. + /// + /// The last two, and the `(n + 3) 2^-54` of the first, depend on no slot and are `drift`: + /// the measured part of `|w - t|`, and `8 (n + 4) 2^-53` for the rest, which comes to + /// `(3n + 9) 2^-53` and so leaves room for a product too small for 32 bits to round in + /// proportion. The half width is therefore `r * length + s * slack + drift`, + /// widened by [`UP`] for its own arithmetic, and both ends are rounded outward to 32 + /// bits. + /// + /// A vector of zeros, stored or searched for, is at distance 1. Its codes or its `unit` + /// are zeros, so the middle of the range is 0, and `drift` is above zero. + pub(super) fn range(&self, slot: &[u8]) -> (f32, f32) { + let scale = f32::from_le_bytes([slot[0], slot[1], slot[2], slot[3]]); + let error = f32::from_le_bytes([slot[4], slot[5], slot[6], slot[7]]); + let middle = f64::from(scale * product(&slot[SLOT_HEADER..], &self.unit)); + let half = f64::from(error) * self.length + f64::from(scale) * self.slack + self.drift; + let half = half * UP; + ( + ((middle - half) as f32).next_down(), + ((middle + half) as f32).next_up(), + ) + } +} diff --git a/v2/contextdb/src/vector/stage.rs b/v2/contextdb/src/vector/stage.rs new file mode 100644 index 00000000..67ea2d1f --- /dev/null +++ b/v2/contextdb/src/vector/stage.rs @@ -0,0 +1,210 @@ +//! What a commit's row changes mean for the vector part, worked out once, inside the commit +//! function, from the complete change set. + +use std::collections::BTreeMap; +use std::collections::btree_map::Entry; + +use super::directory::{self, Partition}; +use super::tally::{position, tuples}; +use crate::memory::{Budget, Reservation}; +use crate::schema::Catalog; +use crate::sql::not_supported; +use crate::store::{Cell, ChangeSet, Key, Read, VectorId, VectorNumber, VectorRecord}; +use crate::types::{Error, Result}; + +/// One staged row whose vector leaves a key tuple, arrives under one, or both. +struct Move { + vector: VectorId, + row: Key, + leaves: Option, + enters: Option, +} + +/// One partition the commit's rows touch. +struct Touched { + /// Its record at the committed moment, `None` when it has none. + was: Option, + /// Its record as the commit leaves it: with no live vector, it has none. + now: Partition, +} + +/// What the commit does to the partitions of one column. +#[derive(Default)] +struct Count { + /// How many it makes less how many it empties. + change: i64, + made: u64, +} + +/// Stages the vector records that the rows staged in `changes` call for. +/// +/// For every staged row of a table with a VECTOR column, the row as the change set's committed +/// moment has it is compared with the row as staged. A vector that is new, changed, or under +/// other partition key values goes on its partition's pending list; the list it was on, if it +/// was on one, lets it go. A partition is made when its first vector arrives and removed, with +/// its blocks, when its last one leaves. A row whose vector and partition key values are as +/// they were stages nothing. +/// +/// Every vector that leaves is taken out before any arrives, so at the limit a commit may +/// move the last vector of one partition into a new one. +/// +/// The commit function calls this under the commit lock, after the change set is laid over the +/// newest moment and every other rule has passed. Refuses with `VectorPartitionLimitExceeded` +/// when a column would end the commit with more partitions than it may have, and stages +/// nothing for a change set that staged no row. +pub(crate) fn stage(changes: &mut ChangeSet, catalog: &Catalog, budget: &Budget) -> Result<()> { + let mut held = budget.reserve(0, "transaction")?; + let moves = moves(changes, catalog, &mut held)?; + if moves.is_empty() { + return Ok(()); + } + let waited = changes.vector_number(VectorNumber::Pending)?; + let mut pending = waited; + let mut touched: BTreeMap<(VectorId, Key), Touched> = BTreeMap::new(); + let mut counts: BTreeMap = BTreeMap::new(); + + for moved in &moves { + let Some(tuple) = &moved.leaves else { continue }; + let left = &mut touch(&mut touched, &*changes, moved.vector, tuple, &mut held)?.now; + let fewer = left.live.checked_sub(1); + left.live = fewer.ok_or_else(|| changes.damaged(directory::SHORT))?; + let waiting = VectorRecord::Pending(moved.vector, left.number, &moved.row); + if directory::kept(&*changes, waiting, |_| ())?.is_some() { + changes.delete_vector(waiting)?; + let fewer = pending.checked_sub(1); + pending = fewer.ok_or_else(|| changes.damaged(directory::SHORT))?; + } + } + // A partition whose last vector left goes, with its blocks. + for ((vector, _), Touched { now: emptied, .. }) in &mut touched { + if emptied.live == 0 { + for block in 0..emptied.blocks { + changes.delete_vector(VectorRecord::Block(*vector, emptied.number, block))?; + } + counts.entry(*vector).or_default().change -= 1; + emptied.blocks = 0; + } + } + for moved in &moves { + let Some(tuple) = &moved.enters else { continue }; + let entered = &mut touch(&mut touched, &*changes, moved.vector, tuple, &mut held)?.now; + if entered.live == 0 { + // A new partition takes the column's next number. + let count = counts.entry(moved.vector).or_default(); + let next = VectorNumber::NextPartition(moved.vector); + // A number is never given twice, and a column has 4,294,967,296 of them. + let number = changes.vector_number(next)? + count.made; + entered.number = u32::try_from(number).map_err(|_| { + not_supported("more than 4,294,967,296 partitions over the life of a VECTOR column") + })?; + (count.change, count.made) = (count.change + 1, count.made + 1); + } + entered.live += 1; + let waiting = VectorRecord::Pending(moved.vector, entered.number, &moved.row); + changes.put_vector(waiting, Vec::new())?; + pending += 1; + } + + for ((vector, tuple), partition) in &touched { + let record = VectorRecord::Partition(*vector, tuple); + if partition.now.live == 0 { + changes.delete_vector(record)?; + } else if partition.was != Some(partition.now) { + changes.put_vector(record, partition.now.encode())?; + } + } + if pending != waited { + changes.put_vector_number(VectorNumber::Pending, pending)?; + } + for (vector, count) in &counts { + let partitions = VectorNumber::Partitions(*vector); + let now = changes + .vector_number(partitions)? + .saturating_add_signed(count.change); + changes.put_vector_number(partitions, now)?; + if count.made == 0 { + continue; + } + let next = VectorNumber::NextPartition(*vector); + changes.put_vector_number(next, changes.vector_number(next)? + count.made)?; + // A statement over the limit was refused when it ran, so this is the refusal only a + // commit can see: another transaction made a partition while this one was open. + if let Some((table, declared)) = catalog.vector(*vector) + && let Some(limit) = declared.limit() + && now > u64::from(limit) + { + return Err(Error::VectorPartitionLimitExceeded { + index: table.vector_ref(declared), + max_partitions: limit, + }); + } + } + Ok(()) +} + +/// The staged rows whose vector leaves a key tuple or arrives under one, in the order of +/// table and row key. A row with a vector before and after moves when its key tuple or its +/// vector changed: the vector it had is loaded to tell, and given back. +fn moves(changes: &ChangeSet, catalog: &Catalog, held: &mut Reservation) -> Result> { + let base = changes.base(); + let mut moves = Vec::new(); + for table in catalog.tables() { + if table.vectors.is_empty() { + continue; + } + let staged = changes.rows_changed(); + for (_, key, after, _) in staged.filter(|staged| staged.0 == table.id) { + let before = base.row(table.id, key)?; + for column in &table.vectors { + let (leaves, enters) = tuples(table, column, before.as_ref(), after, base, held)?; + if leaves == enters { + // The row has a vector under one tuple before and after, or none. + let (Some(row), Some(values), true) = (&before, after, leaves.is_some()) else { + continue; + }; + let at = position(table, &column.column); + let taken = held.bytes(); + let same = match &row.cells[at] { + Cell::Here(vector) => *vector == values[at], + Cell::Beside(beside) => base.load(beside, held)? == values[at], + }; + held.shrink(held.bytes() - taken); + if same { + continue; + } + } + let keys = [Some(key), leaves.as_ref(), enters.as_ref()]; + let keys = keys.into_iter().flatten(); + let bytes: usize = keys.map(|key| key.as_bytes().len()).sum(); + held.grow((size_of::() + bytes) as u64)?; + moves.push(Move { + vector: column.id, + row: key.clone(), + leaves, + enters, + }); + } + } + } + Ok(moves) +} + +/// The partition under a key tuple as the commit has it so far: read from the committed moment +/// the first time, and kept for the length of the call. +fn touch<'a>( + touched: &'a mut BTreeMap<(VectorId, Key), Touched>, + read: &dyn Read, + vector: VectorId, + tuple: &Key, + held: &mut Reservation, +) -> Result<&'a mut Touched> { + Ok(match touched.entry((vector, tuple.clone())) { + Entry::Occupied(known) => known.into_mut(), + Entry::Vacant(first) => { + let was = directory::partition(read, vector, tuple)?; + held.grow(tuple.as_bytes().len() as u64 + 64)?; + let now = was.unwrap_or_default(); + first.insert(Touched { was, now }) + } + }) +} diff --git a/v2/contextdb/src/vector/tally.rs b/v2/contextdb/src/vector/tally.rs new file mode 100644 index 00000000..3e898281 --- /dev/null +++ b/v2/contextdb/src/vector/tally.rs @@ -0,0 +1,229 @@ +//! The partitions a transaction's statements have made and emptied so far, so that the INSERT +//! or UPDATE that would give a column one partition too many is refused when it runs. + +use std::collections::BTreeMap; +use std::collections::btree_map::Entry; + +use super::directory; +use crate::memory::{Budget, Reservation}; +use crate::schema::{Table, VectorColumn}; +use crate::store::{Cell, Key, Read, Row, VectorId, VectorNumber}; +use crate::types::{Error, Result, Value}; + +/// What the rows a transaction has staged mean for the partition count of each VECTOR column +/// that has a PARTITION_KEY: for every key tuple a staged row has put a vector under or taken +/// one from, how many vectors the transaction now has there, and for every column, how many +/// partitions it would have if the transaction were committed as it stands. +/// +/// It is kept in step with the staged rows, row by row, and is exact for the moment the +/// transaction began at together with what the transaction has staged. A partition that +/// another transaction commits meanwhile is not seen here: [`super::stage`] counts again +/// inside the commit function, over the newest moment. +pub(crate) struct Tally { + /// Vectors under each key tuple a statement has touched: what the committed moment has + /// there, changed by what the transaction has staged. + live: BTreeMap<(VectorId, Key), u64>, + /// Partitions of each column a statement has touched, counted the same way. + partitions: BTreeMap, + /// What the statement that is running has noted, in order, so that it can be undone. + noted: Vec, + /// Never read: holds what the three lists above take. + reserved: Reservation, +} + +/// One vector that a statement put under a key tuple or took from it. +struct Note { + place: (VectorId, Key), + /// True for a vector put under the tuple, false for one taken from it. + arrived: bool, + /// True when the tuple had no vector until this one arrived. + began: bool, +} + +impl Note { + /// What a note is counted as while its statement runs. + fn bytes(&self) -> u64 { + self.place.1.as_bytes().len() as u64 + 16 + } +} + +/// Puts one vector under a key tuple, or takes one from it, and keeps the column's number of +/// partitions in step. True when the tuple had none and now has one. +fn shift(live: &mut u64, partitions: &mut u64, arrives: bool) -> bool { + if arrives { + *live += 1; + *partitions += u64::from(*live == 1); + } else { + *live -= 1; + *partitions -= u64::from(*live == 0); + } + arrives && *live == 1 +} + +impl Tally { + /// The tally of a transaction that has staged nothing. + pub(crate) fn new(budget: &Budget) -> Result { + Ok(Tally { + live: BTreeMap::new(), + partitions: BTreeMap::new(), + noted: Vec::new(), + reserved: budget.reserve(0, "transaction")?, + }) + } + + /// Notes one row a statement stages in a table with a VECTOR column: `before` is the row as + /// the transaction had it until now, `None` when there was none, and `after` is the row as + /// staged, `None` for a removal. For each VECTOR column with a PARTITION_KEY, a vector that + /// leaves its key tuple, arrives under one, or moves from one to another changes the count + /// of both tuples; a row whose vector stays under the same tuple, or has none before and + /// after, changes nothing, and its vector is never loaded. + /// + /// `read` is the transaction's change set: a tuple is looked up in the committed moment + /// the first time the transaction touches it, and a key column kept beside the row is + /// loaded from it. `write` calls this for every row, before or after it stages the row. + pub(crate) fn note( + &mut self, + table: &Table, + before: Option<&Row>, + after: Option<&[Value]>, + read: &dyn Read, + ) -> Result<()> { + for column in &table.vectors { + if column.partition_key.is_empty() { + continue; + } + let (left, entered) = tuples(table, column, before, after, read, &mut self.reserved)?; + if left == entered { + continue; + } + for (tuple, arrived) in [(left, false), (entered, true)] { + if let Some(tuple) = tuple { + self.count((column.id, tuple), arrived, read)?; + } + } + } + Ok(()) + } + + /// Counts one vector more under a key tuple, or one fewer, and notes it. The column's + /// partitions and the tuple's vectors are read from the committed moment the first time. + fn count(&mut self, place: (VectorId, Key), arrived: bool, read: &dyn Read) -> Result<()> { + let (vector, tuple) = (place.0, &place.1); + let partitions = match self.partitions.entry(vector) { + Entry::Occupied(counted) => counted.into_mut(), + Entry::Vacant(first) => { + first.insert(read.vector_number(VectorNumber::Partitions(vector))?) + } + }; + let live = match self.live.entry(place.clone()) { + Entry::Occupied(counted) => counted.into_mut(), + Entry::Vacant(first) => { + let partition = directory::partition(read, vector, tuple)?; + self.reserved.grow(tuple.as_bytes().len() as u64 + 64)?; + first.insert(partition.map_or(0, |partition| partition.live)) + } + }; + // A vector leaves a tuple that has one, in a column that has a partition. + if !arrived && (*live == 0 || *partitions == 0) { + return Err(read.damaged(directory::SHORT)); + } + let mut note = Note { + place, + arrived, + began: false, + }; + self.reserved.grow(note.bytes())?; + note.began = shift(live, partitions, arrived); + self.noted.push(note); + Ok(()) + } + + /// The last act of a statement that wrote rows. Refuses with + /// `VectorPartitionLimitExceeded` when the statement put a vector under a key tuple that + /// had none and the column would then have more partitions than its limit, and forgets + /// what the statement noted; otherwise keeps it. + pub(crate) fn settle(&mut self, table: &Table) -> Result<()> { + for column in &table.vectors { + let mut notes = self.noted.iter(); + let began = notes.any(|note| note.place.0 == column.id && note.began); + let partitions = self.partitions.get(&column.id).copied().unwrap_or(0); + if let Some(limit) = column.limit() + && began + && partitions > u64::from(limit) + { + self.forget(); + return Err(Error::VectorPartitionLimitExceeded { + index: table.vector_ref(column), + max_partitions: limit, + }); + } + } + let kept: u64 = self.noted.iter().map(Note::bytes).sum(); + self.reserved.shrink(kept); + self.noted.clear(); + Ok(()) + } + + /// Undoes what the running statement has noted. `database` calls this beside the rewind + /// of the change set when a statement is refused, whatever refused it. + pub(crate) fn forget(&mut self) { + while let Some(note) = self.noted.pop() { + if let Some(live) = self.live.get_mut(¬e.place) + && let Some(partitions) = self.partitions.get_mut(¬e.place.0) + { + shift(live, partitions, !note.arrived); + } + self.reserved.shrink(note.bytes()); + } + } +} + +/// Where a column that a table's declaration names is in a row of the table. +pub(super) fn position(table: &Table, name: &str) -> usize { + let mut columns = table.meta.columns.iter(); + columns + .position(|column| column.name == name) + .expect("a VECTOR column and the columns of its PARTITION_KEY are columns of its table") +} + +/// The key tuple a row's vector was under in a VECTOR column, and the one it is under now. +/// +/// `before` is the row as it was and `after` its values as they are staged; `None` is no row. +/// A row that is not there, and a row whose vector is NULL, has its vector under no tuple. +/// The tuple is the values of the PARTITION_KEY columns as [`Key::from_values`] encodes them, +/// ascending; a column with no PARTITION_KEY has the one tuple of no values. A key column +/// kept beside its row is loaded from `read` through `held`, which is given back what the +/// load took. No vector is loaded. +pub(super) fn tuples( + table: &Table, + column: &VectorColumn, + before: Option<&Row>, + after: Option<&[Value]>, + read: &dyn Read, + held: &mut Reservation, +) -> Result<(Option, Option)> { + let vector = position(table, &column.column); + let keys = column.partition_key.iter(); + let keys: Vec = keys.map(|name| position(table, name)).collect(); + let tuple = |values: Result>| Key::from_values(&values?, &[]).map(Some); + let was = match before { + Some(row) if !matches!(row.cells[vector], Cell::Here(Value::Null)) => { + let taken = held.bytes(); + let values = keys.iter().map(|&at| match &row.cells[at] { + Cell::Here(value) => Ok(value.clone()), + Cell::Beside(beside) => read.load(beside, held), + }); + let was = tuple(values.collect()); + held.shrink(held.bytes() - taken); + was? + } + _ => None, + }; + let is = match after { + Some(values) if values[vector] != Value::Null => { + tuple(Ok(keys.iter().map(|&at| values[at].clone()).collect()))? + } + _ => None, + }; + Ok((was, is)) +} diff --git a/v2/contextdb/src/vector/tests.rs b/v2/contextdb/src/vector/tests.rs new file mode 100644 index 00000000..33218864 --- /dev/null +++ b/v2/contextdb/src/vector/tests.rs @@ -0,0 +1,432 @@ +//! The pure functions of the vector part: the range a slot proves for its vector's similarity, +//! the bytes of a block, and the record of a partition. And what sealing makes of a store whose +//! records do not agree, and what a write and sealing make of a schema record whose VECTOR +//! column is not one: no public call writes such a store, so the last tests here write one +//! through the store itself. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; + +use super::block::{BLOCK_BYTES, BlockReader, BlockWriter, takes}; +use super::directory::Partition; +use super::seal::seal; +use super::slot::{Screen, distance, slot_len, write_slot}; +use super::stage::stage; +use crate::Database; +use crate::memory::Budget; +use crate::schema::Catalog; +use crate::store::{ChangeSet, Key, Store, TableId}; +use crate::types::{Error, Value}; + +/// A fixed scatter of numbers: a xorshift that starts from a seed. +struct Fixed(u64); + +impl Fixed { + fn seeded(seed: u64) -> Fixed { + let spread = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ 0xD6E8_FEB8_6659_FD93; + let mut fixed = Fixed(if spread == 0 { 1 } else { spread }); + for _ in 0..8 { + fixed.unit(); + } + fixed + } + + /// The next number from -1 up to, not including, 1. + fn unit(&mut self) -> f32 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + (self.0 >> 40) as f32 / (1_u64 << 23) as f32 - 1.0 + } +} + +fn vector(seed: u64, dimension: usize) -> Vec { + let mut fixed = Fixed::seeded(seed); + (0..dimension).map(|_| fixed.unit()).collect() +} + +/// The distance as the design writes it: each number widened to 64 bits, the three sums added +/// in the order of the numbers, and 1 for a vector of zeros. +fn written_out(stored: &[f32], query: &[f32]) -> f64 { + let (mut dot, mut aa, mut bb) = (0.0_f64, 0.0_f64, 0.0_f64); + for (x, y) in stored.iter().zip(query) { + let (x, y) = (f64::from(*x), f64::from(*y)); + dot += x * y; + aa += x * x; + bb += y * y; + } + if aa == 0.0 || bb == 0.0 { + 1.0 + } else { + 1.0 - dot / (aa.sqrt() * bb.sqrt()) + } +} + +/// Asserts that the slot written from `stored` gives a range that holds the similarity of +/// `stored` to `query`, one minus their distance, and that the distance is the design's to +/// the bit. +fn the_range_holds(stored: &[f32], query: &[f32], what: &str) { + let found = distance(stored, query); + assert_eq!( + found.to_bits(), + written_out(stored, query).to_bits(), + "{what}: the distance {found:e}" + ); + let mut slot = vec![0_u8; slot_len(stored.len())]; + write_slot(stored, &mut slot); + let (low, high) = Screen::new(query).range(&slot); + let similarity = 1.0 - found; + assert!( + f64::from(low) <= similarity && similarity <= f64::from(high), + "{what}: the similarity {similarity:e} is outside the range from {low:e} to {high:e}" + ); +} + +#[test] +fn the_range_of_a_slot_holds_the_similarity_of_its_vector() { + for dimension in [1, 7, 16, 33, 768] { + // Two hundred pairs: scattered, nearly the same direction, nearly opposite, and of + // lengths a million million apart. + for pair in 0..200_u64 { + let seed = dimension as u64 * 1000 + 2 * pair; + let mut stored = vector(seed, dimension); + let mut query = vector(seed + 1, dimension); + let nudges = query.iter().zip(&stored); + let nudged: Vec = nudges + .map(|(nudge, number)| number + 0.001 * nudge) + .collect(); + match pair % 4 { + 1 => query = nudged, + 2 => query = nudged.iter().map(|number| -number).collect(), + 3 => { + stored.iter_mut().for_each(|number| *number *= 1e-6); + query.iter_mut().for_each(|number| *number *= 1e6); + } + _ => {} + } + the_range_holds( + &stored, + &query, + &format!("dimension {dimension}, pair {pair}"), + ); + } + // A vector of zeros, on either side, is at distance 1; and the largest numbers a + // vector can hold. + let zeros = vec![0.0_f32; dimension]; + let largest = vec![f32::MAX; dimension]; + let other = vector(dimension as u64 * 7919, dimension); + for (stored, query, what) in [ + (&zeros, &other, "a stored vector of zeros"), + (&other, &zeros, "a search for a vector of zeros"), + (&zeros, &zeros, "zeros on both sides"), + (&largest, &other, "a stored vector of the largest number"), + ( + &other, + &largest, + "a search for a vector of the largest number", + ), + (&largest, &largest, "the largest number on both sides"), + (&other, &other, "a vector and itself"), + ] { + the_range_holds(stored, query, &format!("dimension {dimension}, {what}")); + } + assert_eq!(distance(&zeros, &other), 1.0); + assert_eq!(distance(&other, &zeros), 1.0); + } +} + +/// A key of one of several lengths: an INTEGER, or a TEXT of 1 to 40 letters. +fn key(n: i64) -> Key { + let value = if n % 3 == 0 { + Value::Text("k".repeat((n % 40) as usize + 1)) + } else { + Value::Int64(n) + }; + Key::from_values(&[value], &[]).expect("a key of one value") +} + +#[test] +fn a_block_reads_back_every_slot_and_key_and_stays_within_its_bytes() { + const DIMENSION: usize = 16; + // Every slot and key of a block, as they were pushed. + let reads_back = |block: &[u8], pushed: &[(Key, Vec)]| { + let reader = BlockReader::new(block, DIMENSION); + assert_eq!(reader.len(), pushed.len()); + for (at, (key, vector)) in pushed.iter().enumerate() { + let mut slot = vec![0_u8; slot_len(DIMENSION)]; + write_slot(vector, &mut slot); + assert_eq!(reader.slot(at), slot.as_slice(), "slot {at}"); + assert_eq!(reader.key(at), *key, "key {at}"); + } + }; + + // Slots are pushed for as many keys as the writer says its block takes, of more keys than + // a block has room for. + let offered: Vec = (0..(BLOCK_BYTES / slot_len(DIMENSION)) as i64) + .map(key) + .collect(); + let (mut writer, taken) = BlockWriter::new(DIMENSION, None, &offered); + assert!(writer.is_empty()); + let mut pushed: Vec<(Key, Vec)> = Vec::new(); + let refused = loop { + let n = pushed.len() as i64; + let key = key(n); + if pushed.len() == taken { + break key; + } + let vector = vector(n as u64, DIMENSION); + writer.push(&key, &vector); + assert!(!writer.is_empty()); + pushed.push((key, vector)); + }; + let block = writer.finish(); + // A block is its count, the slots, one key end for each slot, and the keys. + let keys: usize = pushed.iter().map(|(key, _)| key.as_bytes().len()).sum(); + assert_eq!( + block.len(), + 4 + pushed.len() * (slot_len(DIMENSION) + 4) + keys + ); + assert_eq!(block[..4], (pushed.len() as u32).to_le_bytes()); + // It is within the most a block may be, and the slot that was refused would have passed it. + assert!(block.len() <= BLOCK_BYTES); + assert!(block.len() + slot_len(DIMENSION) + 4 + refused.as_bytes().len() > BLOCK_BYTES); + reads_back(&block, &pushed); + + // A writer that starts from a block goes on where the block ends. + let (mut first, _) = BlockWriter::new(DIMENSION, None, &offered[..10]); + for (key, vector) in &pushed[..10] { + first.push(key, vector); + } + let first = first.finish(); + reads_back(&first, &pushed[..10]); + let (mut second, room) = BlockWriter::new(DIMENSION, Some(&first), &offered[10..15]); + assert!(!second.is_empty()); + for (at, (key, vector)) in pushed[10..15].iter().enumerate() { + assert!(at < room); + second.push(key, vector); + } + let second = second.finish(); + reads_back(&second, &pushed[..15]); + let (mut at_once, _) = BlockWriter::new(DIMENSION, None, &offered[..15]); + for (key, vector) in &pushed[..15] { + at_once.push(key, vector); + } + assert!( + second == at_once.finish(), + "a block written in two goes is the same block" + ); + + // A writer with no slot takes any key, so that a row is never left without a block. + let long = Key::from_values(&[Value::Text("k".repeat(BLOCK_BYTES))], &[]).expect("a key"); + assert_eq!(takes(DIMENSION, 4, &[long]).0, 1); +} + +#[test] +fn a_partition_record_reads_back() { + let partitions = [ + Partition { + number: 0, + live: 1, + blocks: 0, + }, + Partition { + number: 7, + live: 300, + blocks: 2, + }, + Partition { + number: 0x0102_0304, + live: 0x0506_0708_090A_0B0C, + blocks: 0x0D0E_0F10, + }, + Partition { + number: u32::MAX, + live: u64::MAX, + blocks: u32::MAX, + }, + ]; + for partition in partitions { + // Sixteen bytes: the number, the live vectors and the blocks, low byte first. + let record = partition.encode(); + let mut wanted = partition.number.to_le_bytes().to_vec(); + wanted.extend(partition.live.to_le_bytes()); + wanted.extend(partition.blocks.to_le_bytes()); + assert_eq!(record, wanted); + assert_eq!(Partition::decode(&record), partition); + } +} + +#[test] +fn a_pending_record_of_a_row_with_no_vector_is_reported_as_damage() { + // Which of three pending rows is left with no vector: the first, which a cycle reads to + // find the partition, and a later one, which it reads to write a slot. + for damaged in [1_i64, 2] { + let folder = std::env::temp_dir().join(format!( + "contextdb-pending-without-vector-{damaged}-{}", + std::process::id() + )); + if folder.exists() { + std::fs::remove_dir_all(&folder).expect("clear the folder an earlier run left"); + } + // A sound store: three rows of one partition, each with a vector and each waiting on + // the pending list. + let db = Database::open(&folder).expect("open a fresh store"); + for sql in [ + "CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(4))", + "INSERT INTO docs (id, v) VALUES (1, [1.0, 0.0, 0.0, 0.0]), \ + (2, [0.0, 1.0, 0.0, 0.0]), (3, [0.0, 0.0, 1.0, 0.0])", + ] { + db.execute(sql, &HashMap::new()).expect(sql); + } + db.close().expect("close"); + + // The damage: one row is written again with no vector by the store alone. The vector + // part is not asked what that means, so the row's pending record stays. + let budget = Budget::new(None).expect("a budget"); + let store = Store::open_for_writing(&folder, &budget, None).expect("open the store"); + let view = store.view().expect("the newest commit"); + let catalog = Catalog::load(&view).expect("the tables"); + let table = catalog.table("docs").expect("the table").id; + let mut changes = ChangeSet::new(view, &budget).expect("a change set"); + let key = Key::from_values(&[Value::Int64(damaged)], &[]).expect("a key"); + let row = vec![Value::Int64(damaged), Value::Null]; + changes.put_row(table, key, row, false).expect("stage"); + store.commit(changes).expect("commit the damage"); + + // A cycle over it is told the store is damaged, by name and naming the folder. It + // does not panic, and it leaves nothing staged and nothing counted. + let view = store.view().expect("the newest commit"); + let mut changes = ChangeSet::new(view, &budget).expect("a change set"); + let counted = budget.used(); + match seal(&mut changes, &catalog, &budget, 4096) { + Err(Error::StoreCorrupted { path, .. }) => assert_eq!(path, folder), + other => panic!( + "with no vector in row {damaged}, which has a pending record, a cycle gave \ + {other:?}" + ), + } + assert!(changes.is_empty(), "the refused cycle left records staged"); + assert_eq!(budget.used(), counted, "memory counted after the refusal"); + drop(changes); + drop(store); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); + } +} + +/// What a test of a wrongly named VECTOR column works on: the folder of its store, the store, +/// its budget, and the number of the table `docs`. +type Misnamed = (PathBuf, Store, Budget, TableId); + +/// A store whose schema record names, as the VECTOR column of `docs`, the column `named`. +/// +/// The store is written sound through the public calls: `docs` with the key `id`, an INTEGER +/// column `n` and the VECTOR column `v`, and three rows, each with a vector waiting on the +/// pending list. Then the store alone writes the table's schema record again with one byte +/// changed: the name of the column the record's VECTOR entry is for. No declaration writes such +/// a record. +fn a_vector_entry_named(test: &str, named: u8) -> Misnamed { + let folder = std::env::temp_dir().join(format!("contextdb-{test}-{}", std::process::id())); + if folder.exists() { + std::fs::remove_dir_all(&folder).expect("clear the folder an earlier run left"); + } + let db = Database::open(&folder).expect("open a fresh store"); + for sql in [ + "CREATE TABLE docs (id INTEGER PRIMARY KEY, n INTEGER, v VECTOR(4))", + "INSERT INTO docs (id, n, v) VALUES (1, 10, [1.0, 0.0, 0.0, 0.0]), \ + (2, 20, [0.0, 1.0, 0.0, 0.0]), (3, 30, [0.0, 0.0, 1.0, 0.0])", + ] { + db.execute(sql, &HashMap::new()).expect(sql); + } + db.close().expect("close"); + + let budget = Budget::new(None).expect("a budget"); + let store = Store::open_for_writing(&folder, &budget, None).expect("open the store"); + let view = store.view().expect("the newest commit"); + let table = Catalog::load(&view).expect("the tables"); + let table = table.table("docs").expect("the table").id; + let mut records = view.schema_records().expect("the schema records"); + let mut record = records.pop().expect("the record of the one table"); + assert!(records.is_empty(), "one table has one schema record"); + // A schema record ends with its VECTOR entries, and the one entry here ends with its + // column's name, a text of one byte, and then nine bytes: no PARTITION_KEY column, no + // MAX_PARTITIONS, no SEARCH_MODE. + let name = record.len() - 10; + assert_eq!( + record[name - 4..], + [1, 0, 0, 0, b'v', 0, 0, 0, 0, 0, 0, 0, 0, 0] + ); + record[name] = named; + let mut changes = ChangeSet::new(view, &budget).expect("a change set"); + changes.put_schema("docs", record).expect("stage"); + store.commit(changes).expect("commit the damage"); + (folder, store, budget, table) +} + +/// Asserts that a call over such a store told the store's damage, by name and naming the +/// folder. `what` says which call. +fn told_as_damage(given: crate::types::Result, folder: &Path, what: &str) { + match given { + Err(Error::StoreCorrupted { path, .. }) => assert_eq!(path, folder, "{what}"), + other => panic!("{what} gave {other:?}"), + } +} + +/// A commit that puts a fourth row with a vector into `docs` reaches the vector part, which is +/// asked what the row means for the column. Damage may be told when the tables are read or +/// when the vector part is asked; either way it is damage, and nothing panics. +fn a_write_is_reported_as_damage(test: &str, named: u8, what: &str) { + let (folder, store, budget, table) = a_vector_entry_named(test, named); + let view = store.view().expect("the newest commit"); + let mut changes = ChangeSet::new(view.clone(), &budget).expect("a change set"); + let key = Key::from_values(&[Value::Int64(4)], &[]).expect("a key"); + let numbers = Value::Vector(vec![0.0, 0.0, 0.0, 1.0]); + let row = vec![Value::Int64(4), Value::Int64(40), numbers]; + changes.put_row(table, key, row, true).expect("stage"); + let staged = Catalog::load(&view).and_then(|catalog| stage(&mut changes, &catalog, &budget)); + told_as_damage(staged, &folder, what); + drop(changes); + drop((view, store)); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +/// A sealing cycle over the three pending vectors of `docs`. Damage may be told when the +/// tables are read or by the cycle; either way it is damage, nothing panics, and the cycle +/// leaves nothing staged and nothing counted. +fn a_cycle_is_reported_as_damage(test: &str, named: u8, what: &str) { + let (folder, store, budget, _) = a_vector_entry_named(test, named); + let view = store.view().expect("the newest commit"); + let mut changes = ChangeSet::new(view.clone(), &budget).expect("a change set"); + let counted = budget.used(); + let sealed = + Catalog::load(&view).and_then(|catalog| seal(&mut changes, &catalog, &budget, 4096)); + told_as_damage(sealed, &folder, what); + assert!(changes.is_empty(), "the refused cycle left records staged"); + assert_eq!(budget.used(), counted, "memory counted after the refusal"); + drop(changes); + drop((view, store)); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +#[test] +fn a_write_to_a_vector_column_that_is_no_column_of_its_table_is_reported_as_damage() { + let what = "a commit of a row, with a VECTOR entry for the column `w`, which `docs` has not,"; + a_write_is_reported_as_damage("vector-entry-no-column-write", b'w', what); +} + +#[test] +fn a_cycle_over_a_vector_column_that_is_no_column_of_its_table_is_reported_as_damage() { + let what = "a cycle, with a VECTOR entry for the column `w`, which `docs` has not,"; + a_cycle_is_reported_as_damage("vector-entry-no-column-cycle", b'w', what); +} + +#[test] +fn a_write_to_a_vector_column_that_is_not_of_the_vector_type_is_reported_as_damage() { + let what = "a commit of a row, with a VECTOR entry for the INTEGER column `n`,"; + a_write_is_reported_as_damage("vector-entry-not-vector-write", b'n', what); +} + +#[test] +fn a_cycle_over_a_vector_column_that_is_not_of_the_vector_type_is_reported_as_damage() { + let what = "a cycle, with a VECTOR entry for the INTEGER column `n`,"; + a_cycle_is_reported_as_damage("vector-entry-not-vector-cycle", b'n', what); +} diff --git a/v2/contextdb/src/write/mod.rs b/v2/contextdb/src/write/mod.rs index 4b7a6425..c1e12b74 100644 --- a/v2/contextdb/src/write/mod.rs +++ b/v2/contextdb/src/write/mod.rs @@ -24,6 +24,7 @@ use crate::sql::Plan; use crate::sql::plan::{DeletePlan, InsertPlan, UpdatePlan}; use crate::store::{Cell, ChangeSet, IndexId, Key, Read, View}; use crate::types::{ColumnDefault, Error, Result, SortDirection, Value}; +use crate::vector::Tally; /// Applies one INSERT, UPDATE or DELETE to the change set and returns the number of rows: those /// an INSERT stored, or those the WHERE of an UPDATE or DELETE matched. @@ -36,6 +37,11 @@ use crate::types::{ColumnDefault, Error, Result, SortDirection, Value}; /// a mark on the change set before the call and rewinds to it on refusal. `now` is the /// statement's time in milliseconds: what `DEFAULT NOW()` stores and `NOW()` answers. /// +/// `partitions` is the transaction's tally of vector partitions. Every row staged in a table +/// that has a VECTOR column with a PARTITION_KEY is noted in it, and the statement's last act +/// is to settle it: a statement that would give a column one partition too many is refused +/// there, like any other refusal of this function. +/// /// # Panics /// /// When the plan is not an INSERT, UPDATE or DELETE. @@ -44,13 +50,14 @@ pub(crate) fn apply( params: &HashMap, catalog: &Catalog, changes: &mut ChangeSet, + partitions: &mut Tally, budget: &Budget, now: i64, ) -> Result { match plan { - Plan::Insert(plan) => insert(plan, params, catalog, changes, now), - Plan::Update(plan) => update(plan, params, catalog, changes, budget, now), - Plan::Delete(plan) => delete(plan, params, catalog, changes, budget, now), + Plan::Insert(plan) => insert(plan, params, catalog, changes, partitions, now), + Plan::Update(plan) => update(plan, params, catalog, changes, partitions, budget, now), + Plan::Delete(plan) => delete(plan, params, catalog, changes, partitions, budget, now), other => unreachable!("{other:?} writes no rows"), } } @@ -147,6 +154,7 @@ fn insert( params: &HashMap, catalog: &Catalog, changes: &mut ChangeSet, + partitions: &mut Tally, now: i64, ) -> Result { let table = find(catalog, &plan.table)?; @@ -186,9 +194,10 @@ fn insert( } let row = row.into_iter().enumerate(); let row = row.map(|(at, value)| rules.settle(at, value)); - rules.put(changes, None, row.collect::>()?)?; + rules.put(changes, partitions, None, row.collect::>()?)?; } changes.add_count(table.id, plan.values.len() as i64)?; + rules.settle_partitions(partitions)?; Ok(plan.values.len() as u64) } @@ -197,6 +206,7 @@ fn update( params: &HashMap, catalog: &Catalog, changes: &mut ChangeSet, + partitions: &mut Tally, budget: &Budget, now: i64, ) -> Result { @@ -223,9 +233,10 @@ fn update( // A row that already holds what it is set to is left alone: a statement that changes // nothing stages nothing. if row != *old { - rules.put(changes, Some((key, old)), row)?; + rules.put(changes, partitions, Some((key, old)), row)?; } } + rules.settle_partitions(partitions)?; Ok(found.len() as u64) } @@ -234,6 +245,7 @@ fn delete( params: &HashMap, catalog: &Catalog, changes: &mut ChangeSet, + partitions: &mut Tally, budget: &Budget, now: i64, ) -> Result { @@ -248,6 +260,7 @@ fn delete( now, )?; for (key, row) in &found { + rules.note(partitions, changes, Some(key), None)?; for index in &rules.indexes { changes.delete_index(index.id, &index.entry(row)?, key)?; } @@ -257,6 +270,7 @@ fn delete( if !found.is_empty() { changes.add_count(table.id, -(found.len() as i64))?; } + rules.settle_partitions(partitions)?; Ok(found.len() as u64) } @@ -286,6 +300,9 @@ struct Rules<'a> { /// behind the primary key (`__pk_...`) is left out: it holds none, a row is kept under its /// key. indexes: Vec, + /// The table has a VECTOR column with a PARTITION_KEY: the rows a statement stages in it + /// are noted in the transaction's tally of partitions. + partitioned: bool, } /// One index of a table as a write uses it. @@ -320,13 +337,44 @@ impl<'a> Rules<'a> { .then(|| names.join(", ")), }) }); + let mut vectors = table.vectors.iter(); Ok(Rules { table, key, indexes: indexes.collect::>()?, + partitioned: vectors.any(|vector| !vector.partition_key.is_empty()), }) } + /// Notes in the tally a row the statement is about to stage: `was` is the key the + /// transaction has had the row under until now, `None` for a new row, and `after` the + /// values it is staged with, `None` for a removal. + fn note( + &self, + partitions: &mut Tally, + changes: &ChangeSet, + was: Option<&Key>, + after: Option<&[Value]>, + ) -> Result<()> { + if !self.partitioned { + return Ok(()); + } + let before = match was { + Some(key) => changes.row(self.table.id, key)?, + None => None, + }; + partitions.note(self.table, before.as_ref(), after, changes) + } + + /// The last act of a statement, once every row is staged and every other rule has passed: + /// refuses a statement that gave a column one partition too many. + fn settle_partitions(&self, partitions: &mut Tally) -> Result<()> { + if self.partitioned { + partitions.settle(self.table)?; + } + Ok(()) + } + /// Brings a value written into a column to the column's type, and refuses NULL where the /// column needs a value: a NOT NULL column and the primary key. fn settle(&self, at: usize, value: Value) -> Result { @@ -346,13 +394,15 @@ impl<'a> Rules<'a> { /// /// `was` is the key and the values the row had, for an UPDATE: what the row keeps from /// before is its own and is not looked up, an entry that stays as it was is left alone, and - /// a row whose key changed leaves its old key. + /// a row whose key changed leaves its old key. The row is noted in `partitions` first. fn put( &self, changes: &mut ChangeSet, + partitions: &mut Tally, was: Option<(&Key, &[Value])>, row: Vec, ) -> Result<()> { + self.note(partitions, changes, was.map(|(key, _)| key), Some(&row))?; // No directions: a row's key is encoded ascending. let key = Key::from_values(&pick(&self.key, &row), &[])?; // A new row, or a row whose key changed: the transaction must find no row at this key. diff --git a/v2/contextdb/src/write/values.rs b/v2/contextdb/src/write/values.rs index c1111c5e..0396a1c7 100644 --- a/v2/contextdb/src/write/values.rs +++ b/v2/contextdb/src/write/values.rs @@ -3,13 +3,15 @@ use std::ops::Range; -use crate::types::{ColumnDef, ColumnType, Error, Result, Uuid, Value}; +use crate::query::vector_from_text; +use crate::types::{ColumnDef, ColumnType, Error, Result, Uuid, Value, VectorIndexRef}; /// Brings a value to its column's type, or refuses it with `ColumnTypeMismatch`. /// /// NULL fits every type here; whether the column takes it is the NOT NULL rule's business. The /// only conversions are quoted text into a UUID, TIMESTAMP, JSON or VECTOR column and an INTEGER -/// into a REAL column. +/// into a REAL column. A vector of the column's length with a number that is not finite is +/// refused with `InvalidVectorValue`, however it was written: no search could measure it. pub(super) fn typed(table: &str, column: &ColumnDef, value: Value) -> Result { let converted = match (column.column_type, &value) { (_, Value::Null) @@ -21,7 +23,7 @@ pub(super) fn typed(table: &str, column: &ColumnDef, value: Value) -> Result return Ok(value), (ColumnType::Vector(length), Value::Vector(vector)) if vector.len() == length => { - return Ok(value); + return finite(table, column, value); } (ColumnType::Real, Value::Int64(whole)) => Some(Value::Float64(*whole as f64)), (ColumnType::Uuid, Value::Text(text)) => Uuid::parse_str(text).ok().map(Value::Uuid), @@ -32,12 +34,29 @@ pub(super) fn typed(table: &str, column: &ColumnDef, value: Value) -> Result None, }; - converted.ok_or_else(|| Error::ColumnTypeMismatch { + let converted = converted.ok_or_else(|| Error::ColumnTypeMismatch { table: table.to_owned(), column: column.name.clone(), expected: format!("{:?}", column.column_type).to_uppercase(), actual: kind(&value), - }) + })?; + finite(table, column, converted) +} + +/// The value itself, unless it is a vector that holds a number that is not finite. +fn finite(table: &str, column: &ColumnDef, value: Value) -> Result { + if let Value::Vector(vector) = &value + && let Some(position) = vector.iter().position(|number| !number.is_finite()) + { + return Err(Error::InvalidVectorValue { + index: VectorIndexRef { + table: table.to_owned(), + column: column.name.clone(), + }, + position, + }); + } + Ok(value) } /// What a refusal calls the value it was given. A vector says how long it is, since its length @@ -132,12 +151,3 @@ fn digits(text: &str, positions: Range) -> Option { None } } - -/// Reads the text spelling of a vector, `[0.1, 0.2, 0.3]`. -fn vector_from_text(text: &str) -> Option> { - let inside = text.trim().strip_prefix('[')?.strip_suffix(']')?; - inside - .split(',') - .map(|number| number.trim().parse().ok()) - .collect() -} diff --git a/v2/contextdb/tests/counted_work.rs b/v2/contextdb/tests/counted_work.rs index b32a63d7..db6053a5 100644 --- a/v2/contextdb/tests/counted_work.rs +++ b/v2/contextdb/tests/counted_work.rs @@ -188,3 +188,87 @@ fn a_sort_the_index_already_gives_reads_only_the_limit() { ); assert_eq!(sorted.rows_examined(), 1000); } + +#[test] +fn a_negative_number_in_a_filter_narrows_the_work_as_a_positive_one_does() { + const DIMENSION: usize = 8; + const EACH: i64 = 50; + let folder = support::folder("negative-number-narrows"); + let db = support::open(&folder); + // Two tables of two hundred rows, a quarter of them under each of the numbers -1, 0, 1 + // and 2. In `marks` the number is the PARTITION_KEY of a VECTOR column, and no vector is + // sealed; in `plain` it is a column with an index. + support::run( + &db, + &format!( + "CREATE TABLE marks (id INTEGER PRIMARY KEY, kind INTEGER NOT NULL, \ + v VECTOR({DIMENSION}) PARTITION_KEY (kind))" + ), + ); + support::run( + &db, + "CREATE TABLE plain (id INTEGER PRIMARY KEY, k INTEGER NOT NULL)", + ); + support::run(&db, "CREATE INDEX plain_k ON plain (k)"); + let tx = db.transaction().expect("begin"); + for n in 1..=4 * EACH { + let number = n % 4 - 1; + tx.execute( + "INSERT INTO marks (id, kind, v) VALUES ($id, $kind, $v)", + ¶ms! { "id" => n, "kind" => number, "v" => support::vector(n, DIMENSION) }, + ) + .expect("insert a mark"); + tx.execute( + "INSERT INTO plain (id, k) VALUES ($id, $k)", + ¶ms! { "id" => n, "k" => number }, + ) + .expect("insert a plain row"); + } + tx.commit().expect("commit"); + + // A search of one partition reads that partition's pending rows once and the rows it + // returns once more. Returns the rows it read. + let stored = support::stored(&db, "SELECT id, kind, v FROM marks"); + let query = support::vector(99_000, DIMENSION); + let searched = |number: i64| -> u64 { + let wanted = stored.nearest(|row| row[1] == Value::Int64(number), &query, 10); + assert_eq!(wanted.len(), 10); + let filter = format!("kind = {number}"); + let sql = support::nearest_sql("marks", "v", Some(&filter), 10); + let what = format!("the partition of kind {number}"); + support::assert_nearest(&db, &sql, &query, &wanted, &what).rows_examined() + }; + // A lookup through the index reads the rows it returns. Returns the rows it read. + let looked_up = |number: i64| -> u64 { + let found = support::run( + &db, + &format!("SELECT id FROM plain WHERE k = {number} ORDER BY id"), + ); + let wanted: Vec> = (1..=4 * EACH) + .filter(|n| n % 4 - 1 == number) + .map(|n| vec![Value::Int64(n)]) + .collect(); + support::assert_same_rows( + found.rows(), + &wanted, + &format!("the rows with k = {number}"), + ); + found.rows_examined() + }; + + // A number written without a sign narrows both, and so does zero. + let one_partition = EACH as u64 + 10; + for number in [0, 1, 2] { + assert_eq!(searched(number), one_partition, "WHERE kind = {number}"); + assert_eq!(looked_up(number), EACH as u64, "WHERE k = {number}"); + } + // A number written with a minus sign is the same kind of value and narrows the same. + assert_eq!( + (searched(-1), looked_up(-1)), + (one_partition, EACH as u64), + "the rows read for WHERE kind = -1 by a search of {} rows, and for WHERE k = -1 through \ + the index of a table of {} rows", + 4 * EACH, + 4 * EACH + ); +} diff --git a/v2/contextdb/tests/crash/main.rs b/v2/contextdb/tests/crash/main.rs index 80547beb..37b6395e 100644 --- a/v2/contextdb/tests/crash/main.rs +++ b/v2/contextdb/tests/crash/main.rs @@ -1,6 +1,8 @@ //! Kill and reopen: a writer killed at any storage call of a commit leaves the whole commit or -//! none of it, and the store opens again with nothing to mend by hand. Damage in the file is -//! refused as damage, by name. +//! none of it, and the store opens again with nothing to mend by hand. That holds for a commit +//! of rows with vectors and for the commit of a maintenance cycle that seals vectors: after +//! either, every search finds exactly the vectors of the rows that are there. Damage in the +//! file is refused as damage, by name. //! //! A kill lands at a point the child names. The child's storage door counts the writes and //! syncs of its data file; before call number `stop_at` it announces `stopped` and waits there, @@ -274,11 +276,13 @@ fn a_commit_on_disk_is_reported_and_served_though_the_read_after_it_fails() { Ok(()) } })); + // Under a memory limit of 32 MiB the store's page cache is 8 MiB: a quarter of the limit. + options.memory_limit = Some(32 << 20); let db = support::open_with(&folder, options).expect("open"); support::declare(&db); - // A second table four times the size of the store's 8 MiB page cache, in rows small enough - // to be kept in the table's own pages. Reading all of it leaves the cache full of its - // pages, so the pages a commit writes afterwards have to be read back from the file. + // A second table four times the size of that 8 MiB page cache, in rows small enough to be + // kept in the table's own pages. Reading all of it leaves the cache full of its pages, so + // the pages a commit writes afterwards have to be read back from the file. support::run( &db, "CREATE TABLE padding (id INTEGER PRIMARY KEY, body TEXT)", @@ -370,6 +374,213 @@ fn answers_with_rows_or_store_io(db: &Database, rows: i64, when: &str) { } } +// ---- vectors ---- + +/// The numbers in a vector of the two tests below. +const DIMENSION: usize = 16; + +/// A closed store with `docs`, its VECTOR column partitioned by context: rows 1 to 100 are in +/// blocks and rows 101 to 300 on pending lists, half of each under `c0` and half under `c1`. +fn store_with_vectors(name: &str) -> PathBuf { + let folder = support::folder(name); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, "PARTITION_KEY (ctx)"); + for numbers in [1..=100, 101..=300] { + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, numbers, DIMENSION, |n| format!("c{}", n % 2)); + tx.commit().expect("commit"); + // The first hundred are sealed before the rest arrive. + if support::int(&db, "SELECT COUNT(*) FROM docs") == 100 { + assert_eq!(support::seal_all(&db), 100); + } + } + support::close(db); + folder +} + +/// Asserts that every search of `docs` returns exactly what its rows call for: everything and +/// each context, ten rows and more rows than there are, for three vectors. Returns the rows. +fn searches_agree(db: &Database, state: &str) -> Vec> { + let stored = support::stored(db, "SELECT id, ctx, tag, v FROM docs ORDER BY id"); + let queries = [ + support::vector(98_001, DIMENSION), + support::vector(98_002, DIMENSION), + support::vector(150, DIMENSION), + ]; + for query in &queries { + for ctx in [None, Some("c0"), Some("c1"), Some("new"), Some("newer")] { + let filter = ctx.map(|ctx| format!("ctx = '{ctx}'")); + for k in [10, 1000] { + let passes = |row: &[Value]| ctx.is_none_or(|ctx| row[1] == support::text(ctx)); + let wanted = stored.nearest(passes, query, k); + let sql = support::nearest_sql("docs", "v", filter.as_deref(), k); + support::assert_nearest(db, &sql, query, &wanted, state); + } + } + } + stored.rows().to_vec() +} + +#[test] +fn a_kill_at_any_storage_call_of_a_seal_leaves_every_vector_found() { + if let Some(role) = support::role() { + return one_seal(role); + } + let test = "a_kill_at_any_storage_call_of_a_seal_leaves_every_vector_found"; + // Run once to the end to learn which storage calls belong to the sealing commit. + let folder = store_with_vectors("kill-seal-count"); + let rows = searches_agree(&support::open(&folder), "before the cycle"); + assert_eq!(rows.len(), 300); + let lines = Child::start(test, "sealer stop_at=0", &folder).end(); + let begins = support::last_number(&lines, "sealing").expect("the sealer announces it"); + let calls = support::last_number(&lines, "calls").expect("the sealer announces its calls"); + assert!( + lines.iter().any(|line| line == "sealed 200 0"), + "one cycle was to seal the 200 pending vectors: {lines:?}" + ); + assert!( + calls > begins, + "a cycle that sealed 200 vectors made no write and no sync" + ); + + for k in begins + 1..=calls { + let folder = store_with_vectors(&format!("kill-seal-at-{k}")); + let mut sealer = Child::start(test, &format!("sealer stop_at={k}"), &folder); + let before = sealer.until("stopped"); + sealer.kill(); + let returned = before.iter().any(|line| line.starts_with("sealed ")); + + // The store opens with nothing to mend, every row is there, and every search finds + // exactly the vectors of those rows, whether the seal's commit is there or not. + let db = support::open(&folder); + let state = format!("killed at storage call {k} of a seal"); + assert!( + searches_agree(&db, &state) == rows, + "{state}: the rows of docs" + ); + // The next cycles finish the sealing: all 200 vectors when the killed commit is not + // there, none when it is. A cycle that had returned is there. + let sealed = support::seal_all(&db); + assert!( + sealed == 0 || sealed == 200 && !returned, + "{state}: the next cycles sealed {sealed} vectors; the killed cycle had {}returned", + if returned { "" } else { "not " } + ); + let report = db.run_maintenance_cycle().expect("one more cycle").vector; + assert_eq!( + ( + report.maintenance_vectors_done, + report.maintenance_vectors_remaining + ), + (0, 0), + "{state}: a cycle after the sealing was finished" + ); + assert!( + searches_agree(&db, &state) == rows, + "{state}: the rows after sealing" + ); + } +} + +/// Runs one maintenance cycle. Announces `sealing ` with the storage calls made so far when +/// the cycle is about to begin, `sealed ` when it has returned, and +/// `calls ` once the store is closed. +fn one_seal(role: support::Role) { + let (storage, calls) = stopping_before_call(role.number("stop_at")); + let mut options = DatabaseOpenOptions::default(); + options.storage = Some(storage); + let db = support::open_with(&role.folder, options).expect("open"); + support::say(format!("sealing {}", calls.load(Ordering::SeqCst))); + let report = db.run_maintenance_cycle().expect("a cycle").vector; + support::say(format!( + "sealed {} {}", + report.maintenance_vectors_done, report.maintenance_vectors_remaining + )); + support::close(db); + support::say(format!("calls {}", calls.load(Ordering::SeqCst))); +} + +#[test] +fn a_kill_inside_a_commit_of_vectors_keeps_rows_and_vectors_together() { + if let Some(role) = support::role() { + return one_commit_of_vectors(role); + } + let test = "a_kill_inside_a_commit_of_vectors_keeps_rows_and_vectors_together"; + // The rows before the commit, from a store nothing was done to, and the rows after it, + // from a run of the commit to its end, which also tells which storage calls are its own. + let untouched = store_with_vectors("kill-vectors-before"); + let before = searches_agree(&support::open(&untouched), "before the commit"); + let folder = store_with_vectors("kill-vectors-count"); + let lines = Child::start(test, "writer stop_at=0", &folder).end(); + let begins = support::last_number(&lines, "committing").expect("the writer announces it"); + let calls = support::last_number(&lines, "calls").expect("the writer announces its calls"); + assert!( + calls > begins, + "a commit of sixty vectors made no write and no sync" + ); + let after = searches_agree(&support::open(&folder), "after the commit"); + assert_eq!((before.len(), after.len()), (300, 358)); + + for k in begins + 1..=calls { + let folder = store_with_vectors(&format!("kill-vectors-at-{k}")); + let mut writer = Child::start(test, &format!("writer stop_at={k}"), &folder); + let told = writer.until("stopped"); + writer.kill(); + let returned = told.iter().any(|line| line == "committed"); + + // The rows are those of before the commit or those of after it, and the searches + // answer from the vectors of exactly those rows: never rows of one and vectors of the + // other. A commit that had returned is there. + let db = support::open(&folder); + let state = format!("killed at storage call {k} of a commit of vectors"); + let rows = searches_agree(&db, &state); + assert!( + rows == after || rows == before && !returned, + "{state}: {} rows are there, of neither state or of the wrong one; the commit had \ + {}returned", + rows.len(), + if returned { "" } else { "not " } + ); + // The store goes on from there: what waits is sealed, and the answers stay. + support::seal_all(&db); + assert!( + searches_agree(&db, &state) == rows, + "{state}: the rows after sealing" + ); + } +} + +/// In one transaction adds rows 301 to 360, a third of them under a context that has rows and +/// the rest under two that have none; changes the vector of a sealed row and of a pending row; +/// deletes one of each; and moves a sealed row to one of the new contexts. Announces +/// `committing ` with the storage calls made so far when the commit is about to begin, +/// `committed` when it has returned, and `calls ` once the store is closed. +fn one_commit_of_vectors(role: support::Role) { + let (storage, calls) = stopping_before_call(role.number("stop_at")); + let mut options = DatabaseOpenOptions::default(); + options.storage = Some(storage); + let db = support::open_with(&role.folder, options).expect("open"); + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, 301..=360, DIMENSION, |n| { + ["c0", "new", "newer"][(n % 3) as usize].to_owned() + }); + for id in [5_i64, 150] { + tx.execute( + "UPDATE docs SET v = $v WHERE id = $id", + ¶ms! { "id" => id, "v" => support::vector(98_100 + id, DIMENSION) }, + ) + .expect("change a vector"); + } + support::run(&tx, "DELETE FROM docs WHERE id = 6"); + support::run(&tx, "DELETE FROM docs WHERE id = 151"); + support::run(&tx, "UPDATE docs SET ctx = 'new' WHERE id = 7"); + support::say(format!("committing {}", calls.load(Ordering::SeqCst))); + tx.commit().expect("commit"); + support::say("committed"); + support::close(db); + support::say(format!("calls {}", calls.load(Ordering::SeqCst))); +} + // ---- shared by the tests of this file ---- /// A closed store with the table declared and rows 1 to `rows` committed. diff --git a/v2/contextdb/tests/large_values.rs b/v2/contextdb/tests/large_values.rs new file mode 100644 index 00000000..8dfc170f --- /dev/null +++ b/v2/contextdb/tests/large_values.rs @@ -0,0 +1,752 @@ +//! Large text and JSON values: one that the store keeps beside its row reads back byte for +//! byte whether or not the store keeps it compressed; text that repeats takes less room in +//! the data file than its own bytes and text that does not shrink takes no more than it did; +//! a value longer than 1 MiB and a vector are never compressed; the memory limit holds while +//! a compressed value is read and while one is written; and a damaged length is reported as +//! damage. +//! +//! Nothing here names compression, because nothing in the library does: a test sees it in +//! what a statement returns, in the length of the data file and in the memory counted. + +mod support; + +use std::fs; +use std::path::Path; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use contextdb::serde_json; +use contextdb::{ + Database, DatabaseOpenOptions, Error, ReadSession, ReadSessionOptions, Value, params, +}; + +use support::{Call, Child, Kind, Run}; + +const KIB: usize = 1 << 10; +const MIB: usize = 1 << 20; + +/// The longest text or JSON value the store compresses: 1 MiB. One byte more is stored as it +/// is, whatever it holds. +const LONGEST_COMPRESSED: usize = 1_048_576; + +// ---- shared by the tests of this file ---- + +fn open_limited(folder: &Path, limit: usize) -> Database { + let mut options = DatabaseOpenOptions::default(); + options.memory_limit = Some(limit); + match support::open_with(folder, options) { + Ok(db) => db, + Err(error) => panic!("an open under a limit of {limit} bytes was refused: {error:?}"), + } +} + +fn read_limited(folder: &Path, limit: Option) -> ReadSession { + let mut options = ReadSessionOptions::default(); + options.memory_limit = limit; + match ReadSession::open_with_options(folder, options) { + Ok(reader) => reader, + Err(error) => panic!("a reader under a limit of {limit:?} was refused: {error:?}"), + } +} + +/// The bytes `SHOW MEMORY_LIMIT` reports as used. +fn used(on: &impl Run) -> usize { + let report = support::run(on, "SHOW MEMORY_LIMIT"); + let at = report.columns().iter().position(|name| name == "used"); + match (at, report.rows()) { + (Some(at), [row]) => match row[at] { + Value::Int64(bytes) => bytes as usize, + ref other => panic!("SHOW MEMORY_LIMIT reports {other:?} as used"), + }, + _ => panic!("SHOW MEMORY_LIMIT did not return one row with a column `used`"), + } +} + +/// The length of the store's data file, in bytes. +fn data_file(folder: &Path) -> u64 { + fs::metadata(folder.join("data")) + .expect("the store's data file") + .len() +} + +/// The outcome of a statement that selects large values, with the rows left out of what a +/// failure prints: a megabyte of text does not belong in a test's output. +fn answered(on: &impl Run, sql: &str, id: i64) -> contextdb::Result>> { + on.run(sql, ¶ms! { "id" => id }) + .map(|result| result.rows().to_vec()) +} + +/// The table of every test below but the first: a key and one text. +const NOTES: &str = "CREATE TABLE notes (id INTEGER PRIMARY KEY, body TEXT NOT NULL)"; + +/// Inserts one row of `notes`. +fn insert_note(on: &impl Run, id: i64, body: &str) { + if let Err(error) = on.run( + "INSERT INTO notes (id, body) VALUES ($id, $body)", + ¶ms! { "id" => id, "body" => body }, + ) { + panic!("note {id}, of {} bytes, was refused: {error:?}", body.len()); + } +} + +/// Asserts that the body of one row of `notes` reads back as exactly this text. +fn assert_note_is(on: &impl Run, id: i64, wanted: &str, what: &str) { + let what = format!("{what}: the body of note {id}"); + match answered(on, "SELECT body FROM notes WHERE id = $id", id) { + Ok(rows) => match rows.as_slice() { + [row] => support::assert_text_is(&row[0], wanted, &what), + rows => panic!("{what}: {} rows were returned, not one", rows.len()), + }, + Err(error) => panic!("{what}: the select was refused: {error:?}"), + } +} + +/// A new store whose `notes` holds rows 1 to `rows`, the body of each `len` bytes of `kind`, +/// 500 rows to a transaction, closed. Returns the bytes put in: 8 for each id and the bytes +/// of each body. +fn store_of_notes(folder: &Path, kind: Kind, rows: i64, len: usize) -> u64 { + let db = support::open(folder); + support::run(&db, NOTES); + let mut put_in = 0_u64; + for first in (1..=rows).step_by(500) { + let tx = db.transaction().expect("begin"); + for id in first..(first + 500).min(rows + 1) { + let body = support::large_text(kind, id, len); + put_in += 8 + body.len() as u64; + insert_note(&tx, id, &body); + } + tx.commit().expect("commit five hundred notes"); + } + support::close(db); + put_in +} + +/// Asserts that every row [`store_of_notes`] wrote reads back equal, and that no other row is +/// there. +fn assert_notes_are(on: &impl Run, kind: Kind, rows: i64, len: usize, what: &str) { + assert_eq!( + support::int(on, "SELECT COUNT(*) FROM notes"), + rows, + "{what}" + ); + for id in 1..=rows { + assert_note_is(on, id, &support::large_text(kind, id, len), what); + } +} + +// ---- every kind and size, three ways ---- + +/// The sizes of the first test: the last text that stays in its row and the first that does +/// not, two between, and the first that is too long to be compressed. +const SIZES: [usize; 5] = [1015, 1016, 4000, 70_000, LONGEST_COMPRESSED + 1]; + +/// One row of `docs` for each kind of text at each size. +const DOCS: i64 = (Kind::ALL.len() * SIZES.len()) as i64; + +/// The kind of text row `id` of `docs` is first written with, and the size of its body and +/// of its meta. +fn shape(id: i64) -> (Kind, usize) { + let at = (id - 1) as usize; + (Kind::ALL[at % Kind::ALL.len()], SIZES[at / Kind::ALL.len()]) +} + +/// The body of row `id`: of the row's own kind when it is inserted, `changes` 0, and of the +/// next kind once it has been updated, `changes` 1. +fn body(id: i64, changes: usize) -> String { + let (kind, size) = shape(id); + support::large_text(kind.after(changes), id, size) +} + +/// The meta of row `id`, of the same size as its body: of the kind that shrinks in the rows +/// with an odd id and of the kind that does not in the others, so every size has both. +fn meta(id: i64) -> serde_json::Value { + support::large_json(id % 2 == 1, id, shape(id).1) +} + +/// The rows of `docs` that are there: all of them at first, and all but every third once the +/// bodies have been changed. +fn kept(changes: usize) -> Vec { + (1..=DOCS) + .filter(|id| changes == 0 || id % 3 != 0) + .collect() +} + +/// Asserts that `docs` holds exactly the rows it is to hold, and that the body and the meta +/// of each read back as a `Value::Text` and a `Value::Json` equal to what was put in. Returns +/// how many rows matched. +fn assert_docs_are(on: &impl Run, changes: usize, what: &str) -> u64 { + let kept = kept(changes); + let ids: Vec> = kept.iter().map(|id| vec![Value::Int64(*id)]).collect(); + support::assert_same_rows( + &support::rows(on, "SELECT id FROM docs ORDER BY id"), + &ids, + &format!("{what}: the ids of docs"), + ); + assert_eq!( + support::int(on, "SELECT COUNT(*) FROM docs"), + kept.len() as i64, + "{what}: the kept count of docs" + ); + for id in &kept { + let what = format!("{what}: row {id}, {:?}", shape(*id)); + match answered(on, "SELECT body, meta FROM docs WHERE id = $id", *id) { + Ok(rows) => match rows.as_slice() { + [row] => { + support::assert_text_is(&row[0], &body(*id, changes), &format!("{what}: body")); + support::assert_json_is(&row[1], &meta(*id), &format!("{what}: meta")); + } + rows => panic!("{what}: {} rows were returned, not one", rows.len()), + }, + Err(error) => panic!("{what}: the select was refused: {error:?}"), + } + } + kept.len() as u64 +} + +/// Runs one statement for each row of `docs`, four rows to a transaction. +fn four_to_a_transaction(db: &Database, each: impl Fn(&contextdb::Transaction<'_>, i64)) { + for first in (1..=DOCS).step_by(4) { + let tx = db.transaction().expect("begin"); + for id in first..(first + 4).min(DOCS + 1) { + each(&tx, id); + } + tx.commit().expect("commit four rows"); + } +} + +#[test] +fn large_text_and_json_come_back_byte_for_byte() { + if let Some(role) = support::role() { + return reader_of_docs(&role); + } + let test = "large_text_and_json_come_back_byte_for_byte"; + let folder = support::folder("large-values-come-back"); + // Read through the writer, through the writer after close and reopen, and through a + // reader in another process, which makes the same content from the ids. + let three_ways = |db: Database, changes: usize| -> Database { + assert_docs_are(&db, changes, "through the writer"); + support::close(db); + let db = support::open(&folder); + assert_docs_are(&db, changes, "after close and reopen"); + let role = format!("reader changes={changes}"); + let lines = Child::start(test, &role, &folder).end(); + assert_eq!( + support::last_number(&lines, "matched"), + Some(kept(changes).len() as u64), + "the rows a reader in another process found equal" + ); + db + }; + + let db = support::open(&folder); + support::run( + &db, + "CREATE TABLE docs (id INTEGER PRIMARY KEY, body TEXT, meta JSON)", + ); + four_to_a_transaction(&db, |tx, id| { + let values = params! { "id" => id, "body" => body(id, 0), "meta" => meta(id) }; + if let Err(error) = tx.execute( + "INSERT INTO docs (id, body, meta) VALUES ($id, $body, $meta)", + &values, + ) { + panic!("row {id}, {:?}, was refused: {error:?}", shape(id)); + } + }); + let db = three_ways(db, 0); + + // What a loaded value is counted as does not depend on how it is stored. Rows 13 and 15 + // are alike but for their bodies, 4,000 bytes each: one repeats and is stored compressed, + // the other is scattered and is stored as it is. A result that holds either row holds the + // same number of bytes. + assert_eq!( + (shape(13), shape(15)), + ((Kind::Repeating, 4000), (Kind::Scattered, 4000)) + ); + let held = |id: i64| { + let before = db.memory_used(); + let result = db.execute("SELECT * FROM docs WHERE id = $id", ¶ms! { "id" => id }); + let rows = result.as_ref().map_or(0, |result| result.rows().len()); + assert_eq!(rows, 1, "row {id} was not returned once"); + db.memory_used() - before + }; + let (compressed, plain) = (held(13), held(15)); + assert!( + compressed > 4000, + "a held row of 4,000 bytes was counted as {compressed} bytes" + ); + assert_eq!( + compressed, plain, + "a row whose body is stored compressed and one whose body is not are counted alike" + ); + + // Every body gives way to one of the next kind, so that a value stored compressed is + // replaced by one stored as it is and the other way round; then every third row goes. + four_to_a_transaction(&db, |tx, id| { + let values = params! { "id" => id, "body" => body(id, 1) }; + match tx.execute("UPDATE docs SET body = $body WHERE id = $id", &values) { + Ok(result) => assert_eq!(result.rows_affected(), 1, "the update of row {id}"), + Err(error) => panic!("the update of row {id} was refused: {error:?}"), + } + }); + let tx = db.transaction().expect("begin"); + for id in (3..=DOCS).step_by(3) { + let deleted = tx.execute("DELETE FROM docs WHERE id = $id", ¶ms! { "id" => id }); + match deleted { + Ok(result) => assert_eq!(result.rows_affected(), 1, "the delete of row {id}"), + Err(error) => panic!("the delete of row {id} was refused: {error:?}"), + } + } + tx.commit().expect("commit the deletes"); + assert_eq!(kept(1).len() as i64, DOCS - DOCS / 3); + three_ways(db, 1); +} + +/// Opens the store for reading, checks every row of `docs` against the content made from its +/// id, and announces `matched `. +fn reader_of_docs(role: &support::Role) { + let reader = ReadSession::open(&role.folder).expect("open the store for reading"); + let changes = role.number("changes") as usize; + let matched = assert_docs_are(&reader, changes, "a reader in another process"); + support::say(format!("matched {matched}")); +} + +// ---- the room large values take in the data file ---- + +#[test] +fn a_value_that_does_not_shrink_is_kept_as_it_is() { + // The bytes put in: 1,000 ids of 8 bytes and 1,000 bodies of 4,000 bytes. And the most + // bytes the data file may have: what the same store takes with no compression at all. + const PUT_IN: u64 = 4_008_000; + const RECORDED: u64 = PUT_IN * 5 / 4; // set from the first run + + let folder = support::folder("does-not-shrink"); + // Scattered text does not shrink by an eighth, so every body is stored byte for byte as + // it is. Nothing in the content, its order or the store depends on a clock, a random + // number or a thread, so the file has the same length on every machine and every run. + let put_in = store_of_notes(&folder, Kind::Scattered, 1000, 4000); + assert_eq!(put_in, PUT_IN, "the bytes this test put in"); + let length = data_file(&folder); + + let db = support::open(&folder); + assert_notes_are(&db, Kind::Scattered, 1000, 4000, "after close and reopen"); + assert!( + length <= RECORDED, + "the data file is {length} bytes long; the most recorded for this store is {RECORDED} \ + bytes, and {PUT_IN} bytes were put in" + ); +} + +#[test] +fn a_store_of_repeating_text_is_smaller_than_the_text() { + // The bytes put in: 1,500 ids of 8 bytes and 1,500 bodies of 6,000 bytes. And the most + // bytes the data file may have. + const PUT_IN: u64 = 9_012_000; + const RECORDED: u64 = PUT_IN - 1; // set from the first run + + let folder = support::folder("smaller-than-the-text"); + // Words from a list of 64 repeat, so every body is stored compressed, to about half. The + // length is the same on every machine and every run, as in the test above. + let put_in = store_of_notes(&folder, Kind::Words, 1500, 6000); + assert_eq!(put_in, PUT_IN, "the bytes this test put in"); + let length = data_file(&folder); + + let db = support::open(&folder); + assert_notes_are(&db, Kind::Words, 1500, 6000, "after close and reopen"); + assert!( + length < PUT_IN, + "the data file is {length} bytes long, which is not shorter than the {PUT_IN} bytes of \ + repeating text that were put in" + ); + assert!( + length <= RECORDED, + "the data file is {length} bytes long; the most recorded for this store is {RECORDED} \ + bytes, and {PUT_IN} bytes were put in" + ); +} + +#[test] +fn a_text_longer_than_one_mebibyte_and_a_vector_are_stored_as_they_are() { + // A little over 1 MiB of one sentence repeated. Compressed it would take a few kilobytes. + const OVER: usize = 1_100_000; + // What a reader's limit leaves beside the value: room for the row it comes in, and a + // quarter of what a second copy of the value would need. + const MARGIN: usize = 256 * KIB; + const { assert!(OVER > LONGEST_COMPRESSED && MARGIN < LONGEST_COMPRESSED / 2) }; + + // One byte over the longest value that is compressed, and a good deal over. + for len in [LONGEST_COMPRESSED + 1, OVER] { + // A limit that fits the value once and not twice: the 9 MiB an open sets aside, the + // value's stored record of 9 + len bytes, and the margin. + let once = 9 * MIB + 9 + len + MARGIN; + let folder = support::folder(&format!("longer-than-one-mebibyte-{len}")); + let body = support::large_text(Kind::Repeating, 1, len); + let db = support::open(&folder); + support::run(&db, NOTES); + insert_note(&db, 1, &body); + support::close(db); + // The file holds the value's own bytes: it is at least as long as the value. + let length = data_file(&folder); + assert!( + length >= len as u64, + "the data file is {length} bytes long: too short to hold a text of {len} bytes as \ + it is, so a text longer than {LONGEST_COMPRESSED} bytes was stored compressed" + ); + + // It reads back under a limit that fits it once. `SELECT *` hands the row over as it + // was read; a select list that names the column works its value out as a copy. + let reader = read_limited(&folder, Some(once)); + assert_eq!(used(&reader), 9 * MIB, "set aside by the open"); + match answered(&reader, "SELECT * FROM notes WHERE id = $id", 1) { + Ok(rows) => match rows.as_slice() { + [row] => { + assert_eq!(row[0], Value::Int64(1)); + support::assert_text_is(&row[1], &body, &format!("a text of {len} bytes")); + } + rows => panic!("{} rows were returned, not one", rows.len()), + }, + Err(error) => panic!( + "a text of {len} bytes was refused under a limit that fits it once: {error:?}" + ), + } + assert_eq!( + used(&reader), + 9 * MIB, + "memory in use once the row is dropped" + ); + } + + // A value of exactly 1 MiB of the same text is stored compressed: its data file is too + // short to hold it as it is. It reads back equal all the same. + let folder = support::folder("exactly-one-mebibyte"); + let body = support::large_text(Kind::Repeating, 1, LONGEST_COMPRESSED); + let db = support::open(&folder); + support::run(&db, NOTES); + insert_note(&db, 1, &body); + support::close(db); + let length = data_file(&folder); + assert!( + length < LONGEST_COMPRESSED as u64, + "the data file is {length} bytes long: a repeating text of {LONGEST_COMPRESSED} bytes \ + was stored as it is" + ); + assert_note_is(&read_limited(&folder, None), 1, &body, "exactly 1 MiB"); + + // A vector is never compressed, whatever it holds. A thousand vectors of 1,024 equal + // numbers are 4,096,000 bytes that would shrink to a few kilobytes; each is kept beside + // its row, and the file holds them as they are. + const NUMBERS: usize = 1024; + let folder = support::folder("vectors-as-they-are"); + let db = support::open(&folder); + support::run( + &db, + &format!("CREATE TABLE shapes (id INTEGER PRIMARY KEY, v VECTOR({NUMBERS}))"), + ); + let same = vec![0.5_f32; NUMBERS]; + for first in (1..=1000_i64).step_by(500) { + let tx = db.transaction().expect("begin"); + for id in first..first + 500 { + tx.execute( + "INSERT INTO shapes (id, v) VALUES ($id, $v)", + ¶ms! { "id" => id, "v" => same.clone() }, + ) + .unwrap_or_else(|error| panic!("shape {id} was refused: {error:?}")); + } + tx.commit().expect("commit five hundred shapes"); + } + support::close(db); + let (length, put_in) = (data_file(&folder), 1000 * 4 * NUMBERS as u64); + assert!( + length >= put_in, + "the data file is {length} bytes long: too short to hold {put_in} bytes of vectors as \ + they are, so a vector was stored compressed" + ); + let reader = read_limited(&folder, None); + let last = support::rows(&reader, "SELECT v FROM shapes WHERE id = 1000"); + assert!( + last == [[Value::Vector(same)]], + "the vector of shape 1000 did not read back as it was written" + ); +} + +// ---- the memory limit ---- + +/// A new store whose `notes` holds row 1, with this body, and row 2, with the body `small`, +/// each its own commit; closed. +fn store_of_one_large_note(folder: &Path, body: &str) { + let db = support::open(folder); + support::run(&db, NOTES); + insert_note(&db, 1, body); + insert_note(&db, 2, "small"); + support::close(db); +} + +#[test] +fn a_compressed_value_too_large_for_the_limit_is_refused() { + // The largest value the store compresses: 1 MiB of one sentence repeated. + const LEN: usize = LONGEST_COMPRESSED; + // The limit: the 9 MiB an open sets aside and 64 KiB. Reading the value asks for its + // stored record first, which for this text is a little over 4 KiB (an LZ4 block is never + // under a 255th of what it decodes to) and has room sixteen times over; then for 9 + LEN + // bytes, the value's real size, before any of its bytes is produced, and that is sixteen + // times what is left. So the select is refused after the stored bytes are in hand. + const ROOM: usize = 64 * KIB; + const LIMIT: usize = 9 * MIB + ROOM; + const { assert!(9 + LEN.div_ceil(255) < ROOM / 2 && ROOM < 9 + LEN) }; + + let folder = support::folder("compressed-over-the-limit"); + let body = support::large_text(Kind::Repeating, 1, LEN); + store_of_one_large_note(&folder, &body); + // The value is stored compressed: the file is too short to hold it as it is. + let length = data_file(&folder); + assert!( + length < LEN as u64, + "the data file is {length} bytes long: a repeating text of {LEN} bytes was stored as \ + it is, and this test would show nothing about a compressed value" + ); + + // Through a reader with the limit, then through a writer with it. + let refused = |on: &dyn Run, door: &str| { + let before = used(&on); + assert_eq!(before, 9 * MIB, "{door}: set aside by the open"); + // Asked twice: the first refusal left nothing counted, so the second is the same. + let mut refusals = Vec::new(); + for _ in 0..2 { + match answered(&on, "SELECT body FROM notes WHERE id = $id", 1) { + Err(Error::MemoryBudgetExceeded { + operation, + requested_bytes, + available_bytes, + budget_limit_bytes, + .. + }) => { + assert_eq!(operation, "query", "{door}"); + assert_eq!(budget_limit_bytes, LIMIT as u64, "{door}"); + refusals.push((requested_bytes, available_bytes)); + } + Err(other) => panic!("{door}: the select was refused as {other:?}"), + Ok(rows) => panic!("{door}: the select was answered with {} rows", rows.len()), + } + assert_eq!(used(&on), before, "{door}: memory in use after the refusal"); + } + assert_eq!( + refusals[0], refusals[1], + "{door}: a refused read left memory counted" + ); + // What was refused is the room for the value's real bytes, and by then the statement + // held the stored record. + let (requested, available) = refusals[0]; + assert_eq!( + requested, + 9 + LEN as u64, + "{door}: the bytes that were asked for and refused" + ); + let held = (ROOM as u64).checked_sub(available); + assert!( + held.is_some_and(|held| held >= 9 + (LEN as u64).div_ceil(255)), + "{door}: with {available} of {ROOM} bytes free, the statement did not hold the \ + stored record when it was refused" + ); + // The rest of the table answers in the same session. + assert_note_is(&on, 2, "small", door); + assert_eq!(support::int(&on, "SELECT COUNT(*) FROM notes"), 2, "{door}"); + assert_eq!(used(&on), before, "{door}: memory in use at the end"); + }; + refused(&read_limited(&folder, Some(LIMIT)), "a reader"); + let db = open_limited(&folder, LIMIT); + refused(&db, "the writer"); + assert_eq!(db.memory_used(), 9 * MIB); + support::close(db); + + // The store is unharmed: a reader with no limit returns the value byte for byte. + assert_note_is(&read_limited(&folder, None), 1, &body, "under no limit"); +} + +#[test] +fn a_commit_with_no_room_to_compress_is_written_plain() { + const LEN: usize = LONGEST_COMPRESSED; + // What the compressor works in for a value of LEN bytes: 20 bytes more than 1.1 times + // the value. The commit asks the budget for it and, refused, writes the value as it is. + const BUFFER: usize = LEN + LEN / 10 + 20; + // The room the limit leaves at the commit: enough for whatever else a commit takes, and + // less than half of the compressor's buffer. + const ROOM: usize = 512 * KIB; + const { assert!(BUFFER == 1_153_453 && ROOM + KIB < BUFFER / 2) }; + let body = support::large_text(Kind::Repeating, 1, LEN); + + // With no room. The row is staged under no limit; then the limit is set, through the + // transaction, to what is in use, the row included, rounded up to whole KiB, and ROOM. + let folder = support::folder("no-room-to-compress"); + let db = support::open(&folder); + support::run(&db, NOTES); + let tx = db.transaction().expect("begin"); + insert_note(&tx, 1, &body); + let kib = db.memory_used().div_ceil(KIB) + ROOM / KIB; + support::run(&tx, &format!("SET MEMORY_LIMIT '{kib}K'")); + assert_eq!(db.memory_limit(), Some(kib * KIB), "the limit was taken"); + let free = kib * KIB - db.memory_used(); + assert!( + (ROOM..ROOM + KIB).contains(&free), + "{free} bytes are free at the commit, where {ROOM} were to be" + ); + // Compression never refuses a commit that would be written without it. + if let Err(error) = tx.commit() { + panic!("a commit with {free} bytes free, too few to compress in, was refused: {error:?}"); + } + support::close(db); + // The value went in as it is: the file is longer than the value. + let length = data_file(&folder); + assert!( + length > LEN as u64, + "the data file is {length} bytes long: too short to hold a text of {LEN} bytes as it \ + is, so it was compressed in a buffer the budget had no room for" + ); + // Read by a reader, which has no limit: the store remembers the one set above for its + // next writer. + let reader = read_limited(&folder, None); + assert_note_is(&reader, 1, &body, "written with no room to compress"); + + // With room: the same row, committed the same way with no limit set, is compressed. + let folder = support::folder("room-to-compress"); + let db = support::open(&folder); + support::run(&db, NOTES); + let tx = db.transaction().expect("begin"); + insert_note(&tx, 1, &body); + tx.commit().expect("commit under no limit"); + support::close(db); + let length = data_file(&folder); + assert!( + length < LEN as u64, + "the data file is {length} bytes long: with room to compress, a repeating text of \ + {LEN} bytes was stored as it is" + ); + assert_note_is(&read_limited(&folder, None), 1, &body, "written with room"); +} + +// ---- damage ---- + +#[test] +fn a_damaged_length_of_a_compressed_value_is_reported_as_damage() { + const LEN: usize = LONGEST_COMPRESSED; + // A compressed text record begins with the byte 0x84 and then the length of the value in + // 8 bytes, low byte first. + const BEGINS: [u8; 9] = [0x84, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00, 0x00, 0x00]; + assert_eq!(BEGINS[1..], (LEN as u64).to_le_bytes()); + + let folder = support::folder("damaged-length"); + let body = support::large_text(Kind::Repeating, 1, LEN); + let db = support::open(&folder); + support::run(&db, NOTES); + insert_note(&db, 1, &body); + support::close(db); + + // A length no store holds, one byte more than the value has, and one byte less. + for wrong in [1_u64 << 40, LEN as u64 + 1, LEN as u64 - 1] { + // A storage door that, in every read, puts the wrong length after that 0x84. + let changed = Arc::new(AtomicU64::new(0)); + let mut options = ReadSessionOptions::default(); + options.memory_limit = Some(64 * MIB); + options.storage = Some(support::storage({ + let changed = changed.clone(); + move |call| { + let Call::Read(_, bytes) = call else { + return; + }; + let mut from = 0; + while let Some(found) = bytes[from..] + .windows(BEGINS.len()) + .position(|window| window == BEGINS) + { + let at = from + found; + bytes[at + 1..at + 9].copy_from_slice(&wrong.to_le_bytes()); + changed.fetch_add(1, Ordering::SeqCst); + from = at + BEGINS.len(); + } + } + })); + let reader = ReadSession::open_with_options(&folder, options) + .unwrap_or_else(|error| panic!("the store did not open for reading: {error:?}")); + let before = used(&reader); + + // Damage, by name: never a panic, never the memory limit's refusal, never a value. + // Two checks stand between the changed bytes and an answer, the checksum of the + // page they are in and the record's own length against its stored bytes; whichever + // meets them first, the statement is told the store is damaged. + match answered(&reader, "SELECT body FROM notes WHERE id = $id", 1) { + Err(Error::StoreCorrupted { path, .. }) => assert_eq!(path, folder), + Err(other) => panic!("a recorded length of {wrong} was refused as {other:?}"), + Ok(rows) => panic!( + "a recorded length of {wrong} was answered with {} rows", + rows.len() + ), + } + assert_eq!( + used(&reader), + before, + "memory in use after a recorded length of {wrong}" + ); + // The door did change a length, or this test shows nothing. + assert!( + changed.load(Ordering::SeqCst) > 0, + "no read of the data file held the nine bytes a compressed record of this value \ + begins with, so no length was changed to {wrong}" + ); + } + + // Through an untouched door the value is as it was written. + assert_note_is(&read_limited(&folder, None), 1, &body, "undamaged"); +} + +#[test] +fn a_large_text_selected_by_name_is_held_once() { + // A text of about 900,000 bytes: under the 1 MiB up to which a text may be compressed. + const LEN: usize = 900_000; + // What a reader's limit leaves beside the value: room for the row it comes in and for + // the stored bytes of a compressed value, and well under what a second copy would need. + const MARGIN: usize = 256 * KIB; + const { assert!(LEN < LONGEST_COMPRESSED && MARGIN < LEN / 2 && LEN / 255 < MARGIN / 2) }; + // A limit that fits the value once and not twice: the 9 MiB an open sets aside, the + // value's record of 9 + LEN bytes, and the margin. + let once = 9 * MIB + 9 + LEN + MARGIN; + + // Stored as it is (scattered text does not shrink) and stored compressed (one sentence + // repeated shrinks to a few kilobytes): the statement holds the value once either way. + for (kind, name, as_it_is) in [ + (Kind::Scattered, "plain", true), + (Kind::Repeating, "compressed", false), + ] { + let folder = support::folder(&format!("selected-by-name-{name}")); + let body = support::large_text(kind, 1, LEN); + store_of_one_large_note(&folder, &body); + let length = data_file(&folder); + assert_eq!( + length >= LEN as u64, + as_it_is, + "the data file of a {kind:?} text of {LEN} bytes is {length} bytes long" + ); + + let reader = read_limited(&folder, Some(once)); + assert_eq!(used(&reader), 9 * MIB, "set aside by the open"); + // The select list names the column. The value is one text, whichever way the + // statement asks for it, and a limit with room for it once is room enough. + match answered(&reader, "SELECT body FROM notes WHERE id = $id", 1) { + Ok(rows) => match rows.as_slice() { + [row] => { + assert_eq!(row.len(), 1, "the columns of the row"); + support::assert_text_is(&row[0], &body, &format!("a {kind:?} text by name")); + } + rows => panic!("{} rows were returned, not one", rows.len()), + }, + Err(error) => panic!( + "a {kind:?} text of {LEN} bytes, selected by name, was refused under a limit \ + that fits it once: {error:?}" + ), + } + assert_eq!( + used(&reader), + 9 * MIB, + "memory in use once the row is dropped" + ); + // The small row beside it answers too, and the count. + assert_note_is(&reader, 2, "small", "under the same limit"); + } +} diff --git a/v2/contextdb/tests/memory_ceiling.rs b/v2/contextdb/tests/memory_ceiling.rs index 765b5148..0ea84d8c 100644 --- a/v2/contextdb/tests/memory_ceiling.rs +++ b/v2/contextdb/tests/memory_ceiling.rs @@ -3,13 +3,23 @@ //! //! This test has a file of its own because its program counts every allocation. The count is //! taken in a child, so that it is the count of one store and nothing else. +//! +//! The same count holds a refused read to its word: a value stored compressed whose real size +//! does not fit the reader's limit is refused before its bytes exist. And it holds a sealing +//! cycle to what it reserves: under a limit with exactly the room the cycle's reservations +//! come to, the cycle completes and the process stays within the limit. The cycle is measured +//! with the page cache already full: the limit has room for the whole cache, and a cache that +//! is half empty would hide whatever the cycle held without reserving it. mod support; use std::alloc::{GlobalAlloc, Layout, System}; +use std::fs; use std::sync::atomic::{AtomicUsize, Ordering}; -use contextdb::{DatabaseOpenOptions, params}; +use contextdb::{ + Database, DatabaseOpenOptions, Error, ReadSession, ReadSessionOptions, Value, params, +}; const LIMIT: usize = 64 << 20; const ROWS: i64 = 1_000_000; @@ -159,3 +169,284 @@ fn load_and_scan(role: support::Role) { support::say(format!("before {before}")); support::say(format!("peak {}", ALLOCATED.peak.load(Ordering::Relaxed))); } + +/// The value of the test below: 1 MiB of one sentence repeated, the largest the store +/// compresses. And the reader's limit: the 9 MiB an open sets aside and 64 KiB, which is room +/// for the value's stored record, a little over 4 KiB, and a sixteenth of what the value +/// itself takes. The arithmetic is that of `a_compressed_value_too_large_for_the_limit_is_refused` +/// in `large_values.rs`, which holds the refusal itself. +const VALUE: usize = 1 << 20; +const VALUE_LIMIT: usize = (9 << 20) + (64 << 10); + +#[test] +fn a_refused_compressed_value_is_never_produced() { + if let Some(role) = support::role() { + return refused_select(&role); + } + let folder = support::folder("refused-value-never-produced"); + let db = support::open(&folder); + support::run( + &db, + "CREATE TABLE notes (id INTEGER PRIMARY KEY, body TEXT NOT NULL)", + ); + let large = support::large_text(support::Kind::Repeating, 1, VALUE); + for (id, body) in [(1_i64, large.as_str()), (2, "small")] { + db.execute( + "INSERT INTO notes (id, body) VALUES ($id, $body)", + ¶ms! { "id" => id, "body" => body }, + ) + .unwrap_or_else(|error| panic!("note {id} was refused: {error:?}")); + } + support::close(db); + // The value is stored compressed: the file is too short to hold it as it is. + let length = fs::metadata(folder.join("data")) + .expect("the store's data file") + .len(); + assert!( + length < VALUE as u64, + "the data file is {length} bytes long: a repeating text of {VALUE} bytes was stored as \ + it is, and this test would show nothing about a compressed value" + ); + + let lines = support::Child::start( + "a_refused_compressed_value_is_never_produced", + "reader", + &folder, + ) + .end(); + assert_eq!( + support::last_number(&lines, "refused"), + Some(1), + "the reader's select of the large value was not refused by its limit" + ); + // The value's bytes never existed: from just before the select to just after its refusal + // the most the process had allocated rose by less than the value is long. A read that + // produced the value and then counted it would have taken all of its bytes at once. + let rise = support::last_number(&lines, "rise").expect("the reader announces it"); + assert!( + rise < VALUE as u64, + "while a value of {VALUE} bytes was refused, the process allocated {rise} bytes more \ + than it had before the select" + ); +} + +/// Opens the store for reading under the limit and reads the small row and the count, so that +/// whatever a session sets up with its first statements is there. Then selects the large +/// value, which must be refused by the memory limit, and announces `refused 1` and +/// `rise `: how far the most this process ever had allocated rose above what it had +/// allocated just before the select. +fn refused_select(role: &support::Role) { + let mut options = ReadSessionOptions::default(); + options.memory_limit = Some(VALUE_LIMIT); + let reader = ReadSession::open_with_options(&role.folder, options) + .expect("open the store for reading under the limit"); + assert_eq!(support::int(&reader, "SELECT COUNT(*) FROM notes"), 2); + assert_eq!( + support::rows(&reader, "SELECT body FROM notes WHERE id = 2"), + [[support::text("small")]] + ); + + let before = ALLOCATED.now.load(Ordering::Relaxed); + ALLOCATED.peak.store(before, Ordering::Relaxed); + let outcome = reader.execute("SELECT body FROM notes WHERE id = 1", ¶ms! {}); + let peak = ALLOCATED.peak.load(Ordering::Relaxed); + match outcome { + Err(Error::MemoryBudgetExceeded { .. }) => support::say("refused 1"), + Err(other) => panic!("the select of the large value was refused as {other:?}"), + Ok(result) => panic!( + "the select of the large value was answered with {} rows", + result.rows().len() + ), + } + support::say(format!("rise {}", peak - before)); +} + +// ---- a full cache ---- + +/// Rows of a table beside the one a test measures, each with a text of 900 bytes that does not +/// shrink: 18 MiB, more than twice the 8 MiB cache of a store whose limit is under 32 MiB. +const FILLER_ROWS: i64 = 20_000; +const FILLER_TEXT: usize = 900; + +/// Writes the filler rows, in a table of their own and one commit. +fn write_the_filler(db: &Database) { + support::run( + db, + "CREATE TABLE filler (id INTEGER PRIMARY KEY, body TEXT NOT NULL)", + ); + let tx = db.transaction().expect("begin"); + for id in 1..=FILLER_ROWS { + let body = support::large_text(support::Kind::Scattered, id, FILLER_TEXT); + tx.execute( + "INSERT INTO filler (id, body) VALUES ($id, $body)", + ¶ms! { "id" => id, "body" => body }, + ) + .unwrap_or_else(|error| panic!("filler row {id} was refused: {error:?}")); + } + tx.commit().expect("commit the filler rows"); +} + +/// Reads every filler row's text, so that the cache is full before what is measured begins and +/// nothing is measured with the room an empty cache leaves. +fn fill_the_cache(db: &Database) { + assert_eq!( + support::int(db, "SELECT COUNT(body) FROM filler"), + FILLER_ROWS + ); +} + +// ---- one sealing cycle with exactly the room it reserves ---- + +/// The numbers in a vector of the sealing test, and its rows: as many as one cycle seals, in +/// one partition, with keys of nine bytes. Their slots fill most of one block. +const SEALED_DIMENSION: usize = 40; +const SEALED: usize = 4096; +const KEY: usize = 9; + +/// The bytes of the block one cycle writes for them: its count, and for each vector a slot of +/// eight bytes and one byte a number, the end of its key, and its key. +const BLOCK: usize = 4 + SEALED * (8 + SEALED_DIMENSION + 4 + KEY); +const _: () = assert!(BLOCK == 249_860 && BLOCK <= 262_080); + +/// What the cycle's working memory comes to: the keys of the pending records it collected, +/// each its bytes and 24; one block being built, 262,080 bytes, taken once; the partition +/// record it adds to, the bytes of its key tuple, here none, and 64; and one row with its +/// vector, a row of two values as a statement counts it and the vector's numbers once more. +const WORKING: usize = SEALED * (KEY + 24) + + 262_080 + + 64 + + (size_of::>() + 2 * size_of::() + 2 * 4 * SEALED_DIMENSION); + +/// What the records the cycle stages come to, each its value, 128 bytes and twice its key: the +/// block, whole, under a key of 13 bytes; the removal of each pending record, whose key is nine +/// bytes and the row's key; the partition's entry of 16 bytes under a key of 5; and the +/// store's count of pending records, 128. +const STAGED: usize = + (BLOCK + 128 + 2 * 13) + SEALED * (128 + 2 * (9 + KEY)) + (16 + 128 + 2 * 5) + 128; + +/// A limit that is exactly enough for that one partition: the 9 MiB an open sets aside, and +/// every reservation of one cycle, each once. +const SEALING_LIMIT: usize = (9 << 20) + WORKING + STAGED; + +#[test] +fn a_sealing_cycle_with_exactly_the_room_it_reserves_completes_within_the_limit() { + if let Some(role) = support::role() { + return one_cycle_under_the_limit(&role); + } + let folder = support::folder("sealing-exact-room"); + let db = support::open(&folder); + write_the_filler(&db); + support::run( + &db, + &format!("CREATE TABLE points (id INTEGER PRIMARY KEY, v VECTOR({SEALED_DIMENSION}))"), + ); + let tx = db.transaction().expect("begin"); + for id in 1..=SEALED as i64 { + tx.execute( + "INSERT INTO points (id, v) VALUES ($id, $v)", + ¶ms! { "id" => id, "v" => support::vector(id, SEALED_DIMENSION) }, + ) + .unwrap_or_else(|error| panic!("point {id} was refused: {error:?}")); + } + tx.commit().expect("commit the points"); + support::close(db); + + let lines = support::Child::start( + "a_sealing_cycle_with_exactly_the_room_it_reserves_completes_within_the_limit", + "sealer", + &folder, + ) + .end(); + // The open set 9 MiB aside, so what the limit left the cycle is exactly its reservations. + assert_eq!( + support::last_number(&lines, "set_aside"), + Some(9 << 20), + "set aside by an open under {SEALING_LIMIT} bytes" + ); + let refused: Vec<&String> = lines + .iter() + .filter(|line| line.starts_with("refused ")) + .collect(); + assert!( + refused.is_empty(), + "with {} bytes for its working memory and {} for the records it stages, {} in all, a \ + cycle over one partition of {SEALED} pending vectors was {}", + WORKING, + STAGED, + WORKING + STAGED, + refused[0] + ); + assert_eq!( + ( + support::last_number(&lines, "done"), + support::last_number(&lines, "remaining") + ), + (Some(SEALED as u64), Some(0)), + "what the cycle sealed and left waiting" + ); + assert_eq!( + support::last_number(&lines, "used_after"), + Some(9 << 20), + "memory in use after the cycle" + ); + // What the process really took for the store, from before the open to the end, the + // cycle's commit included, is within the limit. The cycle began with the cache full. + let before = support::last_number(&lines, "before").expect("the sealer announces it"); + let peak = support::last_number(&lines, "peak").expect("the sealer announces it"); + let cached = support::last_number(&lines, "cached").expect("the sealer announces it"); + assert!( + peak <= before + SEALING_LIMIT as u64, + "with {before} bytes allocated before the store was opened, the process reached {peak}: \ + {} bytes on the store's account under a limit of {SEALING_LIMIT}, {} past it, in a \ + cycle over {SEALED} pending vectors, begun with {cached} bytes allocated for a full \ + cache", + peak - before, + peak - before - SEALING_LIMIT as u64 + ); + + // The cycle's commit is there: every vector is in a block, and the nearest are found. + let db = support::open(&folder); + let stored = support::stored(&db, "SELECT id, v FROM points"); + let query = support::vector(94_500, SEALED_DIMENSION); + let wanted = stored.nearest(|_| true, &query, 10); + let sql = support::nearest_sql("points", "v", None, 10); + let result = support::assert_nearest(&db, &sql, &query, &wanted, "after the cycle"); + assert!( + result.rows_examined() < SEALED as u64, + "a search of {SEALED} sealed vectors read {} rows", + result.rows_examined() + ); +} + +/// Opens the store under [`SEALING_LIMIT`], reads every filler row, and runs one maintenance +/// cycle. Announces `set_aside `, what the open counts; `cached `, what this +/// process had allocated more than before the open once the filler rows were read; then +/// `done ` and `remaining `, or `refused `; `used_after `; +/// and, once the store is closed, +/// `before `, what this process had allocated when it opened the store, and +/// `peak `, the most it had allocated at any time after that. +fn one_cycle_under_the_limit(role: &support::Role) { + let before = ALLOCATED.now.load(Ordering::Relaxed); + ALLOCATED.peak.store(before, Ordering::Relaxed); + let mut options = DatabaseOpenOptions::default(); + options.memory_limit = Some(SEALING_LIMIT); + let db = support::open_with(&role.folder, options).expect("open the store under the limit"); + support::say(format!("set_aside {}", db.memory_used())); + fill_the_cache(&db); + let cached = ALLOCATED.now.load(Ordering::Relaxed); + support::say(format!("cached {}", cached.saturating_sub(before))); + match db.run_maintenance_cycle() { + Ok(report) => { + support::say(format!("done {}", report.vector.maintenance_vectors_done)); + support::say(format!( + "remaining {}", + report.vector.maintenance_vectors_remaining + )); + } + Err(error) => support::say(format!("refused {error:?}")), + } + support::say(format!("used_after {}", db.memory_used())); + support::close(db); + support::say(format!("before {before}")); + support::say(format!("peak {}", ALLOCATED.peak.load(Ordering::Relaxed))); +} diff --git a/v2/contextdb/tests/sql/in_lists.sql b/v2/contextdb/tests/sql/in_lists.sql new file mode 100644 index 00000000..7bf8fd02 --- /dev/null +++ b/v2/contextdb/tests/sql/in_lists.sql @@ -0,0 +1,164 @@ +# IN and NOT IN with a list of values, in the WHERE of a SELECT, an UPDATE and a DELETE, as +# standard SQL has them. Row 3 holds NULL in every column but its key and `bin`. + +CREATE TABLE stock (id INTEGER PRIMARY KEY, kind TEXT, qty INTEGER, bin INTEGER, ref UUID); +INSERT INTO stock (id, kind, qty, bin, ref) +VALUES (1, 'bolt', 7, 1, '550e8400-e29b-41d4-a716-446655440001'), + (2, 'nut', 6, 1, '550e8400-e29b-41d4-a716-446655440002'), + (3, NULL, NULL, 2, NULL), + (4, 'bolt', 9, 2, '550e8400-e29b-41d4-a716-446655440004'), + (5, 'washer', 6, 3, NULL); + +# IN is TRUE when the value equals one of the list; NOT IN when it equals none of them. +SELECT id FROM stock WHERE kind IN ('bolt', 'nut') ORDER BY id; +---- +id +1 +2 +4 + +SELECT id FROM stock WHERE kind NOT IN ('bolt', 'nut') ORDER BY id; +---- +id +5 + +SELECT id FROM stock WHERE qty IN (6) ORDER BY id; +---- +id +2 +5 + +SELECT id FROM stock WHERE id IN (4, 1, 404) ORDER BY id; +---- +id +1 +4 + +# It stands wherever a comparison does, beside AND and under NOT. +SELECT id FROM stock WHERE kind IN ('bolt') AND qty NOT IN (7) ORDER BY id; +---- +id +4 + +SELECT id FROM stock WHERE NOT kind IN ('bolt', 'nut') ORDER BY id; +---- +id +5 + +# A NULL on the left equals nothing and differs from nothing, so row 3 is in neither answer. +SELECT id FROM stock WHERE qty IN (6) OR qty NOT IN (6) ORDER BY id; +---- +id +1 +2 +4 +5 + +# A NULL in the list is never equal to anything: IN is TRUE only through another item, and +# NOT IN is never TRUE. +SELECT id FROM stock WHERE qty IN (7, NULL) ORDER BY id; +---- +id +1 + +SELECT id FROM stock WHERE qty NOT IN (7, NULL) ORDER BY id; +---- +id + +# A row value on the left is compared with row values of the same length, column by column. +SELECT id FROM stock WHERE (kind, bin) IN (('bolt', 1), ('nut', 1)) ORDER BY id; +---- +id +1 +2 + +SELECT id FROM stock WHERE (kind, bin) NOT IN (('bolt', 1), ('nut', 2)) ORDER BY id; +---- +id +2 +4 +5 + +SELECT id FROM stock WHERE (kind, bin) IN (('bolt', NULL)) ORDER BY id; +---- +id + +# Two row values are equal when every column is, so one column that differs settles it whatever +# a NULL beside it holds: row 3, (NULL, 2), is not ('bolt', 1). +SELECT id FROM stock WHERE (kind, bin) NOT IN (('bolt', 1)) ORDER BY id; +---- +id +2 +3 +4 +5 + +# A UUID column is compared with the text that spells a UUID, in either case, alone and inside +# a row value. +SELECT id FROM stock WHERE ref IN ('550e8400-e29b-41d4-a716-446655440001', '550E8400-E29B-41D4-A716-446655440004') ORDER BY id; +---- +id +1 +4 + +SELECT id FROM stock WHERE ref NOT IN ('550e8400-e29b-41d4-a716-446655440001') ORDER BY id; +---- +id +2 +4 + +SELECT id FROM stock WHERE (ref, bin) IN (('550e8400-e29b-41d4-a716-446655440002', 1)) ORDER BY id; +---- +id +2 + +# Refused by name: a row value anywhere but beside IN, one of another length than its partner, +# and a list that is a SELECT. +SELECT id FROM stock WHERE (kind, bin) = ('bolt', 1); +---- +error: PlanError + +SELECT id FROM stock WHERE (kind, bin) IN (('bolt', 1, 2)); +---- +error: PlanError + +SELECT id FROM stock WHERE id IN (SELECT id FROM stock); +---- +error: PlanError + +# UPDATE and DELETE find their rows the same way. +UPDATE stock SET qty = 0 WHERE id IN (1, 2); +---- +rows affected: 2 + +UPDATE stock SET bin = 9 WHERE (kind, bin) IN (('bolt', 2), ('washer', 3)); +---- +rows affected: 2 + +UPDATE stock SET qty = 1 WHERE kind NOT IN ('bolt'); +---- +rows affected: 2 + +SELECT id, kind, qty, bin FROM stock ORDER BY id; +---- +id | kind | qty | bin +1 | bolt | 0 | 1 +2 | nut | 1 | 1 +3 | NULL | NULL | 2 +4 | bolt | 9 | 9 +5 | washer | 1 | 9 + +DELETE FROM stock WHERE id IN (2, 404); +---- +rows affected: 1 + +DELETE FROM stock WHERE kind NOT IN ('bolt'); +---- +rows affected: 1 + +SELECT id FROM stock ORDER BY id; +---- +id +1 +3 +4 diff --git a/v2/contextdb/tests/sql/memory_limit.sql b/v2/contextdb/tests/sql/memory_limit.sql index fd112f08..38df9e62 100644 --- a/v2/contextdb/tests/sql/memory_limit.sql +++ b/v2/contextdb/tests/sql/memory_limit.sql @@ -6,6 +6,13 @@ SHOW MEMORY_LIMIT; limit | used | available | startup_ceiling | ceiling_refresh_failures none | * | * | * | * +# A store opened with no limit takes 64 MiB as its first limit. +SET MEMORY_LIMIT 64M; +SHOW MEMORY_LIMIT; +---- +limit | used | available | startup_ceiling | ceiling_refresh_failures +67108864 | * | * | * | * + SET MEMORY_LIMIT 512M; SHOW MEMORY_LIMIT; ---- diff --git a/v2/contextdb/tests/sql/vector_declarations.sql b/v2/contextdb/tests/sql/vector_declarations.sql new file mode 100644 index 00000000..e7767bfc --- /dev/null +++ b/v2/contextdb/tests/sql/vector_declarations.sql @@ -0,0 +1,129 @@ +# What the declaration of a VECTOR column refuses, each by the name of its kind. A refused +# CREATE TABLE declares nothing: after each group the table is still not there. + +# A PARTITION_KEY column is a column of the same table, +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) PARTITION_KEY (nope)); +---- +error: InvalidVectorPartitionDeclaration + +# declared NOT NULL, +CREATE TABLE docs (id INTEGER PRIMARY KEY, tenant TEXT, v VECTOR(3) PARTITION_KEY (tenant)); +---- +error: InvalidVectorPartitionDeclaration + +# of a type whose values are equal or not and nothing between: not REAL, JSON or VECTOR, +CREATE TABLE docs (id INTEGER PRIMARY KEY, weight REAL NOT NULL, v VECTOR(3) PARTITION_KEY (weight)); +---- +error: InvalidVectorPartitionDeclaration + +CREATE TABLE docs (id INTEGER PRIMARY KEY, meta JSON NOT NULL, v VECTOR(3) PARTITION_KEY (meta)); +---- +error: InvalidVectorPartitionDeclaration + +CREATE TABLE docs (id INTEGER PRIMARY KEY, other VECTOR(2) NOT NULL, v VECTOR(3) PARTITION_KEY (other)); +---- +error: InvalidVectorPartitionDeclaration + +# named once, +CREATE TABLE docs (id INTEGER PRIMARY KEY, tenant TEXT NOT NULL, v VECTOR(3) PARTITION_KEY (tenant, tenant)); +---- +error: InvalidVectorPartitionDeclaration + +# and not the vector column itself. +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) PARTITION_KEY (v) NOT NULL); +---- +error: InvalidVectorPartitionDeclaration + +SELECT id FROM docs; +---- +error: TableNotFound + +# MAX_PARTITIONS needs a PARTITION_KEY, and is a whole number from 1 to the largest the engine +# keeps. +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) MAX_PARTITIONS 4); +---- +error: InvalidVectorPartitionDeclaration + +CREATE TABLE docs (id INTEGER PRIMARY KEY, tenant TEXT NOT NULL, v VECTOR(3) PARTITION_KEY (tenant) MAX_PARTITIONS 0); +---- +error: InvalidVectorPartitionDeclaration + +CREATE TABLE docs (id INTEGER PRIMARY KEY, tenant TEXT NOT NULL, v VECTOR(3) PARTITION_KEY (tenant) MAX_PARTITIONS -1); +---- +error: InvalidVectorPartitionDeclaration + +CREATE TABLE docs (id INTEGER PRIMARY KEY, tenant TEXT NOT NULL, v VECTOR(3) PARTITION_KEY (tenant) MAX_PARTITIONS 99999999999); +---- +error: InvalidVectorPartitionDeclaration + +# A vector has at least one number. +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(0)); +---- +error: InvalidVectorPolicyDeclaration + +SELECT id FROM docs; +---- +error: TableNotFound + +# The clauses that describe or steer an index are read by a later version. Each is refused by +# its name, alone or after the clauses this version reads, and never as a syntax error. +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) WITH (quantization = 'SQ8')); +---- +error: PlanError + +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) AUTO_INDEX_AT 10000); +---- +error: PlanError + +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) SEARCH (TARGET_ACCURACY = 95, EXPECTED_VECTORS = 1000000)); +---- +error: PlanError + +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) IVF (LISTS = 64, PROBES = 8)); +---- +error: PlanError + +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) CONSOLIDATION (CHANGE_PERCENT = 25, TOMBSTONE_PERCENT = 5)); +---- +error: PlanError + +CREATE TABLE docs ( + id INTEGER PRIMARY KEY, + tenant TEXT NOT NULL, + v VECTOR(3) PARTITION_KEY (tenant) MAX_PARTITIONS 4 SEARCH_MODE AUTO AUTO_INDEX_AT 10000 +); +---- +error: PlanError + +SELECT id FROM docs; +---- +error: TableNotFound + +# The name was free all along. ALTER is a later version's too: refused, each leaves the column +# as it was declared, with room for one partition and searched INDEXED. +CREATE TABLE docs ( + id INTEGER PRIMARY KEY, + tenant TEXT NOT NULL, + v VECTOR(3) PARTITION_KEY (tenant) MAX_PARTITIONS 1 SEARCH_MODE INDEXED +); +ALTER TABLE docs ALTER COLUMN v SET SEARCH_MODE EXACT; +---- +error: PlanError + +ALTER TABLE docs ALTER COLUMN v SET MAX_PARTITIONS 512; +---- +error: PlanError + +INSERT INTO docs (id, tenant, v) VALUES (1, 'north', [1.0, 0.0, 0.0]); +INSERT INTO docs (id, tenant, v) VALUES (2, 'south', [0.0, 1.0, 0.0]); +---- +error: VectorPartitionLimitExceeded + +SELECT id FROM docs ORDER BY v <=> [1.0, 0.0, 0.0] LIMIT 1; +---- +error: VectorIndexedRouteUnavailable + +SELECT id FROM docs ORDER BY v <=> [1.0, 0.0, 0.0] USE VECTOR EXACT LIMIT 1; +---- +id +1 diff --git a/v2/contextdb/tests/sql/vector_partitions.sql b/v2/contextdb/tests/sql/vector_partitions.sql new file mode 100644 index 00000000..384498e2 --- /dev/null +++ b/v2/contextdb/tests/sql/vector_partitions.sql @@ -0,0 +1,189 @@ +# A VECTOR column with PARTITION_KEY keeps its vectors in one partition for each set of key +# values. A filter on the key columns chooses the partitions a search reads; the answer is still +# the nearest rows the whole filter passes, cut to the LIMIT after every chosen partition is +# compared. From [1, 0]: rows 1 and 3 are at distance 0, row 5 between, rows 2 and 6 at 1, row 4 +# at 2, and row 7 has no vector. The rows are written out of key order. + +CREATE TABLE passages (id INTEGER PRIMARY KEY, a TEXT NOT NULL, b INTEGER NOT NULL, v VECTOR(2) PARTITION_KEY (a, b)); +INSERT INTO passages (id, a, b, v) +VALUES (6, 'y', 1, [0.0, 2.0]), + (3, 'x', 2, [1.0, 0.0]), + (1, 'x', 1, [1.0, 0.0]), + (5, 'y', 1, [1.0, 1.0]), + (2, 'x', 1, [0.0, 1.0]), + (7, 'x', 1, NULL), + (4, 'x', 2, [-1.0, 0.0]); + +# Every key column given: one partition. +SELECT id, v <=> [1.0, 0.0] AS distance FROM passages WHERE a = 'x' AND b = 1 ORDER BY distance LIMIT 5; +---- +id | distance +1 | 0.0 +2 | 1.0 + +# Several partitions, named by a list of row values or by an OR: one answer, nearest first. +SELECT id, v <=> [1.0, 0.0] AS distance FROM passages WHERE (a, b) IN (('x', 1), ('y', 1)) ORDER BY distance LIMIT 5; +---- +id | distance +1 | 0.0 +5 | * +2 | 1.0 +6 | 1.0 + +SELECT id, v <=> [1.0, 0.0] AS distance FROM passages WHERE (a = 'x' AND b = 1) OR (a = 'y' AND b = 1) ORDER BY distance LIMIT 5; +---- +id | distance +1 | 0.0 +5 | * +2 | 1.0 +6 | 1.0 + +# Only the first key column given: every partition under it. Rows 1 and 3 hold the same vector +# in two partitions and come in key order, though row 3 was written first. +SELECT id, v <=> [1.0, 0.0] AS distance FROM passages WHERE a = 'x' ORDER BY distance LIMIT 5; +---- +id | distance +1 | 0.0 +3 | 0.0 +2 | 1.0 +4 | 2.0 + +# No key column given: every partition. +SELECT id, v <=> [1.0, 0.0] AS distance FROM passages ORDER BY distance LIMIT 10; +---- +id | distance +1 | 0.0 +3 | 0.0 +5 | * +2 | 1.0 +6 | 1.0 +4 | 2.0 + +# LIMIT cuts the one answer, not each partition: from [0, 1] the two nearest rows of these two +# partitions are both in ('y', 1). +SELECT id, v <=> [0.0, 1.0] AS distance FROM passages WHERE (a, b) IN (('x', 2), ('y', 1)) ORDER BY distance LIMIT 2; +---- +id | distance +6 | 0.0 +5 | * + +# Key values no row has give no rows, and never the rows of another partition. +SELECT id FROM passages WHERE a = 'z' AND b = 1 ORDER BY v <=> [1.0, 0.0] LIMIT 5; +---- +id + +SELECT id FROM passages WHERE a = 'x' AND b = 3 ORDER BY v <=> [1.0, 0.0] LIMIT 5; +---- +id + +SELECT id FROM passages WHERE (a, b) IN (('z', 9), ('y', 1)) ORDER BY v <=> [1.0, 0.0] LIMIT 5; +---- +id +5 +6 + +# An UPDATE of a key column moves the row to its new partition. +UPDATE passages SET b = 2 WHERE id = 2; +---- +rows affected: 1 + +SELECT id FROM passages WHERE a = 'x' AND b = 1 ORDER BY v <=> [1.0, 0.0] LIMIT 5; +---- +id +1 + +SELECT id FROM passages WHERE a = 'x' AND b = 2 ORDER BY v <=> [1.0, 0.0] LIMIT 5; +---- +id +3 +2 +4 + +# MAX_PARTITIONS is the most partitions that may hold a vector at once. At 2, a row that would +# start a third is refused, by INSERT or by UPDATE, and nothing of it is kept. +CREATE TABLE capped (id INTEGER PRIMARY KEY, a TEXT NOT NULL, v VECTOR(2) PARTITION_KEY (a) MAX_PARTITIONS 2); +INSERT INTO capped (id, a, v) VALUES (1, 'p', [1.0, 0.0]), (2, 'q', [0.0, 1.0]); +INSERT INTO capped (id, a, v) VALUES (3, 'r', [1.0, 0.0]); +---- +error: VectorPartitionLimitExceeded + +SELECT id, a FROM capped ORDER BY id; +---- +id | a +1 | p +2 | q + +# The next row under key values that already hold a vector is taken. +INSERT INTO capped (id, a, v) VALUES (3, 'q', [1.0, 0.0]); +---- +rows affected: 1 + +UPDATE capped SET a = 'r' WHERE id = 3; +---- +error: VectorPartitionLimitExceeded + +# A row with no vector starts no partition; giving it a vector would. +INSERT INTO capped (id, a, v) VALUES (4, 'r', NULL); +---- +rows affected: 1 + +UPDATE capped SET v = [1.0, 0.0] WHERE id = 4; +---- +error: VectorPartitionLimitExceeded + +SELECT id, a, v FROM capped ORDER BY id; +---- +id | a | v +1 | p | [1.0, 0.0] +2 | q | [0.0, 1.0] +3 | q | [1.0, 0.0] +4 | r | NULL + +# Deleting the rows of a partition frees its place, and the third set of key values is taken. +DELETE FROM capped WHERE a = 'p'; +---- +rows affected: 1 + +INSERT INTO capped (id, a, v) VALUES (5, 'r', [0.0, 1.0]); +---- +rows affected: 1 + +UPDATE capped SET v = [1.0, 0.0] WHERE id = 4; +---- +rows affected: 1 + +SELECT id, v <=> [1.0, 0.0] AS distance FROM capped WHERE a = 'r' ORDER BY distance LIMIT 5; +---- +id | distance +4 | 0.0 +5 | 1.0 + +# At the limit, one UPDATE that empties a partition and starts another is taken. +UPDATE capped SET a = 's' WHERE a = 'q'; +---- +rows affected: 2 + +SELECT id FROM capped WHERE a = 's' ORDER BY v <=> [1.0, 0.0] LIMIT 5; +---- +id +3 +2 + +# TEXT key values that differ only in the case of a letter are two partitions: with room for +# two, 'a' and 'A' fill it. +CREATE TABLE cased (id INTEGER PRIMARY KEY, tag TEXT NOT NULL, v VECTOR(2) PARTITION_KEY (tag) MAX_PARTITIONS 2); +INSERT INTO cased (id, tag, v) VALUES (1, 'a', [1.0, 0.0]), (2, 'A', [1.0, 0.0]); +INSERT INTO cased (id, tag, v) VALUES (3, 'b', [1.0, 0.0]); +---- +error: VectorPartitionLimitExceeded + +SELECT id FROM cased WHERE tag = 'A' ORDER BY v <=> [1.0, 0.0] LIMIT 5; +---- +id +2 + +SELECT id FROM cased ORDER BY v <=> [1.0, 0.0] LIMIT 5; +---- +id +1 +2 diff --git a/v2/contextdb/tests/sql/vector_search.sql b/v2/contextdb/tests/sql/vector_search.sql new file mode 100644 index 00000000..8aec8f22 --- /dev/null +++ b/v2/contextdb/tests/sql/vector_search.sql @@ -0,0 +1,177 @@ +# A similarity search: ORDER BY the cosine distance (`<=>`) from a VECTOR column to a vector, with +# a LIMIT. The nearest rows come first, and rows at one distance come in primary-key order. +# From [1, 0, 0, 0]: rows 1 and 2 point the same way (distance 0), row 5 lies between, rows 3 and +# 7 are at a right angle (1), row 8 is all zeros (1, by rule), row 4 points the opposite way (2) +# and row 6 has no vector. The rows are written out of key order. + +CREATE TABLE points (id INTEGER PRIMARY KEY, tag TEXT, v VECTOR(4)); +INSERT INTO points (id, tag, v) +VALUES (7, 'blue', '[0, 0, 3, 0]'), + (4, 'blue', [-1.0, 0.0, 0.0, 0.0]), + (2, 'red', [2.0, 0.0, 0.0, 0.0]), + (8, 'red', [0.0, 0.0, 0.0, 0.0]), + (5, 'red', [0.1, 0.2, 0.0, 0.0]), + (1, 'red', [1.0, 0.0, 0.0, 0.0]), + (6, 'blue', NULL), + (3, 'blue', [0.0, 1.0, 0.0, 0.0]); + +# A vector prints as its numbers in brackets, whichever way it was written, and no vector as NULL. +SELECT id, v FROM points WHERE id >= 5 ORDER BY id; +---- +id | v +5 | [0.1, 0.2, 0.0, 0.0] +6 | NULL +7 | [0.0, 0.0, 3.0, 0.0] +8 | [0.0, 0.0, 0.0, 0.0] + +# The vector searched from is a literal or its quoted text. +SELECT id FROM points ORDER BY v <=> [1.0, 0.0, 0.0, 0.0] LIMIT 3; +---- +id +1 +2 +5 + +SELECT id FROM points ORDER BY v <=> '[1,0,0,0]' LIMIT 3; +---- +id +1 +2 +5 + +# The distance is the same expression in the select list, under the name the statement gives it: +# a REAL, lower for a nearer row. ORDER BY takes that name or the expression, with ASC or without. +SELECT id, v <=> '[1,0,0,0]' AS distance FROM points ORDER BY distance LIMIT 7; +---- +id | distance +1 | 0.0 +2 | 0.0 +5 | * +3 | 1.0 +7 | 1.0 +8 | 1.0 +4 | 2.0 + +SELECT id, v <=> '[1,0,0,0]' AS distance FROM points ORDER BY v <=> '[1,0,0,0]' LIMIT 4; +---- +id | distance +1 | 0.0 +2 | 0.0 +5 | * +3 | 1.0 + +SELECT id, v <=> '[1,0,0,0]' AS distance FROM points ORDER BY v <=> '[1,0,0,0]' ASC LIMIT 4; +---- +id | distance +1 | 0.0 +2 | 0.0 +5 | * +3 | 1.0 + +# One minus the distance is the similarity, worked out like any other arithmetic. +SELECT id, 1 - (v <=> '[1,0,0,0]') AS similarity FROM points ORDER BY v <=> '[1,0,0,0]' LIMIT 7; +---- +id | similarity +1 | 1.0 +2 | 1.0 +5 | * +3 | 0.0 +7 | 0.0 +8 | 0.0 +4 | -1.0 + +# A distance with no name is called what any unnamed expression is called. +SELECT id, v <=> '[1,0,0,0]' FROM points ORDER BY v <=> '[1,0,0,0]' LIMIT 2; +---- +id | expr +1 | 0.0 +2 | 0.0 + +# A search returns the columns the statement asks for and nothing more: no distance from +# SELECT * and none from a list that does not write one. +SELECT * FROM points ORDER BY v <=> '[1,0,0,0]' LIMIT 2; +---- +id | tag | v +1 | red | [1.0, 0.0, 0.0, 0.0] +2 | red | [2.0, 0.0, 0.0, 0.0] + +SELECT id, tag FROM points ORDER BY v <=> '[1,0,0,0]' LIMIT 2; +---- +id | tag +1 | red +2 | red + +# Every row with a vector is at distance 1 from this one, so the answer is in key order. The +# distance from or to a vector of all zeros is 1 as well. +SELECT id, v <=> [0.0, 0.0, 0.0, 1.0] AS distance FROM points ORDER BY distance LIMIT 4; +---- +id | distance +1 | 1.0 +2 | 1.0 +3 | 1.0 +4 | 1.0 + +SELECT id, v <=> [0.0, 0.0, 0.0, 0.0] AS distance FROM points ORDER BY distance LIMIT 2; +---- +id | distance +1 | 1.0 +2 | 1.0 + +# A row with no vector is never returned, so a LIMIT above the table gives the seven rows that +# have one. LIMIT 0 gives none. +SELECT id FROM points ORDER BY v <=> '[1,0,0,0]' LIMIT 100; +---- +id +1 +2 +5 +3 +7 +8 +4 + +SELECT id FROM points ORDER BY v <=> '[1,0,0,0]' LIMIT 0; +---- +id + +# WHERE is applied first: the nearest of the rows it passes, not the rows it passes among the +# nearest. The two nearest rows of the table are red. +SELECT id, v <=> '[1,0,0,0]' AS distance FROM points WHERE tag = 'blue' ORDER BY distance LIMIT 2; +---- +id | distance +3 | 1.0 +7 | 1.0 + +SELECT id, v <=> '[1,0,0,0]' AS distance FROM points WHERE tag = 'blue' AND id > 3 ORDER BY distance LIMIT 5; +---- +id | distance +7 | 1.0 +4 | 2.0 + +SELECT id FROM points WHERE tag = 'green' ORDER BY v <=> '[1,0,0,0]' LIMIT 5; +---- +id + +# Two VECTOR columns of one table, each of its own dimension, are searched each by its own name. +# Row 4 has a vector in one of them only. +CREATE TABLE media (id INTEGER PRIMARY KEY, vector_text VECTOR(2), vector_vision VECTOR(3)); +INSERT INTO media (id, vector_text, vector_vision) +VALUES (1, [1.0, 0.0], [0.0, 0.0, 1.0]), + (2, [0.0, 1.0], [0.0, 1.0, 0.0]), + (3, [-1.0, 0.0], [1.0, 0.0, 0.0]), + (4, NULL, [2.0, 0.0, 0.0]); + +SELECT id, vector_text <=> [1.0, 0.0] AS distance FROM media ORDER BY distance LIMIT 4; +---- +id | distance +1 | 0.0 +2 | 1.0 +3 | 2.0 + +SELECT id, vector_vision <=> [1.0, 0.0, 0.0] AS distance FROM media ORDER BY distance LIMIT 4; +---- +id | distance +3 | 0.0 +4 | 0.0 +1 | 1.0 +2 | 1.0 diff --git a/v2/contextdb/tests/sql/vector_search_refusals.sql b/v2/contextdb/tests/sql/vector_search_refusals.sql new file mode 100644 index 00000000..add273dd --- /dev/null +++ b/v2/contextdb/tests/sql/vector_search_refusals.sql @@ -0,0 +1,165 @@ +# What a similarity search refuses, each by the name of its kind. No refusal changes a row: the +# last statement reads the table back as it was written. + +CREATE TABLE docs (id INTEGER PRIMARY KEY, title TEXT, closeness REAL, v VECTOR(3), w VECTOR(3)); +INSERT INTO docs (id, title, v, w) +VALUES (1, 'one', [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]), (2, 'two', [0.0, 1.0, 0.0], [1.0, 0.0, 0.0]); + +# A search needs a LIMIT: without one it is not supported in this version. The refusal has no +# kind of its own: it is the one every statement of a later version gets, and its sentence +# names a vector search without LIMIT. +SELECT id FROM docs ORDER BY v <=> [1.0, 0.0, 0.0]; +---- +error: PlanError + +SELECT id, v <=> [1.0, 0.0, 0.0] AS distance FROM docs ORDER BY distance; +---- +error: PlanError + +# The statements that describe an index and its partitions belong to a later version too, and +# are refused by name, not as a syntax error. +SHOW VECTOR_INDEXES; +---- +error: PlanError + +SHOW VECTOR_PARTITIONS; +---- +error: PlanError + +# USE VECTOR belongs to a search, not to a plain ORDER BY. +SELECT id FROM docs ORDER BY id USE VECTOR EXACT LIMIT 2; +---- +error: UseVectorRequiresVectorOrder + +# The column left of <=> is a VECTOR column, and the vector has that column's dimension. +SELECT id FROM docs ORDER BY title <=> [1.0, 0.0, 0.0] LIMIT 2; +---- +error: UnknownVectorIndex + +SELECT id FROM docs ORDER BY v <=> [1.0, 0.0] LIMIT 2; +---- +error: VectorIndexDimensionMismatch + +SELECT id FROM docs ORDER BY v <=> '[1,0,0,0]' LIMIT 2; +---- +error: VectorIndexDimensionMismatch + +# This version keeps no index. INDEXED is refused, asked for in the statement or declared on the +# column, and never answered by an exact search under that name; EXACT and AUTO answer. +SELECT id FROM docs ORDER BY v <=> [1.0, 0.0, 0.0] USE VECTOR INDEXED LIMIT 2; +---- +error: VectorIndexedRouteUnavailable + +SELECT id FROM docs ORDER BY v <=> [1.0, 0.0, 0.0] USE VECTOR EXACT LIMIT 2; +---- +id +1 +2 + +CREATE TABLE indexed (id INTEGER PRIMARY KEY, v VECTOR(3) SEARCH_MODE INDEXED); +INSERT INTO indexed (id, v) VALUES (1, [0.0, 1.0, 0.0]), (2, [1.0, 0.0, 0.0]); +SELECT id FROM indexed ORDER BY v <=> [1.0, 0.0, 0.0] LIMIT 2; +---- +error: VectorIndexedRouteUnavailable + +SELECT id FROM indexed ORDER BY v <=> [1.0, 0.0, 0.0] USE VECTOR EXACT LIMIT 2; +---- +id +2 +1 + +SELECT id FROM indexed ORDER BY v <=> [1.0, 0.0, 0.0] USE VECTOR AUTO LIMIT 2; +---- +id +2 +1 + +# <=> in ORDER BY is the whole of vector search: there is no search function, no vector taken +# from another row, and no ranking by a policy in this version. +SELECT vector_search(v, [1.0, 0.0, 0.0], 2) FROM docs; +---- +error: PlanError + +SELECT id FROM docs ORDER BY v <=> ROW_VECTOR('docs', 'v', 1) LIMIT 2; +---- +error: PlanError + +SELECT id FROM docs ORDER BY v <=> [1.0, 0.0, 0.0] USE RANK effective_confidence LIMIT 2; +---- +error: PlanError + +# A distance has a value only for the rows of a search, as that search's one ascending sort key. +# Refused by name: a distance in a WHERE, with a search beside it or without; +SELECT id FROM docs WHERE v <=> [1.0, 0.0, 0.0] < 0.5; +---- +error: PlanError + +SELECT id FROM docs WHERE v <=> [1.0, 0.0, 0.0] < 0.5 ORDER BY v <=> [1.0, 0.0, 0.0] LIMIT 2; +---- +error: PlanError + +# in a select list whose statement does not order by it; +SELECT id, v <=> [1.0, 0.0, 0.0] AS distance FROM docs; +---- +error: PlanError + +SELECT id, v <=> [1.0, 0.0, 0.0] AS distance FROM docs ORDER BY id LIMIT 2; +---- +error: PlanError + +# farthest first, and beside a second sort key; +SELECT id FROM docs ORDER BY v <=> [1.0, 0.0, 0.0] DESC LIMIT 2; +---- +error: PlanError + +SELECT id FROM docs ORDER BY v <=> [1.0, 0.0, 0.0], id LIMIT 2; +---- +error: PlanError + +# in the select list of a search, to another vector or from another column than the ORDER BY's; +SELECT id, v <=> [0.0, 1.0, 0.0] AS other FROM docs ORDER BY v <=> [1.0, 0.0, 0.0] LIMIT 2; +---- +error: PlanError + +SELECT id, w <=> [1.0, 0.0, 0.0] AS other FROM docs ORDER BY v <=> [1.0, 0.0, 0.0] LIMIT 2; +---- +error: PlanError + +# between two columns, and from something that is not a column; +SELECT id FROM docs ORDER BY v <=> w LIMIT 2; +---- +error: PlanError + +SELECT id FROM docs ORDER BY [0.0, 1.0, 0.0] <=> [1.0, 0.0, 0.0] LIMIT 2; +---- +error: PlanError + +# and in a statement that writes. +UPDATE docs SET closeness = v <=> [1.0, 0.0, 0.0]; +---- +error: PlanError + +DELETE FROM docs WHERE v <=> [1.0, 0.0, 0.0] < 0.5; +---- +error: PlanError + +# A search adds no score column. On a table without one the name is a column that is not there; +# on a table with one it is that column's own values. +SELECT id, score FROM docs ORDER BY v <=> [1.0, 0.0, 0.0] LIMIT 2; +---- +error: ColumnNotFound + +CREATE TABLE graded (id INTEGER PRIMARY KEY, score REAL, v VECTOR(3)); +INSERT INTO graded (id, score, v) VALUES (1, 0.25, [0.0, 1.0, 0.0]), (2, 9.5, [1.0, 0.0, 0.0]); +SELECT id, score, v <=> [1.0, 0.0, 0.0] AS distance FROM graded ORDER BY distance LIMIT 2; +---- +id | score | distance +2 | 9.5 | 0.0 +1 | 0.25 | 1.0 + +# Every row of the first table is as it was written. +SELECT id, title, closeness, v, w FROM docs ORDER BY id; +---- +id | title | closeness | v | w +1 | one | NULL | [1.0, 0.0, 0.0] | [0.0, 1.0, 0.0] +2 | two | NULL | [0.0, 1.0, 0.0] | [1.0, 0.0, 0.0] diff --git a/v2/contextdb/tests/store.rs b/v2/contextdb/tests/store.rs index f4a91e55..d8d02c8b 100644 --- a/v2/contextdb/tests/store.rs +++ b/v2/contextdb/tests/store.rs @@ -1,5 +1,5 @@ //! A store is one folder: made with no configuration, the same after close and reopen, and held -//! for writing by one process at a time. +//! for writing by one process at a time, for as long as that process has it open and no longer. mod support; @@ -126,3 +126,81 @@ fn holder(role: support::Role) { } } } + +/// A copy of the descriptor this process holds a store's `lock` file open with: the same open +/// file, as every child process has it that is started while the store is open, until the +/// child replaces itself with its program. The descriptor is found among those the operating +/// system lists for this process, which is how Linux lists them. +#[cfg(target_os = "linux")] +fn copy_of_the_lock(folder: &std::path::Path) -> std::os::fd::OwnedFd { + use std::fs; + use std::os::fd::{BorrowedFd, RawFd}; + + let lock = fs::canonicalize(folder.join("lock")).expect("the store's lock file"); + let mut found: Vec = Vec::new(); + for entry in fs::read_dir("/proc/self/fd").expect("the descriptors of this process") { + let entry = entry.expect("a descriptor of this process"); + // Other tests open and close files meanwhile: a descriptor that has gone is not it. + if fs::read_link(entry.path()).is_ok_and(|open| open == lock) { + let number = entry.file_name().to_string_lossy().parse(); + found.push(number.expect("a descriptor is a number")); + } + } + let [descriptor] = found.as_slice() else { + panic!( + "{} descriptors of this process are of {lock:?}, not one", + found.len() + ); + }; + // SAFETY: the caller has the store open, and the store keeps this descriptor open until + // it is closed; the copy made here is a descriptor of its own. + let held = unsafe { BorrowedFd::borrow_raw(*descriptor) }; + held.try_clone_to_owned() + .expect("copy the descriptor of the lock") +} + +/// A program that starts another process from one thread while another thread closes a store +/// and opens it again meets this for an instant: the child, until it has replaced itself with +/// its program, holds a copy of every file its parent had open, the store's lock among them. +/// Here a child is handed that copy to keep, so that the instant lasts. +#[cfg(target_os = "linux")] +#[test] +fn a_store_closed_here_opens_again_while_a_child_process_holds_a_copy_of_its_lock() { + if support::role().is_some() { + // The child opens nothing. It has the copy it was started with, and waits. + support::say("holding"); + support::hear(); + return; + } + let folder = support::folder("lock-copied-to-a-child"); + let db = support::open(&folder); + support::declare(&db); + support::insert(&db, 1..=3); + let mut heir = support::Child::start_holding( + "a_store_closed_here_opens_again_while_a_child_process_holds_a_copy_of_its_lock", + "heir", + &folder, + copy_of_the_lock(&folder), + ); + heir.expect("holding"); + + // Closing lets go of the store, so that the next writer can open it. No handle in this + // process holds it now, and no other process ever opened it: the next open gets it. + support::close(db); + match support::open_with(&folder, DatabaseOpenOptions::default()) { + Ok(db) => { + support::assert_rows_are(&db, 1..=3); + support::insert(&db, 4..=4); + support::close(db); + } + Err(Error::DatabaseLocked { holder_pid, .. }) => panic!( + "the store this process closed was refused to it as held by process {holder_pid}; \ + this process is {}, and the child that has a copy of the lock is {}", + std::process::id(), + heir.id() + ), + Err(other) => panic!("the store this process closed did not open again: {other:?}"), + } + heir.end(); + support::assert_rows_are(&support::open(&folder), 1..=4); +} diff --git a/v2/contextdb/tests/support/mod.rs b/v2/contextdb/tests/support/mod.rs index a9a202fe..febad0e8 100644 --- a/v2/contextdb/tests/support/mod.rs +++ b/v2/contextdb/tests/support/mod.rs @@ -1,8 +1,9 @@ //! The only shared test code. Each test file says `mod support;`. //! //! It holds what every test needs to stay exact: one store folder per test, the one way a store -//! is opened, a table whose rows are a function of their number, a storage door that lets a test -//! watch the data file, and a second process spoken to over its pipes. +//! is opened, a table whose rows are a function of their number, vectors that are a function of +//! their row's number with the answer every similarity search is held to, a storage door that +//! lets a test watch the data file, and a second process spoken to over its pipes. // Each test program uses a part of this module; the rest would be reported as never used. #![allow(dead_code)] @@ -12,11 +13,13 @@ use std::fmt; use std::fs::{self, File}; use std::io::{self, BufRead, BufReader, Write}; use std::ops::RangeInclusive; +use std::os::fd::OwnedFd; use std::os::unix::fs::FileExt; use std::os::unix::process::ExitStatusExt; use std::path::{Path, PathBuf}; use std::process::{Command, ExitStatus, Stdio}; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::mpsc::{self, Receiver, RecvTimeoutError}; use std::time::Duration; @@ -95,6 +98,14 @@ impl Run for ReadSession { } } +/// A reference to something that runs a statement runs it too, so a test can hand the writer, +/// a transaction or a reader to one function as `&dyn Run`. +impl Run for &R { + fn run(&self, sql: &str, params: &HashMap) -> contextdb::Result { + (**self).run(sql, params) + } +} + /// Runs a statement that must succeed. pub fn run(on: &impl Run, sql: &str) -> QueryResult { match on.run(sql, ¶ms! {}) { @@ -220,6 +231,498 @@ where ); } +// ---- fixed numbers, vectors and text ---- + +/// A generator of fixed numbers: a xorshift seeded with one number, usually a row's. The same +/// seed gives the same numbers on every machine and every run, and neighbouring seeds give +/// numbers that have nothing to do with each other. +pub struct Fixed(u64); + +impl Fixed { + pub fn seeded(seed: i64) -> Fixed { + // Spread the seed over all 64 bits, and never start at zero, where a xorshift stays. + let spread = (seed as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ 0xD6E8_FEB8_6659_FD93; + let mut fixed = Fixed(if spread == 0 { + 0x2545_F491_4F6C_DD1D + } else { + spread + }); + for _ in 0..8 { + fixed.draw(); + } + fixed + } + + /// The next 64 bits. + pub fn draw(&mut self) -> u64 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + self.0 + } + + /// The next number below `bound`. + pub fn below(&mut self, bound: u64) -> u64 { + (self.draw() >> 11) % bound + } + + /// The next number from -1 up to, not including, 1: one of 2^24 evenly spaced values. + pub fn unit(&mut self) -> f32 { + (self.draw() >> 40) as f32 / (1_u64 << 23) as f32 - 1.0 + } +} + +/// The vector of row `n` in a column of this dimension. It is scattered, not a smooth function +/// of `n`: a few rows lie near any vector a test searches with and most lie far from it. +pub fn vector(n: i64, dimension: usize) -> Vec { + let mut fixed = Fixed::seeded(n); + (0..dimension).map(|_| fixed.unit()).collect() +} + +/// `len` lower-case letters, fixed by `seed`. +pub fn letters(seed: i64, len: usize) -> String { + let mut fixed = Fixed::seeded(seed); + (0..len) + .map(|_| char::from(b'a' + fixed.below(26) as u8)) + .collect() +} + +// ---- large text and JSON ---- + +/// What a large text value is made of. The store keeps a text of more than 1,015 bytes beside +/// its row, compressed when that saves an eighth of it; each kind says which way that goes. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Kind { + /// One sentence of 60 bytes, repeated: shrinks to a small part of its size. + Repeating, + /// Words drawn from a list of 64, a space between: shrinks to about half. + Words, + /// Bytes drawn one by one from the 94 printable ASCII characters: does not shrink. + Scattered, + /// Scattered for the first fifteen sixteenths, then one letter: shrinks by a sixteenth, + /// which is less than an eighth. + MostlyScattered, + /// Scattered for the first half, then one letter: shrinks by about half. + HalfScattered, + /// One sentence of two-, three- and four-byte characters, repeated: shrinks. + Wide, +} + +impl Kind { + pub const ALL: [Kind; 6] = [ + Kind::Repeating, + Kind::Words, + Kind::Scattered, + Kind::MostlyScattered, + Kind::HalfScattered, + Kind::Wide, + ]; + + /// The kind `steps` places further on in [`Kind::ALL`], going round at its end. One step + /// turns `Words`, which shrinks by an eighth, into `Scattered`, which does not, and + /// `MostlyScattered`, which does not, into `HalfScattered`, which does. + pub fn after(self, steps: usize) -> Kind { + let at = Kind::ALL + .iter() + .position(|kind| *kind == self) + .expect("every kind is in the list"); + Kind::ALL[(at + steps) % Kind::ALL.len()] + } +} + +/// The sentence of [`Kind::Repeating`]: 60 bytes. +const SENTENCE: &str = "A store keeps what it is given and hands it back as it was. "; + +/// The sentence of [`Kind::Wide`]: characters of one, two, three and four bytes. +const WIDE_SENTENCE: &str = "Grüße aus Zürich für José: 東京の夜は静かで、川は海へ流れる。 🚀🌊✨ "; + +/// The words of [`Kind::Words`]. +const WORDS: [&str; 64] = [ + "about", "after", "again", "always", "answer", "before", "between", "block", "bytes", "cache", + "change", "close", "column", "commit", "count", "data", "every", "file", "first", "folder", + "found", "given", "holds", "index", "into", "keeps", "large", "later", "length", "limit", + "memory", "moment", "never", "number", "once", "open", "order", "other", "page", "place", + "read", "reader", "record", "row", "same", "search", "second", "size", "slice", "small", + "stays", "store", "table", "takes", "text", "then", "there", "three", "under", "value", + "vector", "where", "while", "writer", +]; + +/// One of the 94 printable ASCII characters from `!` to `~`. +fn scattered(fixed: &mut Fixed) -> u8 { + b'!' + fixed.below(94) as u8 +} + +/// Exactly `len` bytes of text of this kind, fixed by `id`: the same on every machine and +/// every run, and different for different ids. +pub fn large_text(kind: Kind, id: i64, len: usize) -> String { + let mut fixed = Fixed::seeded(id); + let mut text = String::with_capacity(len); + match kind { + Kind::Repeating => { + // The sentence over and over, begun at a place that depends on the id. + let from = id.rem_euclid(SENTENCE.len() as i64) as usize; + text.extend(SENTENCE.chars().cycle().skip(from).take(len)); + } + Kind::Words => { + while text.len() < len { + text.push_str(WORDS[fixed.below(64) as usize]); + text.push(' '); + } + // Every word is ASCII, so any byte is a place to cut. + text.truncate(len); + } + Kind::Scattered | Kind::MostlyScattered | Kind::HalfScattered => { + let drawn = match kind { + Kind::Scattered => len, + Kind::MostlyScattered => len - len / 16, + _ => len / 2, + }; + text.extend((0..drawn).map(|_| char::from(scattered(&mut fixed)))); + text.extend(std::iter::repeat_n('e', len - drawn)); + } + Kind::Wide => { + let from = id.rem_euclid(WIDE_SENTENCE.chars().count() as i64) as usize; + for character in WIDE_SENTENCE.chars().cycle().skip(from) { + // Cut at a character: one that would pass the size is left out. + if text.len() + character.len_utf8() > len { + break; + } + text.push(character); + } + text.extend(std::iter::repeat_n('x', len - text.len())); + } + } + assert_eq!(text.len(), len, "a {kind:?} text for row {id}"); + text +} + +/// A JSON value that `serde_json::to_vec` writes in exactly `len` bytes, fixed by `id`; `len` +/// is at least 64. One that `shrinks` is an object holding an array of small objects with the +/// same three keys, and compresses to a small part of its size; one that does not is an +/// object holding one string of scattered text. +pub fn large_json(shrinks: bool, id: i64, len: usize) -> contextdb::serde_json::Value { + use contextdb::serde_json::{Value as Json, json, to_vec}; + let json = if shrinks { + // `{"items":[],"pad":""}` is 21 bytes; every item adds its own bytes and, after the + // first, a comma. Items are added while they fit and `pad` takes up the rest. + let (mut items, mut size) = (Vec::new(), 21); + for number in id.. { + let item = json!({ "id": number, "kind": "note", "ok": true }); + let bytes = to_vec(&item).expect("JSON writes as text").len() + 1; + let bytes = bytes - usize::from(items.is_empty()); + if size + bytes > len { + break; + } + size += bytes; + items.push(item); + } + json!({ "items": items, "pad": "x".repeat(len - size) }) + } else { + // `{"text":""}` is 11 bytes. A quote or a backslash is written as two bytes; where + // only one byte is left, an `x` takes it. + let mut fixed = Fixed::seeded(id); + let (mut text, mut size) = (String::new(), 11); + while size < len { + let character = char::from(scattered(&mut fixed)); + let bytes = if matches!(character, '"' | '\\') { + 2 + } else { + 1 + }; + if size + bytes > len { + text.push('x'); + size += 1; + } else { + text.push(character); + size += bytes; + } + } + json!({ "text": Json::String(text) }) + }; + let written = to_vec(&json).expect("JSON writes as text").len(); + assert_eq!(written, len, "the JSON value for row {id}"); + json +} + +/// Asserts that the bytes read back are the bytes put in. A failure says what was compared, +/// the two lengths and the first place they differ, and never prints the bytes. +pub fn assert_same_bytes(found: &[u8], wanted: &[u8], what: &str) { + if found == wanted { + return; + } + let shared = found.len().min(wanted.len()); + let differs = found + .iter() + .zip(wanted) + .position(|(found, wanted)| found != wanted) + .unwrap_or(shared); + panic!( + "{what}: {} bytes found, {} wanted; they first differ at byte {differs}", + found.len(), + wanted.len() + ); +} + +/// The name of what a value is, for a failure that must not print the value. +fn type_of(value: &Value) -> &'static str { + match value { + Value::Null => "NULL", + Value::Bool(_) => "a BOOLEAN", + Value::Int64(_) => "an INTEGER", + Value::Float64(_) => "a REAL", + Value::Text(_) => "a TEXT", + Value::Uuid(_) => "a UUID", + Value::Timestamp(_) => "a TIMESTAMP", + Value::Json(_) => "a JSON value", + Value::Vector(_) => "a VECTOR", + Value::TxId(_) => "a transaction id", + } +} + +/// Asserts that a value read back is a `Value::Text` of exactly these bytes. +pub fn assert_text_is(found: &Value, wanted: &str, what: &str) { + match found { + Value::Text(found) => assert_same_bytes(found.as_bytes(), wanted.as_bytes(), what), + other => panic!("{what}: read back as {}, not as a TEXT", type_of(other)), + } +} + +/// Asserts that a value read back is a `Value::Json` equal to this one, and that it is written +/// in exactly the bytes this one is written in. +pub fn assert_json_is(found: &Value, wanted: &contextdb::serde_json::Value, what: &str) { + use contextdb::serde_json::to_vec; + match found { + Value::Json(found) => { + let (found_bytes, wanted_bytes) = (to_vec(found), to_vec(wanted)); + assert_same_bytes( + &found_bytes.expect("JSON writes as text"), + &wanted_bytes.expect("JSON writes as text"), + what, + ); + assert!(found == wanted, "{what}: a JSON value that is not equal"); + } + other => panic!("{what}: read back as {}, not as JSON", type_of(other)), + } +} + +// ---- the answer a similarity search is held to ---- + +/// The cosine distance between a stored vector and the vector of a search, worked out the one +/// way the engine is to work it out: each number widened to 64 bits, the dot product and the +/// two squared lengths summed in the order of the numbers, `1 - dot / (sqrt(aa) * sqrt(bb))`, +/// and 1 when either vector is all zeros. +pub fn distance(stored: &[f32], query: &[f32]) -> f64 { + let (mut dot, mut aa, mut bb) = (0.0_f64, 0.0_f64, 0.0_f64); + for (x, y) in stored.iter().zip(query) { + let (x, y) = (f64::from(*x), f64::from(*y)); + dot += x * y; + aa += x * x; + bb += y * y; + } + if aa == 0.0 || bb == 0.0 { + 1.0 + } else { + 1.0 - dot / (aa.sqrt() * bb.sqrt()) + } +} + +/// One row of a search's answer: its primary key and its distance. Two are equal when the keys +/// are equal and the distances are the same to the bit. +#[derive(Clone)] +pub struct Near { + pub key: Value, + pub distance: f64, +} + +impl PartialEq for Near { + fn eq(&self, other: &Near) -> bool { + self.key == other.key && self.distance.to_bits() == other.distance.to_bits() + } +} + +impl fmt::Debug for Near { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "{:?} at {:e} ({:#018x})", + self.key, + self.distance, + self.distance.to_bits() + ) + } +} + +/// The rows of a table as a plain SELECT returned them: the primary key first, then the +/// columns a test filters by, then the VECTOR column last. +pub struct Stored { + rows: Vec>, +} + +/// Reads every row a plain SELECT returns, for the answers of [`Stored::nearest`]. +pub fn stored(on: &impl Run, plain: &str) -> Stored { + Stored { + rows: rows(on, plain), + } +} + +impl Stored { + pub fn rows(&self) -> &[Vec] { + &self.rows + } + + /// The reference answer: of the rows `passes` lets through that have a vector, the `k` + /// least by `(distance, primary key)`, in that order. A search must return exactly these. + pub fn nearest(&self, passes: impl Fn(&[Value]) -> bool, query: &[f32], k: usize) -> Vec { + let mut all = Vec::new(); + for row in self.rows.iter().filter(|row| passes(row)) { + match row.last() { + Some(Value::Vector(vector)) => all.push(Near { + key: row[0].clone(), + distance: distance(vector, query), + }), + Some(Value::Null) => {} + other => panic!("the last column of a stored row is {other:?}, not a vector"), + } + } + all.sort_by(|a, b| { + a.distance + .total_cmp(&b.distance) + .then_with(|| key_order(&a.key, &b.key)) + }); + all.truncate(k); + all + } +} + +/// The order of two primary keys of one type. +fn key_order(a: &Value, b: &Value) -> std::cmp::Ordering { + match (a, b) { + (Value::Int64(a), Value::Int64(b)) => a.cmp(b), + (Value::Text(a), Value::Text(b)) => a.as_bytes().cmp(b.as_bytes()), + other => panic!("no order is written here for the keys {other:?}"), + } +} + +/// The search every test writes unless it says otherwise: the key and the distance of the `k` +/// rows of `table` that pass `filter` and lie nearest `$q` in `column`. +pub fn nearest_sql(table: &str, column: &str, filter: Option<&str>, k: usize) -> String { + let filter = filter.map_or_else(String::new, |filter| format!(" WHERE {filter}")); + format!( + "SELECT id, {column} <=> $q AS distance FROM {table}{filter} ORDER BY distance LIMIT {k}" + ) +} + +/// Runs a statement with `$q` bound to this vector. +pub fn search(on: &impl Run, sql: &str, query: &[f32]) -> contextdb::Result { + on.run(sql, ¶ms! { "q" => query.to_vec() }) +} + +/// What a search returned, read as an answer: the key is the first column and the distance the +/// last, which must be a `Value::Float64`. +pub fn found(result: &QueryResult, sql: &str) -> Vec { + let rows = result.rows().iter(); + rows.map(|row| match (row.first(), row.last()) { + (Some(key), Some(Value::Float64(distance))) if row.len() >= 2 => Near { + key: key.clone(), + distance: *distance, + }, + _ => panic!("{sql}\nreturned the row {row:?}, not a key and a distance that is a REAL"), + }) + .collect() +} + +/// Asserts that a search with `$q` bound to `query` returns exactly `wanted`: the same keys in +/// the same order, each distance the same to the bit. Returns the result, for its counted work. +pub fn assert_nearest( + on: &impl Run, + sql: &str, + query: &[f32], + wanted: &[Near], + what: &str, +) -> QueryResult { + let result = match search(on, sql, query) { + Ok(result) => result, + Err(error) => panic!("{what}: {sql}\nwas refused: {error:?}"), + }; + assert_same_rows(&found(&result, sql), wanted, &format!("{what}: {sql}")); + result +} + +// ---- the table most vector tests search ---- + +/// Every column of `docs`, the key first and the vector last: what [`stored`] reads. +pub const DOCS: &str = "SELECT id, ctx, tag, v FROM docs"; + +/// Declares `docs`, whose row `n` is `(n, , n % 50, )`. `clauses` is what +/// follows `VECTOR(dimension)`: `PARTITION_KEY (ctx)`, `MAX_PARTITIONS 2`, `SEARCH_MODE EXACT` +/// or nothing. +pub fn declare_docs(db: &Database, dimension: usize, clauses: &str) { + run( + db, + &format!( + "CREATE TABLE docs (id INTEGER PRIMARY KEY, ctx TEXT NOT NULL, tag INTEGER NOT NULL, \ + v VECTOR({dimension}) {clauses})" + ), + ); +} + +/// Inserts row `n` of `docs` under this context with this vector, or with no vector for `None`. +pub fn insert_doc( + on: &impl Run, + n: i64, + ctx: &str, + vector: Option>, +) -> contextdb::Result { + on.run( + "INSERT INTO docs (id, ctx, tag, v) VALUES ($id, $ctx, $tag, $v)", + ¶ms! { + "id" => n, "ctx" => ctx, "tag" => n % 50, + "v" => vector.map_or(Value::Null, Value::Vector), + }, + ) +} + +/// Inserts the rows of `docs` with these numbers, one statement each, row `n` under the +/// context `ctx(n)` with the vector [`vector`] gives for `n`. Every one must be taken. +pub fn insert_docs( + on: &impl Run, + numbers: impl IntoIterator, + dimension: usize, + ctx: impl Fn(i64) -> String, +) { + for n in numbers { + if let Err(error) = insert_doc(on, n, &ctx(n), Some(vector(n, dimension))) { + panic!("inserting row {n} of docs was refused: {error:?}"); + } + } +} + +/// A text value, for comparing with a column of a stored row. +pub fn text(value: &str) -> Value { + Value::Text(value.to_owned()) +} + +/// Runs maintenance cycles until no vector waits to be sealed, and returns how many vectors +/// the cycles sealed between them. Nothing else in the tests seals. +pub fn seal_all(db: &Database) -> u64 { + let mut sealed = 0; + loop { + let report = match db.run_maintenance_cycle() { + Ok(report) => report.vector, + Err(error) => panic!("a maintenance cycle was refused: {error:?}"), + }; + sealed += report.maintenance_vectors_done; + if report.maintenance_vectors_remaining == 0 { + return sealed; + } + assert!( + report.maintenance_vectors_done > 0, + "a cycle sealed nothing and left {} vectors waiting", + report.maintenance_vectors_remaining + ); + } +} + // ---- the storage door ---- /// One call the engine makes on its data file, as a test's storage door sees it. @@ -293,6 +796,28 @@ impl StorageBackend for Watched { } } +/// A storage door for a child that is to be stopped inside one statement. While the flag it +/// returns is set it counts the reads of the data file, and at read number `pause_at` it +/// announces `reading` and waits for a line from the parent; `pause_at` 0 never stops. Returns +/// the door, the flag and the count. +pub fn pausing_at_read(pause_at: u64) -> (WrapStorage, Arc, Arc) { + let counting = Arc::new(AtomicBool::new(false)); + let reads = Arc::new(AtomicU64::new(0)); + let door = storage({ + let (counting, reads) = (counting.clone(), reads.clone()); + move |call| { + if matches!(call, Call::Read(..)) + && counting.load(Ordering::SeqCst) + && reads.fetch_add(1, Ordering::SeqCst) + 1 == pause_at + { + say("reading"); + hear(); + } + } + }); + (door, counting, reads) +} + // ---- a second process: the parent's side ---- /// A second process: this test program started again to run one test in a role. It announces @@ -310,12 +835,23 @@ impl Child { /// `if let Some(role) = support::role() { return child(role); }`. `role` is the role's name, /// then any settings as `key=number` words. pub fn start(test: &str, role: &str, folder: &Path) -> Child { + Child::started(test, role, folder, Stdio::inherit()) + } + + /// [`Child::start`] for a child that is handed an open file, which it keeps open for as + /// long as it runs: the file is its standard error, the one channel the roles do not use. + pub fn start_holding(test: &str, role: &str, folder: &Path, held: OwnedFd) -> Child { + Child::started(test, role, folder, Stdio::from(held)) + } + + fn started(test: &str, role: &str, folder: &Path, stderr: Stdio) -> Child { let process = Command::new(std::env::current_exe().expect("the path of this test program")) .args(["--exact", test, "--nocapture"]) .env("CONTEXTDB_TEST_ROLE", role) .env("CONTEXTDB_TEST_FOLDER", folder) .stdin(Stdio::piped()) .stdout(Stdio::piped()) + .stderr(stderr) .spawn() .expect("start this test program a second time"); Child { diff --git a/v2/contextdb/tests/vectors.rs b/v2/contextdb/tests/vectors.rs new file mode 100644 index 00000000..75f380ef --- /dev/null +++ b/v2/contextdb/tests/vectors.rs @@ -0,0 +1,915 @@ +//! Exact similarity search: the rows a search returns are the ones a comparison of every +//! stored vector returns, in the order of `(distance, primary key)`, and the distance it shows +//! is that distance to the bit. This holds whatever has been sealed, changed, deleted or moved, +//! and for every reader from the moment it reads. +//! +//! Every test here works out its own answer from the rows a plain SELECT returns +//! (`support::Stored::nearest`) and holds the search to it. + +mod support; + +use std::collections::BTreeMap; + +use contextdb::{Database, ReadSession, Value, params}; + +use support::{Near, Run}; + +/// The keys of an answer, in its order. +fn ids(answer: &[Near]) -> Vec { + answer + .iter() + .map(|near| match near.key { + Value::Int64(id) => id, + ref other => panic!("the key {other:?} is not an INTEGER"), + }) + .collect() +} + +/// Changes the vector of one row of `docs`, or takes it away for `None`. +fn set_vector(on: &impl Run, id: i64, vector: Option>) { + let updated = on.run( + "UPDATE docs SET v = $v WHERE id = $id", + ¶ms! { "id" => id, "v" => vector.map_or(Value::Null, Value::Vector) }, + ); + match updated { + Ok(result) => assert_eq!(result.rows_affected(), 1, "updating the vector of row {id}"), + Err(error) => panic!("updating the vector of row {id} was refused: {error:?}"), + } +} + +/// The vector of row `n` in the first test: mostly the scattered vector of `n` at one of seven +/// magnitudes, with four rows that hold one vector and a fifth that holds it four times as long +/// (multiplying by four is exact, so its distance is theirs to the bit), two rows of zeros, one +/// row of the largest number an `f32` holds, and one row with no vector. +fn told_apart(n: i64, dimension: usize) -> Option> { + let scaled = |n: i64, by: f32| -> Vec { + let scattered = support::vector(n, dimension).into_iter(); + scattered.map(|number| number * by).collect() + }; + Some(match n { + 5 | 20 | 40 | 60 => scaled(5, 1.0), + 80 => scaled(5, 4.0), + 30 | 31 => vec![0.0; dimension], + 50 => vec![f32::MAX; dimension], + 70 => return None, + _ => scaled(n, 10.0_f32.powi((n % 7) as i32 - 3)), + }) +} + +#[test] +fn exact_search_returns_the_true_nearest_vectors() { + for (dimension, rows) in [(3, 5000_i64), (16, 5000), (33, 600), (384, 300), (768, 300)] { + let folder = support::folder(&format!("exact-nearest-{dimension}")); + let db = support::open(&folder); + support::declare_docs(&db, dimension, ""); + // Not in key order: the even rows from the last down, then the odd rows from the first. + let evens = (1..=rows).rev().filter(|n| n % 2 == 0); + let odds = (1..=rows).filter(|n| n % 2 == 1); + let tx = db.transaction().expect("begin"); + for n in evens.chain(odds) { + if let Err(error) = support::insert_doc(&tx, n, "a", told_apart(n, dimension)) { + panic!("row {n} of dimension {dimension} was refused: {error:?}"); + } + } + tx.commit().expect("commit"); + + // What was inserted is what is stored, number for number. + let stored = support::stored(&db, "SELECT id, ctx, tag, v FROM docs ORDER BY id"); + assert_eq!(stored.rows().len() as i64, rows); + for (n, row) in (1..=rows).zip(stored.rows()) { + let inserted = told_apart(n, dimension).map_or(Value::Null, Value::Vector); + assert!(row[3] == inserted, "row {n} of dimension {dimension}"); + } + + let queries = [ + // The vector of rows 5, 20, 40, 60 and 80: the answer begins with a tie of five. + told_apart(5, dimension).expect("row 5 has a vector"), + // All zeros: every row is at distance 1, so the answer is in key order. + vec![0.0; dimension], + vec![f32::MAX; dimension], + support::vector(100_001, dimension), + support::vector(100_002, dimension), + ]; + // Through the writer or a reader: every search equals the reference from the rows that + // same door reads. + let through = |on: &dyn Run, state: &str| { + let plain = support::stored(&on, support::DOCS); + // Every row with a vector is found, and the row without one never is. + let all = plain.nearest(|_| true, &queries[3], usize::MAX); + assert_eq!(all.len() as i64, rows - 1, "{state}: rows with a vector"); + assert!(!ids(&all).contains(&70), "{state}: the row with no vector"); + // The five rows of one direction come first for that vector, at one distance and + // in key order. + let tie = plain.nearest(|_| true, &queries[0], 5); + assert_eq!( + ids(&tie), + [5, 20, 40, 60, 80], + "{state}: rows of one vector" + ); + let first = tie[0].distance.to_bits(); + assert!(tie.iter().all(|near| near.distance.to_bits() == first)); + // For a vector of zeros every row is at distance 1 and comes in key order. + let zeros = plain.nearest(|_| true, &queries[1], 3); + assert_eq!(ids(&zeros), [1, 2, 3], "{state}: every row at one distance"); + assert!(zeros.iter().all(|near| near.distance == 1.0)); + + for (number, query) in queries.iter().enumerate() { + // One row, ten, and more than the table holds. + for k in [1, 10, rows as usize + 50] { + let wanted = plain.nearest(|_| true, query, k); + let sql = support::nearest_sql("docs", "v", None, k); + let what = format!("dimension {dimension}, {state}, query {number}, k {k}"); + support::assert_nearest(&on, &sql, query, &wanted, &what); + } + } + }; + + through(&db, "unsealed"); + let first = db.run_maintenance_cycle().expect("the first cycle").vector; + assert_eq!( + first.maintenance_vectors_done, + (rows as u64 - 1).min(4096), + "dimension {dimension}: vectors the first cycle sealed" + ); + through(&db, "after one cycle"); + support::seal_all(&db); + through(&db, "after the last cycle"); + support::close(db); + let db = support::open(&folder); + through(&db, "after close and reopen"); + let reader = ReadSession::open(&folder).expect("open the store for reading"); + through(&reader, "through a read session"); + } +} + +#[test] +fn exact_search_rescores_from_stored_vectors() { + const DIMENSION: usize = 16; + let folder = support::folder("exact-rescores"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, "PARTITION_KEY (ctx)"); + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, 1..=400, DIMENSION, |n| format!("c{}", n % 4)); + tx.commit().expect("commit"); + support::seal_all(&db); + + // What each row that is still there holds: its context, and its vector if it has one. + let mut held: BTreeMap>)> = (1..=400) + .map(|n| (n, (n % 4, Some(support::vector(n, DIMENSION))))) + .collect(); + // Vectors that were a row's and are no row's now: a slot written from one of them is the + // nearest thing there is to a search for it, and must give nothing. + let mut left_behind: Vec> = Vec::new(); + let mut steps = support::Fixed::seeded(2); + // Numbers no row was ever given a vector for. + let mut fresh = 10_000_i64; + + for step in 1..=300 { + let at = steps.below(held.len() as u64) as usize; + let id = *held.keys().nth(at).expect("a row that is still there"); + let (ctx, vector) = held.get(&id).expect("the row").clone(); + match steps.below(4) { + 0 => { + fresh += 1; + let new = support::vector(fresh, DIMENSION); + set_vector(&db, id, Some(new.clone())); + left_behind.extend(vector); + held.insert(id, (ctx, Some(new))); + } + 1 => { + support::run(&db, &format!("DELETE FROM docs WHERE id = {id}")); + left_behind.extend(vector); + held.remove(&id); + } + 2 => { + set_vector(&db, id, None); + left_behind.extend(vector); + held.insert(id, (ctx, None)); + } + _ => { + let to = (ctx + 1 + steps.below(3) as i64) % 4; + support::run( + &db, + &format!("UPDATE docs SET ctx = 'c{to}' WHERE id = {id}"), + ); + held.insert(id, (to, vector)); + } + } + if step % 40 == 0 { + db.run_maintenance_cycle().expect("a cycle between changes"); + } + if step % 5 != 0 { + continue; + } + + let stored = support::stored(&db, support::DOCS); + // The rows a plain SELECT returns are the rows this test has left there. + let wanted: Vec<(i64, i64, Option>)> = held + .iter() + .map(|(id, (ctx, vector))| (*id, *ctx, vector.clone())) + .collect(); + let mut read: Vec<(i64, i64, Option>)> = stored + .rows() + .iter() + .map(|row| match row.as_slice() { + [Value::Int64(id), Value::Text(ctx), _, vector] => { + let ctx = ctx[1..].parse().expect("a context is c and a number"); + let vector = match vector { + Value::Vector(vector) => Some(vector.clone()), + _ => None, + }; + (*id, ctx, vector) + } + other => panic!("a row of docs came back as {other:?}"), + }) + .collect(); + read.sort_by_key(|row| row.0); + assert!(read == wanted, "after step {step}: the rows of docs"); + + let mut queries: Vec> = left_behind.iter().rev().take(2).cloned().collect(); + queries.push(support::vector(20_000 + step, DIMENSION)); + for (number, query) in queries.iter().enumerate() { + for k in [1, 10, 1000] { + let what = format!("after step {step}, query {number}"); + let everywhere = stored.nearest(|_| true, query, k); + // No row answers from a vector it no longer has: every distance is the one of + // the vector the row holds now. + for near in &everywhere { + let Value::Int64(id) = near.key else { + panic!("the key {:?}", near.key) + }; + assert!(held.get(&id).is_some_and(|(_, vector)| vector.is_some())); + } + let sql = support::nearest_sql("docs", "v", None, k); + support::assert_nearest(&db, &sql, query, &everywhere, &what); + let in_one = stored.nearest(|row| row[1] == support::text("c1"), query, k); + let sql = support::nearest_sql("docs", "v", Some("ctx = 'c1'"), k); + support::assert_nearest(&db, &sql, query, &in_one, &what); + } + } + } + + // The same after everything pending is sealed, and after the store is opened again. + support::seal_all(&db); + support::close(db); + let db = support::open(&folder); + let stored = support::stored(&db, support::DOCS); + for query in left_behind.iter().rev().take(20) { + for k in [1, 10, 1000] { + let wanted = stored.nearest(|_| true, query, k); + let sql = support::nearest_sql("docs", "v", None, k); + support::assert_nearest(&db, &sql, query, &wanted, "after sealing and reopening"); + } + } +} + +#[test] +fn exact_search_reads_its_own_snapshot() { + const DIMENSION: usize = 16; + let folder = support::folder("exact-own-snapshot"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, "PARTITION_KEY (ctx)"); + let ctx = |n: i64| format!("c{}", n % 2); + // Rows 1 to 100 are in blocks and rows 101 to 200 on pending lists. + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, 1..=100, DIMENSION, ctx); + tx.commit().expect("commit"); + support::seal_all(&db); + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, 101..=200, DIMENSION, ctx); + tx.commit().expect("commit"); + + // The search is for the vector of row 7, which is sealed. Row 160, which is pending, is + // given the same direction, so the committed answer begins with rows 7 and 160. + let query = support::vector(7, DIMENSION); + let doubled: Vec = query.iter().map(|number| number * 2.0).collect(); + set_vector(&db, 160, Some(doubled)); + let reader = ReadSession::open(&folder).expect("open the store for reading"); + + // Holds a search over every partition, and over each one, to the rows `on` itself reads. + fn agrees(on: &impl Run, query: &[f32], what: &str) -> Vec { + let stored = support::stored(on, support::DOCS); + for k in [1, 10, 1000] { + let wanted = stored.nearest(|_| true, query, k); + let sql = support::nearest_sql("docs", "v", None, k); + support::assert_nearest(on, &sql, query, &wanted, what); + for ctx in ["c0", "c1"] { + let wanted = stored.nearest(|row| row[1] == support::text(ctx), query, k); + let filter = format!("ctx = '{ctx}'"); + let sql = support::nearest_sql("docs", "v", Some(&filter), k); + support::assert_nearest(on, &sql, query, &wanted, what); + } + } + stored.nearest(|_| true, query, 1000) + } + // What the transaction below does: adds rows, one of them the nearest there can be; + // changes the vector of a sealed row and of a pending row; deletes one of each. + let change = |tx: &contextdb::Transaction<'_>| { + support::insert_docs(tx, 201..=220, DIMENSION, ctx); + set_vector(tx, 210, Some(query.clone())); + set_vector(tx, 5, Some(support::vector(30_005, DIMENSION))); + set_vector(tx, 150, Some(support::vector(30_150, DIMENSION))); + support::run(tx, "DELETE FROM docs WHERE id = 7"); + support::run(tx, "DELETE FROM docs WHERE id = 160"); + }; + + let committed = agrees(&db, &query, "before any transaction"); + assert_eq!(ids(&committed)[..2], [7, 160]); + assert_eq!(committed.len(), 200); + + // A transaction's searches see what it has done, as its row reads do. + let tx = db.transaction().expect("begin"); + change(&tx); + let inside = agrees(&tx, &query, "inside the transaction"); + assert_eq!(inside.len(), 218); + assert_eq!(ids(&inside)[0], 210, "the row the transaction added"); + for gone in [7, 160] { + assert!(!ids(&inside).contains(&gone), "row {gone} was deleted"); + } + // Nobody else does: the writer's own handle and a reader in the folder answer as before. + assert!(agrees(&db, &query, "beside the open transaction") == committed); + assert!(agrees(&reader, &query, "a reader beside the open transaction") == committed); + for on in [&db as &dyn Run, &reader as &dyn Run] { + let added = on.run("SELECT id FROM docs WHERE id = 210", ¶ms! {}); + assert!(added.expect("look for the added row").rows().is_empty()); + } + + // A rollback leaves the committed answer. + tx.rollback(); + assert!(agrees(&db, &query, "after the rollback") == committed); + assert!(agrees(&reader, &query, "a reader after the rollback") == committed); + + // Committed, the same changes are everyone's, the rows and their vectors together: a + // reader that sees an added row finds it by its vector in the same breath. + let tx = db.transaction().expect("begin"); + change(&tx); + let inside = agrees(&tx, &query, "inside the second transaction"); + tx.commit().expect("commit"); + assert!(agrees(&db, &query, "after the commit") == inside); + assert!(agrees(&reader, &query, "a reader after the commit") == inside); + assert_eq!(support::int(&reader, "SELECT COUNT(*) FROM docs"), 218); + assert_eq!( + support::rows(&reader, "SELECT id FROM docs WHERE id = 210"), + [[Value::Int64(210)]] + ); +} + +#[test] +fn vector_search_filters_before_scoring() { + const DIMENSION: usize = 16; + const ROWS: i64 = 5000; + let folder = support::folder("filters-before-scoring"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, ""); + for first in (1..=ROWS).step_by(1000) { + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, first..first + 1000, DIMENSION, |_| "a".to_owned()); + tx.commit().expect("commit"); + } + let stored = support::stored(&db, support::DOCS); + let tag = |row: &[Value]| match row[2] { + Value::Int64(tag) => tag, + ref other => panic!("the tag {other:?}"), + }; + // A filter that passes every row, half of them, one in fifty, and none. Row `n` has the + // tag `n % 50`. + type Passes = fn(i64) -> bool; + let filters: [(Option<&str>, Passes, usize); 5] = [ + (None, |_| true, 5000), + (Some("tag >= 0"), |_| true, 5000), + (Some("tag < 25"), |tag| tag < 25, 2500), + (Some("tag = 7"), |tag| tag == 7, 100), + (Some("tag = 50"), |_| false, 0), + ]; + let queries = [ + support::vector(40_001, DIMENSION), + support::vector(40_002, DIMENSION), + support::vector(17, DIMENSION), + ]; + // Every search equals the reference; returns the rows each one read. + let check = |state: &str| -> Vec { + let mut examined = Vec::new(); + for (filter, passes, passing) in filters { + for query in &queries { + let all = stored.nearest(|row| passes(tag(row)), query, usize::MAX); + assert_eq!(all.len(), passing, "rows that pass {filter:?}"); + let wanted = &all[..all.len().min(10)]; + let sql = support::nearest_sql("docs", "v", filter, 10); + let what = format!("{state}, filter {filter:?}"); + let result = support::assert_nearest(&db, &sql, query, wanted, &what); + examined.push(result.rows_examined()); + } + } + examined + }; + + check("unsealed"); + assert_eq!(support::seal_all(&db), ROWS as u64); + let examined = check("everything sealed"); + // With every row passing, the one-byte copies rule nearly every row out before it is read: + // a search reads under a tenth of the table. The first six searches are the two filters + // that pass every row, three vectors each. + for read in &examined[..6] { + assert!( + *read < ROWS as u64 / 10, + "a search of {ROWS} sealed rows, all passing, read {read} rows" + ); + } +} + +#[test] +fn equal_distances_come_in_primary_key_order() { + const DIMENSION: usize = 8; + // The first pair of rows with one vector is 3 and 9, the second 4 and 8; the lower key of + // each is under the context x and the higher under y. + let pair = |number: i64| support::vector(50_000 + number, DIMENSION); + let filler = |tx: &contextdb::Transaction<'_>, numbers: std::ops::RangeInclusive| { + support::insert_docs(tx, numbers, DIMENSION, |n| { + let ctx = if n % 2 == 0 { "x" } else { "y" }; + ctx.to_owned() + }); + }; + let insert = |db: &Database, id: i64, ctx: &str, vector: Vec| { + if let Err(error) = support::insert_doc(db, id, ctx, Some(vector)) { + panic!("row {id} was refused: {error:?}"); + } + }; + // The lower key comes first and alone for LIMIT 1, and first of the two for LIMIT 2, + // whichever way both partitions are named. + let lower_first = |db: &Database, query: &[f32], lower: i64, higher: i64, state: &str| { + let stored = support::stored(db, support::DOCS); + for filter in [None, Some("ctx IN ('x', 'y')"), Some("ctx IN ('y', 'x')")] { + for k in [1, 2, 5, 1000] { + let wanted = stored.nearest(|_| true, query, k); + let sql = support::nearest_sql("docs", "v", filter, k); + let what = format!("{state}, {filter:?}"); + let result = support::assert_nearest(db, &sql, query, &wanted, &what); + let found = support::found(&result, &sql); + assert_eq!(ids(&found)[0], lower, "{what}, LIMIT {k}"); + if k >= 2 { + assert_eq!(ids(&found)[1], higher, "{what}, LIMIT {k}"); + assert_eq!( + found[0].distance.to_bits(), + found[1].distance.to_bits(), + "{what}: two rows with one vector are at one distance" + ); + } + } + } + }; + + for lower_goes_in_first in [true, false] { + let name = if lower_goes_in_first { + "lower-first" + } else { + "lower-last" + }; + let folder = support::folder(&format!("equal-distances-{name}")); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, "PARTITION_KEY (ctx)"); + // In the first store the lower key of a pair is inserted first, and its context is + // the first partition; in the second store it is inserted last. + let in_order = |low: (i64, &'static str), high: (i64, &'static str)| { + if lower_goes_in_first { + [low, high] + } else { + [high, low] + } + }; + + // Both rows of the first pair pending. + let [first, second] = in_order((3, "x"), (9, "y")); + insert(&db, first.0, first.1, pair(1)); + insert(&db, second.0, second.1, pair(1)); + let tx = db.transaction().expect("begin"); + filler(&tx, 20..=60); + tx.commit().expect("commit"); + lower_first(&db, &pair(1), 3, 9, &format!("{name}, both pending")); + + // Both sealed. + support::seal_all(&db); + lower_first(&db, &pair(1), 3, 9, &format!("{name}, both sealed")); + + // Of the second pair, the row inserted first is sealed and the other is pending. + let [first, second] = in_order((4, "x"), (8, "y")); + insert(&db, first.0, first.1, pair(2)); + support::seal_all(&db); + insert(&db, second.0, second.1, pair(2)); + let state = format!( + "{name}, row {} sealed and row {} pending", + first.0, second.0 + ); + lower_first(&db, &pair(2), 4, 8, &state); + lower_first(&db, &pair(1), 3, 9, &state); + + // And once the second is sealed too, in a later block. + support::seal_all(&db); + lower_first(&db, &pair(2), 4, 8, &format!("{name}, all sealed")); + } +} + +#[test] +fn a_row_is_in_an_answer_once_after_its_vector_changes() { + const DIMENSION: usize = 16; + let folder = support::folder("in-an-answer-once"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, "PARTITION_KEY (ctx)"); + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, 1..=60, DIMENSION, |n| format!("c{}", n % 3)); + tx.commit().expect("commit"); + support::seal_all(&db); + + // Every vector row 7 has had, and a vector no row has had. + let mut had = vec![support::vector(7, DIMENSION)]; + let elsewhere = support::vector(60_000, DIMENSION); + // With LIMIT above the row count every entry that leads to a row is visited. Row 7 is in + // the answer once, at the distance of the vector it holds now, and no key is there twice. + let once = |on: &dyn Run, had: &[Vec], state: &str| { + let stored = support::stored(&on, support::DOCS); + let now = had.last().expect("row 7 has a vector"); + for query in had.iter().chain([&elsewhere]) { + let wanted = stored.nearest(|_| true, query, 1000); + let sql = support::nearest_sql("docs", "v", None, 1000); + let result = support::assert_nearest(&on, &sql, query, &wanted, state); + let found = support::found(&result, &sql); + let mut keys = ids(&found); + let of_row_7: Vec<&Near> = found + .iter() + .filter(|near| near.key == Value::Int64(7)) + .collect(); + assert_eq!(of_row_7.len(), 1, "{state}: row 7 in one answer"); + assert_eq!( + of_row_7[0].distance.to_bits(), + support::distance(now, query).to_bits(), + "{state}: the distance of row 7" + ); + keys.sort_unstable(); + keys.dedup(); + assert_eq!( + keys.len(), + found.len(), + "{state}: a key twice in one answer" + ); + // A search for one row, with a vector the row no longer has: what a slot left + // behind would answer best. + let sql = support::nearest_sql("docs", "v", None, 1); + support::assert_nearest(&on, &sql, query, &wanted[..1], state); + } + }; + + once(&db, &had, "sealed"); + for change in 1..=2 { + had.push(support::vector(60_000 + change, DIMENSION)); + set_vector(&db, 7, had.last().cloned()); + once(&db, &had, &format!("after change {change} of the vector")); + } + // After the next cycle two slots of row 7 are in blocks: the first and the newest. + assert_eq!( + db.run_maintenance_cycle() + .expect("a cycle") + .vector + .maintenance_vectors_done, + 1 + ); + once(&db, &had, "two slots of the row in blocks"); + + // Under another partition key, searched over all partitions. + support::run(&db, "UPDATE docs SET ctx = 'c0' WHERE id = 7"); + once(&db, &had, "after the row moved to another partition"); + let stored = support::stored(&db, support::DOCS); + for (ctx, there) in [("c0", true), ("c1", false)] { + let now = had.last().expect("row 7 has a vector"); + let wanted = stored.nearest(|row| row[1] == support::text(ctx), now, 1000); + let filter = format!("ctx = '{ctx}'"); + let sql = support::nearest_sql("docs", "v", Some(&filter), 1000); + support::assert_nearest(&db, &sql, now, &wanted, "one partition after the move"); + assert_eq!(ids(&wanted).contains(&7), there, "row 7 under {ctx}"); + } + support::seal_all(&db); + once(&db, &had, "three slots of the row in blocks"); + + // Inside an open transaction that has changed the vector of a sealed row, row 7, and of a + // pending row, row 100, searched through the transaction before it commits. + support::insert_docs(&db, 100..=100, DIMENSION, |_| "c1".to_owned()); + let tx = db.transaction().expect("begin"); + had.push(support::vector(60_003, DIMENSION)); + set_vector(&tx, 7, had.last().cloned()); + set_vector(&tx, 100, Some(support::vector(60_100, DIMENSION))); + once(&tx, &had, "inside a transaction, changed once"); + had.push(support::vector(60_004, DIMENSION)); + set_vector(&tx, 7, had.last().cloned()); + set_vector(&tx, 100, Some(support::vector(60_101, DIMENSION))); + once(&tx, &had, "inside a transaction, changed twice"); + tx.commit().expect("commit"); + once(&db, &had, "after the transaction committed"); + support::seal_all(&db); + once(&db, &had, "after everything was sealed again"); +} + +#[test] +fn the_distance_is_the_expression_in_the_select_list() { + const DIMENSION: usize = 16; + const K: usize = 7; + let folder = support::folder("distance-in-select-list"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, ""); + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, 1..=200, DIMENSION, |_| "a".to_owned()); + tx.commit().expect("commit"); + let query = support::vector(70_000, DIMENSION); + // The name this engine gives a select item that was given none. + let unnamed = support::run(&db, "SELECT id, id + 1 FROM docs LIMIT 1").columns()[1].clone(); + + let check = |state: &str| { + let stored = support::stored(&db, support::DOCS); + let wanted = stored.nearest(|_| true, &query, K); + let distances: Vec = wanted + .iter() + .map(|near| Value::Float64(near.distance)) + .collect(); + let similarities: Vec = wanted + .iter() + .map(|near| Value::Float64(1.0 - near.distance)) + .collect(); + // A statement's rows as its columns, each column a list of values. + let columns_of = |sql: &str, names: &[&str]| -> Vec> { + let result = match support::search(&db, sql, &query) { + Ok(result) => result, + Err(error) => panic!("{state}: {sql}\nwas refused: {error:?}"), + }; + assert_eq!(result.columns(), names, "{state}: the columns of {sql}"); + assert_eq!(result.rows().len(), K, "{state}: the rows of {sql}"); + (0..names.len()) + .map(|at| result.rows().iter().map(|row| row[at].clone()).collect()) + .collect() + }; + // Two lists of values are the same when every REAL is the same to the bit. + let same = |found: &[Value], wanted: &[Value], what: &str| { + let bits = |values: &[Value]| -> Vec { + let values = values.iter(); + values + .map(|value| match value { + Value::Float64(real) => format!("{:#018x}", real.to_bits()), + other => format!("{other:?}"), + }) + .collect() + }; + assert_eq!(bits(found), bits(wanted), "{state}: {what}"); + }; + let keys: Vec = wanted.iter().map(|near| near.key.clone()).collect(); + + // The distance under a name, ordered by the name or by the expression, ASC written + // or not: the same rows in the same order, each with its distance. + for order in ["distance", "v <=> $q", "distance ASC", "v <=> $q ASC"] { + let sql = + format!("SELECT id, v <=> $q AS distance FROM docs ORDER BY {order} LIMIT {K}"); + let columns = columns_of(&sql, &["id", "distance"]); + same(&columns[0], &keys, &sql); + same(&columns[1], &distances, &sql); + } + // Under no name. + let sql = format!("SELECT id, v <=> $q FROM docs ORDER BY v <=> $q LIMIT {K}"); + let columns = columns_of(&sql, &["id", unnamed.as_str()]); + same(&columns[0], &keys, &sql); + same(&columns[1], &distances, &sql); + // Twice in one list, before and after the key. + let sql = format!( + "SELECT v <=> $q AS near, id, v <=> $q AS again FROM docs ORDER BY near LIMIT {K}" + ); + let columns = columns_of(&sql, &["near", "id", "again"]); + same(&columns[0], &distances, &sql); + same(&columns[1], &keys, &sql); + same(&columns[2], &distances, &sql); + // Inside an expression: the similarity is one minus the distance. + let sql = format!( + "SELECT id, 1 - (v <=> $q) AS similarity FROM docs ORDER BY v <=> $q LIMIT {K}" + ); + let columns = columns_of(&sql, &["id", "similarity"]); + same(&columns[0], &keys, &sql); + same(&columns[1], &similarities, &sql); + let sql = format!( + "SELECT id, v <=> $q AS distance, 1 - (v <=> $q) AS similarity FROM docs \ + ORDER BY distance LIMIT {K}" + ); + let columns = columns_of(&sql, &["id", "distance", "similarity"]); + same(&columns[0], &keys, &sql); + same(&columns[1], &distances, &sql); + same(&columns[2], &similarities, &sql); + + // A search that selects no distance returns none. + let sql = format!("SELECT id FROM docs ORDER BY v <=> $q LIMIT {K}"); + same(&columns_of(&sql, &["id"])[0], &keys, &sql); + let sql = format!("SELECT tag, id FROM docs ORDER BY v <=> $q LIMIT {K}"); + same(&columns_of(&sql, &["tag", "id"])[1], &keys, &sql); + // SELECT * returns the table's columns and nothing more: each row whole. + let sql = format!("SELECT * FROM docs ORDER BY v <=> $q LIMIT {K}"); + let columns = columns_of(&sql, &["id", "ctx", "tag", "v"]); + same(&columns[0], &keys, &sql); + for (at, key) in keys.iter().enumerate() { + let row = stored.rows().iter().find(|row| row[0] == *key); + let row = row.expect("a row of the answer is a row of the table"); + let found: Vec = columns.iter().map(|column| column[at].clone()).collect(); + assert!(found == *row, "{state}: row {key:?} of {sql}"); + } + }; + + check("unsealed"); + support::seal_all(&db); + check("sealed"); +} + +// ---- answers where the search has to look twice, and where two entries lead to one row ---- + +/// The vector the two tests below search for: one along the first of eight axes. +const AHEAD: [f32; 8] = [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]; + +/// A vector whose similarity to [`AHEAD`] is `similarity`, as nearly as 32 bits hold it: that +/// much along the first axis and the rest along one of the other seven. The slot a block keeps +/// for such a vector proves its similarity to within one hundredth, so two vectors a tenth +/// apart are never mistaken for one another by a search that only looked at slots. +fn toward(similarity: f32, axis: i64) -> Vec { + let mut vector = vec![0.0_f32; AHEAD.len()]; + vector[0] = similarity; + vector[1 + (axis % 7) as usize] = (1.0 - similarity * similarity).sqrt(); + vector +} + +/// Inserts rows of `docs` under one context: row `first + n` holds the vector at the +/// similarity `nearest - step * n`. +fn insert_toward(on: &impl Run, ctx: &str, first: i64, count: i64, nearest: f32, step: f32) { + for n in 0..count { + let vector = toward(nearest - step * n as f32, first + n); + if let Err(error) = support::insert_doc(on, first + n, ctx, Some(vector)) { + panic!("row {} was refused: {error:?}", first + n); + } + } +} + +#[test] +fn a_search_that_has_to_look_twice_returns_the_true_nearest() { + let folder = support::folder("search-looks-twice"); + let db = support::open(&folder); + support::declare_docs(&db, AHEAD.len(), ""); + // Ten rows very near the vector searched for, ten at about half its similarity and thirty + // far from it, all in one block. + let tx = db.transaction().expect("begin"); + insert_toward(&tx, "near", 1, 10, 0.99, 0.005); + insert_toward(&tx, "mid", 11, 10, 0.50, 0.01); + insert_toward(&tx, "far", 21, 30, 0.10, 0.02); + tx.commit().expect("commit"); + assert_eq!(support::seal_all(&db), 50); + + // The ten nearest rows of the table are the near ones, and the ten nearest of the rest are + // the ones at about half. + let stored = support::stored(&db, support::DOCS); + assert_eq!( + ids(&stored.nearest(|_| true, &AHEAD, 10)), + (1..=10).collect::>() + ); + let not_near = |row: &[Value]| row[1] != support::text("near"); + assert_eq!( + ids(&stored.nearest(not_near, &AHEAD, 10)), + (11..=20).collect::>() + ); + + // A filter lets none of the near rows through. Their slots still say that ten vectors + // lie very near, which rules the slots of every other row out: the first look visits the + // near rows, is told each does not pass, and ends with no row at all. Only a second look, + // at the slots the first one left out, finds the answer. For three rows, ten, and fifteen, + // where the first look ends with some rows and not enough. + let filter = Some("ctx <> 'near'"); + let through = |on: &dyn Run, stored: &support::Stored, state: &str| { + for k in [3, 10, 15] { + let wanted = stored.nearest(not_near, &AHEAD, k); + assert_eq!(wanted.len(), k, "{state}: rows that pass the filter"); + let sql = support::nearest_sql("docs", "v", filter, k); + support::assert_nearest(&on, &sql, &AHEAD, &wanted, state); + } + }; + through(&db, &stored, "near rows filtered out, everything sealed"); + + // Ten rows more, between the far ones and the ones at half, wait on the pending list. They + // are read first and fill an answer of ten, so the first look ends with a full answer of + // the wrong rows: the rows at half are nearer, and their slots were ruled out by the near + // rows' slots. + let tx = db.transaction().expect("begin"); + insert_toward(&tx, "late", 51, 10, 0.30, 0.01); + tx.commit().expect("commit"); + let stored = support::stored(&db, support::DOCS); + let wanted = stored.nearest(not_near, &AHEAD, 10); + assert_eq!(ids(&wanted), (11..=20).collect::>()); + through(&db, &stored, "near rows filtered out, ten rows pending"); + let reader = ReadSession::open(&folder).expect("open the store for reading"); + through(&reader, &stored, "the same through a read session"); + + // The same without a filter once the near rows are deleted: their slots stay in the block, + // and no row answers for them. + assert_eq!( + support::run(&db, "DELETE FROM docs WHERE ctx = 'near'").rows_affected(), + 10 + ); + let stored = support::stored(&db, support::DOCS); + assert_eq!(stored.rows().len(), 50); + for k in [3, 10, 15] { + let wanted = stored.nearest(|_| true, &AHEAD, k); + assert_eq!(ids(&wanted)[..3], [11, 12, 13]); + let sql = support::nearest_sql("docs", "v", None, k); + for on in [&db as &dyn Run, &reader as &dyn Run] { + support::assert_nearest(&on, &sql, &AHEAD, &wanted, "near rows deleted"); + } + } +} + +#[test] +fn a_row_two_entries_lead_to_is_in_the_answer_once_and_pushes_no_row_out() { + // Every search here is for ten rows, and each store is built so that a row of the answer + // is reached a second time when the answer already holds ten rows. Reached twice, the row + // must stay in the answer once, and the tenth row must stay too. + const K: usize = 10; + let nearest_ten = |on: &dyn Run, state: &str| -> (Vec, u64) { + let stored = support::stored(&on, support::DOCS); + let wanted = stored.nearest(|_| true, &AHEAD, K); + assert_eq!(wanted.len(), K, "{state}: the rows of the reference"); + let sql = support::nearest_sql("docs", "v", None, K); + let result = support::assert_nearest(&on, &sql, &AHEAD, &wanted, state); + (ids(&wanted), result.rows_examined()) + }; + + // A row a transaction has staged, whose pending record is still there. Twenty rows wait on + // the pending list; row 20 is the nearest and has the greatest key. A transaction changes + // another column of it, so its search reads row 20 as a staged row, then the twenty + // pending rows in key order, row 20 last of all. + let folder = support::folder("reached-twice-staged-and-pending"); + let db = support::open(&folder); + support::declare_docs(&db, AHEAD.len(), ""); + let tx = db.transaction().expect("begin"); + insert_toward(&tx, "a", 1, 19, 0.88, 0.02); + insert_toward(&tx, "a", 20, 1, 0.95, 0.0); + tx.commit().expect("commit"); + let (committed, examined) = nearest_ten(&db, "twenty pending rows"); + assert_eq!(committed[0], 20); + assert_eq!( + examined, + 20 + K as u64, + "twenty pending rows and ten returned" + ); + let tx = db.transaction().expect("begin"); + support::run(&tx, "UPDATE docs SET tag = 99 WHERE id = 20"); + let (inside, examined) = nearest_ten(&tx, "row 20 staged over its pending record"); + assert_eq!(inside, committed); + // One staged row, twenty pending rows, ten returned: row 20 was reached twice. + assert_eq!(examined, 1 + 20 + K as u64); + tx.rollback(); + + // A row on the pending list whose earlier vector has a slot. Twenty far rows and row 30 + // are in a block. Then nine rows arrive on the pending list, and row 30 is given a vector + // that makes it the third nearest: the ten pending rows are the answer, and they fill it + // before any slot is looked at. The slot row 30 left behind is of a vector nearer than + // the tenth row, so it is visited, and leads to row 30 again. + let folder = support::folder("reached-twice-pending-and-slot"); + let db = support::open(&folder); + support::declare_docs(&db, AHEAD.len(), ""); + let tx = db.transaction().expect("begin"); + insert_toward(&tx, "a", 1, 20, 0.40, 0.02); + insert_toward(&tx, "a", 30, 1, 0.81, 0.0); + tx.commit().expect("commit"); + assert_eq!(support::seal_all(&db), 21); + let tx = db.transaction().expect("begin"); + insert_toward(&tx, "a", 41, 2, 0.90, 0.02); + insert_toward(&tx, "a", 43, 7, 0.84, 0.02); + tx.commit().expect("commit"); + set_vector(&db, 30, Some(toward(0.86, 3))); + let (answer, _) = nearest_ten(&db, "row 30 pending with a slot of its earlier vector"); + assert_eq!(answer, [41, 42, 30, 43, 44, 45, 46, 47, 48, 49]); + let reader = ReadSession::open(&folder).expect("open the store for reading"); + nearest_ten(&reader, "the same through a read session"); + + // Two slots of one vector, one of them left behind. Rows 1 to 8, row 20, row 50 and + // twenty far rows are in a block; rows 20 and 50 hold the same vector, so their slots are + // the same bytes. Then row 50 is given a vector that makes it the nearest. The answer is + // row 50, rows 1 to 8 and row 20: row 20 is its tenth row, and the slot row 50 left + // behind is held back beside the slot of row 20, proving the same range. + let folder = support::folder("reached-twice-slot-of-an-earlier-vector"); + let db = support::open(&folder); + support::declare_docs(&db, AHEAD.len(), ""); + let tx = db.transaction().expect("begin"); + insert_toward(&tx, "a", 1, 8, 0.90, 0.02); + for id in [20, 50] { + if let Err(error) = support::insert_doc(&tx, id, "a", Some(toward(0.60, 3))) { + panic!("row {id} was refused: {error:?}"); + } + } + insert_toward(&tx, "a", 60, 20, 0.40, 0.02); + tx.commit().expect("commit"); + assert_eq!(support::seal_all(&db), 30); + set_vector(&db, 50, Some(toward(0.95, 5))); + let (answer, _) = nearest_ten( + &db, + "row 50 pending, its earlier slot beside the slot of row 20", + ); + assert_eq!(answer, [50, 1, 2, 3, 4, 5, 6, 7, 8, 20]); + // And once row 50's new vector is in the block too. + assert_eq!(support::seal_all(&db), 1); + let (answer, _) = nearest_ten(&db, "two slots of row 50 in the block"); + assert_eq!(answer, [50, 1, 2, 3, 4, 5, 6, 7, 8, 20]); +} diff --git a/v2/contextdb/tests/vectors_memory.rs b/v2/contextdb/tests/vectors_memory.rs new file mode 100644 index 00000000..22c139fc --- /dev/null +++ b/v2/contextdb/tests/vectors_memory.rs @@ -0,0 +1,794 @@ +//! Memory and vectors: a search under a memory limit ends with the whole answer or with a +//! refusal by name, never with fewer rows; a reader with a small limit searches a store whose +//! blocks are several times its page cache; the page cache an open sets aside follows the +//! limit in force; a limit the session's page cache does not fit under is refused; and a +//! maintenance cycle under a tight limit seals what fits and is refused only when one +//! partition's working set does not. + +mod support; + +use std::path::Path; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use contextdb::{ + Database, DatabaseOpenOptions, Error, ReadSession, ReadSessionOptions, Value, params, +}; + +use support::{Call, Run}; + +const MIB: usize = 1 << 20; +const KIB: usize = 1 << 10; + +fn open_limited(folder: &Path, limit: usize) -> Database { + let mut options = DatabaseOpenOptions::default(); + options.memory_limit = Some(limit); + match support::open_with(folder, options) { + Ok(db) => db, + Err(error) => panic!("an open under a limit of {limit} bytes was refused: {error:?}"), + } +} + +fn read_limited(folder: &Path, limit: Option) -> ReadSession { + let mut options = ReadSessionOptions::default(); + options.memory_limit = limit; + match ReadSession::open_with_options(folder, options) { + Ok(reader) => reader, + Err(error) => panic!("a reader under a limit of {limit:?} was refused: {error:?}"), + } +} + +fn reopen(db: Database, folder: &Path) -> Database { + support::close(db); + support::open(folder) +} + +/// One column of what `SHOW MEMORY_LIMIT` reports. +fn shown(on: &impl Run, column: &str) -> Value { + let report = support::run(on, "SHOW MEMORY_LIMIT"); + let at = report + .columns() + .iter() + .position(|name| name == column) + .unwrap_or_else(|| panic!("SHOW MEMORY_LIMIT has no column {column}")); + match report.rows() { + [row] => row[at].clone(), + rows => panic!("SHOW MEMORY_LIMIT returned {} rows", rows.len()), + } +} + +/// The bytes `SHOW MEMORY_LIMIT` reports as used. +fn used(on: &impl Run) -> usize { + match shown(on, "used") { + Value::Int64(bytes) => bytes as usize, + other => panic!("SHOW MEMORY_LIMIT reports {other:?} as used"), + } +} + +#[test] +fn vector_search_ends_complete_or_refused() { + const DIMENSION: usize = 64; + let folder = support::folder("search-complete-or-refused"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, ""); + // 3,600 rows: 3,000 in blocks and 600 on the pending list. + for first in (1..=3600_i64).step_by(600) { + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, first..first + 600, DIMENSION, |_| "a".to_owned()); + tx.commit().expect("commit"); + // Sealed once row 3,000 is in. + if first == 2401 { + assert_eq!(support::seal_all(&db), 3000); + } + } + let stored = support::stored(&db, support::DOCS); + let query = support::vector(96_000, DIMENSION); + // Two searches: ten rows of the half of the table that passes a filter, and two thousand + // rows of all of it. + let half = |row: &[Value]| matches!(row[2], Value::Int64(tag) if tag < 25); + let searches = [ + ( + support::nearest_sql("docs", "v", Some("tag < 25"), 10), + stored.nearest(half, &query, 10), + ), + ( + support::nearest_sql("docs", "v", None, 2000), + stored.nearest(|_| true, &query, 2000), + ), + ]; + assert_eq!((searches[0].1.len(), searches[1].1.len()), (10, 2000)); + let every = stored.nearest(|_| true, &query, 5000); + assert_eq!(every.len(), 3600); + support::close(db); + + // Runs both searches and says which were answered. A search that is answered returns + // every row of the reference; one that is not is refused by the memory limit, by name. + // Either way the memory in use afterwards is what it was before. + let outcomes = |on: &dyn Run, state: &str| -> [bool; 2] { + [0, 1].map(|number| { + let (sql, wanted) = &searches[number]; + let before = used(&on); + let answered = match support::search(&on, sql, &query) { + Ok(result) => { + let found = support::found(&result, sql); + support::assert_same_rows(&found, wanted, &format!("{state}: {sql}")); + true + } + Err(Error::MemoryBudgetExceeded { .. }) => false, + Err(other) => panic!("{state}: {sql}\nwas refused as {other:?}"), + }; + assert_eq!(used(&on), before, "{state}: memory in use after {sql}"); + answered + }) + }; + + // Readers under limits from the smallest that opens a store, where nothing is left for a + // search, to one with room to spare. + let limits = [ + 9 * MIB, + 9 * MIB + 4 * KIB, + 9 * MIB + 16 * KIB, + 9 * MIB + 64 * KIB, + 9 * MIB + 256 * KIB, + 10 * MIB, + 13 * MIB, + 64 * MIB, + ]; + for limit in limits { + let reader = read_limited(&folder, Some(limit)); + let answered = outcomes(&reader, &format!("a reader under {limit} bytes")); + if limit == 9 * MIB { + assert_eq!( + answered, + [false, false], + "searches under a limit with nothing to spare" + ); + } + if limit == 64 * MIB { + assert_eq!(answered, [true, true], "searches under a limit of 64 MiB"); + } + } + // The same through the writer. + let db = open_limited(&folder, 9 * MIB); + assert_eq!(outcomes(&db, "the writer under 9 MiB"), [false, false]); + assert_eq!(db.memory_used(), 9 * MIB); + support::close(db); + let db = open_limited(&folder, 64 * MIB); + assert_eq!(outcomes(&db, "the writer under 64 MiB"), [true, true]); + assert_eq!(db.memory_used(), 17 * MIB); + support::close(db); + + // A cycle under a limit too small for it is refused by name, and nothing was sealed. + let db = open_limited(&folder, 9 * MIB + 64 * KIB); + let before = db.memory_used(); + match db.run_maintenance_cycle() { + Err(Error::MemoryBudgetExceeded { .. }) => {} + other => panic!("a cycle with 64 KiB to work in got {other:?}"), + } + assert_eq!( + db.memory_used(), + before, + "memory in use after the refused cycle" + ); + support::close(db); + // Every vector is still found, the 600 that waited still wait, and the next cycle under + // no limit seals them. + let db = support::open(&folder); + let sql = support::nearest_sql("docs", "v", None, 5000); + let result = support::assert_nearest(&db, &sql, &query, &every, "after the refused cycle"); + assert!( + result.rows_examined() >= 600 + 3600, + "600 pending rows and 3,600 returned rows were read as {} rows", + result.rows_examined() + ); + drop(result); + let report = db + .run_maintenance_cycle() + .expect("a cycle under no limit") + .vector; + assert_eq!( + ( + report.maintenance_vectors_done, + report.maintenance_vectors_remaining + ), + (600, 0) + ); + support::assert_nearest(&db, &sql, &query, &every, "after the cycle that sealed"); +} + +#[test] +fn memory_stays_bounded_as_data_grows_vector() { + const DIMENSION: usize = 768; + // Vectors in the large partition. In a block each takes 8 + 768 bytes of slot, 4 of key + // end and 9 of key, so these are 25 MB of blocks: three times the page cache of 8 MiB + // that the reader's limit gives it, and the fewest that are. + const LARGE: i64 = 31_896; + const SMALL: i64 = 200; + const LIMIT: usize = 16 * MIB; + const CACHE: usize = 8 * MIB; + const BLOCK: u64 = 262_080; + const { assert!(LARGE as u64 * (8 + 768 + 4 + 9) > 3 * CACHE as u64) }; + const { assert!((LARGE as u64 - 1) * (8 + 768 + 4 + 9) <= 3 * CACHE as u64) }; + + let folder = support::folder("memory-bounded-vectors"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, "PARTITION_KEY (ctx)"); + for first in (1..=LARGE + SMALL).step_by(1000) { + let tx = db.transaction().expect("begin"); + let last = (first + 999).min(LARGE + SMALL); + support::insert_docs(&tx, first..=last, DIMENSION, |n| { + String::from(if n <= LARGE { "large" } else { "small" }) + }); + tx.commit().expect("commit"); + } + assert_eq!(support::seal_all(&db), (LARGE + SMALL) as u64); + // The reference for the small partition, from its rows alone. + let small = format!("SELECT id, ctx, tag, v FROM docs WHERE id > {LARGE}"); + let stored = support::stored(&db, &small); + assert_eq!(stored.rows().len() as i64, SMALL); + let query = support::vector(97_000, DIMENSION); + support::close(db); + + // A reader with a limit of 16 MiB, whose storage door counts the bytes it reads. + let read = Arc::new(AtomicU64::new(0)); + let mut options = ReadSessionOptions::default(); + options.memory_limit = Some(LIMIT); + options.storage = Some(support::storage({ + let read = read.clone(); + move |call| { + if let Call::Read(_, bytes) = call { + read.fetch_add(bytes.len() as u64, Ordering::SeqCst); + } + } + })); + let reader = ReadSession::open_with_options(&folder, options) + .expect("open a store larger than the limit"); + // Opening reads no block: less than one block's bytes in all. + let at_open = read.load(Ordering::SeqCst); + assert!( + at_open < BLOCK, + "opening the store read {at_open} bytes of it before any statement" + ); + assert_eq!(used(&reader), 9 * MIB, "set aside by an open under 16 MiB"); + + // It searches the small partition and answers, reading a small part of the file. + for k in [10, 1000] { + let wanted = stored.nearest(|row| row[1] == support::text("small"), &query, k); + let sql = support::nearest_sql("docs", "v", Some("ctx = 'small'"), k); + support::assert_nearest(&reader, &sql, &query, &wanted, "the small partition"); + assert_eq!( + used(&reader), + 9 * MIB, + "memory in use once the answer is dropped" + ); + } + let searching = read.load(Ordering::SeqCst) - at_open; + assert!( + searching < 8 * MIB as u64, + "two searches of a partition of {SMALL} vectors read {searching} bytes of the file" + ); + // The row count of the whole table is still one statement away. + assert_eq!( + support::int(&reader, "SELECT COUNT(*) FROM docs"), + LARGE + SMALL + ); +} + +#[test] +fn the_page_cache_follows_the_limit_in_force() { + let folder = support::folder("page-cache-follows-limit"); + // With no limit an open sets 9 MiB aside: the smallest page cache, 8 MiB, and 1 MiB. + let db = support::open(&folder); + assert_eq!(db.memory_limit(), None); + assert_eq!(db.memory_used(), 9 * MIB, "set aside with no limit"); + assert_eq!(used(&db), 9 * MIB); + support::declare(&db); + support::insert(&db, 1..=10); + support::close(db); + + // Under a limit the cache is a quarter of it and never less than 8 MiB. + for (limit, set_aside) in [(64, 17), (16, 9), (9, 9), (32, 9), (40, 11), (1024, 257)] { + let db = open_limited(&folder, limit * MIB); + assert_eq!(db.memory_limit(), Some(limit * MIB)); + assert_eq!( + db.memory_used(), + set_aside * MIB, + "set aside under a limit of {limit} MiB" + ); + // Under 9 MiB nothing is left for a statement; under the others the store reads. + if limit > 9 { + support::assert_rows_are(&db, 1..=10); + } + support::close(db); + } + + // A reader's cache follows the limit of its own options, and with none it has no limit + // and the smallest cache. + for (limit, set_aside) in [(None, 9), (Some(64), 17), (Some(16), 9), (Some(1024), 257)] { + let reader = read_limited(&folder, limit.map(|limit| limit * MIB)); + assert_eq!( + used(&reader), + set_aside * MIB, + "a reader under {limit:?} MiB" + ); + support::assert_rows_are(&reader, 1..=10); + } + + // A limit the store remembers sizes the cache of the next open that names none. + let db = support::open(&folder); + support::run(&db, "SET MEMORY_LIMIT 512M"); + assert_eq!(db.memory_limit(), Some(512 * MIB)); + assert_eq!( + db.memory_used(), + 9 * MIB, + "the cache of an open keeps its size" + ); + let db = reopen(db, &folder); + assert_eq!(db.memory_limit(), Some(512 * MIB)); + assert_eq!( + db.memory_used(), + 129 * MIB, + "set aside under a remembered limit of 512 MiB" + ); + assert_eq!(used(&db), 129 * MIB); + assert_eq!(shown(&db, "limit"), Value::Int64(512 * MIB as i64)); + assert_eq!(shown(&db, "available"), Value::Int64(383 * MIB as i64)); + // A reader does not take the limit the store remembers. + let reader = read_limited(&folder, None); + assert_eq!(shown(&reader, "limit"), Value::Text("none".to_owned())); + assert_eq!(used(&reader), 9 * MIB); + let reader = read_limited(&folder, Some(16 * MIB)); + assert_eq!(used(&reader), 9 * MIB); + assert_eq!(shown(&reader, "available"), Value::Int64(7 * MIB as i64)); + // An open that names a limit has its own cache and leaves what the store remembers. + support::close(db); + let db = open_limited(&folder, 64 * MIB); + assert_eq!(db.memory_used(), 17 * MIB); + let db = reopen(db, &folder); + assert_eq!(db.memory_limit(), Some(512 * MIB)); + assert_eq!(db.memory_used(), 129 * MIB); +} + +/// Asserts that a limit was refused because the session's page cache does not fit under it: +/// by the memory limit's name, for the operation `set_limit`, in a sentence that names the +/// limit asked for and the smallest limit the session can take, both in bytes, and says to +/// open the store with the limit as an open option. +fn refused_below_the_cache(refused: contextdb::Result<()>, asked: usize, smallest: usize) { + let error = match refused { + Err(error) => error, + Ok(()) => panic!("a limit of {asked} bytes was taken; the smallest is {smallest}"), + }; + let sentence = error.to_string(); + match error { + Error::MemoryBudgetExceeded { operation, .. } => { + assert_eq!(operation, "set_limit", "{sentence}"); + } + other => panic!("a limit of {asked} bytes was refused as {other:?}"), + } + assert!( + sentence.contains(&asked.to_string()), + "the limit asked for: {sentence}" + ); + assert!( + sentence.contains(&smallest.to_string()), + "the smallest limit: {sentence}" + ); + assert!( + sentence.to_lowercase().contains("open the store"), + "the next action, to open the store with the limit as an open option: {sentence}" + ); +} + +#[test] +fn a_limit_below_the_sessions_cache_is_refused_and_one_above_it_is_taken() { + type Set = fn(&Database, usize) -> contextdb::Result<()>; + let in_sql: Set = |db, mib| { + db.execute(&format!("SET MEMORY_LIMIT '{mib}M'"), ¶ms! {}) + .map(|_| ()) + }; + let by_call: Set = |db, mib| db.set_memory_limit(Some(mib * MIB)); + let none = Value::Text("none".to_owned()); + + for (name, set) in [("sql", in_sql), ("call", by_call)] { + // Opened with no limit: 9 MiB are set aside, and that is the smallest limit. + let folder = support::folder(&format!("limit-below-cache-{name}")); + let db = support::open(&folder); + assert_eq!(db.memory_used(), 9 * MIB); + refused_below_the_cache(set(&db, 8), 8 * MIB, 9 * MIB); + refused_below_the_cache(set(&db, 1), MIB, 9 * MIB); + // The limit in force, the memory set aside and the limit the store remembers stay as + // they were. + assert_eq!(db.memory_limit(), None); + assert_eq!(shown(&db, "limit"), none); + assert_eq!(db.memory_used(), 9 * MIB); + let db = reopen(db, &folder); + assert_eq!(db.memory_limit(), None, "a refused limit was remembered"); + assert_eq!(db.memory_used(), 9 * MIB); + + // A limit above it is taken and shown, with 9 MiB still set aside; the next open + // sizes its cache from it. + set(&db, 512).expect("a limit above what is set aside"); + assert_eq!(db.memory_limit(), Some(512 * MIB)); + assert_eq!(shown(&db, "limit"), Value::Int64(512 * MIB as i64)); + assert_eq!(db.memory_used(), 9 * MIB); + assert_eq!(used(&db), 9 * MIB); + let db = reopen(db, &folder); + assert_eq!(db.memory_limit(), Some(512 * MIB)); + assert_eq!(db.memory_used(), 129 * MIB); + // Now 129 MiB is the smallest. + refused_below_the_cache(set(&db, 128), 128 * MIB, 129 * MIB); + assert_eq!(db.memory_limit(), Some(512 * MIB)); + assert_eq!(db.memory_used(), 129 * MIB); + set(&db, 130).expect("a limit just above what is set aside"); + assert_eq!(db.memory_limit(), Some(130 * MIB)); + assert_eq!(db.memory_used(), 129 * MIB); + support::close(db); + + // Opened with options that name 16 MiB: 9 MiB are set aside. 8 MiB is refused and + // 1 GiB is taken, and the next open with no options is under 1 GiB. + let folder = support::folder(&format!("limit-below-small-cache-{name}")); + let db = open_limited(&folder, 16 * MIB); + assert_eq!(db.memory_used(), 9 * MIB); + refused_below_the_cache(set(&db, 8), 8 * MIB, 9 * MIB); + assert_eq!(db.memory_limit(), Some(16 * MIB)); + assert_eq!(db.memory_used(), 9 * MIB); + support::close(db); + assert_eq!( + support::open(&folder).memory_limit(), + None, + "a refused limit was remembered" + ); + let db = open_limited(&folder, 16 * MIB); + set(&db, 1024).expect("a limit above what is set aside"); + assert_eq!(db.memory_limit(), Some(1024 * MIB)); + assert_eq!(db.memory_used(), 9 * MIB); + let db = reopen(db, &folder); + assert_eq!(db.memory_limit(), Some(1024 * MIB)); + assert_eq!(db.memory_used(), 257 * MIB); + // Under the 1 GiB the store remembers the cache is 256 MiB, and 64 MiB is refused. + refused_below_the_cache(set(&db, 64), 64 * MIB, 257 * MIB); + assert_eq!(db.memory_limit(), Some(1024 * MIB)); + } +} + +#[test] +fn a_limit_of_64_mib_is_taken_in_a_store_opened_with_no_limit() { + type Set = fn(&Database, usize) -> contextdb::Result<()>; + let in_sql: Set = |db, mib| { + db.execute(&format!("SET MEMORY_LIMIT {mib}M"), ¶ms! {}) + .map(|_| ()) + }; + let by_call: Set = |db, mib| db.set_memory_limit(Some(mib * MIB)); + let insert = "INSERT INTO big (id, body) VALUES ($id, $body)"; + // Refused by the memory limit, which the refusal names as 64 MiB. + let over_64_mib = + |refused: contextdb::Result, state: &str| match refused { + Err(Error::MemoryBudgetExceeded { + requested_bytes, + available_bytes, + budget_limit_bytes, + .. + }) => { + assert_eq!(budget_limit_bytes, 64 * MIB as u64, "{state}"); + assert!(requested_bytes > available_bytes, "{state}"); + } + Err(other) => panic!("{state}: refused as {other:?}"), + Ok(_) => panic!("{state}: a value larger than what the limit leaves was taken"), + }; + + for (name, set) in [("sql", in_sql), ("call", by_call)] { + // Opened the plain way, with nothing remembered: no limit, and 9 MiB set aside. + let folder = support::folder(&format!("limit-64-in-a-plain-store-{name}")); + let db = support::open(&folder); + assert_eq!(db.memory_limit(), None); + assert_eq!(db.memory_used(), 9 * MIB); + support::run(&db, "CREATE TABLE big (id INTEGER PRIMARY KEY, body TEXT)"); + + // The first limit declared is taken and is in force at once, with the cache this + // session already has. + set(&db, 64).expect("64 MiB in a session opened with no limit"); + assert_eq!(db.memory_limit(), Some(64 * MIB)); + assert_eq!(shown(&db, "limit"), Value::Int64(64 * MIB as i64)); + assert_eq!(db.memory_used(), 9 * MIB); + assert_eq!(used(&db), 9 * MIB); + assert_eq!(shown(&db, "available"), Value::Int64(55 * MIB as i64)); + + // It holds at once for what statements take: a value of 56 MiB is more than the + // 55 MiB the limit leaves and is refused by the memory limit; one of 1 MiB is taken. + let large = "x".repeat(56 * MIB); + over_64_mib( + db.execute(insert, ¶ms! { "id" => 1_i64, "body" => large.as_str() }), + "56 MiB in the session that declared the limit", + ); + assert_eq!( + db.memory_used(), + 9 * MIB, + "a refused statement left memory counted" + ); + let fits = "x".repeat(MIB); + db.execute(insert, ¶ms! { "id" => 2_i64, "body" => fits.as_str() }) + .expect("a value of 1 MiB under a limit of 64 MiB"); + assert_eq!(support::int(&db, "SELECT COUNT(*) FROM big"), 1); + + // The store remembers it: the next plain open is under 64 MiB, with a cache of a + // quarter of it, 16 MiB, and 1 MiB set aside before any statement. + let db = reopen(db, &folder); + assert_eq!(db.memory_limit(), Some(64 * MIB)); + assert_eq!(db.memory_used(), 17 * MIB); + assert_eq!(used(&db), 17 * MIB); + assert_eq!(shown(&db, "limit"), Value::Int64(64 * MIB as i64)); + assert_eq!(shown(&db, "available"), Value::Int64(47 * MIB as i64)); + over_64_mib( + db.execute(insert, ¶ms! { "id" => 1_i64, "body" => large.as_str() }), + "56 MiB after the reopen", + ); + assert_eq!(support::int(&db, "SELECT COUNT(*) FROM big"), 1); + support::close(db); + + // The refusal that stays: a session opened under 1 GiB has a cache of 256 MiB, and + // 64 MiB is refused there with the action, to open the store with the limit. + let folder = support::folder(&format!("limit-64-under-1-gib-{name}")); + let db = open_limited(&folder, 1024 * MIB); + assert_eq!(db.memory_used(), 257 * MIB); + refused_below_the_cache(set(&db, 64), 64 * MIB, 257 * MIB); + assert_eq!(db.memory_limit(), Some(1024 * MIB)); + assert_eq!(db.memory_used(), 257 * MIB); + support::close(db); + assert_eq!( + support::open(&folder).memory_limit(), + None, + "a refused limit was remembered" + ); + } +} + +#[test] +fn a_cycle_under_a_tight_limit_seals_what_fits_and_is_refused_only_when_one_partition_does_not() { + const DIMENSION: usize = 16; + // Forty partitions with a hundred pending vectors each: 4,000 in all, fewer than the + // 4,096 that one cycle with room seals, so only memory cuts a cycle short here. + const PARTITIONS: i64 = 40; + const EACH: i64 = 100; + const ROWS: i64 = PARTITIONS * EACH; + // The bytes of the one block a cycle builds in, which it reserves whole. + const BLOCK: usize = 262_080; + // A limit with room for one partition's working set and not for all forty. One partition + // takes the block being built, the keys of its 100 pending records, and what it stages: + // its block, 100 removals and its entry, about 30 KiB together. Sealing all 4,000 at once + // would stage 4,000 removals at 128 bytes and more each, 512,000 bytes, beside the block + // being built: twice the 393,216 bytes this limit leaves. + const TIGHT: usize = 9 * MIB + 384 * KIB; + const { assert!(BLOCK + 100 * KIB < 384 * KIB && 384 * KIB < BLOCK + ROWS as usize * 128) }; + // A limit under which no partition's working set fits: it leaves half of one block. + const SHORT: usize = 9 * MIB + 128 * KIB; + const { assert!(128 * KIB < BLOCK) }; + + let folder = support::folder("sealing-under-a-tight-limit"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, "PARTITION_KEY (ctx)"); + let partition = |n: i64| format!("p{:02}", n % PARTITIONS); + for first in (1..=ROWS).step_by(1000) { + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, first..first + 1000, DIMENSION, partition); + tx.commit().expect("commit"); + } + support::close(db); + + // The searches are a reader's, which has no limit of its own: ten rows of everything and + // ten of one partition. Returns the rows each search read. + let reader = read_limited(&folder, None); + let query = support::vector(95_000, DIMENSION); + let searches = [None, Some("ctx = 'p07'")]; + let wanted_now = || { + let stored = support::stored(&reader, support::DOCS); + searches.map(|filter| { + let passes = |row: &[Value]| filter.is_none() || row[1] == support::text("p07"); + let wanted = stored.nearest(passes, &query, 10); + assert_eq!(wanted.len(), 10); + wanted + }) + }; + let searched = |wanted: &[Vec; 2], state: &str| -> [u64; 2] { + [0, 1].map(|number| { + let sql = support::nearest_sql("docs", "v", searches[number], 10); + let result = support::assert_nearest(&reader, &sql, &query, &wanted[number], state); + result.rows_examined() + }) + }; + let wanted = wanted_now(); + // Every vector is pending: a search reads each pending row of the partitions it selects + // once, and each row it returns once more. + assert_eq!( + searched(&wanted, "nothing sealed"), + [ROWS as u64 + 10, EACH as u64 + 10] + ); + + // Under the tight limit every cycle seals at least one partition and none is refused. + let db = open_limited(&folder, TIGHT); + let set_aside = db.memory_used(); + let mut remaining = ROWS as u64; + let mut cycles = 0; + while remaining > 0 { + cycles += 1; + assert!( + cycles <= PARTITIONS, + "{remaining} vectors still wait after {PARTITIONS} cycles over {PARTITIONS} partitions" + ); + let report = match db.run_maintenance_cycle() { + Ok(report) => report.vector, + Err(error) => panic!( + "cycle {cycles}, with {remaining} vectors waiting in partitions of {EACH}, was \ + refused: {error:?}" + ), + }; + let done = report.maintenance_vectors_done; + assert!( + done >= EACH as u64 && done <= remaining, + "cycle {cycles} sealed {done} of {remaining} waiting vectors: at least one \ + partition's {EACH} were to be sealed" + ); + assert_eq!( + report.maintenance_vectors_remaining, + remaining - done, + "what cycle {cycles} reports as still waiting" + ); + if cycles == 1 { + assert!( + done < ROWS as u64, + "the first cycle sealed all {ROWS} vectors in {} bytes of room", + TIGHT - 9 * MIB + ); + } + assert_eq!( + db.memory_used(), + set_aside, + "memory in use after cycle {cycles}" + ); + remaining -= done; + // The same rows at the same distances after every cycle. + searched( + &wanted, + &format!("after cycle {cycles} under the tight limit"), + ); + } + let position = db.committed_watermark(); + let report = db + .run_maintenance_cycle() + .expect("a cycle with nothing to seal"); + assert_eq!( + ( + report.vector.maintenance_vectors_done, + report.vector.maintenance_vectors_remaining + ), + (0, 0) + ); + assert_eq!(db.committed_watermark(), position); + support::close(db); + + // Two hundred more vectors come to wait in one partition. + let db = support::open(&folder); + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, ROWS + 1..=ROWS + 200, DIMENSION, |_| "p07".to_owned()); + tx.commit().expect("commit"); + support::close(db); + let wanted = wanted_now(); + let read_before = searched(&wanted, "two hundred waiting"); + + // Under a limit that leaves less than the block a cycle builds in, the working set of + // that one partition does not fit, and the cycle is refused by name. Nothing was staged: + // no commit was made and nothing is counted. + let db = open_limited(&folder, SHORT); + let (set_aside, position) = (db.memory_used(), db.committed_watermark()); + match db.run_maintenance_cycle() { + Err(Error::MemoryBudgetExceeded { + operation, + budget_limit_bytes, + .. + }) => { + assert_eq!(operation, "maintenance"); + assert_eq!(budget_limit_bytes, SHORT as u64); + } + other => panic!("a cycle with half a block of room got {other:?}"), + } + assert_eq!( + db.memory_used(), + set_aside, + "memory in use after the refusal" + ); + assert_eq!( + db.committed_watermark(), + position, + "a refused cycle committed" + ); + // The store is unharmed: the same search answers the same and reads the same rows, the + // two hundred that waited among them. + assert_eq!( + searched(&wanted, "after the refused cycle"), + read_before, + "the rows the searches read after the refused cycle" + ); + + // With the limit raised the cycle goes through, and all two hundred were still waiting. + db.set_memory_limit(Some(64 * MIB)) + .expect("raise the limit above what the session has set aside"); + let report = db + .run_maintenance_cycle() + .expect("a cycle under the raised limit") + .vector; + assert_eq!( + ( + report.maintenance_vectors_done, + report.maintenance_vectors_remaining + ), + (200, 0) + ); + searched(&wanted, "after the cycle under the raised limit"); +} + +#[test] +fn the_memory_of_a_search_is_given_back_when_its_result_is_dropped() { + const DIMENSION: usize = 16; + let folder = support::folder("search-gives-memory-back"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, "PARTITION_KEY (ctx)"); + // Two partitions: 300 rows in blocks, then 100 on pending lists. + let ctx = |n: i64| format!("c{}", n % 2); + for (numbers, sealed) in [(1..=300, true), (301..=400, false)] { + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, numbers, DIMENSION, ctx); + tx.commit().expect("commit"); + if sealed { + assert_eq!(support::seal_all(&db), 300); + } + } + let query = support::vector(96_500, DIMENSION); + let reader = read_limited(&folder, Some(16 * MIB)); + // A transaction that has staged rows of its own: its searches read those too. + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, 401..=420, DIMENSION, ctx); + + // Every partition, one partition, a filter on another column, and more rows than there + // are; with only the key selected, with the distance, and with every column. + let searches = [ + support::nearest_sql("docs", "v", None, 10), + support::nearest_sql("docs", "v", Some("ctx = 'c1'"), 10), + support::nearest_sql("docs", "v", Some("tag < 25"), 50), + support::nearest_sql("docs", "v", None, 1000), + "SELECT id FROM docs ORDER BY v <=> $q LIMIT 10".to_owned(), + "SELECT * FROM docs ORDER BY v <=> $q LIMIT 100".to_owned(), + ]; + let doors = [ + (&db as &dyn Run, "the writer"), + (&tx as &dyn Run, "a transaction"), + (&reader as &dyn Run, "a reader"), + ]; + for (on, door) in doors { + for sql in &searches { + let before = used(&on); + let result = match support::search(&on, sql, &query) { + Ok(result) => result, + Err(error) => panic!("through {door}: {sql}\nwas refused: {error:?}"), + }; + assert!( + result.rows().len() >= 10, + "through {door}: the rows of {sql}" + ); + // While the result lives its rows are counted, so the figure does move. + let held = used(&on); + assert!( + held > before, + "through {door}: with the result of {sql}\nin hand, {held} bytes are in use, \ + and {before} were before it" + ); + drop(result); + assert_eq!( + used(&on), + before, + "through {door}: memory in use once the result of {sql}\nis dropped" + ); + } + } + // What the transaction holds is its own rows, before and after: it commits them. + tx.commit().expect("commit"); + assert_eq!(support::int(&db, "SELECT COUNT(*) FROM docs"), 420); +} diff --git a/v2/contextdb/tests/vectors_partitions.rs b/v2/contextdb/tests/vectors_partitions.rs new file mode 100644 index 00000000..e36fba63 --- /dev/null +++ b/v2/contextdb/tests/vectors_partitions.rs @@ -0,0 +1,697 @@ +//! The partitions of a VECTOR column: a search over one, several or all of them gives one +//! answer; two key tuples are one partition only when every value is equal as SQL `=` has it; +//! and a column's partition limit counts the partitions that hold a live vector, refuses the +//! statement that would pass it and changes nothing when it does. + +mod support; + +use contextdb::{Database, Error, ReadSession, ReadSessionOptions, Value, VectorIndexRef, params}; + +use support::{Child, Near, Run}; + +/// A table whose VECTOR column is partitioned by two columns. +const NOTES: &str = "SELECT id, ctx, kind, v FROM notes"; + +fn declare_notes(db: &Database, dimension: usize, clauses: &str) { + support::run( + db, + &format!( + "CREATE TABLE notes (id INTEGER PRIMARY KEY, ctx TEXT NOT NULL, \ + kind INTEGER NOT NULL, v VECTOR({dimension}) PARTITION_KEY (ctx, kind) {clauses})" + ), + ); +} + +fn insert_note(on: &impl Run, id: i64, ctx: &str, kind: i64, vector: Vec) { + let inserted = on.run( + "INSERT INTO notes (id, ctx, kind, v) VALUES ($id, $ctx, $kind, $v)", + ¶ms! { "id" => id, "ctx" => ctx, "kind" => kind, "v" => vector }, + ); + if let Err(error) = inserted { + panic!("note {id} under ({ctx:?}, {kind}) was refused: {error:?}"); + } +} + +/// True for a row of `notes` under this key tuple. +fn under(row: &[Value], ctx: &str, kind: i64) -> bool { + row[1] == support::text(ctx) && row[2] == Value::Int64(kind) +} + +fn ids(answer: &[Near]) -> Vec { + answer + .iter() + .map(|near| match near.key { + Value::Int64(id) => id, + ref other => panic!("the key {other:?} is not an INTEGER"), + }) + .collect() +} + +#[test] +fn multi_partition_search_returns_one_global_answer() { + const DIMENSION: usize = 16; + let folder = support::folder("multi-partition-search"); + let db = support::open(&folder); + // Six partitions of a hundred notes each: three contexts, two kinds. + declare_notes(&db, DIMENSION, ""); + let ctx = |n: i64| ["a", "b", "c"][(n % 3) as usize]; + let kind = |n: i64| (n / 3) % 2; + let tx = db.transaction().expect("begin"); + for n in 1..=600 { + insert_note(&tx, n, ctx(n), kind(n), support::vector(n, DIMENSION)); + } + tx.commit().expect("commit"); + // And a table partitioned by one column: three partitions of a hundred. + support::declare_docs(&db, DIMENSION, "PARTITION_KEY (ctx)"); + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, 1..=300, DIMENSION, |n| ctx(n).to_owned()); + tx.commit().expect("commit"); + + let queries = [ + support::vector(80_001, DIMENSION), + support::vector(80_002, DIMENSION), + ]; + // Each scope as a filter, the rows it lets through, and how many partitions it selects. + type Scope = ( + &'static str, + Option<&'static str>, + fn(&[Value]) -> bool, + u64, + ); + let of_notes: [Scope; 9] = [ + ( + "one", + Some("ctx = 'a' AND kind = 1"), + |row| under(row, "a", 1), + 1, + ), + ( + "several, as a list of row values", + Some("(ctx, kind) IN (('a', 1), ('b', 0))"), + |row| under(row, "a", 1) || under(row, "b", 0), + 2, + ), + ( + "the same list the other way round", + Some("(ctx, kind) IN (('b', 0), ('a', 1))"), + |row| under(row, "a", 1) || under(row, "b", 0), + 2, + ), + ( + "several, as an OR", + Some("(ctx = 'a' AND kind = 1) OR (ctx = 'b' AND kind = 0)"), + |row| under(row, "a", 1) || under(row, "b", 0), + 2, + ), + ("all", None, |_| true, 6), + ( + "every partition of one first key value", + Some("ctx = 'a'"), + |row| row[1] == support::text("a"), + 2, + ), + ( + "a tuple no row has", + Some("ctx = 'zz' AND kind = 7"), + |_| false, + 0, + ), + ( + "a tuple with a first value rows have and a second none has", + Some("ctx = 'a' AND kind = 7"), + |_| false, + 0, + ), + ( + "one tuple rows have and one no row has", + Some("(ctx, kind) IN (('a', 1), ('zz', 7))"), + |row| under(row, "a", 1), + 1, + ), + ]; + let of_docs: [Scope; 5] = [ + ( + "one", + Some("ctx = 'b'"), + |row| row[1] == support::text("b"), + 1, + ), + ( + "several, as a list", + Some("ctx IN ('a', 'b')"), + |row| row[1] != support::text("c"), + 2, + ), + ( + "the same list the other way round", + Some("ctx IN ('b', 'a')"), + |row| row[1] != support::text("c"), + 2, + ), + ("all", None, |_| true, 3), + ("a key no row has", Some("ctx IN ('zz')"), |_| false, 0), + ]; + // Every scope of both tables equals the reference, for ten rows and for more rows than + // there are. `pending` says how many vectors of each partition are on its pending list: + // then a search reads exactly those of the partitions it selects, and each row it returns + // once more. + let check = |pending: Option, state: &str| { + for (table, plain, scopes) in [ + ("notes", NOTES, of_notes.as_slice()), + ("docs", support::DOCS, of_docs.as_slice()), + ] { + let stored = support::stored(&db, plain); + for (name, filter, passes, partitions) in scopes { + for query in &queries { + for k in [10, 1000] { + let wanted = stored.nearest(passes, query, k); + let sql = support::nearest_sql(table, "v", *filter, k); + let what = format!("{state}, {table}, {name}"); + let result = support::assert_nearest(&db, &sql, query, &wanted, &what); + if let Some(pending) = pending { + assert_eq!( + result.rows_examined(), + partitions * pending + wanted.len() as u64, + "{what}: rows read by {sql}" + ); + } + } + } + } + } + }; + + check(Some(100), "unsealed"); + support::seal_all(&db); + check(None, "sealed"); + + // LIMIT cuts after the whole scope: three notes of the direction searched for go into the + // partition that was made last, the one of note 6, and a search of everything with LIMIT 3 + // returns them. Multiplying by two and by four is exact, so the three are at one distance. + let near = &queries[0]; + for (id, times) in [(601, 1.0_f32), (602, 2.0), (603, 4.0)] { + let vector = near.iter().map(|number| number * times).collect(); + insert_note(&db, id, ctx(6), kind(6), vector); + } + for state in ["the three pending", "the three sealed"] { + let stored = support::stored(&db, NOTES); + let wanted = stored.nearest(|_| true, near, 3); + assert_eq!(ids(&wanted), [601, 602, 603]); + let sql = support::nearest_sql("notes", "v", None, 3); + support::assert_nearest(&db, &sql, near, &wanted, state); + support::seal_all(&db); + } + + // With everything sealed, a search of one partition gives the same rows for the same + // counted work while other partitions grow, pending or sealed, and while new ones are made. + let one = support::nearest_sql("notes", "v", Some("ctx = 'c' AND kind = 0"), 10); + let measure = |state: &str| -> Vec<(Vec, u64)> { + let stored = support::stored(&db, NOTES); + let measured = queries.iter().map(|query| { + let wanted = stored.nearest(|row| under(row, "c", 0), query, 10); + let result = support::assert_nearest(&db, &one, query, &wanted, state); + (wanted, result.rows_examined()) + }); + measured.collect() + }; + let before = measure("before other partitions grow"); + let tx = db.transaction().expect("begin"); + for n in 1001..=1500 { + insert_note(&tx, n, "b", 0, support::vector(n, DIMENSION)); + } + tx.commit().expect("commit"); + assert!(measure("another partition grew, its new rows pending") == before); + support::seal_all(&db); + assert!(measure("another partition grew, its new rows sealed") == before); + for n in 1501..=1600 { + insert_note(&db, n, "d", n % 5, support::vector(n, DIMENSION)); + } + assert!(measure("five more partitions were made") == before); + check(None, "at the end"); +} + +#[test] +fn vector_partition_identity_is_typed_sql_equality() { + const DIMENSION: usize = 8; + let folder = support::folder("partition-identity"); + let db = support::open(&folder); + declare_notes(&db, DIMENSION, "MAX_PARTITIONS 5"); + // Five key tuples that a careless comparison takes for fewer: three texts that differ in + // the case of a letter or in a space at the end, and two tuples that print alike when + // their values are written one after the other. + let tuples: [(&str, i64); 5] = [("ab", 1), ("Ab", 1), ("ab ", 1), ("1", 11), ("11", 1)]; + let tx = db.transaction().expect("begin"); + for n in 1..=100 { + let (ctx, kind) = tuples[(n % 5) as usize]; + insert_note(&tx, n, ctx, kind, support::vector(n, DIMENSION)); + } + tx.commit().expect("commit"); + // They are five partitions and the limit is five: a sixth tuple is refused. + let sixth = db.execute( + "INSERT INTO notes (id, ctx, kind, v) VALUES ($id, $ctx, $kind, $v)", + ¶ms! { + "id" => 101_i64, "ctx" => "AB", "kind" => 1_i64, + "v" => support::vector(101, DIMENSION), + }, + ); + match sixth { + Err(Error::VectorPartitionLimitExceeded { + max_partitions: 5, .. + }) => {} + other => panic!("a sixth key tuple under MAX_PARTITIONS 5 got {other:?}"), + } + + let queries = [ + support::vector(81_001, DIMENSION), + support::vector(81_002, DIMENSION), + ]; + // A search of one tuple returns that tuple's rows and no other's. While nothing is sealed + // it reads the twenty vectors of its own partition and each row it returns once more. + let each_alone = |on: &dyn Run, pending: Option, state: &str| { + let stored = support::stored(&on, NOTES); + for (ctx, kind) in tuples { + let filter = format!("ctx = '{ctx}' AND kind = {kind}"); + for query in &queries { + let wanted = stored.nearest(|row| under(row, ctx, kind), query, 1000); + let sql = support::nearest_sql("notes", "v", Some(&filter), 1000); + let what = format!("{state}, ({ctx:?}, {kind})"); + let result = support::assert_nearest(&on, &sql, query, &wanted, &what); + if let Some(pending) = pending { + assert_eq!(wanted.len(), 20, "{what}"); + assert_eq!(result.rows_examined(), pending + 20, "{what}"); + } + } + } + }; + each_alone(&db, Some(20), "unsealed"); + support::seal_all(&db); + each_alone(&db, None, "sealed"); + + // An UPDATE of a key column moves the row's vector. Row 5 is under ("ab", 1). + let found_under = |on: &dyn Run, ctx: &str, state: &str| -> bool { + let stored = support::stored(&on, NOTES); + let query = support::vector(5, DIMENSION); + let wanted = stored.nearest(|row| under(row, ctx, 1), &query, 1000); + let filter = format!("ctx = '{ctx}' AND kind = 1"); + let sql = support::nearest_sql("notes", "v", Some(&filter), 1000); + support::assert_nearest(&on, &sql, &query, &wanted, state); + ids(&wanted).contains(&5) + }; + assert!(found_under(&db, "ab", "before the move")); + assert!(!found_under(&db, "Ab", "before the move")); + let reader = ReadSession::open(&folder).expect("open the store for reading"); + let tx = db.transaction().expect("begin"); + support::run(&tx, "UPDATE notes SET ctx = 'Ab' WHERE id = 5"); + // Inside the transaction it is under the new key and not under the old. + assert!(found_under(&tx, "Ab", "inside the moving transaction")); + assert!(!found_under(&tx, "ab", "inside the moving transaction")); + // For everyone else it has not moved. + for on in [&db as &dyn Run, &reader as &dyn Run] { + assert!(found_under(on, "ab", "beside the moving transaction")); + assert!(!found_under(on, "Ab", "beside the moving transaction")); + } + // The commit moves the row and its vector together. + tx.commit().expect("commit"); + for on in [&db as &dyn Run, &reader as &dyn Run] { + assert!(found_under(on, "Ab", "after the move")); + assert!(!found_under(on, "ab", "after the move")); + } + each_alone(&db, None, "after the move, the moved vector pending"); + support::seal_all(&db); + each_alone(&db, None, "after the move, everything sealed"); + // Back again, in a statement of its own. + support::run(&db, "UPDATE notes SET ctx = 'ab' WHERE id = 5"); + assert!(found_under(&db, "ab", "after the move back")); + assert!(!found_under(&db, "Ab", "after the move back")); + each_alone(&db, None, "after the move back"); +} + +/// The column the limit refusals of the next test name. +fn docs_v() -> VectorIndexRef { + VectorIndexRef { + table: "docs".to_owned(), + column: "v".to_owned(), + } +} + +/// Asserts that a statement or a commit was refused because `docs.v` may have two partitions. +fn refused_at_two(outcome: contextdb::Result, what: &str) { + match outcome { + Err(Error::VectorPartitionLimitExceeded { + index, + max_partitions, + }) => { + assert_eq!(index, docs_v(), "{what}: the column the refusal names"); + assert_eq!(max_partitions, 2, "{what}: the limit the refusal names"); + } + other => panic!("{what}: was to be refused at the partition limit and got {other:?}"), + } +} + +/// The rows of `docs`, the commit position, and the answer of a search of everything and of +/// each context a test uses: what a refused statement must leave exactly as it was. +fn state_of(db: &Database) -> (Vec>, contextdb::TxId, Vec>) { + let rows = support::rows(db, "SELECT id, ctx, tag, v FROM docs ORDER BY id"); + let stored = support::stored(db, support::DOCS); + let query = support::vector(82_000, 8); + let mut answers = Vec::new(); + let all = stored.nearest(|_| true, &query, 1000); + let sql = support::nearest_sql("docs", "v", None, 1000); + support::assert_nearest(db, &sql, &query, &all, "the state of docs"); + answers.push(all); + for ctx in ["a", "b", "c", "d", "e", "m", "n", "z"] { + let wanted = stored.nearest(|row| row[1] == support::text(ctx), &query, 1000); + let filter = format!("ctx = '{ctx}'"); + let sql = support::nearest_sql("docs", "v", Some(&filter), 1000); + support::assert_nearest(db, &sql, &query, &wanted, "the state of docs"); + answers.push(wanted); + } + (rows, db.committed_watermark(), answers) +} + +fn insert_one(on: &impl Run, id: i64, ctx: &str) -> contextdb::Result<()> { + support::insert_doc(on, id, ctx, Some(support::vector(id, 8))).map(|_| ()) +} + +#[test] +fn vector_partition_limit_counts_and_refuses() { + let folder = support::folder("partition-limit"); + let db = support::open(&folder); + + // With no MAX_PARTITIONS written the limit is 256. + support::run( + &db, + "CREATE TABLE wide (id INTEGER PRIMARY KEY, ctx INTEGER NOT NULL, \ + v VECTOR(4) PARTITION_KEY (ctx))", + ); + let wide = "INSERT INTO wide (id, ctx, v) VALUES ($id, $ctx, $v)"; + let tx = db.transaction().expect("begin"); + for n in 1..=256_i64 { + tx.execute( + wide, + ¶ms! { "id" => n, "ctx" => n, "v" => support::vector(n, 4) }, + ) + .unwrap_or_else(|error| panic!("partition {n} of 256 was refused: {error:?}")); + } + tx.commit().expect("commit 256 partitions"); + let past = db.execute( + wide, + ¶ms! { "id" => 257_i64, "ctx" => 257_i64, "v" => support::vector(257, 4) }, + ); + match past { + Err(Error::VectorPartitionLimitExceeded { + index, + max_partitions: 256, + }) => assert_eq!((index.table.as_str(), index.column.as_str()), ("wide", "v")), + other => panic!("partition 257 with no MAX_PARTITIONS written got {other:?}"), + } + db.execute( + wide, + ¶ms! { "id" => 258_i64, "ctx" => 256_i64, "v" => support::vector(258, 4) }, + ) + .expect("a second row under an existing partition"); + assert_eq!(support::int(&db, "SELECT COUNT(*) FROM wide"), 257); + + // A column with MAX_PARTITIONS 2, and two partitions of ten rows. + support::declare_docs(&db, 8, "PARTITION_KEY (ctx) MAX_PARTITIONS 2"); + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, 1..=10, 8, |_| "a".to_owned()); + support::insert_docs(&tx, 11..=20, 8, |_| "b".to_owned()); + tx.commit().expect("commit"); + let two = state_of(&db); + assert_eq!(two.0.len(), 20); + + // The INSERT of a third key tuple is refused, and so is the UPDATE that moves a row to + // one while its partition keeps other rows. Row, vector and partitions are as before. + refused_at_two(insert_one(&db, 21, "c"), "the INSERT of a third tuple"); + assert!(state_of(&db) == two); + let moved = db.execute("UPDATE docs SET ctx = 'c' WHERE id = 1", ¶ms! {}); + refused_at_two(moved, "the UPDATE that moves a row to a third tuple"); + assert!(state_of(&db) == two); + // A multi-row INSERT whose second row would pass the limit stores neither row; nor does + // one whose first row would. + for (first, second) in [("a", "c"), ("c", "a")] { + let both = db.execute( + "INSERT INTO docs (id, ctx, tag, v) VALUES (22, $first, 0, $v), (23, $second, 0, $v)", + ¶ms! { "first" => first, "second" => second, "v" => support::vector(22, 8) }, + ); + refused_at_two(both, "an INSERT of two rows, one under a third tuple"); + assert!(state_of(&db) == two); + } + + // Inside a transaction the refusal is the statement's: what the transaction did before is + // still there, a later statement under an existing tuple succeeds, and it commits. + let tx = db.transaction().expect("begin"); + insert_one(&tx, 30, "a").expect("a row under an existing tuple"); + refused_at_two( + insert_one(&tx, 31, "c"), + "a third tuple inside a transaction", + ); + assert_eq!(support::int(&tx, "SELECT COUNT(*) FROM docs"), 21); + insert_one(&tx, 32, "b").expect("a later statement under an existing tuple"); + tx.commit() + .expect("a transaction commits after one of its statements was refused"); + assert_eq!( + support::rows(&db, "SELECT id FROM docs WHERE id >= 30 ORDER BY id"), + [[Value::Int64(30)], [Value::Int64(32)]] + ); + + // In one transaction at the limit: deleting a partition's last vector and then inserting + // under a new tuple succeeds, and inserting under the emptied tuple after that is refused. + let tx = db.transaction().expect("begin"); + support::run(&tx, "DELETE FROM docs WHERE ctx = 'b'"); + insert_one(&tx, 40, "c").expect("a new tuple once a partition was emptied"); + refused_at_two(insert_one(&tx, 41, "b"), "the emptied tuple, now a third"); + insert_one(&tx, 42, "c").expect("the transaction is still usable"); + tx.commit().expect("commit"); + // The partitions are a and c now. + refused_at_two(insert_one(&db, 43, "b"), "a third tuple after the commit"); + let a_and_c = state_of(&db); + assert!(a_and_c.2[2].is_empty() && a_and_c.2[3].len() == 2); + + // The one refusal only a commit sees. Deleting c leaves room for one partition. A + // transaction starts d while there is room; another statement starts e meanwhile; the + // transaction's commit is refused with the same kind and nothing of it is stored. + support::run(&db, "DELETE FROM docs WHERE ctx = 'c'"); + let tx = db.transaction().expect("begin"); + insert_one(&tx, 50, "d").expect("a second partition while there is room"); + insert_one(&tx, 51, "a").expect("and a row under an existing tuple"); + insert_one(&db, 52, "e").expect("another statement takes the room"); + let a_and_e = state_of(&db); + refused_at_two( + tx.commit(), + "the commit of a partition there is no room for now", + ); + assert!(state_of(&db) == a_and_e); + assert!(support::rows(&db, "SELECT id FROM docs WHERE id IN (50, 51)").is_empty()); + + // A NULL vector takes no partition: at the limit a row with none goes in under a third + // tuple, and giving it a vector is what is refused. + db.execute( + "INSERT INTO docs (id, ctx, tag) VALUES (60, 'z', 0)", + ¶ms! {}, + ) + .expect("a row with no vector under a third tuple"); + let with_vector = "UPDATE docs SET v = $v WHERE id = 60"; + let vector = params! { "v" => support::vector(60, 8) }; + refused_at_two( + db.execute(with_vector, &vector), + "a vector for the row under z", + ); + refused_at_two( + insert_one(&db, 61, "n"), + "a third tuple beside the row with no vector", + ); + // Taking the only vector of e away frees its partition, and the row under z takes it. + support::run(&db, "UPDATE docs SET v = NULL WHERE ctx = 'e'"); + db.execute(with_vector, &vector) + .expect("a vector for the row under z once e holds none"); + refused_at_two( + db.execute("UPDATE docs SET v = $v WHERE id = 52", &vector), + "e again", + ); + + // Deleting a partition's last vector frees it. + support::run(&db, "DELETE FROM docs WHERE ctx = 'z'"); + insert_one(&db, 70, "m").expect("a new tuple once z was deleted"); + refused_at_two(insert_one(&db, 71, "n"), "a third tuple beside a and m"); + + // At the limit, one UPDATE that moves the last vector of a partition to a new key + // succeeds; one that moves a vector out of a partition that keeps others is refused. + let before = state_of(&db); + let moved = db.execute("UPDATE docs SET ctx = 'n' WHERE id = 1", ¶ms! {}); + refused_at_two(moved, "moving one of the vectors of a to a third tuple"); + assert!(state_of(&db) == before); + support::run(&db, "UPDATE docs SET ctx = 'n' WHERE id = 70"); + let after = state_of(&db); + // Under m there is nothing now, and under n the one row. + assert!(after.2[6].is_empty()); + assert_eq!(ids(&after.2[7]), [70]); + refused_at_two(insert_one(&db, 72, "m"), "m, which is a third tuple again"); + + // After sealing everything, an UPDATE of a column that is neither the vector nor a key + // column leaves the next cycle with nothing to do and the searches' counted work as it was. + support::seal_all(&db); + let sealed = state_of(&db); + let work = |db: &Database| -> Vec { + let query = support::vector(82_001, 8); + let filters = [None, Some("ctx = 'a'"), Some("ctx = 'n'")]; + let searches = filters.iter().map(|filter| { + let sql = support::nearest_sql("docs", "v", *filter, 5); + let result = support::search(db, &sql, &query).expect("a search"); + result.rows_examined() + }); + searches.collect() + }; + let read_before = work(&db); + let updated = support::run(&db, "UPDATE docs SET tag = tag + 100"); + assert_eq!(updated.rows_affected(), sealed.0.len() as u64); + let position = db.committed_watermark(); + let cycle = db.run_maintenance_cycle().expect("a cycle").vector; + assert_eq!( + ( + cycle.maintenance_vectors_done, + cycle.maintenance_vectors_remaining + ), + (0, 0), + "a cycle after an UPDATE that changed no vector and no key" + ); + assert_eq!( + db.committed_watermark(), + position, + "a cycle with nothing to do" + ); + assert_eq!(work(&db), read_before); + assert!(state_of(&db).2 == sealed.2); +} + +/// Rows in each of the two partitions of the next test. +const HELD: i64 = 600; + +/// The search the child of the next test is stopped inside: both partitions, every row. +fn of_both() -> String { + support::nearest_sql("docs", "v", Some("ctx IN ('a', 'b')"), 100_000) +} + +/// An answer as lines a child announces: ` `. +fn lines(word: &str, answer: &[Near]) -> Vec { + let answer = answer.iter(); + answer + .map(|near| format!("{word} {:?} {:016x}", near.key, near.distance.to_bits())) + .collect() +} + +#[test] +fn the_partition_cap_counts_live_partitions_while_a_reader_holds_an_earlier_state() { + if let Some(role) = support::role() { + return searcher(&role); + } + let test = "the_partition_cap_counts_live_partitions_while_a_reader_holds_an_earlier_state"; + const DIMENSION: usize = 64; + let folder = support::folder("cap-beside-a-reader"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, "PARTITION_KEY (ctx) MAX_PARTITIONS 2"); + let ctx = |n: i64| String::from(if n <= HELD { "a" } else { "b" }); + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, 1..=2 * HELD, DIMENSION, ctx); + tx.commit().expect("commit"); + support::seal_all(&db); + let query = support::vector(83_000, DIMENSION); + let began_at = support::stored(&db, support::DOCS).nearest(|_| true, &query, usize::MAX); + assert_eq!(began_at.len() as i64, 2 * HELD); + + // Run the search once to its end in a child, to learn how many storage reads it makes. + let mut told = Child::start(test, "searcher pause_at=0", &folder).end(); + let reads = support::last_number(&told, "reads").expect("the searcher announces its reads"); + told.retain(|line| line.starts_with("row ")); + assert!( + told == lines("row", &began_at), + "the rows of a search nothing stopped" + ); + assert!( + reads > 1, + "a search of every row made {reads} storage reads" + ); + // A second child, stopped halfway through the storage reads of the same search. + let halfway = reads.div_ceil(2); + let mut reader = Child::start(test, &format!("searcher pause_at={halfway}"), &folder); + reader.expect("reading"); + + // At the cap, the writer deletes every row of one partition and inserts a row under a + // third key tuple. It is not refused: the partition the stopped reader still sees holds + // no live vector, so it does not count. + support::run(&db, "DELETE FROM docs WHERE ctx = 'b'"); + if let Err(error) = support::insert_doc(&db, 5000, "c", Some(support::vector(5000, DIMENSION))) + { + panic!("a third key tuple, after the second was emptied, was refused: {error:?}"); + } + // The cap still holds for the partitions that are live now. + match support::insert_doc(&db, 5001, "d", Some(support::vector(5001, DIMENSION))) { + Err(Error::VectorPartitionLimitExceeded { + max_partitions: 2, .. + }) => {} + other => panic!("a third live partition got {other:?}"), + } + // More commits and a seal, so that pages the deleted partition gave up come up for reuse. + for first in (5100..5700_i64).step_by(100) { + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, first..first + 100, DIMENSION, |_| "c".to_owned()); + tx.commit().expect("commit"); + } + support::seal_all(&db); + let now = support::stored(&db, support::DOCS).nearest(|_| true, &query, usize::MAX); + assert_eq!(now.len() as i64, HELD + 601); + + // Let go, the reader's search answers with the rows of the commit it began at, the + // removed partition's rows among them; its next search sees the store as it is now. + reader.tell("go"); + let searched = reader.until("searched"); + assert!( + searched == lines("row", &began_at), + "the rows of the stopped search" + ); + assert!( + searched + .iter() + .any(|line| line.starts_with(&format!("row Int64({}) ", 2 * HELD))) + ); + let mut after = reader.until("done"); + after.retain(|line| line.starts_with("now ")); + assert!( + after == lines("now", &now), + "the rows of the reader's next search" + ); + reader.end(); +} + +/// Opens the store for reading and searches both partitions for every row, announcing each +/// row of the answer as `row `, then `searched` and `reads `, the storage reads +/// the search made. At read number `pause_at` of that search it announces `reading` and waits +/// for a line. Then it searches every partition again and announces each row as +/// `now `, then `done`. +fn searcher(role: &support::Role) { + let (door, searching, reads) = support::pausing_at_read(role.number("pause_at")); + let mut options = ReadSessionOptions::default(); + options.storage = Some(door); + let reader = + ReadSession::open_with_options(&role.folder, options).expect("open the store for reading"); + let query = support::vector(83_000, 64); + + searching.store(true, std::sync::atomic::Ordering::SeqCst); + let sql = of_both(); + let result = support::search(&reader, &sql, &query).expect("a reader's search is not refused"); + searching.store(false, std::sync::atomic::Ordering::SeqCst); + for line in lines("row", &support::found(&result, &sql)) { + support::say(line); + } + support::say("searched"); + let reads = reads.load(std::sync::atomic::Ordering::SeqCst); + support::say(format!("reads {reads}")); + + let sql = support::nearest_sql("docs", "v", None, 100_000); + let result = support::search(&reader, &sql, &query).expect("the reader's next search"); + for line in lines("now", &support::found(&result, &sql)) { + support::say(line); + } + support::say("done"); +} diff --git a/v2/contextdb/tests/vectors_statements.rs b/v2/contextdb/tests/vectors_statements.rs new file mode 100644 index 00000000..8c78d0ea --- /dev/null +++ b/v2/contextdb/tests/vectors_statements.rs @@ -0,0 +1,1176 @@ +//! What a statement may say about a VECTOR column, through every door: the search mode a +//! request runs in, the refusal of a search that asks for an index, every other refusal of a +//! similarity search by its kind, and the three clauses of a VECTOR column's declaration. + +mod support; + +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use contextdb::{ + Database, Error, ReadSession, ReadSessionOptions, Value, VectorIndexRef, + VectorPartitionDeclarationIssue, VectorPolicyDeclarationIssue, params, +}; + +use support::{Call, Near, Run}; + +const DIMENSION: usize = 8; + +/// How the refusal of a statement that only a later version runs ends. +const LATER_VERSION: &str = "is not supported in this version of the library"; + +fn column(table: &str, column: &str) -> VectorIndexRef { + VectorIndexRef { + table: table.to_owned(), + column: column.to_owned(), + } +} + +/// Declares a table like `docs` under another name, with these clauses after `VECTOR(8)`. +fn declare(db: &Database, table: &str, clauses: &str) { + support::run( + db, + &format!( + "CREATE TABLE {table} (id INTEGER PRIMARY KEY, ctx TEXT NOT NULL, \ + tag INTEGER NOT NULL, v VECTOR({DIMENSION}) {clauses})" + ), + ); +} + +/// Inserts rows with these numbers: row `n` under the context `c`. +fn fill(on: &impl Run, table: &str, numbers: std::ops::RangeInclusive) { + for n in numbers { + let inserted = on.run( + &format!("INSERT INTO {table} (id, ctx, tag, v) VALUES ($id, $ctx, $tag, $v)"), + ¶ms! { + "id" => n, "ctx" => format!("c{}", n % 3), "tag" => n % 50, + "v" => support::vector(n, DIMENSION), + }, + ); + if let Err(error) = inserted { + panic!("row {n} of {table} was refused: {error:?}"); + } + } +} + +/// A search of `table` with this filter and this clause between ORDER BY and LIMIT. +fn search_sql(table: &str, filter: Option<&str>, clause: &str, k: usize) -> String { + let filter = filter.map_or_else(String::new, |filter| format!(" WHERE {filter}")); + format!( + "SELECT id, v <=> $q AS distance FROM {table}{filter} ORDER BY distance {clause} LIMIT {k}" + ) +} + +/// Asserts that a search was refused because it asked for an index, naming this column. +fn refused_as_indexed(outcome: contextdb::Result, table: &str, sql: &str) { + match outcome { + Err(Error::VectorIndexedRouteUnavailable { index }) => { + assert_eq!(index, column(table, "v"), "the column named for {sql}"); + } + other => panic!("{sql}\nasked for an index and got {other:?}"), + } +} + +/// True for a sentence that names both the function `ROW_VECTOR`, in any letter case, and the +/// `<=>` it was written beside. Every test of that refusal holds it to both. +fn names_row_vector_beside_a_distance(sentence: &str) -> bool { + sentence.to_uppercase().contains("ROW_VECTOR") && sentence.contains("<=>") +} + +#[test] +fn vector_search_mode_routes_each_request() { + let folder = support::folder("search-mode-routes"); + let db = support::open(&folder); + // Four tables with the same rows: no SEARCH_MODE, and each of the three. + let tables = [ + ("plain", "PARTITION_KEY (ctx)"), + ("auto", "PARTITION_KEY (ctx) SEARCH_MODE AUTO"), + ("exact", "PARTITION_KEY (ctx) SEARCH_MODE EXACT"), + ("indexed", "PARTITION_KEY (ctx) SEARCH_MODE INDEXED"), + ]; + for (table, clauses) in tables { + declare(&db, table, clauses); + } + // Rows 1 to 80 in blocks, rows 81 to 120 on pending lists. + for (table, _) in tables { + let tx = db.transaction().expect("begin"); + fill(&tx, table, 1..=80); + tx.commit().expect("commit"); + } + support::seal_all(&db); + for (table, _) in tables { + let tx = db.transaction().expect("begin"); + fill(&tx, table, 81..=120); + tx.commit().expect("commit"); + } + let query = support::vector(90_000, DIMENSION); + type Scope = (Option<&'static str>, fn(&[Value]) -> bool); + let scopes: [Scope; 3] = [ + (None, |_| true), + (Some("ctx = 'c1'"), |row| row[1] == support::text("c1")), + (Some("ctx IN ('c0', 'c2')"), |row| { + row[1] != support::text("c1") + }), + ]; + + // Through one door: for every table and scope, what each way of writing the mode gives. + let through = |on: &dyn Run, door: &str| { + for (table, _) in tables { + let stored = support::stored(&on, &format!("SELECT id, ctx, tag, v FROM {table}")); + for (filter, passes) in scopes { + let wanted = stored.nearest(passes, &query, 10); + let answer = |clause: &str| -> u64 { + let sql = search_sql(table, filter, clause, 10); + let what = format!("through {door}"); + let result = support::assert_nearest(&on, &sql, &query, &wanted, &what); + result.rows_examined() + }; + let refused = |clause: &str| { + let sql = search_sql(table, filter, clause, 10); + refused_as_indexed(support::search(&on, &sql, &query), table, &sql); + }; + // USE VECTOR overrides the column's mode both ways: EXACT and AUTO answer on + // every table, the one declared INDEXED too, and INDEXED is refused on every + // table, the one declared EXACT too. AUTO and EXACT are one search: the same + // rows for the same counted work. + let exact = answer("USE VECTOR EXACT"); + assert_eq!( + answer("USE VECTOR AUTO"), + exact, + "AUTO and EXACT on {table}" + ); + assert_eq!(answer("use vector exact"), exact, "{table}"); + refused("USE VECTOR INDEXED"); + // With no clause the column's mode is in force, and no mode is AUTO. + if table == "indexed" { + refused(""); + } else { + assert_eq!(answer(""), exact, "no clause on {table}"); + } + } + } + }; + through(&db, "the writer"); + let tx = db.transaction().expect("begin"); + through(&tx, "a transaction"); + tx.rollback(); + let reader = ReadSession::open(&folder).expect("open the store for reading"); + through(&reader, "a read session"); + + // A table of one row and a column declared SEARCH_MODE EXACT, with nothing else set. + support::run( + &db, + "CREATE TABLE two (id INTEGER PRIMARY KEY, v VECTOR(2) SEARCH_MODE EXACT)", + ); + support::run(&db, "INSERT INTO two (id, v) VALUES (1, [1.0, 0.0])"); + let found = db + .execute( + "SELECT id, v <=> $q AS distance FROM two ORDER BY distance LIMIT 5", + ¶ms! { "q" => vec![1.0_f32, 0.0] }, + ) + .expect("search the table of one row"); + let wanted = Near { + key: Value::Int64(1), + distance: support::distance(&[1.0, 0.0], &[1.0, 0.0]), + }; + assert_eq!(support::found(&found, "the table of one row"), [wanted]); +} + +#[test] +fn indexed_search_is_refused_at_every_door() { + let folder = support::folder("indexed-refused"); + let db = support::open(&folder); + declare(&db, "plain", "PARTITION_KEY (ctx)"); + declare(&db, "indexed", "PARTITION_KEY (ctx) SEARCH_MODE INDEXED"); + declare(&db, "empty", "SEARCH_MODE INDEXED"); + // Two thousand pending rows in each: a search that read rows would read them all. + for table in ["plain", "indexed"] { + let tx = db.transaction().expect("begin"); + fill(&tx, table, 1..=2000); + tx.commit().expect("commit"); + } + let query = support::vector(91_000, DIMENSION); + // Every way a search asks for an index: in the statement, by the column's declaration, + // both, on a table with no row, and for a partition no row is in. + let asking: Vec<(&str, String)> = vec![ + ("plain", search_sql("plain", None, "USE VECTOR INDEXED", 5)), + ( + "plain", + search_sql("plain", Some("ctx = 'c1'"), "USE VECTOR INDEXED", 5), + ), + ( + "plain", + search_sql("plain", Some("ctx = 'zz'"), "USE VECTOR INDEXED", 5), + ), + ("indexed", search_sql("indexed", None, "", 5)), + ("indexed", search_sql("indexed", Some("ctx = 'zz'"), "", 5)), + ( + "indexed", + search_sql("indexed", None, "USE VECTOR INDEXED", 5), + ), + ( + "indexed", + "SELECT * FROM indexed ORDER BY v <=> $q LIMIT 1".to_owned(), + ), + ("empty", search_sql("empty", None, "", 5)), + ("empty", search_sql("empty", None, "USE VECTOR INDEXED", 5)), + ]; + let before = ( + db.committed_watermark(), + support::int(&db, "SELECT COUNT(*) FROM indexed"), + ); + let reader = ReadSession::open(&folder).expect("open the store for reading"); + let tx = db.transaction().expect("begin"); + for (table, sql) in &asking { + // Never an answer, through the writer, a transaction and a reader alike. + refused_as_indexed(support::search(&db, sql, &query), table, sql); + refused_as_indexed(support::search(&tx, sql, &query), table, sql); + refused_as_indexed(support::search(&reader, sql, &query), table, sql); + } + // The refused transaction is as usable as before, and nothing changed. + fill(&tx, "indexed", 2001..=2001); + tx.rollback(); + assert_eq!( + ( + db.committed_watermark(), + support::int(&db, "SELECT COUNT(*) FROM indexed") + ), + before + ); + // The same column answers when the statement asks for the exact search. + let stored = support::stored(&db, "SELECT id, ctx, tag, v FROM indexed"); + let wanted = stored.nearest(|_| true, &query, 5); + let sql = search_sql("indexed", None, "USE VECTOR EXACT", 5); + support::assert_nearest(&reader, &sql, &query, &wanted, "EXACT on an INDEXED column"); + support::close(db); + + // At once, with no row read. Each statement runs first in a session of its own, whose + // storage door counts the reads of the data file. The refused search reads no more of the + // file than a count of the whole table, which reads no row; a search that answers reads + // more than that count, so the comparison can tell. + let reads_of = |sql: &str| -> (u64, contextdb::Result) { + let reads = Arc::new(AtomicU64::new(0)); + let mut options = ReadSessionOptions::default(); + options.storage = Some(support::storage({ + let reads = reads.clone(); + move |call| { + if matches!(call, Call::Read(..)) { + reads.fetch_add(1, Ordering::SeqCst); + } + } + })); + let session = ReadSession::open_with_options(&folder, options).expect("open for reading"); + let opened = reads.load(Ordering::SeqCst); + let outcome = support::search(&session, sql, &query); + (reads.load(Ordering::SeqCst) - opened, outcome) + }; + let (count_reads, count) = reads_of("SELECT COUNT(*) FROM indexed"); + let count = count.expect("count the table"); + assert_eq!(count.rows(), [[Value::Int64(2000)]]); + assert_eq!(count.rows_examined(), 0); + for (table, sql) in asking.iter().filter(|(table, _)| *table == "indexed") { + let (reads, outcome) = reads_of(sql); + refused_as_indexed(outcome, table, sql); + assert!( + reads <= count_reads, + "{sql}\nwas refused after {reads} storage reads; a count that reads no row made \ + {count_reads}" + ); + } + let (exact_reads, exact) = reads_of(&sql); + assert_eq!(exact.expect("the exact search").rows().len(), 5); + assert!( + exact_reads > count_reads, + "an exact search of 2,000 pending rows made {exact_reads} storage reads, a count \ + {count_reads}" + ); +} + +#[test] +fn vector_query_errors_are_typed() { + let folder = support::folder("vector-query-errors"); + let db = support::open(&folder); + support::run( + &db, + "CREATE TABLE docs (id INTEGER PRIMARY KEY, ctx TEXT NOT NULL, r REAL, v VECTOR(4), \ + w VECTOR(4))", + ); + for n in 1..=20_i64 { + db.execute( + "INSERT INTO docs (id, ctx, r, v, w) VALUES ($id, 'a', 0.5, $v, $w)", + ¶ms! { + "id" => n, "v" => support::vector(n, 4), "w" => support::vector(100 + n, 4), + }, + ) + .expect("insert a row"); + } + // Half of them sealed. + support::seal_all(&db); + support::run(&db, "UPDATE docs SET r = 1.5 WHERE id > 10"); + for n in 21..=30_i64 { + db.execute( + "INSERT INTO docs (id, ctx, v, w) VALUES ($id, 'a', $v, $w)", + ¶ms! { + "id" => n, "v" => support::vector(n, 4), "w" => support::vector(100 + n, 4), + }, + ) + .expect("insert a row"); + } + let reader = ReadSession::open(&folder).expect("open the store for reading"); + let state = |db: &Database| { + let rows = support::rows(db, "SELECT id, ctx, r, v, w FROM docs ORDER BY id"); + (rows, db.committed_watermark()) + }; + let before = state(&db); + + let good = support::vector(92_000, 4); + let with = |index: usize, number: f32| -> Vec { + let mut vector = good.clone(); + vector[index] = number; + vector + }; + // The refusal of a statement run with `$q` and `$v` bound to these vectors. A SELECT is + // run through the writer, a transaction and a reader, and must be refused alike. + let refusal = |sql: &str, q: &[f32], v: &[f32]| -> Error { + let bound = params! { "q" => q.to_vec(), "v" => v.to_vec(), "p" => vec![0.5_f32; 4] }; + let kind = |error: &Error| format!("{error:?}"); + let from_writer = match db.execute(sql, &bound) { + Err(error) => error, + Ok(result) => panic!("{sql}\nwas not refused: {result:?}"), + }; + if sql.starts_with("SELECT") { + let tx = db.transaction().expect("begin"); + for from in [tx.execute(sql, &bound), reader.execute(sql, &bound)] { + match from { + Err(error) => assert_eq!(kind(&error), kind(&from_writer), "{sql}"), + Ok(result) => panic!("{sql}\nwas not refused at every door: {result:?}"), + } + } + } + from_writer + }; + let plan_error = |sql: &str, names: &[&str]| match refusal(sql, &good, &good) { + Error::PlanError(text) => { + for name in names { + assert!( + text.contains(name), + "{sql}\nwas refused without naming {name}: {text}" + ); + } + } + other => panic!("{sql}\nwas refused as {other:?}, not as a plan error"), + }; + + // A search with no LIMIT has no kind of its own: it is refused as every statement of a + // later version is, in a sentence that names it. + for sql in [ + "SELECT id FROM docs ORDER BY v <=> $q", + "SELECT id, v <=> $q AS distance FROM docs ORDER BY distance", + "SELECT id FROM docs WHERE ctx = 'a' ORDER BY v <=> $q USE VECTOR EXACT", + ] { + match refusal(sql, &good, &good) { + Error::PlanError(text) => assert!( + text.contains("vector search without LIMIT") && text.ends_with(LATER_VERSION), + "{sql}\nwas refused in the sentence: {text}" + ), + other => panic!("{sql}\nwas refused as {other:?}, not as a plan error"), + } + } + // USE VECTOR in a SELECT that does not order by a distance. + for sql in [ + "SELECT id FROM docs ORDER BY id USE VECTOR EXACT LIMIT 3", + "SELECT id FROM docs USE VECTOR AUTO LIMIT 3", + "SELECT id FROM docs WHERE ctx = 'a' ORDER BY r, id USE VECTOR INDEXED LIMIT 3", + ] { + let refused = refusal(sql, &good, &good); + assert!( + matches!(refused, Error::UseVectorRequiresVectorOrder), + "{sql}\ngot {refused:?}" + ); + } + // A distance from a column that is not a VECTOR column. + for (sql, not_a_vector) in [ + ("SELECT id FROM docs ORDER BY ctx <=> $q LIMIT 3", "ctx"), + ( + "SELECT id, r <=> $q AS distance FROM docs ORDER BY distance LIMIT 3", + "r", + ), + ] { + match refusal(sql, &good, &good) { + Error::UnknownVectorIndex { index } => assert_eq!(index, column("docs", not_a_vector)), + other => panic!("{sql}\ngot {other:?}"), + } + } + // A vector of another length than the column's. + let search = "SELECT id, v <=> $q AS distance FROM docs ORDER BY distance LIMIT 3"; + for length in [1_usize, 3, 5, 768] { + match refusal(search, &vec![0.5; length], &good) { + Error::VectorIndexDimensionMismatch { + index, + expected, + actual, + } => { + assert_eq!(index, column("docs", "v")); + assert_eq!((expected, actual), (4, length)); + } + other => panic!("a search with a vector of {length} numbers got {other:?}"), + } + } + // A number that is not finite, in the vector of a search and in a vector that is written. + let not_finite = [ + (0, f32::NAN), + (2, f32::NAN), + (3, f32::INFINITY), + (1, f32::NEG_INFINITY), + ]; + for (position, number) in not_finite { + let bad = with(position, number); + let cases = [ + (search, &bad, &good, "v"), + ( + "SELECT id FROM docs ORDER BY w <=> $q LIMIT 3", + &bad, + &good, + "w", + ), + ( + "INSERT INTO docs (id, ctx, v) VALUES (99, 'a', $v)", + &good, + &bad, + "v", + ), + ( + "INSERT INTO docs (id, ctx, v, w) VALUES (99, 'a', $q, $v)", + &good, + &bad, + "w", + ), + ("UPDATE docs SET v = $v WHERE id = 1", &good, &bad, "v"), + ("UPDATE docs SET w = $v WHERE id >= 1", &good, &bad, "w"), + ]; + for (sql, q, v, named) in cases { + match refusal(sql, q, v) { + Error::InvalidVectorValue { + index, + position: at, + } => { + assert_eq!(index, column("docs", named), "{sql}"); + assert_eq!(at, position, "{sql}: where {number} stands"); + } + other => panic!("{sql}\nwith {number} at {position} got {other:?}"), + } + } + } + + // ROW_VECTOR is a function this version does not have: the refusal names the function and + // the `<=>` it stands beside. + let sql = "SELECT id FROM docs ORDER BY v <=> ROW_VECTOR('docs', 'v', 1) LIMIT 3"; + match refusal(sql, &good, &good) { + Error::PlanError(text) => assert!( + names_row_vector_beside_a_distance(&text), + "{sql}\nwas refused without naming both ROW_VECTOR and <=>: {text}" + ), + other => panic!("{sql}\ngot {other:?}"), + } + // A distance anywhere but as the one ORDER BY key of a search, ascending, and in that + // search's select list. + for sql in [ + // In a WHERE. + "SELECT id FROM docs WHERE v <=> $q < 0.5", + "SELECT id FROM docs WHERE v <=> $q < 0.5 ORDER BY id LIMIT 3", + "SELECT id FROM docs WHERE v <=> $q < 0.5 ORDER BY v <=> $q LIMIT 3", + "DELETE FROM docs WHERE v <=> $q < 0.5", + // In a select list with no ORDER BY on it. + "SELECT id, v <=> $q AS distance FROM docs", + "SELECT id, v <=> $q AS distance FROM docs ORDER BY id LIMIT 3", + // With DESC. + "SELECT id FROM docs ORDER BY v <=> $q DESC LIMIT 3", + "SELECT id, v <=> $q AS distance FROM docs ORDER BY distance DESC LIMIT 3", + // Beside a second ORDER BY item. + "SELECT id FROM docs ORDER BY v <=> $q, id LIMIT 3", + "SELECT id FROM docs ORDER BY id, v <=> $q LIMIT 3", + // With a left side that is not a column. + "SELECT id FROM docs ORDER BY $q <=> v LIMIT 3", + // In a write. + "UPDATE docs SET r = v <=> $q WHERE id = 1", + "INSERT INTO docs (id, ctx, r) VALUES (99, 'a', $v <=> $q)", + ] { + plan_error(sql, &["<=>"]); + } + // A second distance that is not the ORDER BY's own: of another column, or to another + // vector. The refusal names `<=>` and the column, as `column `: the letter alone is + // in nearly every sentence. + plan_error( + "SELECT id, w <=> $q AS other FROM docs ORDER BY v <=> $q LIMIT 3", + &["<=>", "column w"], + ); + plan_error( + "SELECT id, v <=> $p AS other FROM docs ORDER BY v <=> $q LIMIT 3", + &["<=>", "column v"], + ); + plan_error( + "SELECT id, v <=> $q AS distance, 1 - (w <=> $q) AS other FROM docs \ + ORDER BY distance LIMIT 3", + &["<=>", "column w"], + ); + + // There is no score column: on a table with no such column it is a name like any other. + for sql in [ + "SELECT id, score FROM docs ORDER BY v <=> $q LIMIT 3", + "SELECT id, v <=> $q AS distance, score FROM docs ORDER BY distance LIMIT 3", + ] { + match refusal(sql, &good, &good) { + Error::ColumnNotFound { table, column } => { + assert_eq!( + (table.as_str(), column.as_str()), + ("docs", "score"), + "{sql}" + ); + } + other => panic!("{sql}\ngot {other:?}"), + } + } + + // Each of them left the store as it was, and the column still answers. + assert!(state(&db) == before); + let stored = support::stored(&db, "SELECT id, ctx, r, v FROM docs"); + let wanted = stored.nearest(|_| true, &good, 3); + support::assert_nearest(&db, search, &good, &wanted, "after every refusal"); +} + +#[test] +fn what_the_parser_reads_and_this_version_does_not_run_is_refused_when_planned() { + let folder = support::folder("refused-when-planned"); + let db = support::open(&folder); + support::run( + &db, + "CREATE TABLE pairs (id INTEGER PRIMARY KEY, a INTEGER, b INTEGER, v VECTOR(3), \ + w VECTOR(3))", + ); + support::run( + &db, + "INSERT INTO pairs (id, a, b, v, w) VALUES (1, 1, 2, [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]), \ + (2, 3, 4, [0.0, 1.0, 0.0], [1.0, 0.0, 0.0])", + ); + let reader = ReadSession::open(&folder).expect("open the store for reading"); + let state = |db: &Database| { + let rows = support::rows(db, "SELECT id, a, b, v, w FROM pairs ORDER BY id"); + (rows, db.committed_watermark()) + }; + let before = state(&db); + + // The sentence a statement is refused in. It is a plan error, never a syntax error, and + // the same through the writer, a transaction and a reader. + let refused_when_planned = |sql: &str| -> String { + let tx = db.transaction().expect("begin"); + let none = params! {}; + let doors = [ + db.execute(sql, &none), + tx.execute(sql, &none), + reader.execute(sql, &none), + ]; + let sentences = doors.map(|outcome| match outcome { + Err(Error::PlanError(text)) => text, + Err(other) => panic!("{sql}\nwas refused as {other:?}, not as a plan error"), + Ok(result) => panic!("{sql}\nwas not refused: {result:?}"), + }); + assert!( + sentences[0] == sentences[1] && sentences[0] == sentences[2], + "{sql}\nwas refused in three sentences: {sentences:?}" + ); + let [sentence, ..] = sentences; + sentence + }; + + // A row value is read wherever an expression is, and refused by name anywhere but beside + // IN: compared with `=`, and in a select list. + for sql in [ + "SELECT id FROM pairs WHERE (a, b) = (1, 2)", + "SELECT (a, b) FROM pairs", + ] { + let sentence = refused_when_planned(sql); + assert!( + sentence.to_lowercase().contains("row value"), + "{sql}\nwas refused without naming the row value: {sentence}" + ); + } + // Anything is read on the right of `<=>`, and anything but a vector literal, its quoted + // text or a bound name is refused naming `<=>`: a column, and ROW_VECTOR(...), whose + // refusal names the function too. + let sql = "SELECT id FROM pairs ORDER BY v <=> w LIMIT 3"; + let sentence = refused_when_planned(sql); + assert!( + sentence.contains("<=>"), + "{sql}\nwas refused without naming <=>: {sentence}" + ); + let sql = "SELECT id FROM pairs ORDER BY v <=> ROW_VECTOR('pairs', 'v', 1) LIMIT 3"; + let sentence = refused_when_planned(sql); + assert!( + names_row_vector_beside_a_distance(&sentence), + "{sql}\nwas refused without naming both ROW_VECTOR and <=>: {sentence}" + ); + // The two statements that describe an index and its partitions, and a search with no + // LIMIT, are refused as every statement of a later version is, each by its name. + for (sql, name) in [ + ("SHOW VECTOR_INDEXES", "VECTOR_INDEXES"), + ("SHOW VECTOR_PARTITIONS", "VECTOR_PARTITIONS"), + ( + "SELECT id FROM pairs ORDER BY v <=> [1.0, 0.0, 0.0]", + "vector search without LIMIT", + ), + ] { + let sentence = refused_when_planned(sql); + assert!( + sentence.contains(name) && sentence.ends_with(LATER_VERSION), + "{sql}\nwas refused in the sentence: {sentence}" + ); + } + + // Each of them left the store as it was, and the column still answers. + assert!(state(&db) == before); + let stored = support::stored(&db, "SELECT id, a, b, v FROM pairs"); + let query = [1.0_f32, 0.0, 0.0]; + let wanted = stored.nearest(|_| true, &query, 2); + let sql = support::nearest_sql("pairs", "v", None, 2); + support::assert_nearest(&db, &sql, &query, &wanted, "after every refusal"); +} + +/// The words of a declaration's line for one column: what stands between spaces, commas and +/// parentheses. +fn words(declaration: &str, column: &str) -> Vec { + let line = declaration + .lines() + .find(|line| line.trim_start().starts_with(&format!("{column} "))) + .unwrap_or_else(|| panic!("no line for the column {column} in\n{declaration}")); + line.split(|c: char| c.is_whitespace() || "(),".contains(c)) + .filter(|word| !word.is_empty()) + .map(str::to_owned) + .collect() +} + +#[test] +fn vector_column_declarations_are_checked_stored_and_printed_back() { + let folder = support::folder("vector-declarations"); + let db = support::open(&folder); + let fresh = db.committed_watermark(); + + // Each clause that cannot stand, by its kind, naming the column; and no table declared. + use VectorPartitionDeclarationIssue as Issue; + let refused: [(&str, Issue); 14] = [ + ( + "v VECTOR(4) PARTITION_KEY (nowhere)", + Issue::UnknownPartitionKeyColumn, + ), + ( + "v VECTOR(4) PARTITION_KEY (ctx, nowhere)", + Issue::UnknownPartitionKeyColumn, + ), + ( + "v VECTOR(4) PARTITION_KEY (loose)", + Issue::NullablePartitionKeyColumn, + ), + ( + "v VECTOR(4) PARTITION_KEY (ctx, loose)", + Issue::NullablePartitionKeyColumn, + ), + ( + "v VECTOR(4) PARTITION_KEY (r)", + Issue::UnsupportedPartitionKeyColumnType, + ), + ( + "v VECTOR(4) PARTITION_KEY (doc)", + Issue::UnsupportedPartitionKeyColumnType, + ), + ( + "v VECTOR(4) PARTITION_KEY (other)", + Issue::UnsupportedPartitionKeyColumnType, + ), + ( + "v VECTOR(4) PARTITION_KEY (ctx, ctx)", + Issue::DuplicatePartitionKeyColumn, + ), + ( + "v VECTOR(4) PARTITION_KEY (v) NOT NULL", + Issue::VectorColumnInPartitionKey, + ), + ( + "v VECTOR(4) MAX_PARTITIONS 4", + Issue::MaxPartitionsWithoutPartitionKey, + ), + ( + "v VECTOR(4) PARTITION_KEY (ctx) MAX_PARTITIONS 0", + Issue::MaxPartitionsNotPositive, + ), + ( + "v VECTOR(4) PARTITION_KEY (ctx) MAX_PARTITIONS -1", + Issue::MaxPartitionsNotPositive, + ), + ( + "v VECTOR(4) PARTITION_KEY (ctx) MAX_PARTITIONS 4294967296", + Issue::MaxPartitionsOutOfRange, + ), + ( + "v VECTOR(4) PARTITION_KEY (ctx) MAX_PARTITIONS 99999999999999999999999999999999", + Issue::MaxPartitionsOutOfRange, + ), + ]; + let table = |vector_column: &str| { + format!( + "CREATE TABLE t (id INTEGER PRIMARY KEY, ctx TEXT NOT NULL, kind INTEGER NOT NULL, \ + loose TEXT, r REAL NOT NULL, doc JSON NOT NULL, other VECTOR(4) NOT NULL, \ + {vector_column})" + ) + }; + for (vector_column, wanted) in refused { + let sql = table(vector_column); + match support::refusal(&db, &sql) { + Error::InvalidVectorPartitionDeclaration { index, issue } => { + assert_eq!(index, column("t", "v"), "{sql}"); + assert_eq!(issue, wanted, "{sql}"); + } + other => panic!("{sql}\nwas refused as {other:?}"), + } + assert!(db.table_names().is_empty(), "{sql}\ndeclared a table"); + } + for vector_column in [ + "v VECTOR(0)", + "v VECTOR(0) PARTITION_KEY (ctx) MAX_PARTITIONS 2", + ] { + let sql = table(vector_column); + match support::refusal(&db, &sql) { + Error::InvalidVectorPolicyDeclaration { index, issue } => { + assert_eq!(index, column("t", "v"), "{sql}"); + assert_eq!( + issue, + VectorPolicyDeclarationIssue::DimensionsNotPositive, + "{sql}" + ); + } + other => panic!("{sql}\nwas refused as {other:?}"), + } + assert!(db.table_names().is_empty(), "{sql}\ndeclared a table"); + } + assert_eq!( + db.committed_watermark(), + fresh, + "a refused declaration committed" + ); + + // What stands: each column with the words its declaration prints back for it. The three + // clauses come in one order; MAX_PARTITIONS is printed for every column with a + // PARTITION_KEY, as 256 when none was written; SEARCH_MODE only when it was written. + let stands: [(&str, &str, &[&str]); 7] = [ + ("bare", "VECTOR(4)", &["bare", "VECTOR", "4"]), + ( + "keyed", + "VECTOR(4) PARTITION_KEY (ctx)", + &[ + "keyed", + "VECTOR", + "4", + "PARTITION_KEY", + "ctx", + "MAX_PARTITIONS", + "256", + ], + ), + ( + "two_keys", + "VECTOR(4) PARTITION_KEY (kind, ctx) MAX_PARTITIONS 4294967295", + &[ + "two_keys", + "VECTOR", + "4", + "PARTITION_KEY", + "kind", + "ctx", + "MAX_PARTITIONS", + "4294967295", + ], + ), + ( + "capped", + "VECTOR(4) PARTITION_KEY (ctx) MAX_PARTITIONS 2 SEARCH_MODE AUTO", + &[ + "capped", + "VECTOR", + "4", + "PARTITION_KEY", + "ctx", + "MAX_PARTITIONS", + "2", + "SEARCH_MODE", + "AUTO", + ], + ), + ( + "strict", + "VECTOR(768) partition_key (ctx) search_mode exact NOT NULL", + &[ + "strict", + "VECTOR", + "768", + "PARTITION_KEY", + "ctx", + "MAX_PARTITIONS", + "256", + "SEARCH_MODE", + "EXACT", + "NOT", + "NULL", + ], + ), + ( + "by_index", + "VECTOR(3) SEARCH_MODE INDEXED", + &["by_index", "VECTOR", "3", "SEARCH_MODE", "INDEXED"], + ), + ("one", "VECTOR(1) MAX_PARTITIONS 1 PARTITION_KEY (ctx)", &[]), + ]; + // The last of them writes two clauses the wrong way round, which is not read. + let (out_of_order, standing) = stands.split_last().expect("seven columns"); + let wrong = format!( + "CREATE TABLE t (id INTEGER PRIMARY KEY, ctx TEXT NOT NULL, {} {})", + out_of_order.0, out_of_order.1 + ); + assert!( + matches!( + support::refusal(&db, &wrong), + Error::ParseError(_) | Error::PlanError(_) + ), + "{wrong}" + ); + assert!(db.table_names().is_empty()); + + let columns: Vec = standing + .iter() + .map(|(name, clauses, _)| format!("{name} {clauses}")) + .collect(); + support::run( + &db, + &format!( + "CREATE TABLE t (id INTEGER PRIMARY KEY, ctx TEXT NOT NULL, kind INTEGER NOT NULL, {})", + columns.join(", ") + ), + ); + let printed = db.declaration("t").expect("the declaration of t"); + let check = |declaration: &str, state: &str| { + for (name, _, wanted) in standing { + assert_eq!( + words(declaration, name), + *wanted, + "{state}: the line of {name}" + ); + } + }; + check(&printed, "as declared"); + + // The limit that was printed is the limit that holds: `capped` takes two partitions. + let insert = |db: &Database, id: i64, ctx: &str| { + db.execute( + "INSERT INTO t (id, ctx, kind, capped, strict) VALUES ($id, $ctx, 0, $v, $e)", + ¶ms! { + "id" => id, "ctx" => ctx, "v" => support::vector(id, 4), + "e" => support::vector(id, 768), + }, + ) + }; + let holds = |db: &Database, state: &str| { + for (id, ctx) in [(1, "a"), (2, "b")] { + if let Err(error) = insert(db, id, ctx) { + panic!("{state}: partition {ctx} of two was refused: {error:?}"); + } + } + match insert(db, 3, "c") { + Err(Error::VectorPartitionLimitExceeded { + index, + max_partitions, + }) => { + assert_eq!(index, column("t", "capped"), "{state}"); + assert_eq!(max_partitions, 2, "{state}"); + } + other => panic!("{state}: a third partition of `capped` got {other:?}"), + } + // And the mode that was printed is the mode in force. + let sql = "SELECT id FROM t ORDER BY by_index <=> $q LIMIT 1"; + match db.execute(sql, ¶ms! { "q" => vec![0.5_f32; 3] }) { + Err(Error::VectorIndexedRouteUnavailable { index }) => { + assert_eq!(index, column("t", "by_index"), "{state}"); + } + other => panic!("{state}: {sql}\ngot {other:?}"), + } + }; + holds(&db, "as declared"); + + // After reopen the same text, through the writer and through a reader. + support::close(db); + let db = support::open(&folder); + assert_eq!( + db.declaration("t").expect("the declaration"), + printed, + "after reopen" + ); + let reader = ReadSession::open(&folder).expect("open the store for reading"); + assert_eq!( + reader.declaration("t").expect("a reader's declaration"), + printed + ); + + // The printed text declares the same table in a new store. + let folder = support::folder("vector-declarations-again"); + let again = support::open(&folder); + for statement in printed.split_inclusive(";\n") { + support::run(&again, statement); + } + assert_eq!( + again.declaration("t").expect("the declaration"), + printed, + "declared again" + ); + check( + &again.declaration("t").expect("the declaration"), + "declared again", + ); + holds(&again, "declared again"); +} + +/// How a statement was refused through the writer, a transaction and a reader, with `$q` bound +/// to this vector: it must be refused at all three, and alike. The transaction is rolled back. +fn refused_at_every_door(db: &Database, reader: &ReadSession, sql: &str, query: &[f32]) -> Error { + let bound = params! { "q" => query.to_vec() }; + let tx = db.transaction().expect("begin"); + let doors = [ + ("the writer", db.execute(sql, &bound)), + ("a transaction", tx.execute(sql, &bound)), + ("a reader", reader.execute(sql, &bound)), + ]; + let refusals = doors.map(|(door, outcome)| match outcome { + Err(error) => error, + Ok(result) => panic!("{sql}\nwas not refused through {door}: {result:?}"), + }); + let kinds: Vec = refusals.iter().map(|error| format!("{error:?}")).collect(); + assert!( + kinds[0] == kinds[1] && kinds[0] == kinds[2], + "{sql}\nwas refused in three ways: {kinds:?}" + ); + let [from_writer, ..] = refusals; + from_writer +} + +#[test] +fn a_second_sort_key_after_a_distance_is_refused_as_a_later_form_whose_order_by_is_there() { + let folder = support::folder("second-sort-key"); + let db = support::open(&folder); + declare(&db, "docs", ""); + fill(&db, "docs", 1..=5); + let reader = ReadSession::open(&folder).expect("open the store for reading"); + let query = support::vector(93_000, DIMENSION); + let before = ( + support::rows(&db, "SELECT id, ctx, tag, v FROM docs ORDER BY id"), + db.committed_watermark(), + ); + + // What a statement is told whose distance is in no ORDER BY at all: there the ORDER BY is + // what is missing, and the sentence may say so. + let no_order_by = "SELECT id, v <=> $q AS distance FROM docs LIMIT 3"; + let missing = match refused_at_every_door(&db, &reader, no_order_by, &query) { + Error::PlanError(text) => text, + other => panic!("{no_order_by}\nwas refused as {other:?}"), + }; + + // A search with a second sort key after its distance has its ORDER BY. It is a form of a + // later version: refused when planned, as every such form is, naming `<=>`, and not in + // the sentence of a statement that lacks the ORDER BY. That holds with USE VECTOR too, + // which is not told that it "requires ORDER BY ... <=>" by a statement that wrote one. + let mut wrong = Vec::new(); + for sql in [ + "SELECT id FROM docs ORDER BY v <=> $q, id LIMIT 3", + "SELECT id, v <=> $q AS distance FROM docs ORDER BY distance, id LIMIT 3", + "SELECT id FROM docs ORDER BY v <=> $q, id USE VECTOR EXACT LIMIT 3", + ] { + match refused_at_every_door(&db, &reader, sql, &query) { + Error::PlanError(text) => { + if !(text.contains("<=>") && text.ends_with(LATER_VERSION)) { + wrong.push(format!( + "{sql}\nwas not refused as a form of a later version that names <=>: {text}" + )); + } + if text == missing { + wrong.push(format!( + "{sql}\nwas told what a statement with no ORDER BY on its distance is \ + told: {text}" + )); + } + } + other => wrong.push(format!( + "{sql}\nwas refused as {other:?} ({other}), not as a plan error" + )), + } + } + assert!(wrong.is_empty(), "{}", wrong.join("\n\n")); + + // Each of them left the store as it was, and the search without the second key answers. + let after = ( + support::rows(&db, "SELECT id, ctx, tag, v FROM docs ORDER BY id"), + db.committed_watermark(), + ); + assert!(after == before, "a refused statement changed the store"); + let stored = support::stored(&db, "SELECT id, ctx, tag, v FROM docs"); + let wanted = stored.nearest(|_| true, &query, 3); + let sql = support::nearest_sql("docs", "v", None, 3); + support::assert_nearest(&db, &sql, &query, &wanted, "after every refusal"); +} + +#[test] +fn a_distance_on_a_column_of_another_table_is_a_column_that_is_not_found() { + let folder = support::folder("distance-wrongly-qualified"); + let db = support::open(&folder); + declare(&db, "docs", ""); + declare(&db, "other", ""); + fill(&db, "docs", 1..=5); + fill(&db, "other", 1..=5); + let reader = ReadSession::open(&folder).expect("open the store for reading"); + let query = support::vector(94_000, DIMENSION); + + // Under the name of a table that is declared and is not the statement's, and under a name + // no table has. + let mut wrong = Vec::new(); + for qualifier in ["other", "nowhere"] { + // What any column written under that name is refused as: not found, naming the + // statement's table and the column. + let plain = format!("SELECT {qualifier}.v FROM docs LIMIT 3"); + let not_found = refused_at_every_door(&db, &reader, &plain, &query); + match ¬_found { + Error::ColumnNotFound { table, column } => { + assert_eq!((table.as_str(), column.as_str()), ("docs", "v"), "{plain}"); + } + other => panic!("{plain}\nwas refused as {other:?}"), + } + // The same column on the left of `<=>` is refused the same way: as the distance the + // search is ordered by, under its select-list name or written in the ORDER BY, and as + // a second distance beside it. + for sql in [ + format!( + "SELECT id, {qualifier}.v <=> $q AS distance FROM docs ORDER BY distance LIMIT 3" + ), + format!("SELECT id FROM docs ORDER BY {qualifier}.v <=> $q LIMIT 3"), + format!("SELECT id, {qualifier}.v <=> $q AS far FROM docs ORDER BY v <=> $q LIMIT 3"), + ] { + let refused = refused_at_every_door(&db, &reader, &sql, &query); + if format!("{refused:?}") != format!("{not_found:?}") { + wrong.push(format!( + "{sql}\nwas refused as {refused:?}, not as {not_found:?}" + )); + } + } + } + assert!(wrong.is_empty(), "{}", wrong.join("\n\n")); + + // Under its own table's name the column is the column, and the search answers. + let stored = support::stored(&db, "SELECT id, ctx, tag, v FROM docs"); + let wanted = stored.nearest(|_| true, &query, 3); + for sql in [ + "SELECT id, docs.v <=> $q AS distance FROM docs ORDER BY distance LIMIT 3", + "SELECT id, docs.v <=> $q AS distance FROM docs ORDER BY docs.v <=> $q LIMIT 3", + ] { + support::assert_nearest( + &db, + sql, + &query, + &wanted, + "the column under its table's name", + ); + support::assert_nearest(&reader, sql, &query, &wanted, "through a reader"); + } +} + +#[test] +fn a_distance_on_a_table_that_does_not_exist_is_a_table_that_is_not_found() { + let folder = support::folder("distance-no-such-table"); + let db = support::open(&folder); + declare(&db, "other", ""); + fill(&db, "other", 1..=5); + let reader = ReadSession::open(&folder).expect("open the store for reading"); + let query = support::vector(94_100, DIMENSION); + + // What any statement on a table that is not there is told. + let plain = "SELECT other.v FROM missing LIMIT 3"; + match refused_at_every_door(&db, &reader, plain, &query) { + Error::TableNotFound(table) => assert_eq!(table, "missing", "{plain}"), + other => panic!("{plain}\nwas refused as {other:?}"), + } + // A distance on a column of a table that is there changes nothing about that: the table + // of the statement is still the first thing that is wrong with it. + let mut wrong = Vec::new(); + for sql in [ + "SELECT id FROM missing ORDER BY other.v <=> $q LIMIT 3", + "SELECT other.v <=> $q AS d FROM missing ORDER BY d LIMIT 3", + ] { + match refused_at_every_door(&db, &reader, sql, &query) { + Error::TableNotFound(table) if table == "missing" => {} + other => wrong.push(format!( + "{sql}\nwas refused as {other:?}, not as the table missing not found" + )), + } + } + assert!(wrong.is_empty(), "{}", wrong.join("\n\n")); +} + +#[test] +fn a_distance_in_the_where_of_a_write_is_refused_on_a_table_with_no_row() { + let folder = support::folder("distance-in-a-write-no-row"); + let db = support::open(&folder); + support::run( + &db, + "CREATE TABLE docs (id INTEGER PRIMARY KEY, r REAL, v VECTOR(4))", + ); + let bound = params! { "q" => vec![0.5_f32; 4] }; + let writes = [ + "UPDATE docs SET r = 1.5 WHERE v <=> $q < 0.5", + "DELETE FROM docs WHERE v <=> $q < 0.5", + ]; + + // With no row in the table and with one: the refusal of a statement does not depend on + // what its table holds. Through the writer and inside a transaction, which stays usable. + let mut wrong = Vec::new(); + for rows in [0_i64, 1] { + if rows == 1 { + support::run( + &db, + "INSERT INTO docs (id, r, v) VALUES (1, 0.25, [1.0, 0.0, 0.0, 0.0])", + ); + } + let position = db.committed_watermark(); + let tx = db.transaction().expect("begin"); + for sql in writes { + let doors = [ + ("the writer", db.execute(sql, &bound)), + ("a transaction", tx.execute(sql, &bound)), + ]; + for (door, outcome) in doors { + match outcome { + Err(Error::PlanError(text)) if text.contains("<=>") => {} + Err(other) => wrong.push(format!( + "{sql}\nwith {rows} rows, through {door}, was refused as {other:?}, not \ + as a plan error that names <=>" + )), + Ok(result) => wrong.push(format!( + "{sql}\nwith {rows} rows, through {door}, was answered with {} rows \ + affected, not refused", + result.rows_affected() + )), + } + } + } + // The transaction is as usable as before, and nothing was written. + assert_eq!(support::int(&tx, "SELECT COUNT(*) FROM docs"), rows); + tx.rollback(); + assert_eq!(db.committed_watermark(), position, "with {rows} rows"); + assert_eq!(support::int(&db, "SELECT COUNT(*) FROM docs"), rows); + } + assert!(wrong.is_empty(), "{}", wrong.join("\n\n")); + assert_eq!( + support::rows(&db, "SELECT id, r FROM docs"), + [[Value::Int64(1), Value::Float64(0.25)]] + ); +} diff --git a/v2/contextdb/tests/vectors_upkeep.rs b/v2/contextdb/tests/vectors_upkeep.rs new file mode 100644 index 00000000..17a1d4fc --- /dev/null +++ b/v2/contextdb/tests/vectors_upkeep.rs @@ -0,0 +1,423 @@ +//! Sealing: vectors move from pending lists into blocks only when a maintenance cycle is +//! called, a bounded number in each cycle; a reader in another process searches while the +//! writer seals and gets the answer of one commit every time; and a small store of fixed +//! content stays within the size that was recorded for it. + +mod support; + +use std::fs; +use std::sync::atomic::Ordering; + +use contextdb::{ + Database, MaintenanceReport, QueryResult, ReadSession, ReadSessionOptions, params, +}; + +use support::{Child, Run}; + +/// What one cycle reports: the vectors it sealed and the vectors still waiting. +fn cycle(db: &Database) -> (u64, u64) { + let report = db + .run_maintenance_cycle() + .expect("a maintenance cycle") + .vector; + ( + report.maintenance_vectors_done, + report.maintenance_vectors_remaining, + ) +} + +#[test] +fn sealing_advances_only_when_asked_and_in_bounded_cycles() { + const DIMENSION: usize = 8; + let folder = support::folder("sealing-in-cycles"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, "PARTITION_KEY (ctx)"); + // Five thousand vectors: three thousand under a, the partition made first, and two + // thousand under b. + for first in (1..=5000_i64).step_by(1000) { + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, first..first + 1000, DIMENSION, |n| { + String::from(if n <= 3000 { "a" } else { "b" }) + }); + tx.commit().expect("commit"); + } + let queries = [ + support::vector(94_001, DIMENSION), + support::vector(94_002, DIMENSION), + ]; + // Ten rows of everything, of a and of b: each equals the reference for both vectors. + // Returns the rows the three searches read for the first vector. + let examined = |on: &dyn Run, state: &str| -> [u64; 3] { + let stored = support::stored(&on, support::DOCS); + let scopes = [ + (None, None), + (Some("ctx = 'a'"), Some("a")), + (Some("ctx = 'b'"), Some("b")), + ]; + scopes.map(|(filter, ctx): (Option<&str>, Option<&str>)| { + let passes = + |row: &[contextdb::Value]| ctx.is_none_or(|ctx| row[1] == support::text(ctx)); + let sql = support::nearest_sql("docs", "v", filter, 10); + let mut read = Vec::new(); + for query in &queries { + let wanted = stored.nearest(passes, query, 10); + assert_eq!(wanted.len(), 10, "{state}"); + let result = support::assert_nearest(&on, &sql, query, &wanted, state); + read.push(result.rows_examined()); + } + read[0] + }) + }; + + // Before any cycle every vector is pending: a search reads each pending row of the + // partitions it selects once, and each row it returns once more. + assert_eq!(examined(&db, "nothing sealed"), [5010, 3010, 2010]); + // Nothing seals unless a cycle is called: not a search, not a close, not an open. + support::close(db); + let db = support::open(&folder); + assert_eq!( + examined(&db, "nothing sealed, after reopen"), + [5010, 3010, 2010] + ); + let reader = ReadSession::open(&folder).expect("open the store for reading"); + assert_eq!( + examined(&reader, "nothing sealed, a reader"), + [5010, 3010, 2010] + ); + + // The first cycle seals 4,096 vectors in one commit: the first in the order of partition + // and row key, so all of a and the first 1,096 of b. 904 of b stay pending. + let unsealed = db.committed_watermark(); + assert_eq!(cycle(&db), (4096, 904)); + let after_first = db.committed_watermark(); + assert!(after_first > unsealed, "a cycle that sealed made no commit"); + let [everything, a, b] = examined(&db, "after the first cycle"); + assert!( + (914..5000).contains(&everything), + "everything read {everything} rows" + ); + assert!( + (10..914).contains(&a), + "a, which has nothing pending, read {a} rows" + ); + assert!( + (914..2010).contains(&b), + "b, which has 904 pending, read {b} rows" + ); + assert_eq!( + examined(&reader, "after the first cycle, a reader"), + [everything, a, b] + ); + + // The second seals the rest, and the third has nothing to do and makes no commit. + assert_eq!(cycle(&db), (904, 0)); + let after_second = db.committed_watermark(); + assert!( + after_second > after_first, + "a cycle that sealed made no commit" + ); + assert_eq!(cycle(&db), (0, 0)); + assert_eq!( + db.committed_watermark(), + after_second, + "a cycle with nothing to seal" + ); + let sealed = examined(&db, "after the last cycle"); + for (read, pending) in sealed.into_iter().zip([5000, 3000, 2000]) { + assert!( + (10..pending).contains(&read), + "a sealed search read {read} rows" + ); + } + assert_eq!(examined(&reader, "after the last cycle, a reader"), sealed); + + // A cycle reports exactly the vectors that came to wait since the one before. + let position = |db: &Database| db.committed_watermark(); + // Ten new rows. + support::insert_docs(&db, 5001..=5010, DIMENSION, |_| "a".to_owned()); + assert_eq!(cycle(&db), (10, 0)); + // A sealed row whose vector changes, twice, waits once. + for fresh in [94_101, 94_102] { + db.execute( + "UPDATE docs SET v = $v WHERE id = 1", + ¶ms! { "v" => support::vector(fresh, DIMENSION) }, + ) + .expect("change a vector"); + } + assert_eq!(cycle(&db), (1, 0)); + // Five new rows, of which two are deleted and one loses its vector before the cycle. + support::insert_docs(&db, 5011..=5015, DIMENSION, |_| "b".to_owned()); + support::run(&db, "DELETE FROM docs WHERE id IN (5011, 5012)"); + support::run(&db, "UPDATE docs SET v = NULL WHERE id = 5013"); + assert_eq!(cycle(&db), (2, 0)); + // A row with no vector, and a sealed row that is deleted, leave nothing to seal. + support::insert_doc(&db, 5016, "a", None).expect("a row with no vector"); + support::run(&db, "DELETE FROM docs WHERE id = 2"); + let before = position(&db); + assert_eq!(cycle(&db), (0, 0)); + assert_eq!(position(&db), before); + // A sealed row that moves to the other partition waits there. + support::run(&db, "UPDATE docs SET ctx = 'b' WHERE id = 3"); + assert_eq!(cycle(&db), (1, 0)); + examined(&db, "at the end"); + + // A store with no VECTOR column, and a store with no table: zeros, and no commit. + let folder = support::folder("sealing-nothing"); + let db = support::open(&folder); + let before = position(&db); + let report = db + .run_maintenance_cycle() + .expect("a cycle of an empty store"); + assert_eq!(report, MaintenanceReport::default()); + assert_eq!(position(&db), before); + support::declare(&db); + support::insert(&db, 1..=10); + let before = position(&db); + let report = db + .run_maintenance_cycle() + .expect("a cycle of a store with no vectors"); + assert_eq!( + ( + report.vector.maintenance_vectors_done, + report.vector.maintenance_vectors_remaining + ), + (0, 0) + ); + assert_eq!(position(&db), before); + support::assert_rows_are(&db, 1..=10); +} + +/// The dimension, the vector and the statement of the searches in the next test. +const SEARCHED: usize = 16; + +fn searched_for() -> Vec { + support::vector(95_000, SEARCHED) +} + +fn search_of_ten() -> String { + support::nearest_sql("docs", "v", None, 10) +} + +/// A search's answer and its counted work as one line: the rows it read, then each key with +/// the bits of its distance. Two searches of one commit give the same line. +fn told(result: &QueryResult) -> String { + let rows: Vec = support::found(result, "the search of ten") + .iter() + .map(|near| format!("{:?}:{:016x}", near.key, near.distance.to_bits())) + .collect(); + format!("{} {}", result.rows_examined(), rows.join(",")) +} + +#[test] +fn a_reader_in_another_process_searches_while_the_writer_seals() { + if let Some(role) = support::role() { + return match role.name.as_str() { + "looping" => looping(&role), + "stopped" => stopped(&role), + other => panic!("no role {other:?} in this test"), + }; + } + let test = "a_reader_in_another_process_searches_while_the_writer_seals"; + let folder = support::folder("reader-beside-sealing"); + let db = support::open(&folder); + support::declare_docs(&db, SEARCHED, "PARTITION_KEY (ctx)"); + let insert = |numbers: std::ops::Range| { + let tx = db.transaction().expect("begin"); + support::insert_docs(&tx, numbers, SEARCHED, |n| format!("c{}", n % 2)); + tx.commit().expect("commit"); + }; + for first in (1..6001).step_by(1000) { + insert(first..first + 1000); + } + // The line of the newest commit: the search equals the reference there, and this is what + // it reads and returns. + let newest = || -> String { + let stored = support::stored(&db, support::DOCS); + let wanted = stored.nearest(|_| true, &searched_for(), 10); + let sql = search_of_ten(); + told(&support::assert_nearest( + &db, + &sql, + &searched_for(), + &wanted, + "the writer", + )) + }; + let mut commits = vec![newest()]; + + // One reader is stopped halfway through the storage reads of a search of the first commit, + // where every vector is pending. Where halfway is comes from a run of the same search to + // its end. + let lines = Child::start(test, "stopped pause_at=0", &folder).end(); + assert_eq!(lines.first(), commits.first(), "a search nothing stopped"); + let reads = support::last_number(&lines, "reads").expect("the reader announces its reads"); + assert!( + reads > 1, + "a search of 6,000 pending rows made {reads} storage reads" + ); + let mut held = Child::start( + test, + &format!("stopped pause_at={}", reads.div_ceil(2)), + &folder, + ); + held.expect("reading"); + + // A second reader searches in a loop, and announces each answer that differs from its + // last. Every line it announces must be the line of a commit, and it must come to the + // newest one: it is never refused and never told to try again. + let mut looping = Child::start(test, "looping", &folder); + let mut follow = |commits: &[String]| { + let newest = commits.last().expect("a commit"); + loop { + let Some(line) = looping.line() else { + panic!("the looping reader ended; it had still to reach\n{newest}") + }; + assert!( + commits.contains(&line), + "the looping reader announced an answer that is no commit's:\n{line}" + ); + if line == *newest { + return; + } + } + }; + follow(&commits); + + // The writer seals cycle by cycle and inserts and changes between. Each step is one + // commit that changes what a search reads, so the reader can be seen to arrive at it. + let steps: [&dyn Fn(); 7] = [ + &|| assert_eq!(cycle(&db), (4096, 1904)), + &|| insert(6001..6501), + &|| assert_eq!(cycle(&db), (2404, 0)), + &|| insert(6501..7001), + &|| assert_eq!(cycle(&db), (500, 0)), + &|| { + let tx = db.transaction().expect("begin"); + for id in 1..=300_i64 { + tx.execute( + "UPDATE docs SET v = $v WHERE id = $id", + ¶ms! { "id" => id, "v" => support::vector(95_100 + id, SEARCHED) }, + ) + .expect("change a vector"); + } + tx.commit().expect("commit"); + }, + &|| assert_eq!(cycle(&db), (300, 0)), + ]; + for step in steps { + step(); + let line = newest(); + assert!( + Some(&line) != commits.last(), + "a step left the search's line as it was" + ); + commits.push(line); + follow(&commits); + } + + // The reader that was stopped inside its search through all of it answers with the rows + // of the commit it began at, for the work that commit took; its next search is of the + // newest commit. + held.tell("go"); + assert_eq!( + held.line().as_ref(), + commits.first(), + "the search that was stopped" + ); + assert!(held.line().is_some_and(|line| line.starts_with("reads "))); + assert_eq!( + held.line().as_ref(), + commits.last(), + "the stopped reader's next search" + ); + held.end(); + looping.kill(); +} + +/// Opens the store for reading and searches without end, announcing an answer's line each +/// time it differs from the one before. The parent kills it. +fn looping(role: &support::Role) { + let reader = ReadSession::open(&role.folder).expect("open the store for reading"); + let (sql, query) = (search_of_ten(), searched_for()); + let mut last = String::new(); + loop { + let result = match support::search(&reader, &sql, &query) { + Ok(result) => result, + Err(error) => panic!("a reader was refused while the writer sealed: {error:?}"), + }; + let line = told(&result); + if line != last { + support::say(&line); + last = line; + } + } +} + +/// Opens the store for reading and searches once. At storage read number `pause_at` of that +/// search it announces `reading` and waits for a line; `pause_at` 0 never stops. Then it +/// announces the answer's line and `reads `, the storage reads the search made. Then it +/// searches once more and announces that answer's line. +fn stopped(role: &support::Role) { + let (door, searching, reads) = support::pausing_at_read(role.number("pause_at")); + let mut options = ReadSessionOptions::default(); + options.storage = Some(door); + let reader = + ReadSession::open_with_options(&role.folder, options).expect("open the store for reading"); + let (sql, query) = (search_of_ten(), searched_for()); + + searching.store(true, Ordering::SeqCst); + let result = support::search(&reader, &sql, &query).expect("a reader's search is not refused"); + searching.store(false, Ordering::SeqCst); + support::say(told(&result)); + support::say(format!("reads {}", reads.load(Ordering::SeqCst))); + let next = support::search(&reader, &sql, &query).expect("the reader's next search"); + support::say(told(&next)); +} + +#[test] +fn a_small_fixed_store_stays_within_its_recorded_size() { + // The bytes put in: 1,500 ids of 8 bytes, 1,500 vectors of 1,536 bytes and 4,800,000 + // bytes of text. And the most bytes the data file may have. + const PUT_IN: u64 = 7_116_000; + // The length is of a store that keeps a large text compressed when that saves an eighth. + // `body` is letters drawn one by one, which do not shrink by an eighth: stored as it is. + const RECORDED: u64 = 2 * PUT_IN; // set from the first run + + let folder = support::folder("small-fixed-store"); + let db = support::open(&folder); + support::run( + &db, + "CREATE TABLE notes (id INTEGER PRIMARY KEY, body TEXT NOT NULL, embedding VECTOR(384))", + ); + // Nothing in the content, its order or the store depends on a clock, a random number or a + // thread, so the file has the same length on every machine and every run. + let mut put_in = 0_u64; + for first in (1..=1500_i64).step_by(500) { + let tx = db.transaction().expect("begin"); + for id in first..first + 500 { + // 200, 1,700, 3,200, 4,700 or 6,200 lower-case letters: one size that stays in + // its row and four that do not. + let body = support::letters(id, 200 + 1500 * (id % 5) as usize); + let embedding = support::vector(id, 384); + put_in += 8 + body.len() as u64 + 4 * embedding.len() as u64; + tx.execute( + "INSERT INTO notes (id, body, embedding) VALUES ($id, $body, $embedding)", + ¶ms! { "id" => id, "body" => body, "embedding" => embedding }, + ) + .unwrap_or_else(|error| panic!("note {id} was refused: {error:?}")); + } + tx.commit().expect("commit five hundred notes"); + } + assert_eq!(put_in, PUT_IN, "the bytes this test put in"); + assert_eq!(support::seal_all(&db), 1500); + support::close(db); + + let length = fs::metadata(folder.join("data")) + .expect("the store's data file") + .len(); + assert!( + length <= RECORDED, + "the data file is {length} bytes long; the most recorded for this store is {RECORDED} \ + bytes, and {PUT_IN} bytes were put in" + ); +} diff --git a/v2/numbers.sh b/v2/numbers.sh index 9d4f7933..ce54d13c 100755 --- a/v2/numbers.sh +++ b/v2/numbers.sh @@ -9,6 +9,18 @@ # builds in release mode, works in a folder of its own under $TMPDIR and removes that folder # when it is done. Packages cargo has downloaded before are not downloaded again, so the two # build times leave the download out. +# +# Three groups of lines are taken on data the script is given, each in a variable of the +# environment, and are printed as `unavailable` without it: +# +# SESSIONS=:: text files loaded one row a line: the +# first names of the list, each a +# path under the folder +# PASSAGES=: a dump of search passages and the answers +# PASSAGES_LARGE=: recorded for it; a second, larger one +# +# The head of contextdb/examples/passages.rs says what a dump and its answers hold. A store's +# size is the bytes of every file in its folder after it was closed. set -u cd "$(dirname "$0")/.." || exit 1 scratch=$(mktemp -d) || exit 1 @@ -39,11 +51,42 @@ else echo "storage_library_changed_lines unavailable -" fi -# The README's program at 20,000 rows of 1.6 KB, and a million rows under a 64 MiB limit. -for what in readme million; do - cargo run --quiet --release --locked --manifest-path v2/Cargo.toml -p contextdb --example measure -- "$what" "$scratch/$what" || - echo "${what}_measurements unavailable -" +# The measuring programs are built as any program that uses the library is: in release mode +# with nothing added. This line says what compiler flags the environment adds, so that a +# number from a build made for one processor is not read as an ordinary build's. +flags="${RUSTFLAGS:-}${CARGO_ENCODED_RUSTFLAGS:+ $CARGO_ENCODED_RUSTFLAGS}${CARGO_BUILD_RUSTFLAGS:+ $CARGO_BUILD_RUSTFLAGS}" +echo "find_build_rustflags $(printf '%s' "${flags:-none}" | tr -s ' \037' '__') -" +example() { + cargo run --quiet --release --locked --manifest-path v2/Cargo.toml -p contextdb --example "$@" +} + +# 20,000 notes of 1.6 KB, a million rows under a 64 MiB limit, and 20,000 vectors of 384 +# numbers with the 10 nearest of 100 others. +for what in readme million vectors; do + example measure -- "$what" "$scratch/$what" || echo "${what}_measurements unavailable -" + rm -rf -- "$scratch/$what" +done + +# Text files, one row a line: what was put in beside what the store takes. +if [ -n "${SESSIONS:-}" ]; then + IFS=: read -r corpus list count <<< "$SESSIONS" + example measure -- sessions "$corpus" "$list" "$count" "$scratch/sessions" || + echo "sessions_measurements unavailable -" + rm -rf -- "$scratch/sessions" +else + echo "sessions_measurements unavailable -" +fi + +# Search passages: load, seal, search through SQL, and the answers beside the recorded ones. +given=0 +for input in "${PASSAGES:-}" "${PASSAGES_LARGE:-}"; do + [ -n "$input" ] || continue + given=1 + example passages -- "${input%%:*}" "${input#*:}" "$scratch/passages" || + echo "passages_measurements unavailable -" + rm -rf -- "$scratch/passages" done +[ "$given" = 1 ] || echo "passages_measurements unavailable -" # From a fresh clone to a green `cargo test`: nothing built before. start=$SECONDS @@ -53,9 +96,9 @@ else echo "fresh_clone_to_green_test unavailable -" fi -# The first build of a program that embeds the library: the README's own program. +# The first build of a program that embeds the library: the README's first program. mkdir -p "$scratch/embed/src" -sed -n '/^```rust/,/^```$/p' v2/contextdb/README.md | sed '1d;$d' > "$scratch/embed/src/main.rs" +awk '/^```rust/ { block++; next } /^```$/ && block == 1 { exit } block == 1' v2/contextdb/README.md > "$scratch/embed/src/main.rs" printf '[package]\nname = "embed"\nversion = "0.0.0"\nedition = "2024"\n\n[dependencies]\ncontextdb = { path = "%s/v2/contextdb" }\n' "$PWD" > "$scratch/embed/Cargo.toml" start=$SECONDS if CARGO_TARGET_DIR="$scratch/embed-target" cargo build --release --manifest-path "$scratch/embed/Cargo.toml" > "$scratch/embed.log" 2>&1; then diff --git a/v2/redb/README.md b/v2/redb/README.md index 34270d8b..522dd137 100644 --- a/v2/redb/README.md +++ b/v2/redb/README.md @@ -30,6 +30,17 @@ Checking pages when they are read: pages they read into the page cache, so the cache only holds pages that were checked or that were written through the same handle. +Keeping a file close to the size of what it holds: + +- A leaf written in key order is left full. When a leaf must split because a pair was inserted + after its last pair, the first leaf keeps every pair that was there and the second starts with + the new pair; upstream divides every leaf where half its bytes are, which leaves a table that + is written in key order with half-empty leaves. Every other split is still in half. +- The file grows in small steps. Once its usable bytes pass 1 MiB the file grows by a sixteenth + of them, and never by less than twice the allocation that needed the room; upstream doubles + the file while it has one region and adds a whole region at a time after that. A file under + 1 MiB still doubles, and a region is still filled before the next one is started. + Everything below is upstream's README. # redb diff --git a/v2/redb/src/tree_store/btree_base.rs b/v2/redb/src/tree_store/btree_base.rs index 25b5b3e3..fc7fe86e 100644 --- a/v2/redb/src/tree_store/btree_base.rs +++ b/v2/redb/src/tree_store/btree_base.rs @@ -728,7 +728,14 @@ impl<'a, 'b> LeafBuilder<'a, 'b> { required_size > self.mem.get_page_size() && self.pairs.len() > 1 } - pub(super) fn build_split<'txn>(self) -> Result<(PageMut<'txn>, &'a [u8], PageMut<'txn>)> { + // Divides the pairs between two pages where half their bytes are. When `appended` is set, + // the last pair was added after every pair the leaf already held: the first page then keeps + // all of those and the second starts with the new pair, so that a tree written in key order + // is left with full leaves instead of half-empty ones. + pub(super) fn build_split<'txn>( + self, + appended: bool, + ) -> Result<(PageMut<'txn>, &'a [u8], PageMut<'txn>)> { let total_size = self.total_key_bytes + self.total_value_bytes; let mut division = 0; let mut first_split_key_bytes = 0; @@ -737,7 +744,7 @@ impl<'a, 'b> LeafBuilder<'a, 'b> { first_split_key_bytes += key.len(); first_split_value_bytes += value.len(); division += 1; - if first_split_key_bytes + first_split_value_bytes >= total_size / 2 { + if !appended && first_split_key_bytes + first_split_value_bytes >= total_size / 2 { break; } } diff --git a/v2/redb/src/tree_store/btree_mutator.rs b/v2/redb/src/tree_store/btree_mutator.rs index 44f2ee6d..9ac67cac 100644 --- a/v2/redb/src/tree_store/btree_mutator.rs +++ b/v2/redb/src/tree_store/btree_mutator.rs @@ -417,7 +417,8 @@ impl<'a, 'b, K: Key, V: Value> MutateHelper<'a, 'b, K, V> { old_value: existing_value, } } else { - let (new_page1, split_key, new_page2) = builder.build_split()?; + let appended = !found && accessor.num_pairs() == position; + let (new_page1, split_key, new_page2) = builder.build_split(appended)?; let split_key = split_key.to_vec(); let page_number = page.get_page_number(); let existing_value = if found { @@ -899,7 +900,8 @@ impl<'a, 'b, K: Key, V: Value> MutateHelper<'a, 'b, K, V> { .push_all_except(&partial_child_accessor, Some(deleted_pair)); } if child_builder.should_split() { - let (new_page1, split_key, new_page2) = child_builder.build_split()?; + let (new_page1, split_key, new_page2) = + child_builder.build_split(false)?; builder.push_key(split_key); builder.push_child(new_page1.get_page_number(), DEFERRED); builder.push_child(new_page2.get_page_number(), DEFERRED); diff --git a/v2/redb/src/tree_store/page_store/page_manager.rs b/v2/redb/src/tree_store/page_store/page_manager.rs index 341f64ce..f0c3e12e 100644 --- a/v2/redb/src/tree_store/page_store/page_manager.rs +++ b/v2/redb/src/tree_store/page_store/page_manager.rs @@ -1194,24 +1194,34 @@ impl TransactionalMemory { 2u64.pow(required_order_allocation.into()) * u64::from(state.header.page_size()); let max_region_size = u64::from(state.header.layout().full_region_layout().num_pages()) * u64::from(state.header.page_size()); + // The file grows in small steps, so that it stays close to the size of what it holds: + // by a sixteenth of its usable bytes once they pass 1 MiB, and never by less than twice + // the allocation that needs the room. A smaller file doubles. + let usable = layout.usable_bytes(); + let step = max( + if usable < (1 << 20) { + usable + } else { + usable / 16 + }, + required_growth * 2, + ); let next_desired_size = if layout.num_full_regions() > 0 { if let Some(trailing) = layout.trailing_region_layout() { - if 2 * required_growth < max_region_size - trailing.usable_bytes() { - // Fill out the trailing region - layout.usable_bytes() + (max_region_size - trailing.usable_bytes()) + let room = max_region_size - trailing.usable_bytes(); + if 2 * required_growth < room { + // Grow the trailing region, at most until it is full + usable + min(step, room) } else { - // Fill out trailing & Grow by 1 region - layout.usable_bytes() + 2 * max_region_size - trailing.usable_bytes() + // Fill out trailing & start the next region + usable + room + min(step, max_region_size) } } else { - // Grow by 1 region - layout.usable_bytes() + max_region_size + // Start the next region + usable + min(step, max_region_size) } } else { - max( - layout.usable_bytes() * 2, - layout.usable_bytes() + required_growth * 2, - ) + usable + step }; let new_layout = DatabaseLayout::calculate( next_desired_size, diff --git a/v2/redb/tests/contextdb_changes.rs b/v2/redb/tests/contextdb_changes.rs index 42dcb9e6..897528e4 100644 --- a/v2/redb/tests/contextdb_changes.rs +++ b/v2/redb/tests/contextdb_changes.rs @@ -368,3 +368,116 @@ fn a_page_with_a_wrong_checksum_is_refused_when_read() { Ok(_) => panic!("a damaged page was answered with a wrong value"), } } + +/// The table of the two tests below: the slices of large values, each under its place. +const SLICES: TableDefinition = TableDefinition::new("slices"); + +/// The 243 bytes kept under `place`. +fn slice(place: u64) -> Vec { + format!("slice {place:08} ") + .into_bytes() + .into_iter() + .cycle() + .take(243) + .collect() +} + +/// Writes the slices of these places in this order, 500 to a commit. +fn commit_slices(db: &Database, places: &[u64]) { + for batch in places.chunks(500) { + let write = db.begin_write().unwrap(); + { + let mut table = write.open_table(SLICES).unwrap(); + for place in batch { + table.insert(*place, slice(*place).as_slice()).unwrap(); + } + } + write.commit().unwrap(); + } +} + +/// Asserts that the table holds exactly the slices of the places below `count`, each whole. +fn assert_slices_are(db: &Database, count: u64) { + let read = db.begin_read().unwrap(); + let table = read.open_table(SLICES).unwrap(); + let mut next = 0; + for entry in table.iter().unwrap() { + let (place, value) = entry.unwrap(); + assert_eq!( + place.value(), + next, + "the places come in order, none missing" + ); + assert!( + value.value() == slice(next).as_slice(), + "the slice at {next}" + ); + next += 1; + } + assert_eq!(next, count, "slices read back"); +} + +#[test] +fn a_table_written_in_key_order_has_full_leaves() { + let path = folder("a-table-written-in-key-order").join("data"); + let db = writer(&path); + let in_order: Vec = (0..16_000).collect(); + commit_slices(&db, &in_order); + + // A leaf is 4 bytes and, for each slice, 8 of key, 4 of length and 243 of value: sixteen + // slices take 4,084 of a page's 4,096 bytes. Written in key order, every leaf is left + // full, so 16,000 slices are in exactly 1,000 leaves. + let write = db.begin_write().unwrap(); + let stats = write.stats().unwrap(); + write.abort().unwrap(); + assert_eq!(stats.page_size(), 4096); + assert_eq!( + stats.leaf_pages(), + 1000, + "leaf pages of 16,000 slices in key order" + ); + assert_slices_are(&db, 16_000); + + // A table written in an order far from key order still reads back every slice. 7,919 and + // 16,000 share no factor, so the places below are every place once. + let path = folder("a-table-written-in-scattered-key-order").join("data"); + let db = writer(&path); + let scattered: Vec = (0..16_000_u64).map(|n| n * 7919 % 16_000).collect(); + commit_slices(&db, &scattered); + assert_slices_are(&db, 16_000); + drop(db); + assert_slices_are(&writer(&path), 16_000); +} + +#[test] +fn a_file_is_at_most_a_sixteenth_longer_than_its_pages() { + let path = folder("a-file-is-at-most-a-sixteenth-longer").join("data"); + let db = writer(&path); + let in_order: Vec = (0..40_000).collect(); + commit_slices(&db, &in_order); + + let write = db.begin_write().unwrap(); + let stats = write.stats().unwrap(); + write.abort().unwrap(); + let pages = stats.allocated_pages() * stats.page_size() as u64; + assert!( + pages > 40_000 * 243, + "40,000 slices of 243 bytes are in {pages} bytes of pages" + ); + // Seventeen sixteenths of the pages, and 64 KiB. + let most = pages * 17 / 16 + (64 << 10); + let length = std::fs::metadata(&path).unwrap().len(); + assert!( + length <= most, + "after the last commit the file is {length} bytes long; its pages are {pages} bytes, \ + and it may be {most}" + ); + drop(db); + let length = std::fs::metadata(&path).unwrap().len(); + assert!( + length <= most, + "after the database was dropped the file is {length} bytes long; its pages are \ + {pages} bytes, and it may be {most}" + ); + assert_slices_are(&writer(&path), 40_000); +}