From 718698cda4a67f18f026d40663d2086f4c11d346 Mon Sep 17 00:00:00 2001 From: Lucid Programmer Date: Thu, 12 Nov 2026 00:02:35 +0530 Subject: [PATCH 1/2] contextdb: add a search index for vector columns with exact, indexed and automatic search modes --- .github/workflows/check.yml | 4 +- v2/AGENTS.md | 46 +- v2/README.md | 7 + v2/contextdb-cli/src/main.rs | 3 +- v2/contextdb-cli/src/output.rs | 26 + v2/contextdb-cli/src/shell.rs | 25 +- v2/contextdb-cli/tests/shell.rs | 279 +++- v2/contextdb/README.md | 384 ++++- v2/contextdb/examples/measure.rs | 355 ++++- v2/contextdb/examples/passages.rs | 806 ++++++++--- v2/contextdb/examples/support/mod.rs | 639 +++++++- v2/contextdb/src/database/mod.rs | 104 +- v2/contextdb/src/query/inspect.rs | 423 ++++++ v2/contextdb/src/query/mod.rs | 21 + v2/contextdb/src/query/search.rs | 111 +- v2/contextdb/src/read/mod.rs | 13 + v2/contextdb/src/schema/check.rs | 59 +- v2/contextdb/src/schema/mod.rs | 76 + v2/contextdb/src/schema/record.rs | 26 +- v2/contextdb/src/schema/render.rs | 23 +- v2/contextdb/src/schema/tests.rs | 57 + v2/contextdb/src/sql/ast.rs | 27 + v2/contextdb/src/sql/grammar.pest | 35 +- v2/contextdb/src/sql/parser/mod.rs | 43 +- v2/contextdb/src/sql/parser/tests.rs | 120 +- v2/contextdb/src/sql/plan.rs | 14 +- v2/contextdb/src/sql/planner.rs | 10 + v2/contextdb/src/store/disk.rs | 37 +- v2/contextdb/src/store/keys.rs | 106 +- v2/contextdb/src/store/mod.rs | 64 +- v2/contextdb/src/store/record.rs | 53 +- v2/contextdb/src/store/tests.rs | 200 ++- v2/contextdb/src/types/mod.rs | 40 +- v2/contextdb/src/upkeep/mod.rs | 31 +- v2/contextdb/src/vector/block.rs | 9 +- v2/contextdb/src/vector/build.rs | 976 +++++++++++++ v2/contextdb/src/vector/directory.rs | 151 +- v2/contextdb/src/vector/grouping.rs | 242 ++++ v2/contextdb/src/vector/inspect.rs | 161 +++ v2/contextdb/src/vector/mod.rs | 36 +- v2/contextdb/src/vector/pack.rs | 534 +++++++ v2/contextdb/src/vector/seal.rs | 98 +- v2/contextdb/src/vector/search.rs | 1107 ++++++++++++-- v2/contextdb/src/vector/sketch.rs | 218 +++ v2/contextdb/src/vector/slot.rs | 80 +- v2/contextdb/src/vector/stage.rs | 92 +- v2/contextdb/src/vector/tests.rs | 1282 ++++++++++++++++- v2/contextdb/src/vector/train.rs | 161 +++ v2/contextdb/src/write/mod.rs | 55 +- v2/contextdb/tests/crash/main.rs | 130 ++ v2/contextdb/tests/memory_ceiling.rs | 88 ++ v2/contextdb/tests/readers.rs | 122 ++ .../tests/sql/vector_declarations.sql | 19 +- v2/contextdb/tests/sql/vector_index_small.sql | 70 + v2/contextdb/tests/sql/vector_inspection.sql | 63 + v2/contextdb/tests/sql/vector_policy.sql | 100 ++ .../tests/sql/vector_search_refusals.sql | 22 +- v2/contextdb/tests/store.rs | 26 + v2/contextdb/tests/support/mod.rs | 251 +++- v2/contextdb/tests/transactions.rs | 70 + v2/contextdb/tests/vectors_index.rs | 828 +++++++++++ v2/contextdb/tests/vectors_inspection.rs | 388 +++++ v2/contextdb/tests/vectors_memory.rs | 78 + v2/contextdb/tests/vectors_reduced.rs | 129 ++ v2/contextdb/tests/vectors_statements.rs | 476 +++++- v2/contextdb/tests/vectors_upkeep.rs | 346 ++++- v2/numbers.sh | 69 +- 67 files changed, 11912 insertions(+), 832 deletions(-) create mode 100644 v2/contextdb/src/query/inspect.rs create mode 100644 v2/contextdb/src/vector/build.rs create mode 100644 v2/contextdb/src/vector/grouping.rs create mode 100644 v2/contextdb/src/vector/inspect.rs create mode 100644 v2/contextdb/src/vector/pack.rs create mode 100644 v2/contextdb/src/vector/sketch.rs create mode 100644 v2/contextdb/src/vector/train.rs create mode 100644 v2/contextdb/tests/sql/vector_index_small.sql create mode 100644 v2/contextdb/tests/sql/vector_inspection.sql create mode 100644 v2/contextdb/tests/sql/vector_policy.sql create mode 100644 v2/contextdb/tests/vectors_index.rs create mode 100644 v2/contextdb/tests/vectors_inspection.rs create mode 100644 v2/contextdb/tests/vectors_reduced.rs diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 60c96adc..6e3a9949 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -20,8 +20,8 @@ env: # 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 + SHIPPED_LINES: 15500 + TESTS: 260 jobs: lint: diff --git a/v2/AGENTS.md b/v2/AGENTS.md index ed6ac801..2664c8a5 100644 --- a/v2/AGENTS.md +++ b/v2/AGENTS.md @@ -23,22 +23,28 @@ The library is twelve parts, one folder each under `contextdb/src/`: | `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, 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 | +| `query` | answers a SELECT from one table; holds the one expression evaluator; counts the rows it reads and a search's distance evaluations; answers `SHOW VECTOR_INDEXES`, `SHOW VECTOR_PARTITIONS` and `explain` from the vector part's small records | 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 | +| `database` | the handle a store is opened for writing with, its transactions, the one commit function; `run_maintenance_cycle` (the sealing commit, then one step for every partition with work, each worked out with no lock held and committed through the one commit function); `explain` | has a second way to commit | +| `read` | a session that reads a store without holding it for writing, from any process; its `explain` | 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, the grouping of a partition (its centres, a one-bit sketch of every slot, the partition's mean), search over them (every slot, or the sketches of the nearest groups and of the tail, then the slots those keep), building a grouping in steps and putting it in force, sealing, and what inspection reads | 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; reads a clock, takes the commit lock, or starts a thread | +| `upkeep` | the sealing of one cycle, at most 4,096 waiting vectors moved into blocks in one commit (the 4,096 is `SEAL_ROWS` in `upkeep/mod.rs`), and the report of a cycle; `database` drives the grouping steps that follow it | 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. +- `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; `Reduced`, the one place that turns a vector into one byte a number for a column declared `SQ8`. - `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. +- `grouping.rs`: the grouping of a partition (its slots sorted by the centre each is nearest, packed into full blocks, with the centres, the place where each group begins and a sketch of every slot) and the partition's mean; written once by the build that made it and never changed. +- `sketch.rs`: one vector as a sketch keeps it (a bit for each number and a factor, taken from a reference: the centre of its group, or the partition's mean for a tail block), and the range a search gets from a sketch; the counting loop. +- `train.rs`: finding centres, the arithmetic of a build over slots alone; nothing there reads a store. +- `build.rs`: when a grouping is due, the record of how far a build has come, and the steps that find centres and say which group each slot goes to. +- `pack.rs`: the last steps of a build: packing the slots of each group into the blocks of the new grouping with a sketch of each, measuring how far a search scans and how wide a sketch's estimate is taken to be, and putting the grouping in force. +- `search.rs`: the one search, and the plan it runs: scores the rows in no block exactly, screens the slots of the selected partitions (every slot of a partition searched whole; of one searched by its nearest groups, the slots its sketches keep), 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. +- `inspect.rs`: what inspection reads of a VECTOR column, its partitions as their records describe them; no block, no centre and no row. ## Where the parts meet @@ -56,11 +62,14 @@ A part reaches another only through what that part makes visible: the `pub(crate | `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` | +| `VectorColumn` (a VECTOR column's number, partition key, limit, mode, quantization, `AUTO_INDEX_AT` and consolidation, with the numbers in force), 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 | +| `search`, `plan`, `Nearest`, `Route`, `distance` | `vector` | `query` calls `search` with the partitions the filter selects, the mode in force, the column's declaration and whether the statement has a filter on rows, 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; `search` answers nothing at all only when INDEXED is refused (a selected partition is not usable, or a filter on rows leaves the index nothing to choose among), which `query` turns into the refusal. `plan` settles what a search does with each partition, from small records alone, and `search` runs exactly that plan, so `explain` cannot disagree with it | | `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` | +| `waiting`, `prepare`, `stage_step`, `stage_refusal`, `remaining` (and the values they return) | `vector` | `database` lists the partitions with work, has a step worked out with no lock held, and stages it through the one commit function; a step the memory limit refuses is recorded, not returned | +| `partitions`, `Inspected`, `State` | `vector` | `query` answers the two SHOW statements and `explain` with them | +| `Reduced` | `vector` | `write` gives a column declared `SQ8` the vector `Reduced` makes of the one written, and hands the store its scale and codes with the row (`store::Compact`) | | `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 | @@ -79,12 +88,26 @@ 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` | +| 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`; a grouping, centres, a build, assigned groups, a sketch record or a partition entry that disagrees with itself or with its partition, likewise: `a_grouping_that_disagrees_with_its_partition_is_damage`, `centres_that_do_not_come_to_their_groups_are_damage`, `a_build_whose_cursor_is_past_its_blocks_is_damage`, `assigned_groups_that_do_not_match_their_block_are_damage`, `a_sketch_record_that_does_not_match_its_block_is_damage`, `a_partition_entry_that_disagrees_with_itself_is_damage`; a reduced vector's record of another length: `contextdb/src/store/tests.rs`: `a_reduced_record_of_another_length_is_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` | +| Every row with a vector is found by an exact search, whether its vector waits on a list, is in a block written since the last grouping or is in the partition's grouping once and not as a stale slot, and after a kill inside a commit, a seal or a build | `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`, `a_grouping_holds_every_live_vector_once_and_no_stale_slot`, `vectors_sealed_while_a_first_build_starts_are_found_after_the_switch`; `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_kill_at_any_storage_call_of_a_build_leaves_every_vector_found`, `a_killed_build_goes_on_from_its_record` | | 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 search under INDEXED or AUTO returns only rows an exact search of the same statement could return, each with its true distance and none twice; it is never answered short: it returns `LIMIT` rows whenever the selected partitions hold that many rows the filter passes | `contextdb/tests/vectors_index.rs`: `indexed_is_served_at_every_size`, `a_search_that_widens_returns_no_row_twice`, `a_nearer_allowed_row_in_a_group_taken_later_is_in_the_answer`, `a_search_with_a_filter_on_rows_widens_until_it_is_whole` | +| A partition with no grouping, a grouping of one group, or sketches that cannot tell its vectors apart is searched as an exact search searches it; below its size a column pays for no index | `contextdb/tests/vectors_index.rs`: `a_column_below_its_size_pays_nothing_and_is_searched_whole`, `a_grouping_whose_sketches_cannot_tell_its_vectors_apart_answers_as_exact` | +| On the fixed vectors of the tests an indexed search counts fewer distance evaluations than an exact one (on other data it can count more, each an eighth of the bytes), reads sketches before blocks, and never costs more row visits than an exact one | `contextdb/tests/vectors_index.rs`: `an_indexed_search_does_fewer_distance_evaluations_than_exact`, `an_indexed_search_reads_sketches_before_blocks`, `an_indexed_search_that_takes_every_group_does_the_exact_evaluations` | +| INDEXED is refused by name, before the filter would be asked of the whole table, when a partition is far behind its writes or a filter leaves the index nothing to choose among; AUTO is never refused for a grouping and answers exactly where one is behind or damaged | `contextdb/tests/vectors_index.rs`: `indexed_is_refused_whole_when_one_partition_is_behind`, `an_indexed_search_short_of_its_limit_is_refused_before_the_whole_table_is_read`, `an_indexed_search_whose_first_pass_takes_every_slot_is_refused_before_a_row_is_asked`, `auto_answers_exactly_where_a_grouping_is_behind_or_damaged`; `contextdb/src/vector/tests.rs`: `auto_answers_exactly_where_a_grouping_is_damaged` | +| AUTO is exact until the selected partitions hold `AUTO_INDEX_AT` vectors, then indexed in each partition that has a grouping and is not behind | `contextdb/tests/vectors_index.rs`: `auto_uses_groupings_from_auto_index_at_on`, `auto_searches_each_partition_by_what_its_grouping_saves` | +| A write after a grouping is found at once, before any cycle and after sealing | `contextdb/tests/vectors_index.rs`: `writes_after_a_grouping_are_found_at_once` | +| A grouping is built in bounded cycles and switched in by one commit; the old grouping serves until then, and a write or search runs between any two cycles | `contextdb/tests/vectors_upkeep.rs`: `a_grouping_is_built_in_bounded_cycles_and_switched_in_one`; `contextdb/tests/transactions.rs`: `a_statement_commits_while_a_cycle_works_out_a_step` | +| Every block of a grouping, and every block sealed after the partition's mean exists, has a sketch for each slot; the index takes the bytes the layout says | `contextdb/tests/vectors_upkeep.rs`: `a_grouping_keeps_a_sketch_of_every_slot_and_measures_its_width`, `sealing_after_a_grouping_keeps_a_sketch_of_every_slot` | +| A partition is grouped again only when its changes reach the declared share | `contextdb/tests/vectors_upkeep.rs`: `a_partition_is_grouped_again_only_when_its_changes_reach_the_declared_share` | +| A partition that loses its last vector leaves no record of its grouping or its build | `contextdb/src/vector/tests.rs`: `a_partition_that_empties_leaves_no_record_of_its_grouping_or_its_build` | +| A statement reads one grouping whole, the one in force at its commit, in another process too; a new process searches a built grouping at once, and a grouping is there after reopen | `contextdb/tests/readers.rs`: `a_reader_sees_one_grouping_whole_while_the_writer_builds_and_switches`, `a_new_process_searches_a_built_grouping_at_once`; `contextdb/tests/store.rs`: `a_grouping_is_there_after_close_and_reopen` | +| A build or an indexed search under the memory limit stays under it; a step the limit refuses is recorded as a stall and the old grouping serves; a commit with no room for its vector records changes nothing | `contextdb/tests/memory_ceiling.rs`: `an_indexed_search_and_a_build_stay_under_the_limit`; `contextdb/tests/vectors_upkeep.rs`: `a_cycle_under_a_tight_limit_records_the_stall_and_the_old_grouping_serves`; `contextdb/tests/vectors_memory.rs`: `a_commit_with_no_room_for_its_vector_records_changes_nothing`; `contextdb/tests/vectors_index.rs`: `a_search_with_no_room_for_its_sketches_is_refused_by_the_memory_limit` | +| Inspection and `explain` read no block and change nothing; each partition is described as its records say | `contextdb/tests/vectors_inspection.rs`: `inspection_and_explain_read_no_block_and_change_nothing`, `show_vector_indexes_describes_every_vector_column`, `show_vector_partitions_lists_partitions_by_key_with_limit_and_offset`, `explain_names_how_each_partition_is_searched` | +| A column declared `SQ8` keeps the vector rounded to one byte a number, reads it back number for number, ranks by it, and is grouped and searched as any other | `contextdb/tests/vectors_reduced.rs`: `a_reduced_column_keeps_the_nearest_one_byte_vector`, `exact_on_a_reduced_column_ranks_by_the_kept_vectors`, `a_reduced_column_is_grouped_and_searched_as_any_other`; `contextdb/src/store/tests.rs`: `a_reduced_vector_reads_back_number_for_number_in_a_quarter_of_the_room` | | 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` | @@ -96,6 +119,7 @@ Every path in this guide is written from `v2/`. |---|---|---| | 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 | +| When or how a partition is grouped (the size at which a column is grouped, the share of change that groups it again, the steps or the numbers a build measures) | `contextdb/src/vector/build.rs` (`due`, `usable`, the steps), `contextdb/src/vector/pack.rs` (packing and measuring), `contextdb/src/vector/train.rs` (centres); the numbers in force from `VectorColumn` in `contextdb/src/schema/`. A new record for a grouping is a `VectorRecord` kind in `contextdb/src/store/`, which keeps it and never reads it. A change that moves what a build measures changes how a search uses it in `contextdb/src/vector/search.rs`, in the same change | `contextdb/tests/vectors_upkeep.rs` for what a build does and when, `contextdb/tests/vectors_index.rs` for what a search returns and does, `contextdb/tests/crash/main.rs` for a kill inside a build; the cycle counts and sizes the tests record change only with the design they state | | 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 | diff --git a/v2/README.md b/v2/README.md index d4022b40..d97ad9dc 100644 --- a/v2/README.md +++ b/v2/README.md @@ -64,3 +64,10 @@ UPDATE_EXPECTED=1 cargo test --manifest-path v2/Cargo.toml -p contextdb --test s `v2/numbers.sh` prints the numbers in the library's README for the commit it is run at. It only measures: it never passes or fails, and the check does not run it. + +The lines about vector search are taken by the programs in `contextdb/examples/`, the generated +vectors always and the groups of lines about real search passages only on data given in the +environment variables named at the head of the script (without them each of those lines prints +`unavailable`). Every time of a search is taken by one thread, in a read session under the memory +limit that a line of the output states; `NUMBERS_LARGE=1` adds five million generated vectors and +takes the better part of an hour. Run it on a machine that is doing nothing else. diff --git a/v2/contextdb-cli/src/main.rs b/v2/contextdb-cli/src/main.rs index 40a0c493..cb84062c 100644 --- a/v2/contextdb-cli/src/main.rs +++ b/v2/contextdb-cli/src/main.rs @@ -67,13 +67,14 @@ MEMORY: `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 SELECT or any SHOW statement {"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":...}} + .explain {"explain":{"lines":[...]}} .help {"help":[...]} a refusal {"error":{"class":...,"message":...,"detail":{"kind":...}}} on the error stream; "detail" is left out when the diff --git a/v2/contextdb-cli/src/output.rs b/v2/contextdb-cli/src/output.rs index 35870b3c..31b202f3 100644 --- a/v2/contextdb-cli/src/output.rs +++ b/v2/contextdb-cli/src/output.rs @@ -40,6 +40,32 @@ impl Output { } } + /// The rows of an inspection statement: one `name: value` to a line, with an empty line + /// between rows. + pub fn fields(self, result: &QueryResult) { + if self.json { + return self.result(result); + } + for (at, row) in result.rows().iter().enumerate() { + if at > 0 { + println!(); + } + for (name, value) in result.columns().iter().zip(row) { + println!("{name}: {}", text(value)); + } + } + } + + /// What `.explain` answers: its lines. + pub fn explain(self, lines: &str) { + if self.json { + let lines: Vec<&str> = lines.lines().collect(); + println!("{}", json!({ "explain": { "lines": lines } })); + } else { + println!("{lines}"); + } + } + /// A statement that returns no rows succeeded and changed this many. pub fn rows_affected(self, rows: u64) { if self.json { diff --git a/v2/contextdb-cli/src/shell.rs b/v2/contextdb-cli/src/shell.rs index cdee5824..ca7dca57 100644 --- a/v2/contextdb-cli/src/shell.rs +++ b/v2/contextdb-cli/src/shell.rs @@ -18,11 +18,13 @@ pub enum Store { Write(Database), } -const HELP: [&str; 6] = [ +const HELP: [&str; 8] = [ "; Run a SQL statement. It may span lines and ends at its `;`.", "BEGIN; ... COMMIT; Run statements as one transaction; ROLLBACK; undoes them, except SET MEMORY_LIMIT. Needs --write.", ".tables \\dt List the tables.", ".schema
\\d
Show a table's declaration as SQL.", + ".explain Show how a similarity search would be answered, without running it.", + "SHOW VECTOR_INDEXES; SHOW VECTOR_PARTITIONS; Describe the VECTOR columns and their partitions.", ".help \\? Show this list.", ".quit \\q End the session.", ]; @@ -151,6 +153,15 @@ impl<'a> Shell<'a> { Store::Write(db) => db.declaration(table), } .map(|declaration| self.output.schema(table, &declaration)), + (Some(".explain"), Some(_), _) => { + let sql = line[".explain".len()..].trim(); + let sql = sql.strip_suffix(';').unwrap_or(sql); + match self.store { + Store::Read(reader) => reader.explain(sql), + Store::Write(db) => db.explain(sql), + } + .map(|lines| self.output.explain(&lines)) + } (Some(".help" | "\\?"), None, _) => { self.output.help(&HELP); Ok(()) @@ -186,7 +197,17 @@ impl<'a> Shell<'a> { Ok(()) }) } else { - self.execute(sql).map(|result| self.output.result(&result)) + // The rows of the two inspection statements are too wide for a table. + let mut words = sql.split_whitespace().map(str::to_ascii_uppercase); + let inspection = words.next().as_deref() == Some("SHOW") + && words.next().is_some_and(|word| word.starts_with("VECTOR_")); + self.execute(sql).map(|result| { + if inspection { + self.output.fields(&result); + } else { + self.output.result(&result); + } + }) }; if let Err(error) = ran { self.refused = true; diff --git a/v2/contextdb-cli/tests/shell.rs b/v2/contextdb-cli/tests/shell.rs index 69001312..4829643b 100644 --- a/v2/contextdb-cli/tests/shell.rs +++ b/v2/contextdb-cli/tests/shell.rs @@ -894,8 +894,8 @@ VALUES (1, 'a', [1.0, 0.0, 0.0]), (2, 'a', [0.0, 1.0, 0.0]), (3, 'b', [-1.0, 0.0 } #[test] -fn indexed_search_is_refused_by_name_through_the_program() { - let dir = folder("indexed-search-is-refused"); +fn an_indexed_search_runs_through_the_program() { + let dir = folder("an-indexed-search-runs"); let store = dir.join("store"); let made = shell( &dir, @@ -905,7 +905,8 @@ fn indexed_search_is_refused_by_name_through_the_program() { ); assert_eq!(made.code, Some(0), "{}", made.stderr); - // The same search asked for INDEXED and then for EXACT. + // The same search asked for INDEXED and then for EXACT: no maintenance cycle has grouped + // these rows, so INDEXED searches every one of them and both give the same table. let search = |mode: &str| { format!( "SELECT id FROM passages ORDER BY embedding <=> [1.0, 0.0, 0.0] USE VECTOR {mode} \ @@ -913,55 +914,259 @@ fn indexed_search_is_refused_by_name_through_the_program() { ) }; let input = format!("{}{}", search("INDEXED"), search("EXACT")); - let exact = "+----+\n| id |\n+----+\n| 2 |\n| 1 |\n+----+\n(2 rows)\n"; + let table = "+----+\n| id |\n+----+\n| 2 |\n| 1 |\n+----+\n(2 rows)\n"; + let document = json!({"result": {"columns": ["id"], "rows": [{"id": 2}, {"id": 1}]}}); - // Reading or writing, this version keeps no index and says so. + // Reading or writing, for a person and for a program. 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:?}"); + assert_eq!(run.code, Some(0), "{mode:?}: {}", run.stderr); + assert_eq!(run.errors(), Vec::<&str>::new(), "{mode:?}"); + assert_eq!(run.stdout, format!("{table}{table}"), "{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!(run.code, Some(0), "{mode:?}: {}", run.stderr); assert_eq!( documents(&run.stdout), - [json!({"result": {"columns": ["id"], "rows": [{"id": 2}, {"id": 1}]}})], + [document.clone(), document.clone()], "{mode:?}" ); } } +/// What the two tests below search and inspect: a column with a partition for each `kind`, two +/// vectors in each, and no maintenance cycle run. +const KINDS: &str = "\ +CREATE TABLE notes (id INTEGER PRIMARY KEY, kind TEXT NOT NULL, v VECTOR(3) PARTITION_KEY (kind)); +INSERT INTO notes (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]); +"; + +/// The fields of `SHOW VECTOR_INDEXES`, in the order they are printed. +const INDEX_FIELDS: [&str; 31] = [ + "table", + "column", + "dimension", + "quantization", + "vector_count", + "bytes", + "partition_key_columns", + "max_partitions", + "live_partitions", + "search_mode", + "declared_auto_index_at", + "effective_auto_index_at", + "declared_consolidation_mode", + "declared_consolidation_change_percent", + "declared_consolidation_tombstone_percent", + "effective_consolidation_mode", + "effective_consolidation_change_percent", + "effective_consolidation_tombstone_percent", + "pending_inserts", + "tombstones", + "durable_vector_bytes", + "durable_index_bytes", + "indexed_available", + "maintenance_state", + "maintenance_vectors_total", + "maintenance_vectors_done", + "maintenance_vectors_remaining", + "unavailable_partitions", + "stalled_partitions", + "broad_route", + "recovery_action", +]; + +/// The fields of `SHOW VECTOR_PARTITIONS`, in the order they are printed. +const PARTITION_FIELDS: [&str; 18] = [ + "table", + "column", + "partition_key", + "live_rows", + "pending_inserts", + "tombstones", + "durable_vector_bytes", + "durable_index_bytes", + "indexed_available", + "availability_reason", + "maintenance_state", + "maintenance_reason", + "maintenance_vectors_total", + "maintenance_vectors_done", + "maintenance_vectors_remaining", + "lists", + "probes", + "recovery_action", +]; + +#[test] +fn vector_inspection_prints_one_field_to_a_line() { + let dir = folder("vector-inspection-prints"); + let store = dir.join("store"); + let more = "CREATE TABLE other (id INTEGER PRIMARY KEY, w VECTOR(2));\n"; + let made = shell(&dir, &[arg(&store), "--write"], &format!("{KINDS}{more}")); + assert_eq!(made.code, Some(0), "{}", made.stderr); + + // Each statement with the fields it prints and, for each row, the values of three of them. + let indexes = [["notes", "v", "3"], ["other", "w", "2"]]; + let partitions = [ + ["notes", "v", r#"{"kind":"a"}"#], + ["notes", "v", r#"{"kind":"b"}"#], + ["other", "w", "{}"], + ]; + let cases = [ + ("SHOW VECTOR_INDEXES;\n", &INDEX_FIELDS[..], &indexes[..]), + ( + "SHOW VECTOR_PARTITIONS;\n", + &PARTITION_FIELDS[..], + &partitions[..], + ), + ]; + for (statement, fields, rows) in cases { + // For a person: rows with an empty line between them, each `name: value` lines in the + // statement's order, and no line too long for a terminal. + let run = shell(&dir, &[arg(&store)], statement); + assert_eq!(run.code, Some(0), "{}", run.stderr); + assert!( + run.stdout.lines().all(|line| line.chars().count() <= 120), + "{}", + run.stdout + ); + let printed: Vec> = run + .stdout + .split("\n\n") + .map(|row| { + row.lines() + .filter(|line| !line.starts_with('(')) + .map(|line| { + line.split_once(": ") + .unwrap_or_else(|| panic!("not `name: value`: {line:?}")) + }) + .collect::>() + }) + .filter(|row| !row.is_empty()) + .collect(); + assert_eq!(printed.len(), rows.len(), "{}", run.stdout); + for (row, first) in printed.iter().zip(rows) { + let names: Vec<&str> = row.iter().map(|(name, _)| *name).collect(); + assert_eq!(names, fields, "{}", run.stdout); + let values: Vec<&str> = row[..3].iter().map(|(_, value)| *value).collect(); + assert_eq!(values, first, "{}", run.stdout); + } + + // For a program: the usual result document, its columns the same names. + let run = shell(&dir, &[arg(&store), "--json"], statement); + assert_eq!(run.code, Some(0), "{}", run.stderr); + let results = documents(&run.stdout); + assert_eq!(results.len(), 1, "{}", run.stdout); + assert_eq!(results[0]["result"]["columns"], json!(fields)); + let found = results[0]["result"]["rows"].as_array().unwrap(); + assert_eq!(found.len(), rows.len(), "{}", run.stdout); + assert_eq!(found[0]["table"], "notes", "{}", run.stdout); + } +} + +#[test] +fn explain_prints_how_each_partition_is_searched() { + let dir = folder("explain-prints"); + let store = dir.join("store"); + let made = shell(&dir, &[arg(&store), "--write"], KINDS); + assert_eq!(made.code, Some(0), "{}", made.stderr); + + // Four vectors in two partitions, far below the 100,000 from which AUTO uses a grouping, + // so the search asked for under AUTO is the exact one, and none grouped. A line ending in + // `*` is checked by its name alone: with no grouping the partition has no count of groups + // to print. + let lines = [ + "requested_mode: AUTO", + "resolved_mode: EXACT", + "effective_auto_index_at: 100000", + "aggregate_allowed_vectors: 4", + r#"partition_key_columns: ["kind"]"#, + "partitions: 2", + r#"partition: {"kind":"a"}"#, + " route: exact", + " lists: *", + " probes: *", + " live_rows: 2", + " outside_grouping: 2", + r#"partition: {"kind":"b"}"#, + " route: exact", + " lists: *", + " probes: *", + " live_rows: 2", + " outside_grouping: 2", + "refusal: none", + ]; + let same = |printed: &[&str]| { + printed.len() == lines.len() + && lines.iter().zip(printed).all(|(expected, line)| { + expected + .strip_suffix('*') + .map_or(expected == line, |name| line.starts_with(name)) + }) + }; + let search = ".explain SELECT id FROM notes ORDER BY v <=> '[1,0,0]' LIMIT 3\n"; + + // The lines and nothing else: the search is not run, so no table of rows is printed. + let run = shell(&dir, &[arg(&store)], search); + assert_eq!(run.code, Some(0), "{}", run.stderr); + assert!( + same(&run.stdout.lines().collect::>()), + "{}", + run.stdout + ); + + // Any other statement is refused in one line, and the session goes on to the next. + let other = ".explain SELECT id FROM notes ORDER BY id\n"; + let next = "SELECT id FROM notes WHERE id = 5;\n"; + let run = shell(&dir, &[arg(&store)], &format!("{other}{next}")); + assert_eq!(run.code, Some(1), "{}", run.stderr); + let refusals: Vec<&str> = run.stderr.lines().collect(); + assert_eq!(refusals.len(), 1, "{}", run.stderr); + assert!( + refusals[0].starts_with("Error: ") + && refusals[0].contains( + "explaining this statement is not supported in this version of the library" + ), + "{}", + run.stderr + ); + assert_eq!(tables(&run.stdout), [[["id"], ["5"]]], "{}", run.stdout); + + // For a program: one document that holds the same lines. + let run = shell(&dir, &[arg(&store), "--json"], search); + assert_eq!(run.code, Some(0), "{}", run.stderr); + let results = documents(&run.stdout); + assert_eq!(results.len(), 1, "{}", run.stdout); + let printed: Vec<&str> = results[0]["explain"]["lines"] + .as_array() + .unwrap_or_else(|| panic!("no explain.lines: {}", results[0])) + .iter() + .map(|line| line.as_str().unwrap()) + .collect(); + assert!(same(&printed), "{}", run.stdout); +} + +#[test] +fn help_lists_explain() { + let dir = folder("help-lists-explain"); + let store = dir.join("store"); + let made = shell(&dir, &[arg(&store), "--write"], CREATE_NOTES); + assert_eq!(made.code, Some(0), "{}", made.stderr); + let run = shell(&dir, &[arg(&store)], ".help\n"); + assert_eq!(run.code, Some(0), "{}", run.stderr); + assert!( + run.stdout.lines().any(|line| line.starts_with(".explain")), + "{}", + run.stdout + ); +} + #[test] fn a_read_shell_searches_while_a_writer_holds_the_store() { let dir = folder("a-read-shell-searches-while-a-writer-holds"); diff --git a/v2/contextdb/README.md b/v2/contextdb/README.md index 42032c29..b209699d 100644 --- a/v2/contextdb/README.md +++ b/v2/contextdb/README.md @@ -10,14 +10,15 @@ 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; +- vector search: the rows nearest a vector by cosine distance, each with its distance, over + vectors that a column may keep in partitions; exact on request, and through an index that + `run_maintenance_cycle` builds and brings up to date once a column holds many vectors; - 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: 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 +What it does not have yet: 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: @@ -65,9 +66,9 @@ fn main() -> contextdb::Result<()> { ``` 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 +calls `db.run_maintenance_cycle()?` again and again, as the paragraph on upkeep under Vector search says, because +this version does not do it by itself; a search before that still finds every vector, and reads +more to do it. 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. @@ -76,8 +77,11 @@ 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; +SHOW VECTOR_INDEXES; +SHOW VECTOR_PARTITIONS [FOR table.column] [LIMIT count [OFFSET start]]; ``` - Every table has exactly one `PRIMARY KEY` column. A default is a literal or `NOW()`. @@ -102,12 +108,13 @@ SET MEMORY_LIMIT 512M; SET MEMORY_LIMIT NONE; SHOW MEMORY_LIMIT; 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. + returns. The clauses of a `VECTOR` column are described under Vector search (quantization, modes, upkeep), + and the two `SHOW` statements there too, under inspection. - `--` and `/* */` comments are read past. 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: +columns; and the vector clauses `IVF`, `SEARCH (...)` and `WITH (quantization = 'SQ4')`. 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 @@ -124,7 +131,7 @@ 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` finite numbers, each kept as an `f32`, written `[0.1, 0.2, 0.3]`; searched with `<=>` | `Value::Vector` | +| `VECTOR(n)` | a vector of exactly `n` finite numbers, each kept as an `f32` (a column declared `SQ8` keeps a rounded vector, see the paragraph on quantization under Vector search), 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 @@ -151,8 +158,10 @@ SELECT id, embedding <=> $q AS distance FROM notes ORDER BY distance LIMIT 10; 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. +- How much of the column a search compares with `$q` is the search's mode: `EXACT` compares + every stored vector and returns the true nearest rows; `INDEXED` chooses among them through + the column's index; `AUTO`, the default, is `EXACT` until a column is large and then uses the + index (the paragraph on modes below). 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`, @@ -160,13 +169,35 @@ in the select list of a statement that does not order by it, in an `INSERT`, `UP `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 +**Upkeep.** 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, +vectors not yet in a block, which a search compares one by one by reading their rows. Once a +column is large enough (the paragraph on modes below) the engine also groups a partition's vectors: +it finds the centres of groups of near vectors, files each vector under its nearest centre, and +keeps a one-bit sketch of every vector, an eighth of the size of its one-byte copy. A grouping +is made once and replaced whole; a vector written after it waits in blocks of its own, each +with its sketches, until the partition is grouped again. + +`Database::run_maintenance_cycle` does this work in bounded steps, and nothing in this version +calls it on its own. One call first seals: it moves at most 4,096 waiting vectors into blocks +in one commit. Then it takes one step for each partition that has a grouping due or being +built, each step in a commit of its own: begin, ten rounds that find the centres, file the +vectors under them, pack the groups into blocks, measure (twice), and put the new grouping in +force. A partition is grouped again when the vectors written, removed or changed since reach +`CHANGE_PERCENT` of its grouping (20 when left out) or the removed or changed ones reach +`TOMBSTONE_PERCENT` (10); `CONSOLIDATION NONE` never groups it again. The call returns a +`MaintenanceReport`, and its `vector` field says what is left: +`maintenance_vectors_done` and `maintenance_vectors_remaining` (vectors), `remaining_partitions` +(partitions with a grouping due or being built) and `stalled_partitions` (those whose last step +the memory limit refused; a further cycle does nothing for them until the limit is raised, and +with no limit declared none stall). A host runs cycles while `remaining_partitions` is above +`stalled_partitions` or `maintenance_vectors_remaining` is not zero, as the program in the paragraph on the index below does. + +A step is worked out with no lock held, so a write waits only for a step's commit and a reader +not at all, and a search made between two cycles finds every vector: it sees the grouping in +force when the statement began, whole, or the one before it, whole, never a mix. An exact +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 @@ -186,15 +217,258 @@ column and its limit; the statement changes nothing, and inside a transaction th 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`. +`USE VECTOR`, written after the `ORDER BY` item, says it for one statement. The mode in force is +the statement's, else the column's, else `AUTO`. + +- `EXACT` compares the search vector with every stored vector of every partition the statement + selects: the grouping's blocks, the blocks written since and the waiting list. It returns the + true nearest rows, and reads no centre and no sketch. +- `INDEXED` searches the waiting vectors of each selected partition and, of the rest, only + the vectors the index chooses (the paragraph on the index below). A partition with no grouping, + or with a grouping of one group, has nothing to choose among, and is searched as `EXACT` + searches it. `INDEXED` is refused, before anything is read, when a selected partition's + writes are far ahead of upkeep (for a partition with no grouping or a grouping of one group, + the vectors outside its grouping number at least twice `AUTO_INDEX_AT`; for one with more + groups, at least twice what makes it due to be grouped again, which under + `CONSOLIDATION NONE` is never), and when a filter on rows leaves it nothing to choose among, as the + paragraph on filters below says. +- `AUTO`, which is the default, adds up the vectors of the partitions the statement selects. + Below `AUTO_INDEX_AT` it is `EXACT`. From there on each selected partition that has a + grouping and is not far behind is searched as `INDEXED` searches it, and every other + partition as `EXACT` does. `AUTO` is never refused for a grouping. + +`AUTO_INDEX_AT` is 100,000 when not written. Until a column holds that many vectors in all its +partitions together (or declares `SEARCH_MODE INDEXED`), no centre, sketch or mean is +written, and `durable_index_bytes` is 0 unless removed or changed rows (10% by default) have +given a partition a grouping of one group, which takes 28 bytes. In a column that has reached it, a partition is grouped once it +holds a hundredth of it, 1,000 vectors when nothing is written, so a store of many middling +partitions is indexed too; a partition of a few hundred vectors is not cut into groups. A +declared `AUTO_INDEX_AT` of any size is served. A column that should never be grouped declares +one it will not reach. + +**The index.** An indexed search of a partition that has groups compares the search vector with +the centres of its groups and takes the groups whose centre is nearly as near as the nearest +centre; then it estimates from the sketches of those groups' vectors, and of the vectors +written since the grouping, which of them can be among the nearest; only those go on to the +one-byte copy and then to the row. So every row in an indexed answer has its true distance, +the answer is in the order of those distances, no row comes twice, and a search is never +answered short: it returns `LIMIT` rows whenever the selected partitions hold that many rows +the statement's filter passes. What a sketch can cost is a row that an exact search would have +put among the nearest. + +An indexed search returns rows that an exact search would rank among its nearest, each with +its true distance. How far from the nearest group it scans and how far it trusts a sketch +are measured when the index is built, on questions made from the stored vectors: half are +stored vectors themselves, half lie between two stored vectors, and a sketch counts as kept +only if it would be kept both as a grouped row and as a row written after the build. That is +a measure over those questions, not a promise for any one question: a question unlike both +kinds, rows written after the build that are unlike the rows the build saw, and a `LIMIT` +other than ten are not what the measure was taken on. `USE VECTOR EXACT` or `SEARCH_MODE EXACT` +is exact. + +A search with a filter on rows uses the index too. The filter is asked of each row the search +visits, and the search widens until it holds as many allowed rows as were asked for: a filter +that passes most rows costs what the unfiltered search costs, and one that passes few costs +more, up to the work of `EXACT`. `INDEXED` is refused by name only when a filter passes too few +rows for the index to choose among, because being sure of the answer would mean asking the filter +of the whole table under the index's name; `AUTO` does that work and returns every row there +is. ```text -no vector index ships in this version, so notes.embedding cannot be searched INDEXED; use EXACT or AUTO +notes.embedding cannot be searched INDEXED for this statement now: its index is behind its writes, or the statement's filter leaves the index nothing to choose among; run a maintenance cycle, or use EXACT or AUTO ``` +`QueryResult::distance_evaluations` counts the work of a similarity search: the centres it +compared with the search vector, the sketches it estimated, the one-byte copies it screened and +the rows whose distance it worked out. It is an exact count, the same on every run over the +same store, and 0 for any other statement. An indexed search usually counts far fewer than an +exact one, and can count more when it leaves few groups out, because it estimates a sketch of +every vector it keeps; each of those counts an eighth of the bytes. `Database::explain`, and +`ReadSession::explain`, return what a similarity search would do, without running it. This +program builds a store of 2,000 vectors, runs upkeep as above, and searches it both ways: + +```rust,no_run +use contextdb::{params, Database}; + +// A number from -1 to 1 for each seed, the same on every machine. +fn spread(seed: u64) -> f32 { + let mut x = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15); + x ^= x >> 29; + x = x.wrapping_mul(0xBF58_476D_1CE4_E5B9); + x ^= x >> 32; + (x % 2001) as f32 / 1000.0 - 1.0 +} + +// Eight numbers for each id: the topic's (fifty topics) and a little of the id's own. +fn vector(id: u64) -> Vec { + (0..8u64) + .map(|j| spread((id % 50) * 8 + j + 1) + 0.15 * spread(id * 8 + j + 100_000)) + .collect() +} + +fn main() -> contextdb::Result<()> { + let db = Database::open("./notes-index")?; + db.execute( + "CREATE TABLE notes (id INTEGER PRIMARY KEY, embedding VECTOR(8) AUTO_INDEX_AT 1000)", + ¶ms! {}, + )?; + let tx = db.transaction()?; + for id in 0..2000u64 { + tx.execute( + "INSERT INTO notes (id, embedding) VALUES ($id, $embedding)", + ¶ms! { "id" => id as i64, "embedding" => vector(id) }, + )?; + } + tx.commit()?; + + // Upkeep: seal the new vectors and group them, a bounded step at a time. + let mut cycles = 0; + loop { + let report = db.run_maintenance_cycle()?.vector; + cycles += 1; + if report.remaining_partitions <= report.stalled_partitions + && report.maintenance_vectors_remaining == 0 + { + break; + } + } + println!("{cycles} cycles"); + + let mut ids = Vec::new(); + for mode in ["EXACT", "INDEXED"] { + let sql = format!( + "SELECT id FROM notes ORDER BY embedding <=> $q USE VECTOR {mode} LIMIT 10" + ); + let found = db.execute(&sql, ¶ms! { "q" => vector(5000) })?; + println!("{mode}: {} distance evaluations", found.distance_evaluations()); + ids.push(found.rows().to_vec()); + } + println!("the same ten rows: {}", ids[0] == ids[1]); + println!("{}", db.explain("SELECT id FROM notes ORDER BY embedding <=> $q LIMIT 10")?); + db.close() +} +``` + +It prints, for the first run (`./notes-index` must not exist): + +```text +16 cycles +EXACT: 2019 distance evaluations +INDEXED: 225 distance evaluations +the same ten rows: true +requested_mode: AUTO +resolved_mode: AUTO +effective_auto_index_at: 1000 +aggregate_allowed_vectors: 2000 +partition_key_columns: [] +partitions: 1 +partition: {} + route: indexed + lists: 45 + probes: 8 + live_rows: 2000 + outside_grouping: 0 +refusal: none +``` + +`lists` is the groups of the grouping in force and `probes` the most of them a search scanned +for the questions its build made, `outside_grouping` the slots in blocks written since (stale ones +included) plus the vectors still waiting, and `route` is `exact` or `indexed`; `refusal` is `none` or, for INDEXED over a partition that is +behind its writes, the sentence the search would be refused with (the refusal for a filter is +known only when the search runs). A statement that names a partition with a `$name` has nothing to select +with, since `explain` binds nothing: it lists every partition of the column and says so, +`partitions: every partition ($t is not bound)`. `explain` of a statement that is not a +similarity search is refused as a later version's. + +**Quantization.** `WITH (quantization = 'SQ8')` on a column keeps each vector as one scale and +one whole number from -127 to 127 for each of its numbers, so a vector of `n` numbers takes +`n + 13` bytes in place of `4n + 9`: in the two stores of real passages under Numbers below a passage +takes about 2.0 KB in place of 4.4 and 4.9, and the store 0.45 times the room. The cost is in the answers. The vector the row keeps is the rounded one, and it is the +one a `SELECT` returns and every search, `EXACT` included, ranks by (`[0.1, 0.2, 0.3]` is +read back as `[0.0992126, 0.2007874, 0.3]`); and on those stores, searched exactly, the ten nearest rows of an +`SQ8` column differ from the full-precision column's on 12 to 21 questions in 100, by whether the +search covers one, several or all of its partitions. It is an option a table chooses and never a default: a +column keeps `F32` unless it says otherwise. `SQ4` is refused by name. + +**Inspection.** `SHOW VECTOR_INDEXES` returns one row for each `VECTOR` column of the store, +in the order of table name and column, and `SHOW VECTOR_PARTITIONS` one row for each partition +of the column after `FOR`, or of every `VECTOR` column, in the order of table, column and key +tuple; `LIMIT` takes the first rows and `OFFSET`, only after `LIMIT`, skips some. Both read +the vector part's small records and no block, no centre and no row, change nothing, and work in +a writer, a transaction and a read session alike. The shell prints their rows one field to a +line, with an empty line between rows. For the store above, `contextdb ./notes-index` prints +these (the lines of `SHOW VECTOR_INDEXES` that tell the column's declaration, size and +state, with `…` for the lines left out): + +```text +table: notes +column: embedding +dimension: 8 +quantization: F32 +vector_count: 2000 +… +search_mode: AUTO +declared_auto_index_at: 1000 +effective_auto_index_at: 1000 +… +durable_vector_bytes: 32000 +durable_index_bytes: 7856 +indexed_available: true +maintenance_state: idle +… +recovery_action: none +``` + +and for `SHOW VECTOR_PARTITIONS`: + +```text +table: notes +column: embedding +partition_key: {} +live_rows: 2000 +pending_inserts: 0 +tombstones: 0 +durable_vector_bytes: 32000 +durable_index_bytes: 7856 +indexed_available: true +availability_reason: none +maintenance_state: idle +maintenance_reason: none +maintenance_vectors_total: 0 +maintenance_vectors_done: 0 +maintenance_vectors_remaining: 0 +lists: 45 +probes: 8 +recovery_action: none +``` + +`.explain ` in the shell prints what `Database::explain` returns; `.explain SELECT id FROM notes +ORDER BY embedding <=> $q LIMIT 10` prints: + +```text +requested_mode: AUTO +resolved_mode: AUTO +effective_auto_index_at: 1000 +aggregate_allowed_vectors: 2000 +partition_key_columns: [] +partitions: 1 +partition: {} + route: indexed + lists: 45 + probes: 8 + live_rows: 2000 + outside_grouping: 0 +refusal: none +``` + +Inspection tells the groups a partition has (`lists`) and the most groups a search scanned +while the index was measured (`probes`); it does not print the measured distance or width. A +column with no `PARTITION_KEY` has one row in `SHOW VECTOR_PARTITIONS`, with `partition_key` +`{}`, and a partitioned column with no vector has none. A partition's `maintenance_state` is `idle`, `pending` (a grouping is due and not begun), +`building` or `stalled`, and its `maintenance_reason` says why upkeep has work: `first_grouping`, +`changes`, `tombstones`, `damaged` or `memory_limit`. For `pending` and `building` the +`recovery_action` is `run a maintenance cycle`, and for `stalled` it is `raise the memory limit`. + ## Binding values `$name` in a statement is the entry `"name"` of the map passed with it. `params!` writes the @@ -350,7 +624,9 @@ contextdb ./notes --write read and write it, creating it if it is missing A statement ends at `;`. `BEGIN;`, `COMMIT;` and `ROLLBACK;` run a transaction, which needs `--write`; `.tables` lists -the tables, `.schema
` prints a table's declaration and `.help` lists the commands. With +the tables, `.schema
` prints a table's declaration, `.explain ` prints how a +similarity search would be answered without running it (anything else it refuses like a statement +of a later version), and `.help` lists the commands. With statements piped in there is no prompt, and a refusal does not stop what comes after it. A refused statement or command is one line on the error stream, `Error: ` and then the refusal's sentence, or the program's own sentence where one of its options is the way on; a line that @@ -368,22 +644,56 @@ options, the exit statuses and the JSON forms. Measured at one commit by `v2/numbers.sh`, which prints them again on any machine: a release build on a 16-core AMD Ryzen 9 9955HX with 64 GB of memory and an NVMe drive, running Linux and -nothing else. +nothing else. Every search time is the 95th percentile of 100 questions asked one after the +other by one thread, in a read session under a 1 GiB memory limit unless a row says 64 MiB, in +an ordinary release build with no processor flag (the counting loop uses the processor's +bit-count instruction, which this processor has). Recall is the share of an exact search's ten +nearest rows that the search returned, averaged over the 100 questions, and the next number is +how many of the 100 questions got all ten. | What | Measured | |---|---| -| 20,000 rows of 1.6 KB, one commit per row: time to insert | 11.0 s | +| 20,000 rows of 1.6 KB, one commit per row: time to insert | 11.1 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 | +| Open that store | 0.561 ms | +| `SELECT COUNT(*)` over it | 0.030 ms, 0 rows read | +| A million rows under a 64 MiB limit: time to load | 10.5 s | +| Scan the million rows | 236.582 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 | +| Most memory the process held, as the operating system reports it | 21.5 MiB | +| First build of a program that uses the library | 8 s, 47 packages, a 3.2 MiB program | +| From a fresh clone to a passing `cargo test` | 3.2 min | +| Ten nearest of 20,000 vectors, sealed, through SQL, exact: 95th percentile | 1.548 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 of Rust and grammar in the library and the program | 15,440 | +| Tests | 253 | | 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 | + +Searches for the ten nearest, indexed beside exact from the same run. Where a store is below +`AUTO_INDEX_AT` nothing is grouped and an indexed search is the exact one. + +| Search | Indexed | Exact | Recall, questions with all ten | +|---|---|---|---| +| 10,000 generated vectors, one partition, 100 groups (`AUTO_INDEX_AT 1000`) | 0.541 ms | 1.185 ms | 0.9960, 96 of 100 | +| 100,008 real passages, eight partitions, 884 groups: one partition | 0.682 ms | 1.411 ms | 1.0000, 100 of 100 | +| the same: several partitions | 1.537 ms | 3.841 ms | 0.9990, 99 of 100 | +| the same: all eight | 4.052 ms | 10.880 ms | 0.9990, 99 of 100 | +| the same: all eight, under 64 MiB | 11.383 ms | not taken | 0.9990, 99 of 100 | +| the same: all eight, a filter that passes 992 of the 100,008 rows | 32.411 ms | 38.453 ms | 1.0000, 100 of 100 | +| the same, 9,996 more passages written after the build and sealed | 3.749 ms | 10.781 ms | 1.0000, 100 of 100 | +| 1,000,000 generated vectors of 768 numbers around 300 hidden centres, one partition, 1,000 groups | 1.109 ms | 180.929 ms | 1.0000, 100 of 100 | +| 1,000,000 generated camera sightings, one partition, 1,000 groups | 30.348 ms | 175.409 ms | 1.0000, 100 of 100 | +| 11,588 real passages, eight partitions, below the size: one / several / all | 0.303 / 0.724 / 1.440 ms | 0.302 / 0.724 / 1.440 ms | 1.0000, 100 of 100 each | + +Getting a store to those groupings takes `run_maintenance_cycle` calls: 34 cycles and 1.8 s for +the 10,000, 56 cycles and 26.9 s for the 100,008 passages, 1,150 cycles and 362.8 s for the +million. The index takes 129, 125 and 101 bytes a vector there. Right after a build a store's +file is 1.20, 1.13 and 1.31 times the size it had sealed, which is more than the index holds, +because the pages the build freed stay in the file in this version; a store that is never +grouped stays at 1.00 times (the 11,588 passages). + +With `WITH (quantization = 'SQ8')` the same two stores of real passages take 0.45 times the room +(a passage 1,968 bytes in place of 4,411 in the 11,588, and 1,992 in place of 4,946 in the +100,008). Searched exactly, the same ten rows come back as from the full-precision store on 88, +84 and 79 questions of 100 in the 11,588 (one, several, all partitions) and on 86, 83 and 85 in +the 100,008. diff --git a/v2/contextdb/examples/measure.rs b/v2/contextdb/examples/measure.rs index f59e3875..e1fbded1 100644 --- a/v2/contextdb/examples/measure.rs +++ b/v2/contextdb/examples/measure.rs @@ -14,6 +14,23 @@ //! 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. +//! - `measure index ` measures similarity search under each mode on a million generated +//! vectors of 768 numbers in one partition, a row being an INTEGER key and the vector, and +//! then on a million generated sightings: lines `synthetic_1000000_...` and +//! `sightings_1000000_...`. `measure synthetic ` and `measure sightings +//! ` do one of the two at another size. The first are the index pilot's vectors, to +//! the number: a xorshift from `0x9E37_79B9_7F4A_7C15`, 300 hidden centres, each vector a +//! centre plus 0.06 of noise a number, and 100 questions with four thirds of that noise. The +//! second stand in for a camera store: three classes, a subject for every 2,500 rows with +//! eight directions of its own for its appearance, tracks of ten sightings whose appearance +//! drifts, a little noise on every number; 100 questions, eighty new sightings of stored +//! subjects and twenty of subjects never stored. Each case loads two stores, 500 rows to a +//! commit: one whose column declares `AUTO_INDEX_AT 1000000000`, so that nothing is grouped, +//! sealed and searched under EXACT; and one that declares nothing, with every maintenance +//! cycle run, opened and searched once in a new process (`measure first ...`, which the +//! program starts itself as) and then searched, in one read session under 1 GiB, under +//! EXACT, INDEXED and AUTO in turn: every question once after 5 searches of each mode that +//! are not timed. On sightings a hundred stored vectors are then each searched for by itself. //! //! 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. @@ -23,12 +40,16 @@ mod support; +use std::ffi::OsStr; use std::fs::File; use std::io::{BufRead, BufReader}; use std::path::Path; use contextdb::{Database, DatabaseOpenOptions, Result, Uuid, params}; -use support::{bytes_on_disk, line, mebibytes, milliseconds, percentile, ratio, seconds, timed}; +use support::{ + EXACT, Fallible, READER_LIMIT, Values, bytes_on_disk, line, mebibytes, milliseconds, + percentile, ratio, seconds, timed, +}; #[global_allocator] static ALLOCATOR: support::Counting = support::Counting; @@ -266,29 +287,347 @@ fn sessions(corpus: &Path, list: &Path, count: usize, folder: &Path) -> Result<( Ok(()) } +/// The index pilot's generator of numbers: a xorshift, and a number near a normal one made of +/// six of its draws. +struct Random(u64); + +impl Random { + fn next(&mut self) -> u64 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + self.0 + } + + fn below(&mut self, n: usize) -> usize { + (self.next() % n as u64) as usize + } + + fn g(&mut self) -> f32 { + let s: f32 = (0..6) + .map(|_| (self.next() >> 40) as f32 / (1u64 << 24) as f32) + .sum(); + (s - 3.0) * 1.414 + } +} + +const DIMENSIONS: usize = 768; + +/// What a case's rows are handed to, one vector after another. +type Each<'a> = &'a mut dyn FnMut(Vec) -> Fallible<()>; + +/// The index pilot's vectors, to the number: hands `each` every row in order and answers the +/// 100 questions, which the generator makes after the rows. +fn synthetic(rows: usize, each: Each<'_>) -> Fallible>> { + fn norm(v: &mut [f32]) { + let n = v.iter().map(|x| x * x).sum::().sqrt(); + v.iter_mut().for_each(|x| *x /= n); + } + let mut rng = Random(0x9E37_79B9_7F4A_7C15); + let (nc, noise) = (300_u64, 0.06_f32); + let mut centers = Vec::new(); + for _ in 0..nc { + let mut c: Vec = (0..DIMENSIONS).map(|_| rng.g()).collect(); + norm(&mut c); + centers.push(c); + } + for _ in 0..rows { + let c = ¢ers[(rng.next() % nc) as usize]; + let mut v: Vec = c.iter().map(|x| x + noise * rng.g()).collect(); + norm(&mut v); + each(v)?; + } + let mut questions = Vec::new(); + for _ in 0..100 { + let c = ¢ers[(rng.next() % nc) as usize]; + let mut q: Vec = c.iter().map(|x| x + noise * 4.0 / 3.0 * rng.g()).collect(); + norm(&mut q); + questions.push(q); + } + Ok(questions) +} + +/// Generated sightings, to the number as the second measurement of the index made them. A +/// subject is a direction of its own beside its class's; its appearances vary along eight +/// directions of its own; a track is ten sightings in a row whose appearance drifts; every +/// sighting has a little noise on every number. +struct Sightings { + subjects: usize, + classes: Vec>, + bases: Vec>, + axes: Vec>, + /// The generator's state after the directions, from which the questions are made. + after: u64, +} + +impl Sightings { + const AXES: usize = 8; + const TRACK: usize = 10; + + fn unit(v: &mut [f32]) { + let n = v.iter().map(|x| f64::from(*x) * f64::from(*x)); + let n = n.sum::().sqrt(); + if n > 0.0 { + for x in v { + *x = (f64::from(*x) / n) as f32; + } + } + } + + fn new(rows: usize) -> Sightings { + let subjects = (rows / 2500).max(8); + let mut random = Random(0x51A7_1265_0BAD_CAFE); + let mut directions = |count: usize| -> Vec> { + let direction = |_| { + let mut c: Vec = (0..DIMENSIONS).map(|_| random.g()).collect(); + Sightings::unit(&mut c); + c + }; + (0..count).map(direction).collect() + }; + Sightings { + subjects, + classes: directions(3), + bases: directions(subjects + 20), + axes: directions((subjects + 20) * Sightings::AXES), + after: random.0, + } + } + + /// One sighting of a subject with this appearance. + fn make(&self, subject: usize, appearance: &[f32; 8], random: &mut Random) -> Vec { + let class = &self.classes[subject % 3]; + let base = &self.bases[subject]; + let mut out: Vec = (0..DIMENSIONS) + .map(|d| { + 0.5 * class[d] + 0.6 * base[d] + 0.25 * random.g() / (DIMENSIONS as f32).sqrt() + }) + .collect(); + for (j, a) in appearance.iter().enumerate() { + let axis = &self.axes[subject * Sightings::AXES + j]; + let weight = 0.5 * a / (Sightings::AXES as f32).sqrt(); + for (x, along) in out.iter_mut().zip(axis) { + *x += weight * along; + } + } + Sightings::unit(&mut out); + out + } + + /// The first `count` sightings of one track. + fn track(&self, track: usize, count: usize) -> Vec> { + let seed = (track as u64 + 1).wrapping_mul(0x2545_F491_4F6C_DD1D); + let mut random = Random(0x51A7_1265_0BAD_CAFE ^ seed); + random.next(); + let subject = random.below(self.subjects); + let mut appearance = [0f32; 8]; + for a in &mut appearance { + *a = random.g(); + } + let mut sightings = Vec::with_capacity(count); + for _ in 0..count { + sightings.push(self.make(subject, &appearance, &mut random)); + for a in &mut appearance { + *a = 0.9 * *a + 0.436 * random.g(); + } + } + sightings + } + + /// Hands `each` the first `rows` sightings, track after track. + fn rows(&self, rows: usize, each: Each<'_>) -> Fallible<()> { + for track in 0..rows.div_ceil(Sightings::TRACK) { + let count = Sightings::TRACK.min(rows - track * Sightings::TRACK); + for sighting in self.track(track, count) { + each(sighting)?; + } + } + Ok(()) + } + + /// The 100 questions: eighty new sightings of stored subjects, twenty of subjects never + /// stored. + fn questions(&self) -> Vec> { + let mut random = Random(self.after); + let mut questions = Vec::new(); + for j in 0..100 { + let subject = if j % 5 == 4 { + self.subjects + j / 5 + } else { + random.below(self.subjects) + }; + let mut appearance = [0f32; 8]; + for a in &mut appearance { + *a = random.g(); + } + questions.push(self.make(subject, &appearance, &mut random)); + } + questions + } +} + +/// The search of a generated case: what follows `ORDER BY distance`, and the vector. +fn nearest(tail: &str, question: &[f32]) -> (String, Values) { + let sql = + format!("SELECT id, embedding <=> $query AS distance FROM items ORDER BY distance{tail}"); + (sql, params! { "query" => question.to_vec() }) +} + +/// A new store of the generated cases' table, its vector column declared with `clauses`, and +/// every row that `rows` hands over inserted, 500 to a commit, under the keys 0, 1, 2 and so +/// on. Answers the store and what `rows` answers. +fn filled( + store: &Path, + clauses: &str, + rows: &dyn Fn(Each<'_>) -> Fallible, +) -> Fallible<(Database, T)> { + let db = Database::open(store)?; + db.execute( + &format!("CREATE TABLE items (id INTEGER PRIMARY KEY, embedding VECTOR(768){clauses})"), + ¶ms! {}, + )?; + let mut tx = db.transaction()?; + let mut id = 0_i64; + let answered = rows(&mut |vector| { + tx.execute( + "INSERT INTO items (id, embedding) VALUES ($id, $embedding)", + ¶ms! { "id" => id, "embedding" => vector }, + )?; + id += 1; + if id % 500 == 0 { + std::mem::replace(&mut tx, db.transaction()?).commit()?; + } + Ok(()) + })?; + tx.commit()?; + Ok((db, answered)) +} + +/// One generated case: the two stores and every line taken on them. `rows` hands over the +/// case's rows and answers its questions; `stored` are rows to search for by themselves, each +/// with its key. +fn index( + prefix: &str, + folder: &Path, + count: usize, + rows: &dyn Fn(Each<'_>) -> Fallible>>, + stored: &[(u128, Vec)], +) -> Fallible<()> { + support::must_be_empty(folder); + std::fs::create_dir_all(folder)?; + let store = folder.join("sealed"); + let ((db, questions), load) = { + let (filled, load) = timed(|| filled(&store, " AUTO_INDEX_AT 1000000000", rows)); + (filled?, load) + }; + line(&format!("{prefix}load"), seconds(load), "s"); + let sealed = support::sealed(prefix, &db, &store)?; + let reader = support::session(&store, READER_LIMIT)?; + let before = support::ask(&reader, &nearest, &[EXACT], &questions, 1)?; + let name = format!("{prefix}find_exact_before"); + support::setting(&name, READER_LIMIT); + support::times(&name, &before[0]); + support::counted(&name, &before[0]); + drop(reader); + + let store = folder.join("built"); + let (db, _) = filled(&store, "", rows)?; + let held = count as u64 * (8 + 4 * DIMENSIONS as u64); + let parts = support::built(prefix, &db, &store, (held, sealed, DIMENSIONS))?; + let question = folder.join("question.f32"); + let bytes = questions[0].iter().flat_map(|number| number.to_le_bytes()); + std::fs::write(&question, bytes.collect::>())?; + let first = [OsStr::new("first"), store.as_os_str(), OsStr::new(prefix)]; + support::in_new_process(&[first[0], first[1], first[2], question.as_os_str()])?; + let reader = support::session(&store, READER_LIMIT)?; + let find = format!("{prefix}find"); + let every = support::selected(&parts, &[]); + support::modes(&find, &reader, &nearest, &questions, (1, &every[..]))?; + if !stored.is_empty() { + let name = format!("{find}_own_grouped_found"); + support::own(&name, &reader, &nearest, stored)?; + } + Ok(()) +} + +/// The index pilot's vectors in one partition. +fn synthetic_case(rows: usize, folder: &Path) -> Fallible<()> { + let prefix = format!("synthetic_{rows}_"); + index(&prefix, folder, rows, &|each| synthetic(rows, each), &[]) +} + +/// Generated sightings in one partition, and a hundred of them, evenly spread, each searched +/// for by itself. +fn sightings_case(rows: usize, folder: &Path) -> Fallible<()> { + let sightings = Sightings::new(rows); + let mut stored = Vec::new(); + for row in (0..100).map(|i| i * rows / 100) { + let (track, place) = (row / Sightings::TRACK, row % Sightings::TRACK); + if let Some(vector) = sightings.track(track, place + 1).pop() { + stored.push((row as u128, vector)); + } + } + let load = |each: Each<'_>| { + sightings.rows(rows, each)?; + Ok(sightings.questions()) + }; + index(&format!("sightings_{rows}_"), folder, rows, &load, &stored) +} + +/// What the program does when it starts itself: opens the built store and runs one search for +/// the vector the file holds. +fn first(store: &Path, prefix: &str, question: &Path) -> Fallible<()> { + let question = std::fs::read(question)?; + let (numbers, _) = question.as_chunks::<4>(); + let question: Vec = numbers + .iter() + .map(|four| f32::from_le_bytes(*four)) + .collect(); + let names = (format!("{prefix}open"), format!("{prefix}find_auto_first")); + support::first( + store, + (&names.0, &names.1), + &nearest(support::AUTO, &question), + ) +} + fn usage() -> ! { eprintln!( - "usage: measure readme|million|vectors \n \ - measure sessions " + "usage: measure readme|million|vectors|index \n \ + measure sessions \n \ + measure synthetic|sightings " ); std::process::exit(2); } -fn main() -> Result<()> { +fn main() -> Fallible<()> { let arguments: Vec = std::env::args().skip(1).collect(); let arguments: Vec<&str> = arguments.iter().map(String::as_str).collect(); + if let ["first", store, prefix, question] = arguments.as_slice() { + return first(Path::new(store), prefix, Path::new(question)); + } 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), + ["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), + Ok(count) => sessions(Path::new(corpus), Path::new(list), count, folder)?, Err(_) => usage(), }, + ["index", _] => { + synthetic_case(1_000_000, &folder.join("synthetic"))?; + sightings_case(1_000_000, &folder.join("sightings"))?; + } + [which @ ("synthetic" | "sightings"), rows, _] => match (*which, rows.parse()) { + ("synthetic", Ok(rows)) => synthetic_case(rows, folder)?, + (_, Ok(rows)) => sightings_case(rows, folder)?, + _ => usage(), + }, _ => usage(), } + Ok(()) } diff --git a/v2/contextdb/examples/passages.rs b/v2/contextdb/examples/passages.rs index 2b27bf7b..55c1a8b7 100644 --- a/v2/contextdb/examples/passages.rs +++ b/v2/contextdb/examples/passages.rs @@ -1,7 +1,15 @@ -//! Loads a dump of search passages into a store and measures searches on them, printing one -//! number per line as `name value unit`. +//! Loads a dump of search passages into stores and measures searches on them under each mode, +//! printing one number per line as `name value unit`. //! -//! `passages ` +//! The program is called in three ways: +//! +//! - `passages [sq8][,tail]`: every passage of the +//! dump, declared as Context Graph declares its table. Lines `passages__...`. +//! - `passages single `: the first 10,000 passages, in dump order, of +//! the dump's (Context, generation) pair of index 1, in a table whose vector column also +//! declares `AUTO_INDEX_AT 1000`. Lines `single__...`. +//! - `passages first single|store `: what the program starts +//! itself as, to open a built store in a new process and run one search. //! //! 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 @@ -11,66 +19,82 @@ //! `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. +//! rounded to 32 bits. The folder must be missing or empty; each store is a folder inside it. +//! +//! What it does, in order, each store loaded 500 passages to a commit with no memory limit +//! and every read session opened under 1 GiB unless a line's name says another limit: //! -//! What it does, in order: +//! 1. `sealed`: a store whose column declares `AUTO_INDEX_AT 1000000000`, so that no partition +//! is ever grouped. One search of every question with every vector still waiting; cycles +//! until nothing waits; its size; EXACT in each scope (`exact_before`). +//! 2. `built`: the store as its case declares it. Cycles until nothing remains; what +//! inspection says of the index; its size; in a new process, opening it and one search; +//! then, in one session, each scope under EXACT, INDEXED and AUTO in turn within each of +//! five rounds, the two indexed modes compared with EXACT of the same run; the answers +//! beside the recorded ones; INDEXED again under 64 MiB; the statement with one condition +//! more that passes one passage in a hundred, under EXACT and INDEXED; and a hundred stored +//! vectors each searched for by itself. +//! 3. with `sq8`: a store whose column is declared `WITH (quantization = 'SQ8')` and never +//! grouped, searched under EXACT, its answers beside the first store's. +//! 4. with `tail`: a store built from all but the last tenth of each partition's passages, in +//! dump order, which are written after the build and sealed, so that they stay in each +//! partition's tail; EXACT and INDEXED over all of it, and a hundred of the tail's vectors +//! each searched for by itself. //! -//! 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. +//! The scopes are one Context (the first), several (the first three) and all of them; the +//! second way of calling has the one scope of its pair and leaves the scope out of its names. +//! The search is Context Graph's statement as it sends it, its two flags included. The added +//! condition is on the column `id`: sixteen ranges of ids, evenly spread through the ids in +//! order, that hold a hundredth of the passages together. //! -//! 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. +//! A time named `p50` or `p95` is the middle of the five rounds' values, each the time at +//! position `round((count - 1) * share)` of a round's times in order; `p95_spread` is the +//! largest of the five less the smallest. "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::collections::BTreeSet; +use std::ffi::OsStr; use std::fs::File; -use std::io::{BufReader, Read}; -use std::path::Path; -use std::time::Duration; +use std::io::{BufReader, Read, Seek, SeekFrom}; +use std::path::{Path, PathBuf}; 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}; +use contextdb::{Database, Uuid, Value, params}; +use support::{ + AUTO, Asked, EXACT, Fallible, INDEXED, READER_LIMIT, Values, against, ask, counted, line, + mebibytes, ratio, seconds, selected, session, setting, timed, times, work, +}; #[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"]; +const ROUNDS: usize = 5; +/// The clause of a column no partition of which is ever grouped. +const NEVER_GROUPED: &str = " AUTO_INDEX_AT 1000000000"; + +struct Passage { + id: Uuid, + pair: usize, + flags: u32, +} -/// What the dump says beside its passages. +/// A dump: its passages without their vectors, and what it says beside them. struct Dump { + folder: PathBuf, /// 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>, + passages: Vec, } fn floats(bytes: &[u8]) -> Vec { @@ -94,76 +118,207 @@ fn read_dump(folder: &Path) -> Fallible { } let questions = floats(&std::fs::read(folder.join("queries.f32"))?); let questions = questions.chunks_exact(DIMENSIONS).map(<[f32]>::to_vec); + let mut passages = Vec::new(); + for record in std::fs::read(folder.join("rows.bin"))?.chunks_exact(36) { + let number = |at: usize| { + u32::from_le_bytes([record[at], record[at + 1], record[at + 2], record[at + 3]]) + }; + passages.push(Passage { + id: Uuid::from_slice(&record[8..24])?, + pair: number(28) as usize, + flags: number(32), + }); + } Ok(Dump { + folder: folder.to_owned(), contexts, pairs, questions: questions.collect(), + passages, }) } -/// 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; +impl Dump { + /// The Context and generation of a pair. + fn pair(&self, pair: usize) -> (Uuid, Uuid) { + (self.contexts[self.pairs[pair].0], self.pairs[pair].1) + } + + /// The active pairs of these Contexts, each once, in the order of their ids. + fn active(&self, contexts: &[usize]) -> Vec<(Uuid, Uuid)> { + let active = self.pairs.iter().filter(|pair| pair.2); + let wanted = active.filter(|pair| contexts.contains(&pair.0)); + let pairs: BTreeSet<(Uuid, Uuid)> = + wanted.map(|pair| (self.contexts[pair.0], pair.1)).collect(); + pairs.into_iter().collect() + } + + /// As many as `count` of these passages that have a vector, evenly spread, each with its + /// id as a number and its vector. + fn stored(&self, places: &[usize], count: usize) -> Fallible)>> { + let with = places.iter().copied(); + let with: Vec = with + .filter(|at| self.passages[*at].flags & 4 != 0) + .collect(); + let count = count.min(with.len()); + let mut vectors = File::open(self.folder.join("vectors.f32"))?; + let mut bytes = vec![0u8; 4 * DIMENSIONS]; + let mut stored = Vec::with_capacity(count); + for at in (0..count).map(|i| with[i * with.len() / count]) { + vectors.seek(SeekFrom::Start((at * 4 * DIMENSIONS) as u64))?; + vectors.read_exact(&mut bytes)?; + stored.push((self.passages[at].id.as_u128(), floats(&bytes))); + } + Ok(stored) + } +} + +/// What is loaded and how it is declared and searched. +struct Case<'a> { + dump: &'a Dump, + /// What every line's name begins with. + prefix: String, + /// The passages loaded, by their place in the dump, in its order. + loaded: Vec, + /// What the vector column declares after `SEARCH_MODE AUTO`. + declared: &'static str, + /// Each scope: what its lines' names hold after `find_`, and the pairs it selects. The + /// last is the widest. + scopes: Vec<(&'static str, Vec<(Uuid, Uuid)>)>, + /// The folder of recorded answers, for a case that has them. + answers: Option<&'a Path>, +} + +impl<'a> Case<'a> { + /// Every passage of the dump, as Context Graph declares its table. + fn store(dump: &'a Dump, answers: Option<&'a Path>) -> Case<'a> { + let every: Vec = (0..dump.contexts.len()).collect(); + Case { + dump, + prefix: format!("passages_{}_", dump.passages.len()), + loaded: (0..dump.passages.len()).collect(), + declared: "", + scopes: vec![ + ("one_", dump.active(&[0])), + ("several_", dump.active(&[0, 1, 2])), + ("all_", dump.active(&every)), + ], + answers, + } + } + + /// The first 10,000 passages of the pair of index 1, in a column that declares + /// `AUTO_INDEX_AT 1000`. + fn single(dump: &'a Dump) -> Case<'a> { + let places = 0..dump.passages.len(); + let of_pair = places.filter(|at| dump.passages[*at].pair == 1); + let loaded: Vec = of_pair.take(10_000).collect(); + Case { + dump, + prefix: format!("single_{}_", loaded.len()), + loaded, + declared: " AUTO_INDEX_AT 1000", + scopes: vec![("", vec![dump.pair(1)])], + answers: None, + } + } + + /// The name of a scope's lines up to the mode: `find_`, or `find`. + fn find(&self, scope: &str) -> String { + format!("{}find_{scope}", self.prefix) + .trim_end_matches('_') + .to_owned() + } + + /// The name of one line of the case. + fn name(&self, what: &str) -> String { + format!("{}{what}", self.prefix) + } + + /// The widest scope. + fn widest(&self) -> &(&'static str, Vec<(Uuid, Uuid)>) { + &self.scopes[self.scopes.len() - 1] + } +} + +/// Opens a new store with Context Graph's table, its vector column declared `with` this +/// quantization clause and with these clauses after `SEARCH_MODE AUTO`. +fn create(store: &Path, with: &str, clauses: &str) -> Fallible { + let db = Database::open(store)?; + db.execute( + &format!( + "CREATE TABLE event_search_passages ( + id UUID PRIMARY KEY, + context_id UUID NOT NULL, + generation_id UUID NOT NULL, + vector_text VECTOR(768){with} PARTITION_KEY (context_id, generation_id) + MAX_PARTITIONS 4096 SEARCH_MODE AUTO{clauses}, + is_current BOOLEAN NOT NULL, + is_searchable BOOLEAN NOT NULL DEFAULT true + )" + ), + ¶ms! {}, + )?; + Ok(db) +} + +/// Inserts these passages of the dump, given by their places in rising order. Returns the +/// bytes held. +fn load(db: &Database, dump: &Dump, places: &[usize]) -> Fallible { + let mut vectors = BufReader::new(File::open(dump.folder.join("vectors.f32"))?); + let mut vector = vec![0u8; 4 * DIMENSIONS]; + let (mut held, mut next) = (0, 0); let mut tx = db.transaction()?; - for passage in 1..=passages as usize { - rows.read_exact(&mut record)?; + for (count, at) in places.iter().enumerate() { + vectors.seek_relative(((at - next) * 4 * DIMENSIONS) as i64)?; 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 { + next = at + 1; + let passage = &dump.passages[*at]; + let (context, generation) = dump.pair(passage.pair); + let has_vector = passage.flags & 4 != 0; + let embedding = if has_vector { Value::Vector(floats(&vector)) + } else { + Value::Null }; - held += 3 * 16 - + 2 - + if flags & 4 == 0 { - 0 - } else { - 4 * DIMENSIONS as u64 - }; + held += 3 * 16 + 2 + if has_vector { 4 * DIMENSIONS as u64 } else { 0 }; 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], + "id" => passage.id, + "context" => context, "generation" => generation, "vector" => embedding, - "current" => flags & 1 != 0, - "searchable" => flags & 2 != 0, + "current" => passage.flags & 1 != 0, + "searchable" => passage.flags & 2 != 0, }, )?; - if passage % PER_COMMIT == 0 { + if (count + 1) % PER_COMMIT == 0 { std::mem::replace(&mut tx, db.transaction()?).commit()?; } } tx.commit()?; - Ok((passages, held)) + Ok(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) { +/// Context Graph's search statement for these pairs, with its values: `tail` is what follows +/// `ORDER BY distance`, and `among`, when it holds any, adds the condition that the passage's +/// id lies in one of these ranges, at or after a range's first id and before its second. +fn statement( + pairs: &[(Uuid, Uuid)], + among: &[(Uuid, Uuid)], + tail: &str, + question: &[f32], +) -> (String, Values) { 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)); + for (i, (context, generation)) in pairs.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 { @@ -171,65 +326,45 @@ fn statement(dump: &Dump, scope: &[usize], question: &[f32]) -> (String, HashMap } else { format!("(context_id, generation_id) IN ({})", names.join(", ")) }; + let mut ranges = Vec::new(); + for (i, (from, before)) in among.iter().enumerate() { + values.insert(format!("id_from_{i}"), Value::Uuid(*from)); + values.insert(format!("id_before_{i}"), Value::Uuid(*before)); + ranges.push(format!("(id >= $id_from_{i} AND id < $id_before_{i})")); + } + let condition = if ranges.is_empty() { + String::new() + } else { + format!(" AND ({})", ranges.join(" OR ")) + }; 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" + AND is_current = $is_current AND is_searchable = $is_searchable{condition} \ + ORDER BY distance{tail}" ); (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() +/// The key tuples of these pairs as the store writes them: for each of the two ids the tag +/// of a UUID and its bytes. +fn tuples(pairs: &[(Uuid, Uuid)]) -> Vec> { + let tuple = |(context, generation): &(Uuid, Uuid)| { + [ + &[5][..], + &context.as_bytes()[..], + &[5], + &generation.as_bytes()[..], + ] + .concat() + }; + pairs.iter().map(tuple).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)> { +fn compare(answers: &[Vec<(u128, 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)?; @@ -242,12 +377,13 @@ fn compare(answers: &[Vec<(Uuid, f64)>], recorded: &Path) -> Fallible<(usize, us }; let ids = ids.iter().map(|id| id.as_str().unwrap_or_default().parse()); let ids: Vec = ids.collect::>()?; + let ids: Vec = ids.iter().map(Uuid::as_u128).collect(); let answer = answers .get(question as usize) .ok_or("a question out of range")?; - let ours: Vec = answer.iter().map(|(id, _)| *id).collect(); + let ours: Vec = answer.iter().map(|(id, _)| *id).collect(); same_order += usize::from(ours == ids); - let as_set = |ids: &[Uuid]| ids.iter().copied().collect::>(); + let as_set = |ids: &[u128]| 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) { @@ -259,109 +395,353 @@ fn compare(answers: &[Vec<(Uuid, f64)>], recorded: &Path) -> Fallible<(usize, us 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! {}, - )?; +/// The first store: never grouped. Returns the bytes held, the store's bytes, and EXACT's +/// answers in each scope. +fn sealed(case: &Case<'_>, store: &Path) -> Fallible<(u64, u64, Vec)> { + let dump = case.dump; + let db = create(store, "", NEVER_GROUPED)?; support::take_peak_allocated(); - let (loaded, load_time) = timed(|| load(&db, dump_folder, &dump)); - let (passages, held) = loaded?; + let (held, load_time) = timed(|| load(&db, dump, &case.loaded)); + let held = held?; + line(&case.name("load"), seconds(load_time), "s"); 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", + &case.name("load_peak_allocated"), + mebibytes(load_peak), + "MiB", ); - 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"); + + let (scope, pairs) = case.widest(); + let find = |tail: &str, question: &[f32]| statement(pairs, &[], tail, question); + let unsealed = session(store, READER_LIMIT)?; + let unsealed = ask(&unsealed, &find, &[AUTO], &dump.questions, 1)?; + let name = format!("{}find_unsealed_{scope}", case.prefix); + let name = name.trim_end_matches('_'); + setting(name, READER_LIMIT); + times(name, &unsealed[0]); + counted(name, &unsealed[0]); + + let bytes = support::sealed(&case.prefix, &db, store)?; + let reader = session(store, READER_LIMIT)?; + let mut exact = Vec::new(); + for (scope, pairs) in &case.scopes { + let find = |tail: &str, question: &[f32]| statement(pairs, &[], tail, question); + let mut asked = ask(&reader, &find, &[EXACT], &dump.questions, ROUNDS)?; + let name = format!("{}_exact_before", case.find(scope)); + setting(&name, READER_LIMIT); + times(&name, &asked[0]); + counted(&name, &asked[0]); + exact.append(&mut asked); + } + Ok((held, bytes, exact)) +} + +/// Ranges of ids: the first id of each and the first id after it. +type Ranges = Vec<(Uuid, Uuid)>; + +/// Sixteen ranges of ids, evenly spread through the loaded passages' ids in order, that hold +/// a hundredth of the passages together: each range's first id and the first id after it, +/// how many passages they hold, and in how many pairs those lie. `None` for fewer than 1,600 +/// passages. +fn hundredth(case: &Case<'_>) -> Option<(Ranges, usize, usize)> { + const RANGES: usize = 16; + let passages = case.loaded.iter().map(|at| &case.dump.passages[*at]); + let mut ids: Vec<(Uuid, usize)> = passages.map(|passage| (passage.id, passage.pair)).collect(); + ids.sort_unstable(); + let each = ids.len() / (100 * RANGES); + let (mut ranges, mut pairs) = (Vec::new(), BTreeSet::new()); + for range in 0..RANGES { + let from = (2 * range + 1) * ids.len() / (2 * RANGES) - each / 2; + let inside = ids.get(from..from + each).filter(|_| each > 0)?; + pairs.extend(inside.iter().map(|(_, pair)| *pair)); + ranges.push((inside[0].0, ids.get(from + each)?.0)); + } + Some((ranges, each * RANGES, pairs.len())) +} + +/// The second store, declared as the case says, and every line taken on it. +fn built(case: &Case<'_>, folder: &Path, held: u64, sealed: u64) -> Fallible<()> { + let (dump, store) = (case.dump, folder.join("built")); + let db = create(&store, "", case.declared)?; + load(&db, dump, &case.loaded)?; + let parts = support::built(&case.prefix, &db, &store, (held, sealed, DIMENSIONS))?; + let kind = if case.answers.is_some() { + "store" + } else { + "single" + }; + let first = [OsStr::new("first"), OsStr::new(kind)]; + support::in_new_process(&[ + first[0], + first[1], + dump.folder.as_os_str(), + store.as_os_str(), + ])?; 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 reader = session(&store, READER_LIMIT)?; + let mut asked = Vec::new(); + for (scope, pairs) in &case.scopes { + let find = |tail: &str, question: &[f32]| statement(pairs, &[], tail, question); + let parts = selected(&parts, &tuples(pairs)); + let how = (ROUNDS, &parts[..]); + asked.push(support::modes( + &case.find(scope), + &reader, + &find, + &dump.questions, + how, + )?); } - let all = &unlimited[2]; - line(&name("find_all_first"), milliseconds(all.first), "ms"); + if let Some(answers) = case.answers { + let mut gap = 0.0; + for (mode, at) in [("", 0), ("auto_", 2)] { + for (what, order) in [("same_ten", false), ("same_order", true)] { + for ((scope, _), asked) in case.scopes.iter().zip(&asked) { + let scope = scope.trim_end_matches('_'); + let recorded = answers.join(format!("engine-{scope}.jsonl")); + let (same_ten, same_order, far) = compare(&asked[at].answers, &recorded)?; + let same = if order { same_order } else { same_ten }; + line(&case.name(&format!("{mode}{what}_{scope}")), same, "of_100"); + gap = f64::max(gap, far); + } + } + } + line(&case.name("distance_gap_max"), gap, "-"); + + // INDEXED again where the blocks do not stay in the cache. + let small = session(&store, 64 << 20)?; + for ((scope, pairs), asked) in case.scopes.iter().zip(&asked) { + let find = |tail: &str, question: &[f32]| statement(pairs, &[], tail, question); + let under = ask(&small, &find, &[INDEXED], &dump.questions, ROUNDS)?; + let name = format!("{}_indexed_under_64m", case.find(scope)); + setting(&name, 64 << 20); + times(&name, &under[0]); + counted(&name, &under[0]); + against(&name, &under[0], &asked[0]); + } + + // The statement with a condition that passes a passage in a hundred. INDEXED is + // refused where the filter leaves the index nothing to choose among, as in a store + // none of whose partitions is grouped: EXACT is then timed alone. + let (scope, pairs) = case.widest(); + let find = case.find(scope); + if let Some((among, passing, in_pairs)) = hundredth(case) { + let name = format!("{find}_filtered"); + line(&format!("{name}_rows_passing"), passing, "rows"); + let in_pairs = (format!("{name}_pairs_with_passing_rows"), in_pairs); + line(&in_pairs.0, in_pairs.1, "pairs"); + let filtered = |tail: &str, question: &[f32]| statement(pairs, &among, tail, question); + let both = ask( + &reader, + &filtered, + &[EXACT, INDEXED], + &dump.questions, + ROUNDS, + ); + let asked = match both { + Ok(asked) => asked, + Err(error) => { + eprintln!("{find}_indexed_filtered: {error}"); + ask(&reader, &filtered, &[EXACT], &dump.questions, ROUNDS)? + } + }; + setting(&name, READER_LIMIT); + let name = format!("{find}_exact_filtered"); + times(&name, &asked[0]); + counted(&name, &asked[0]); + let name = format!("{find}_indexed_filtered"); + match asked.get(1) { + Some(indexed) => { + times(&name, indexed); + counted(&name, indexed); + against(&name, indexed, &asked[0]); + } + None => line(&name, "refused", "-"), + } + } else { + line(&format!("{find}_filtered"), "unavailable", "-"); + } + } + + let (scope, pairs) = case.widest(); + let find = |tail: &str, question: &[f32]| statement(pairs, &[], tail, question); + let name = format!("{}_own_grouped_found", case.find(scope)); + support::own(&name, &reader, &find, &dump.stored(&case.loaded, 100)?)?; + let find_peak = support::take_peak_allocated(); line( - &name("find_all_max"), - milliseconds(percentile(&all.times, 1.0)), - "ms", + &case.name("find_peak_allocated"), + mebibytes(find_peak), + "MiB", ); - 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"); + Ok(()) +} + +/// A store with a tail: built from all but the last tenth of each pair's passages, which are +/// written after the build and sealed. +fn tail(case: &Case<'_>, store: &Path) -> Fallible<()> { + let dump = case.dump; + let mut of_pair = vec![0_usize; dump.pairs.len()]; + for at in &case.loaded { + of_pair[dump.passages[*at].pair] += 1; + } + let mut seen = vec![0_usize; dump.pairs.len()]; + let (mut first, mut last) = (Vec::new(), Vec::new()); + for at in &case.loaded { + let pair = dump.passages[*at].pair; + seen[pair] += 1; + if seen[pair] > of_pair[pair] - of_pair[pair] / 10 { + last.push(*at); + } else { + first.push(*at); + } + } + let db = create(store, "", case.declared)?; + load(&db, dump, &first)?; + support::upkeep(&db)?; + load(&db, dump, &last)?; + let cycles = support::upkeep(&db)?; + db.close()?; + line(&case.name("tail_passages"), last.len(), "passages"); + line(&case.name("tail_cycles"), cycles, "cycles"); + + let parts = support::parts(store, DIMENSIONS)?; + let (scope, pairs) = case.widest(); + let parts = selected(&parts, &tuples(pairs)); + let find = |tail: &str, question: &[f32]| statement(pairs, &[], tail, question); + let reader = session(store, READER_LIMIT)?; + let asked = ask(&reader, &find, &[EXACT, INDEXED], &dump.questions, ROUNDS)?; + let scope_find = case.find(scope); + setting(&format!("{scope_find}_tail"), READER_LIMIT); + let name = format!("{scope_find}_exact_tail"); + times(&name, &asked[0]); + counted(&name, &asked[0]); + if let Some(answers) = case.answers { + let scope = scope.trim_end_matches('_'); + let recorded = answers.join(format!("engine-{scope}.jsonl")); + let same_ten = compare(&asked[0].answers, &recorded)?.0; + line(&format!("{name}_same_ten"), same_ten, "of_100"); + } + let name = format!("{scope_find}_indexed_tail"); + times(&name, &asked[1]); + counted(&name, &asked[1]); + against(&name, &asked[1], &asked[0]); + work(&name, &asked[1], &parts, &dump.questions); + let name = format!("{scope_find}_own_tail_found"); + support::own(&name, &reader, &find, &dump.stored(&last, 100)?) +} + +/// A store that keeps only the reduced vectors, never grouped, searched under EXACT beside +/// the first store's answers under EXACT. +fn reduced(case: &Case<'_>, store: &Path, sealed: u64, exact: &[Asked]) -> Fallible<()> { + let dump = case.dump; + let db = create(store, " WITH (quantization = 'SQ8')", NEVER_GROUPED)?; + load(&db, dump, &case.loaded)?; + support::upkeep(&db)?; + db.close()?; + let reader = session(store, READER_LIMIT)?; + let mut compared = Vec::new(); + for ((_, pairs), exact) in case.scopes.iter().zip(exact) { + let find = |tail: &str, question: &[f32]| statement(pairs, &[], tail, question); + let asked = ask(&reader, &find, &[EXACT], &dump.questions, 1)?; + let (mut same_ten, mut same_order, mut shared) = (0, 0, 0); + for (ours, theirs) in asked[0].answers.iter().zip(&exact.answers) { + let ids = |answer: &[(u128, f64)]| answer.iter().map(|row| row.0).collect::>(); + let (ours, theirs) = (ids(ours), ids(theirs)); + let found = theirs.iter().filter(|id| ours.contains(id)).count(); + shared += found; + same_ten += usize::from(found == theirs.len() && ours.len() == theirs.len()); + same_order += usize::from(ours == theirs); + } + compared.push((same_ten, same_order, shared)); } - for (what, share) in [("p50", 0.5), ("p95", 0.95), ("max", 1.0)] { - let rows = percentile(&all.examined, share); + let scopes = || { + case.scopes + .iter() + .map(|scope| scope.0.trim_end_matches('_')) + }; + for (scope, (same_ten, _, _)) in scopes().zip(&compared) { line( - &name(&format!("find_all_rows_examined_{what}")), - rows, - "rows", + &case.name(&format!("sq8_same_ten_{scope}")), + same_ten, + "of_100", ); } + for (scope, (_, same_order, _)) in scopes().zip(&compared) { + line( + &case.name(&format!("sq8_same_order_{scope}")), + same_order, + "of_100", + ); + } + for (scope, (_, _, shared)) in scopes().zip(&compared) { + line( + &case.name(&format!("sq8_answers_shared_{scope}")), + shared, + "of_1000", + ); + } + let bytes = support::bytes_on_disk(store)?; + let passages = case.loaded.len().max(1) as u64; + line(&case.name("sq8_store_size"), mebibytes(bytes), "MiB"); + line( + &case.name("sq8_bytes_per_passage"), + bytes / passages, + "bytes", + ); + line(&case.name("sq8_size_over_f32"), ratio(bytes, sealed), "x"); + Ok(()) +} - 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)?); +/// Every store of a case, each a folder inside `folder`. +fn measure(case: &Case<'_>, folder: &Path, extras: &str) -> Fallible<()> { + support::must_be_empty(folder); + std::fs::create_dir_all(folder)?; + if case.declared == " AUTO_INDEX_AT 1000" { + line(&case.name("auto_index_at"), 1000, "vectors"); } - for (scope, (same_ten, _, _)) in SCOPES.iter().zip(&compared) { - line(&name(&format!("same_ten_{scope}")), same_ten, "of_100"); + let (held, sealed_bytes, exact) = sealed(case, &folder.join("sealed"))?; + built(case, folder, held, sealed_bytes)?; + if extras.split(',').any(|extra| extra == "sq8") { + reduced(case, &folder.join("sq8"), sealed_bytes, &exact)?; } - for (scope, (_, same_order, _)) in SCOPES.iter().zip(&compared) { - line(&name(&format!("same_order_{scope}")), same_order, "of_100"); + if extras.split(',').any(|extra| extra == "tail") { + tail(case, &folder.join("tail"))?; } - 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(()) } + +fn main() -> Fallible<()> { + let arguments: Vec = std::env::args().skip(1).collect(); + let arguments: Vec<&str> = arguments.iter().map(String::as_str).collect(); + match arguments.as_slice() { + ["first", kind, dump, store] => { + let dump = read_dump(Path::new(dump))?; + let case = if *kind == "single" { + Case::single(&dump) + } else { + Case::store(&dump, None) + }; + let (scope, pairs) = case.widest(); + let first = format!("{}_auto_first", case.find(scope)); + let search = statement(pairs, &[], AUTO, &dump.questions[0]); + support::first(Path::new(store), (&case.name("open"), &first), &search) + } + ["single", dump, folder] => { + let dump = read_dump(Path::new(dump))?; + measure(&Case::single(&dump), Path::new(folder), "") + } + [dump, answers, folder] | [dump, answers, folder, _] => { + let dump = read_dump(Path::new(dump))?; + let case = Case::store(&dump, Some(Path::new(answers))); + measure(&case, Path::new(folder), arguments.get(3).unwrap_or(&"")) + } + _ => { + eprintln!( + "usage: passages [sq8][,tail]\n \ + passages single " + ); + std::process::exit(2); + } + } +} diff --git a/v2/contextdb/examples/support/mod.rs b/v2/contextdb/examples/support/mod.rs index 19e57b98..f738f46a 100644 --- a/v2/contextdb/examples/support/mod.rs +++ b/v2/contextdb/examples/support/mod.rs @@ -1,12 +1,45 @@ //! 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. +//! the size of a folder, the `name value unit` lines they print, and the measuring of a +//! similarity search under each mode beside what the search worked through. +//! +//! One thing here reads a store's data file itself and not through the library: [`parts`], +//! which reads the vector part's partition, grouping, centre and sketch records with the +//! storage library, because the library's inspection gives a grouping's group count and +//! nothing of its gap or its width, and counts a search's work as one number. The layout it +//! reads is the one `contextdb/src/store/keys.rs` and `contextdb/src/vector/` write; a change +//! there changes it here, and a record it cannot read stops the program. use std::alloc::{GlobalAlloc, Layout, System}; +use std::collections::HashMap; +use std::ffi::OsStr; use std::fmt::Display; use std::path::Path; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::{Duration, Instant}; +use contextdb::{Database, QueryResult, ReadSession, ReadSessionOptions, Value}; +use redb::{ReadOnlyDatabase, ReadableDatabase, TableDefinition}; + +pub type Fallible = Result>; + +/// The values a statement is run with, by name. +pub type Values = HashMap; + +/// What makes the search statement of a case: given what follows `ORDER BY distance` and the +/// vector searched for, the statement and its values. +pub type Statement<'a> = &'a dyn Fn(&str, &[f32]) -> (String, Values); + +/// The memory limit a read session is opened under: what Context Graph gives a reader. +pub const READER_LIMIT: usize = 1 << 30; + +/// Searches of each mode that are run and not timed before the timed ones. +pub const UNTIMED: usize = 5; + +/// What follows `ORDER BY distance` under each mode, for the ten nearest. +pub const EXACT: &str = " USE VECTOR EXACT LIMIT 10"; +pub const INDEXED: &str = " USE VECTOR INDEXED LIMIT 10"; +pub const AUTO: &str = " LIMIT 10"; + /// The system allocator, counting the bytes handed out and not yet given back. A program makes /// it its global allocator. pub struct Counting; @@ -103,3 +136,607 @@ pub fn must_be_empty(folder: &Path) { std::process::exit(2); } } + +/// The sum of the numbers over how many there are, to one place. +fn mean(numbers: impl Iterator) -> String { + let (mut sum, mut count) = (0_i64, 0_i64); + for number in numbers { + sum += number; + count += 1; + } + format!("{:.1}", sum as f64 / count.max(1) as f64) +} + +// --------------------------------------------------------------------------------------------- +// A store built for a case + +/// Runs maintenance cycles until one reports nothing remaining, and answers how many it ran: +/// no vector waiting, and no partition with a grouping due or being built that is not stalled. +pub fn upkeep(db: &Database) -> Fallible { + for cycle in 1..=10_000_000 { + let report = db.run_maintenance_cycle()?.vector; + if report.maintenance_vectors_remaining == 0 + && report.remaining_partitions <= report.stalled_partitions + { + if report.stalled_partitions > 0 { + return Err("a partition's build is stalled for memory".into()); + } + return Ok(cycle); + } + } + Err("ten million maintenance cycles left work remaining".into()) +} + +/// The first store of a case, in which no partition is ever grouped: runs its cycles, closes +/// it and prints `seal`, `seal_cycles` and `store_size_sealed`. Answers the store's bytes. +pub fn sealed(prefix: &str, db: &Database, store: &Path) -> Fallible { + let (cycles, time) = timed(|| upkeep(db)); + let cycles = cycles?; + db.close()?; + let bytes = bytes_on_disk(store)?; + line(&format!("{prefix}seal"), seconds(time), "s"); + line(&format!("{prefix}seal_cycles"), cycles, "cycles"); + line( + &format!("{prefix}store_size_sealed"), + mebibytes(bytes), + "MiB", + ); + Ok(bytes) +} + +/// One whole number of the first row of an inspection statement, by the name of its column, +/// or the sum of it over every row. NULL counts as nothing. +fn inspected(found: &QueryResult, field: &str) -> Fallible { + let at = found.columns().iter().position(|name| name == field); + let at = at.ok_or("an inspection statement has no such field")?; + let mut sum = 0; + for row in found.rows() { + if let Value::Int64(number) = row[at] { + sum += number as u64; + } + } + Ok(sum) +} + +/// The second store of a case, declared as the case says: runs cycles until nothing remains, +/// closes it and prints `upkeep`, `upkeep_cycles`, `lists`, `probes`, the index's bytes, the +/// store's size beside what was `held` and beside the first store's bytes, and the least and +/// the most gap and width of its groupings. Answers its partitions as their records describe +/// them. +pub fn built( + prefix: &str, + db: &Database, + store: &Path, + (held, sealed, dimension): (u64, u64, usize), +) -> Fallible> { + let (cycles, time) = timed(|| upkeep(db)); + let cycles = cycles?; + let name = |what: &str| format!("{prefix}{what}"); + line(&name("upkeep"), seconds(time), "s"); + line(&name("upkeep_cycles"), cycles, "cycles"); + let partitions = db.execute("SHOW VECTOR_PARTITIONS", &Values::new())?; + line(&name("lists"), inspected(&partitions, "lists")?, "groups"); + line(&name("probes"), inspected(&partitions, "probes")?, "groups"); + let indexes = db.execute("SHOW VECTOR_INDEXES", &Values::new())?; + let index_bytes = inspected(&indexes, "durable_index_bytes")?; + let vectors = inspected(&indexes, "vector_count")?; + line(&name("index_bytes"), index_bytes, "bytes"); + let each = index_bytes / vectors.max(1); + line(&name("index_bytes_per_vector"), each, "bytes"); + db.close()?; + let bytes = bytes_on_disk(store)?; + line(&name("store_size"), mebibytes(bytes), "MiB"); + line(&name("size_over_held"), ratio(bytes, held), "x"); + line(&name("size_over_sealed"), ratio(bytes, sealed), "x"); + line(&name("bytes_per_passage"), bytes / vectors.max(1), "bytes"); + let parts = parts(store, dimension)?; + let grouped = parts.iter().filter(|part| part.groups.is_some()).count(); + line(&name("partitions"), parts.len(), "partitions"); + line(&name("partitions_grouped"), grouped, "partitions"); + measured(prefix, &parts.iter().collect::>()); + Ok(parts) +} + +/// Starts this program again with these arguments and waits for it: what a new process +/// measures, it prints itself. +pub fn in_new_process(arguments: &[&OsStr]) -> Fallible<()> { + let program = std::env::current_exe()?; + let status = std::process::Command::new(program) + .args(arguments) + .status()?; + if status.success() { + Ok(()) + } else { + Err("the new process did not finish".into()) + } +} + +/// What a new process does first: opens a read session on a built store under the reader's +/// limit and runs one search, and prints the time of each under these two names. +pub fn first(store: &Path, names: (&str, &str), (sql, values): &(String, Values)) -> Fallible<()> { + let (opened, open) = timed(|| session(store, READER_LIMIT)); + let opened = opened?; + let (found, find) = timed(|| opened.execute(sql, values)); + found?; + line(names.0, milliseconds(open), "ms"); + line(names.1, milliseconds(find), "ms"); + Ok(()) +} + +// --------------------------------------------------------------------------------------------- +// Searches under each mode + +/// A read session on a store under a memory limit. +pub fn session(store: &Path, limit: usize) -> Fallible { + let mut options = ReadSessionOptions::default(); + options.memory_limit = Some(limit); + Ok(ReadSession::open_with_options(store, options)?) +} + +/// What the questions gave under one mode: for each round the middle and the 95th percentile +/// of the questions' times, and from the first round, for each question, the search's counted +/// work and its answer (each row's id as a number, and its distance). +#[derive(Default)] +pub struct Asked { + pub p50: Vec, + pub p95: Vec, + pub evaluations: Vec, + pub examined: Vec, + pub answers: Vec>, +} + +/// The id and the distance of every row of a search's answer: the id is its first column, a +/// UUID or an INTEGER, and the distance its last. +fn answer(found: &QueryResult) -> Fallible> { + let rows = found + .rows() + .iter() + .map(|row| match (row.first(), row.last()) { + (Some(Value::Uuid(id)), Some(Value::Float64(distance))) => { + Ok((id.as_u128(), *distance)) + } + (Some(Value::Int64(id)), Some(Value::Float64(distance))) => { + Ok((*id as u128, *distance)) + } + _ => Err("a search row is not an id and a distance".into()), + }); + rows.collect() +} + +/// Asks every question under each of `tails` (what follows `ORDER BY distance`: a mode and a +/// limit), in one session: first [`UNTIMED`] searches under each that are not timed, then +/// `rounds` rounds, and within each round every question under the first tail, then under +/// the next. Each `execute` call is timed whole. +pub fn ask( + session: &ReadSession, + statement: Statement<'_>, + tails: &[&str], + questions: &[Vec], + rounds: usize, +) -> Fallible> { + let search = |tail: &str, question: &[f32]| { + let (sql, values) = statement(tail, question); + timed(|| session.execute(&sql, &values)) + }; + for tail in tails { + for question in questions.iter().take(UNTIMED) { + search(tail, question).0?; + } + } + let mut asked: Vec = tails.iter().map(|_| Asked::default()).collect(); + for round in 0..rounds { + for (tail, asked) in tails.iter().zip(&mut asked) { + let mut times = Vec::with_capacity(questions.len()); + for question in questions { + let (found, time) = search(tail, question); + let found = found?; + times.push(time); + if round == 0 { + asked.evaluations.push(found.distance_evaluations()); + asked.examined.push(found.rows_examined()); + asked.answers.push(answer(&found)?); + } + } + asked.p50.push(percentile(×, 0.5)); + asked.p95.push(percentile(×, 0.95)); + } + } + Ok(asked) +} + +/// Prints the memory limit and the count of threads the lines after it were taken under. +pub fn setting(name: &str, limit: usize) { + let value = format!("memory_limit_{limit}_bytes_threads_1"); + line(&format!("{name}_session"), value, "-"); +} + +/// Prints `_p50` and `_p95`, each the middle of the rounds' values, and with more rounds than +/// one `_p95_spread`, the largest of the rounds' 95th percentiles less the smallest. +pub fn times(name: &str, asked: &Asked) { + let p50 = milliseconds(percentile(&asked.p50, 0.5)); + line(&format!("{name}_p50"), p50, "ms"); + let p95 = milliseconds(percentile(&asked.p95, 0.5)); + line(&format!("{name}_p95"), p95, "ms"); + if asked.p95.len() > 1 { + let spread = percentile(&asked.p95, 1.0) - percentile(&asked.p95, 0.0); + line(&format!("{name}_p95_spread"), milliseconds(spread), "ms"); + } +} + +/// Prints `_distance_evaluations_mean` and `_rows_examined_p50`. +pub fn counted(name: &str, asked: &Asked) { + let evaluations = mean(asked.evaluations.iter().map(|count| *count as i64)); + let name = |what: &str| format!("{name}_{what}"); + line( + &name("distance_evaluations_mean"), + evaluations, + "evaluations", + ); + let rows = percentile(&asked.examined, 0.5); + line(&name("rows_examined_p50"), rows, "rows"); +} + +/// Prints `_recall` (the share of `exact`'s answers, over all questions, that this mode +/// returned), `_all_ten` (questions whose answer holds every row of `exact`'s) and +/// `_not_fewer_evaluations` (questions on which it did as many evaluations as `exact` or +/// more). +pub fn against(name: &str, asked: &Asked, exact: &Asked) { + let (mut shared, mut wanted, mut all) = (0, 0, 0); + for (ours, theirs) in asked.answers.iter().zip(&exact.answers) { + let found = theirs + .iter() + .filter(|(id, _)| ours.iter().any(|row| row.0 == *id)); + let found = found.count(); + shared += found; + wanted += theirs.len(); + all += usize::from(found == theirs.len()); + } + let recall = format!("{:.4}", shared as f64 / wanted.max(1) as f64); + let of = format!("of_{}", exact.answers.len()); + let pairs = asked.evaluations.iter().zip(&exact.evaluations); + let not_fewer = pairs.filter(|(ours, theirs)| ours >= theirs).count(); + line(&format!("{name}_recall"), recall, "-"); + line(&format!("{name}_all_ten"), all, &of); + line(&format!("{name}_not_fewer_evaluations"), not_fewer, &of); +} + +/// Asks the questions under EXACT, INDEXED and AUTO in turn within each round, in one session +/// under the reader's limit, and prints every line of the three for `name`: the times, the +/// counted work, the two indexed modes against EXACT of this run, and under INDEXED what the +/// price is read from ([`work`]). `parts` are the partitions the statement selects. Answers +/// what the three gave, in that order. +pub fn modes( + name: &str, + session: &ReadSession, + statement: Statement<'_>, + questions: &[Vec], + (rounds, parts): (usize, &[&Part]), +) -> Fallible> { + let asked = ask( + session, + statement, + &[EXACT, INDEXED, AUTO], + questions, + rounds, + )?; + setting(name, READER_LIMIT); + for (mode, of) in ["exact", "indexed", "auto"].iter().zip(&asked) { + let name = format!("{name}_{mode}"); + times(&name, of); + counted(&name, of); + if *mode != "exact" { + against(&name, of, &asked[0]); + } + if *mode == "indexed" { + work(&name, of, parts, questions); + } + } + Ok(asked) +} + +/// How many of these stored vectors come back first when each is searched for by itself, +/// under INDEXED and beside it under EXACT, with `LIMIT 1` and with `LIMIT 10`. Under EXACT a +/// vector that another row holds too may come back second. +pub fn own( + name: &str, + session: &ReadSession, + statement: Statement<'_>, + stored: &[(u128, Vec)], +) -> Fallible<()> { + setting(name, READER_LIMIT); + for (mode, clause) in [("", "INDEXED"), ("_exact", "EXACT")] { + for limit in [1, 10] { + let tail = format!(" USE VECTOR {clause} LIMIT {limit}"); + let mut found = 0; + for (id, vector) in stored { + let (sql, values) = statement(&tail, vector); + let rows = answer(&session.execute(&sql, &values)?)?; + found += usize::from(rows.first().is_some_and(|row| row.0 == *id)); + } + let name = format!("{name}{mode}_limit_{limit}"); + line(&name, found, &format!("of_{}", stored.len())); + } + } + Ok(()) +} + +// --------------------------------------------------------------------------------------------- +// What a search's price is read from: the groupings as their records hold them + +const VECTORS: TableDefinition<&[u8], &[u8]> = TableDefinition::new("vectors"); +const BLOCKS: TableDefinition<&[u8], &[u8]> = TableDefinition::new("vector_blocks"); + +/// One partition of a VECTOR column as its records describe it. +pub struct Part { + /// Its key tuple as the store writes it: for each value a tag and the value's bytes. + tuple: Vec, + /// The slots of its grouping and of its tail together. + slots: u64, + /// The slots of its tail. + tail: u64, + /// Its grouping, when it has one of more groups than one. + groups: Option, +} + +struct Groups { + /// How far a search scans, in ten-thousandths of a product. + gap: u32, + /// How wide a search takes a sketch, in quarters of its factor; 0 when it reads none. + width: u32, + /// Every centre, one after another. + centres: Vec, + /// The slots of each group. + sizes: Vec, +} + +fn short(bytes: &[u8], at: usize) -> u32 { + u32::from_le_bytes(std::array::from_fn(|byte| bytes[at + byte])) +} + +fn long(bytes: &[u8], at: usize) -> u64 { + u64::from_le_bytes(std::array::from_fn(|byte| bytes[at + byte])) +} + +/// The partitions of the one VECTOR column of a closed store, read from its data file. +pub fn parts(store: &Path, dimension: usize) -> Fallible> { + let data = ReadOnlyDatabase::open(store.join("data"))?; + let read = data.begin_read()?; + let (vectors, blocks) = (read.open_table(VECTORS)?, read.open_table(BLOCKS)?); + let mut parts = Vec::new(); + // A partition's record is kept under the byte 1, its column and its key tuple. + for record in vectors.range::<&[u8]>(&[1_u8][..]..&[2_u8][..])? { + let (key, entry) = record?; + let (key, entry) = (key.value(), entry.value()); + if entry.len() != 73 || key.len() < 5 { + return Err("a partition record is not as this program reads it".into()); + } + let (tail, flags) = (long(entry, 24), entry[40]); + let (number, blocks_held, groups) = (short(entry, 41), short(entry, 45), short(entry, 49)); + let grouped = if flags & 1 == 0 { 0 } else { long(entry, 53) }; + let mut part = Part { + tuple: key[5..].to_vec(), + slots: grouped + tail, + tail, + groups: None, + }; + if flags & 1 != 0 && groups > 1 { + // The grouping's record: the gap, the width, the probes, then where each group + // begins and where the last ends, a block and a slot each. + let under = [&key[1..5], &number.to_be_bytes()[..]].concat(); + let record = vectors.get([&[3_u8][..], &under].concat().as_slice())?; + let record = record.ok_or("a grouping has no record")?; + let record = record.value(); + if record.len() != 20 + 8 * groups as usize { + return Err("a grouping record is not as this program reads it".into()); + } + // How many slots lie before each block: a block's count opens its sketch record. + let mut before = vec![0_u64]; + for block in 0..blocks_held { + let key = [&[7_u8][..], &under, &block.to_be_bytes()].concat(); + let sketches = blocks.get(key.as_slice())?; + let sketches = sketches.ok_or("a block of a grouping has no sketches")?; + before.push(before[block as usize] + u64::from(short(sketches.value(), 0))); + } + let place = |group: usize| { + let (block, slot) = (short(record, 12 + 8 * group), short(record, 16 + 8 * group)); + before + .get(block as usize) + .map(|slots| slots + u64::from(slot)) + }; + let mut sizes = Vec::with_capacity(groups as usize); + for group in 0..groups as usize { + let (Some(from), Some(to)) = (place(group), place(group + 1)) else { + return Err("a grouping's places are past its blocks".into()); + }; + sizes.push(to - from); + } + let mut centres = Vec::with_capacity(groups as usize * dimension); + let from = [&[4_u8][..], &under].concat(); + let to = [&[4_u8][..], &key[1..5], &(number + 1).to_be_bytes()].concat(); + for record in blocks.range::<&[u8]>(from.as_slice()..to.as_slice())? { + let (_, value) = record?; + let (numbers, _) = value.value().as_chunks::<4>(); + centres.extend(numbers.iter().map(|number| f32::from_le_bytes(*number))); + } + if centres.len() != groups as usize * dimension { + return Err("a grouping's centres are not as many as its groups".into()); + } + part.groups = Some(Groups { + gap: short(record, 0), + width: short(record, 4), + centres, + sizes, + }); + } + parts.push(part); + } + Ok(parts) +} + +/// Those of `parts` whose key tuple is one of `tuples`, or every one when none is given. +pub fn selected<'a>(parts: &'a [Part], tuples: &[Vec]) -> Vec<&'a Part> { + let wanted = |part: &&Part| tuples.is_empty() || tuples.contains(&part.tuple); + parts.iter().filter(wanted).collect() +} + +/// Prints the least and the most gap and width among the groupings of these partitions: the +/// gap as a difference of products, the width as a multiple of a sketch's factor. `none` when +/// no partition of them has a grouping of more groups than one. +fn measured(name: &str, parts: &[&Part]) { + let groupings = || parts.iter().filter_map(|part| part.groups.as_ref()); + let gaps = groupings().map(|groups| groups.gap); + let widths = groupings().map(|groups| groups.width); + let name = |what: &str| format!("{name}{what}"); + let (Some(gap_least), Some(gap_most)) = (gaps.clone().min(), gaps.max()) else { + for what in ["gap_least", "gap_most", "width_least", "width_most"] { + line(&name(what), "none", "-"); + } + return; + }; + let gap = |gap: u32| format!("{:.4}", f64::from(gap) / 10_000.0); + line(&name("gap_least"), gap(gap_least), "product"); + line(&name("gap_most"), gap(gap_most), "product"); + let width = |width: Option| format!("{:.2}", f64::from(width.unwrap_or(0)) / 4.0); + line(&name("width_least"), width(widths.clone().min()), "factors"); + line(&name("width_most"), width(widths.max()), "factors"); +} + +/// The numbers of one vector multiplied by another's and added up as the library adds them: +/// in 32 bits, sixteen sums side by side. +fn product(one: &[f32], other: &[f32]) -> f32 { + let mut lanes = [0.0_f32; 16]; + let (one_runs, one_rest) = one.as_chunks::<16>(); + let (other_runs, other_rest) = other.as_chunks::<16>(); + for (one, other) in one_runs.iter().zip(other_runs) { + for ((lane, one), other) in lanes.iter_mut().zip(one).zip(other) { + *lane += one * other; + } + } + for ((lane, one), other) in lanes.iter_mut().zip(one_rest).zip(other_rest) { + *lane += one * other; + } + lanes.iter().sum() +} + +/// What an unfiltered search under INDEXED does with these partitions for one question, worked +/// out from their records as `vector/search.rs` works it out: the centres and means it +/// compares, the groups inside the gap, the slots of those groups and of the tails beside +/// them, how many of those slots it estimates a sketch of, and the slots it screens with no +/// sketch read (a partition with no groups, and a grouping of width 0). +#[derive(Default)] +struct Reached { + centres: i64, + groups: i64, + slots: i64, + sketches: i64, + unsketched: i64, +} + +fn reached(parts: &[&Part], question: &[f32]) -> Reached { + let length = question.iter().map(|x| f64::from(*x) * f64::from(*x)); + let length = length.sum::().sqrt(); + let unit = question.iter().map(|x| (f64::from(*x) / length) as f32); + let unit: Vec = unit.collect(); + let mut reached = Reached::default(); + for part in parts { + let Some(groups) = &part.groups else { + reached.unsketched += part.slots as i64; + continue; + }; + let centres = groups.centres.chunks(unit.len()); + let products: Vec = centres.map(|centre| product(centre, &unit)).collect(); + let best = products.iter().copied().fold(f32::MIN, f32::max); + let apart = + |product: f32| ((f64::from(best) - f64::from(product)) * 10_000.0).min(20_000.0); + let mut slots = part.tail as i64; + for (size, product) in groups.sizes.iter().zip(&products) { + if apart(*product) - f64::from(groups.gap) <= 0.0 { + reached.groups += 1; + slots += *size as i64; + } + } + reached.centres += products.len() as i64; + reached.slots += slots; + if groups.width == 0 { + reached.unsketched += slots; + } else { + // The partition's mean is compared too. + reached.centres += 1; + reached.sketches += slots; + } + } + reached +} + +/// Prints, for the unfiltered searches of `asked` under INDEXED over `parts`, what their price +/// is read from: the least and the most gap and width of the groupings, and at the mean +/// question the passes a search made, the centres compared, the groups inside the gap, the +/// slots of those groups and of the tails, the sketches estimated, the slots screened and the +/// rows asked for. +/// +/// The library counts one number a search, `distance_evaluations`: centres and means compared, +/// sketches estimated, slots screened and rows asked for. The rows asked for are the rows the +/// statement examined less the rows it returned. The rest is worked out here. A search makes +/// one pass, and another whenever the last row of its answer lies below what its sketches +/// ended with; every pass compares the centres and estimates the sketches of the groups +/// inside the gap again, and a slot is screened once. So with `c` centres and `s` sketches a +/// pass ([`reached`]), a search of `p` passes counts from `p (c + s)` up to `p (c + s) + s`, +/// ranges that do not meet: the passes are the count over `c + s`, rounded down, and the slots +/// screened are what is left. `_work_out_of_range` counts the questions on which what is left +/// is fewer than the slots screened with no sketch or more than those and every sketched +/// slot: there a later pass took in more groups than the first, or this program and the +/// library disagree about a group at the gap's edge, and the three numbers are not exact. +pub fn work(name: &str, asked: &Asked, parts: &[&Part], questions: &[Vec]) { + measured(&format!("{name}_"), parts); + let mut reaches = Vec::with_capacity(questions.len()); + let (mut asked_for, mut screened, mut passes) = (Vec::new(), Vec::new(), Vec::new()); + let (mut out, mut again) = (0, 0); + for (at, question) in questions.iter().enumerate() { + let mut reached = reached(parts, question); + let rows = asked.examined[at] as i64 - asked.answers[at].len() as i64; + let counted = asked.evaluations[at] as i64 - rows; + let a_pass = reached.centres + reached.sketches; + let made = match reached.sketches { + 0 => 1, + _ => (counted / (a_pass + reached.unsketched)).max(1), + }; + let left = counted - made * a_pass; + let most = reached.unsketched + reached.sketches; + out += usize::from(left < reached.unsketched || left > most); + again += usize::from(made > 1); + reached.centres *= made; + reached.sketches *= made; + asked_for.push(rows); + screened.push(left); + passes.push(made); + reaches.push(reached); + } + let of = |what: fn(&Reached) -> i64| mean(reaches.iter().map(what)); + let name = |what: &str| format!("{name}_{what}"); + let hundred = format!("of_{}", questions.len()); + line(&name("passes_mean"), mean(passes.into_iter()), "passes"); + line(&name("searches_of_more_passes"), again, &hundred); + line(&name("centres_compared_mean"), of(|r| r.centres), "centres"); + line(&name("groups_scanned_mean"), of(|r| r.groups), "groups"); + line( + &name("slots_in_scanned_groups_mean"), + of(|r| r.slots), + "slots", + ); + line( + &name("sketches_estimated_mean"), + of(|r| r.sketches), + "sketches", + ); + line( + &name("slots_screened_mean"), + mean(screened.into_iter()), + "slots", + ); + line( + &name("rows_asked_mean"), + mean(asked_for.into_iter()), + "rows", + ); + line(&name("work_out_of_range"), out, &hundred); +} diff --git a/v2/contextdb/src/database/mod.rs b/v2/contextdb/src/database/mod.rs index 939ed370..b21f800e 100644 --- a/v2/contextdb/src/database/mod.rs +++ b/v2/contextdb/src/database/mod.rs @@ -1,8 +1,9 @@ //! database: the handle a store is opened for writing with, its transactions, and the one //! commit function. //! -//! It never has a second way to commit: a statement run on its own, a transaction and a change -//! of the memory limit each build a change set and hand it to the same function. +//! It never has a second way to commit: a statement run on its own, a transaction, a change of +//! the memory limit and every step of a maintenance cycle each build a change set and hand it +//! to the same function. #[cfg(test)] mod tests; @@ -496,38 +497,99 @@ impl Database { locked(&self.must_be_open().newest).commit } + /// What a statement would do as of the newest commit, without running it, as text of one + /// `name: value` to a line. For a similarity search it names the mode asked for and the one + /// in force, and for each partition the filter selects how it is searched: every vector of + /// it, or the groups of vectors nearest the one searched for. Nothing is read but the + /// store's small records, and nothing changes. A transaction has no such call. + /// + /// # Errors + /// + /// Whatever refuses the statement before it reads a row, by its own kind; + /// [`crate::Error::StoreIo`] or [`crate::Error::StoreCorrupted`] when the store cannot be + /// read; [`crate::Error::DatabaseClosed`] on a handle that was closed. + pub fn explain(&self, sql: &str) -> Result { + let plan = planned(sql)?; + let open = self.opened()?; + let (view, catalog) = open.newest()?; + let now = system::now_millis(&*open.clock); + query::explain(&plan, &catalog, &view, &open.budget, now) + } + /// 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. + /// A cycle first moves at most 4,096 vectors from the pending lists of VECTOR columns into + /// blocks, in one commit. Then every partition that has a grouping due or being built is + /// advanced by one step, each step in a commit of its own: a build begins, moves its centres + /// once over a sample of at most 64 vectors a group, gives the vectors of whole blocks their + /// groups in at most 4,194,304 comparisons, packs whole groups of at most 4,096 vectors + /// together into new blocks, measures at most 4,194,304 vectors or sketches against its + /// questions, or puts the new grouping in force. A step that goes past its bound does so + /// by one block, one group or one question. A store with nothing to do 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. + /// Searches give the same rows before and after a cycle; they read fewer rows after it. + /// + /// Under a memory limit with no room to seal everything, a cycle seals the partitions it has + /// room for, at least one. A step of a build that has no room is not an error: the cycle + /// records the bytes the step needed, reports the partition as stalled, and goes on. /// /// 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. + /// calls this after it has written, and again as [`crate::VectorMaintenanceReport`] says. + /// A step is worked out with no lock held, so other calls on the handle wait only for a + /// step's commit, as they wait for any other commit, and never for its arithmetic; 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. + /// plus 24), the partition's record, and one row with its vector at a time; then nothing + /// was sealed. [`crate::Error::StoreIo`] or [`crate::Error::StoreCorrupted`] when the store + /// cannot be read or written; [`crate::Error::DatabaseClosed`] on a handle that was closed. + /// The commits made before the error stay, 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)?; + let mut report = { + 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)?; + report + }; + let (view, catalog) = open.newest()?; + let mut listed = open.budget.reserve(0, "maintenance")?; + for work in vector::waiting(&view, &catalog, &mut listed)? { + // Worked out over the newest commit with no lock held; staged under the commit + // lock, where the vector part checks that the build has not moved since. + let (view, catalog) = open.newest()?; + let step = vector::prepare(&work, &view, &catalog, &open.budget); + let held = locked(&open.committing); + let mut staged = open.begin()?; + let done = step.and_then(|step| match step { + Some(step) => vector::stage_step(step, &mut staged.changes), + None => Ok(0), + }); + match done { + Ok(done) => report.vector.maintenance_vectors_done += done, + // A step that does not fit, worked out or staged, is recorded in a change set + // of its own and is no error of the cycle. + Err(Error::MemoryBudgetExceeded { + requested_bytes, .. + }) => { + staged = open.begin()?; + vector::stage_refusal(&work, requested_bytes, &mut staged.changes)?; + } + Err(other) => return Err(other), + } + open.commit(&held, staged)?; + } + let (view, catalog) = open.newest()?; + let left = vector::remaining(&view, &catalog)?; + report.vector.maintenance_vectors_remaining += left.vectors; + report.vector.remaining_partitions = left.partitions; + report.vector.stalled_partitions = left.stalled; Ok(report) } diff --git a/v2/contextdb/src/query/inspect.rs b/v2/contextdb/src/query/inspect.rs new file mode 100644 index 00000000..5f12b0bb --- /dev/null +++ b/v2/contextdb/src/query/inspect.rs @@ -0,0 +1,423 @@ +//! Inspection of VECTOR columns: what `SHOW VECTOR_INDEXES` and `SHOW VECTOR_PARTITIONS` return +//! and what explaining a similarity search says. All three read the vector part's small records +//! and no block, no centre and no row, and none changes anything. + +use std::cell::Cell; +use std::collections::HashMap; + +use serde_json::json; + +use super::QueryResult; +use super::search::Scope; +use crate::memory::Budget; +use crate::schema::{Catalog, Consolidation, Quantization, Table, VectorColumn}; +use crate::sql::{Plan, not_supported}; +use crate::store::{Key, Read, VectorNumber}; +use crate::types::{Error, Result, Value, VectorIndexRef, VectorSearchMode}; +use crate::vector::{self, Inspected, Nearest, Route, State}; + +const INDEX_FIELDS: [&str; 31] = [ + "table", + "column", + "dimension", + "quantization", + "vector_count", + "bytes", + "partition_key_columns", + "max_partitions", + "live_partitions", + "search_mode", + "declared_auto_index_at", + "effective_auto_index_at", + "declared_consolidation_mode", + "declared_consolidation_change_percent", + "declared_consolidation_tombstone_percent", + "effective_consolidation_mode", + "effective_consolidation_change_percent", + "effective_consolidation_tombstone_percent", + "pending_inserts", + "tombstones", + "durable_vector_bytes", + "durable_index_bytes", + "indexed_available", + "maintenance_state", + "maintenance_vectors_total", + "maintenance_vectors_done", + "maintenance_vectors_remaining", + "unavailable_partitions", + "stalled_partitions", + "broad_route", + "recovery_action", +]; + +const PARTITION_FIELDS: [&str; 18] = [ + "table", + "column", + "partition_key", + "live_rows", + "pending_inserts", + "tombstones", + "durable_vector_bytes", + "durable_index_bytes", + "indexed_available", + "availability_reason", + "maintenance_state", + "maintenance_reason", + "maintenance_vectors_total", + "maintenance_vectors_done", + "maintenance_vectors_remaining", + "lists", + "probes", + "recovery_action", +]; + +/// Answers `SHOW VECTOR_INDEXES` and `SHOW VECTOR_PARTITIONS`; `None` for any other plan. +/// +/// `SHOW VECTOR_INDEXES` gives one row for each VECTOR column of the store, in the order of +/// table name and column; `SHOW VECTOR_PARTITIONS` one row for each partition, of the column +/// after `FOR` or of every column, in the order of table, column and key tuple, after `OFFSET` +/// rows and at most `LIMIT`. +pub(super) fn show( + plan: &Plan, + catalog: &Catalog, + read: &dyn Read, + budget: &Budget, +) -> Result> { + // The partitions read are held until their rows are made. + let held = &mut budget.reserve(0, "query")?; + let (fields, rows) = match plan { + Plan::ShowVectorIndexes => { + let mut rows = Vec::new(); + for (table, column) in columns(catalog, None)? { + let parts = + vector::partitions(read, column, table.dimension(column), 0, None, held)?; + rows.push(index_row(table, column, &parts)); + } + (&INDEX_FIELDS[..], rows) + } + Plan::ShowVectorPartitions { + column, + limit, + offset, + } => { + let (mut skip, mut left, mut rows) = (*offset, *limit, Vec::new()); + for (table, column) in columns(catalog, column.as_ref())? { + let n = table.dimension(column); + let parts = vector::partitions(read, column, n, skip, left, held)?; + // A column with no PARTITION_KEY always has its one partition. + let has = read.vector_number(VectorNumber::Partitions(column.id))?; + skip = skip.saturating_sub(has.max(u64::from(column.partition_key.is_empty()))); + left = left.map(|left| left - parts.len() as u64); + let shown = parts.iter(); + rows.extend(shown.map(|part| partition_row(table, column, part))); + } + (&PARTITION_FIELDS[..], rows) + } + _ => return Ok(None), + }; + let fields = fields.iter().map(|name| (*name).to_owned()).collect(); + Ok(Some(QueryResult::from_rows(fields, rows))) +} + +/// The VECTOR column after `FOR`, or every VECTOR column in the order of table name and +/// column. +fn columns<'a>( + catalog: &'a Catalog, + named: Option<&(String, String)>, +) -> Result> { + if let Some((table, column)) = named { + let table = catalog + .table(table) + .ok_or_else(|| Error::TableNotFound(table.clone()))?; + let index = VectorIndexRef { + table: table.name.clone(), + column: column.clone(), + }; + let column = table.vector(column); + return Ok(vec![( + table, + column.ok_or(Error::UnknownVectorIndex { index })?, + )]); + } + let mut found = Vec::new(); + for table in catalog.tables() { + let mut vectors: Vec<_> = table.vectors.iter().map(|column| (table, column)).collect(); + vectors.sort_by(|left, right| left.1.column.cmp(&right.1.column)); + found.extend(vectors); + } + Ok(found) +} + +fn whole(number: u64) -> Value { + Value::Int64(i64::try_from(number).unwrap_or(i64::MAX)) +} + +fn text(word: &str) -> Value { + Value::Text(word.to_owned()) +} + +fn mode_name(mode: VectorSearchMode) -> &'static str { + match mode { + VectorSearchMode::Auto => "AUTO", + VectorSearchMode::Exact => "EXACT", + VectorSearchMode::Indexed => "INDEXED", + } +} + +fn state_name(state: State) -> &'static str { + match state { + State::Idle => "idle", + State::Pending => "pending", + State::Building => "building", + State::Stalled => "stalled", + } +} + +/// What brings a partition, or a column, out of this state. +fn recovery(state: State) -> &'static str { + match state { + State::Idle => "none", + State::Pending | State::Building => "run a maintenance cycle", + State::Stalled => "raise the memory limit", + } +} + +/// The names of the PARTITION_KEY columns as a JSON array. +fn key_columns(column: &VectorColumn) -> String { + json!(column.partition_key).to_string() +} + +/// A partition's key tuple as a JSON object of the key columns, in declared order, and their +/// values. +fn key_object(column: &VectorColumn, tuple: &Key) -> String { + let values = tuple.values().unwrap_or_default(); + let members = column + .partition_key + .iter() + .zip(&values) + .map(|(name, value)| { + let value = match value { + Value::Bool(value) => json!(value), + Value::Int64(value) | Value::Timestamp(value) => json!(value), + Value::Float64(value) => json!(value), + Value::Text(value) => json!(value), + Value::Uuid(value) => json!(value.to_string()), + _ => serde_json::Value::Null, + }; + format!("{}:{value}", json!(name)) + }); + format!("{{{}}}", members.collect::>().join(",")) +} + +/// The row of `SHOW VECTOR_INDEXES` for one column, from all its partitions. +fn index_row(table: &Table, column: &VectorColumn, parts: &[Inspected]) -> Vec { + let sum = |of: fn(&Inspected) -> u64| parts.iter().map(of).sum::(); + let count = |of: fn(&Inspected) -> bool| parts.iter().filter(|part| of(part)).count() as u64; + let optional = |number: Option| number.map_or(Value::Null, whole); + let percent = |number: Option| optional(number.map(u64::from)); + let (declared, change, tombstone) = match column.consolidation { + None => (Value::Null, None, None), + Some(Consolidation::None) => (text("NONE"), None, None), + Some(Consolidation::Percent { change, tombstone }) => (text("PERCENT"), change, tombstone), + }; + let in_force = column.consolidation(); + let state = parts.iter().map(|part| part.state).max(); + let state = state.unwrap_or(State::Idle); + let (slot_bytes, index_bytes) = (sum(|part| part.slot_bytes), sum(|part| part.index_bytes)); + let (total, done) = (sum(|part| part.total), sum(|part| part.done)); + let unavailable = count(|part| !part.usable); + vec![ + text(&table.name), + text(&column.column), + whole(table.dimension(column) as u64), + text(match column.quantization { + Quantization::F32 => "F32", + Quantization::Sq8 => "SQ8", + }), + whole(sum(|part| part.live)), + whole(slot_bytes + index_bytes), + text(&key_columns(column)), + optional(column.limit().map(u64::from)), + whole(count(|part| part.live > 0)), + text(mode_name(column.mode())), + optional(column.auto_index_at), + whole(column.index_at()), + declared, + percent(change), + percent(tombstone), + text(in_force.map_or("NONE", |_| "PERCENT")), + percent(in_force.map(|(change, _)| change)), + percent(in_force.map(|(_, tombstone)| tombstone)), + whole(sum(|part| part.pending)), + whole(sum(|part| part.stale)), + whole(slot_bytes), + whole(index_bytes), + Value::Bool(unavailable == 0), + text(state_name(state)), + whole(total), + whole(done), + whole(total - done), + whole(unavailable), + whole(count(|part| part.state == State::Stalled)), + text("fan_out"), + text(recovery(state)), + ] +} + +/// The row of `SHOW VECTOR_PARTITIONS` for one partition of a column. +fn partition_row(table: &Table, column: &VectorColumn, part: &Inspected) -> Vec { + // Why the grouping is due or being built: the slots rows have left behind, when they + // alone are enough; else what has changed since. + let reason = match part.state { + _ if part.damage.is_some() => "damaged", + State::Idle => "none", + State::Stalled => "memory_limit", + _ if part.groups == 0 => "first_grouping", + _ if part.left_behind => "tombstones", + _ => "changes", + }; + let grouping = |number: u32| match part.groups { + 0 => Value::Null, + _ => whole(u64::from(number)), + }; + vec![ + text(&table.name), + text(&column.column), + text(&key_object(column, &part.tuple)), + whole(part.live), + whole(part.pending), + whole(part.stale), + whole(part.slot_bytes), + whole(part.index_bytes), + Value::Bool(part.usable), + text(if part.usable { "none" } else { "behind_writes" }), + text(state_name(part.state)), + text(reason), + whole(part.total), + whole(part.done), + whole(part.total - part.done), + grouping(part.groups), + grouping(part.probes), + text(recovery(part.state)), + ] +} + +/// What a statement would do, as text of one `name: value` to a line. For a similarity search: +/// the mode asked for and the one in force, the column's `AUTO_INDEX_AT`, how many partitions +/// the filter selects, and for each of them how it is searched; then what refuses the search, +/// or `none`. The search is not run. A partition named through a `$name` has no value here: +/// every partition is then listed, and the line of their count says so. +pub(crate) fn explain( + plan: &Plan, + catalog: &Catalog, + read: &dyn Read, + budget: &Budget, + now: i64, +) -> Result { + let plan = match plan { + Plan::Project { input, .. } => &**input, + plan => plan, + }; + let Plan::VectorSearch { + table, + column, + k, + filter, + mode, + .. + } = plan + else { + return Err(not_supported("explaining this statement")); + }; + let table = catalog + .table(table) + .ok_or_else(|| Error::TableNotFound(table.clone()))?; + let index = VectorIndexRef { + table: table.name.clone(), + column: column.clone(), + }; + let Some(vector) = table.vector(column) else { + return Err(Error::UnknownVectorIndex { index }); + }; + let mode = mode.unwrap_or_else(|| vector.mode()); + let key = vector.partition_key.iter(); + let key = key.filter_map(|name| table.meta.columns.iter().position(|c| c.name == *name)); + let scope = Scope { + table, + key: key.collect(), + params: &HashMap::new(), + now, + unbound: Cell::default(), + }; + let mut keys = scope.keys(filter.as_ref())?; + let every = scope.unbound.take(); + if every.is_some() { + keys.clear(); + } + // The list is let go with the text made from it. + let mut reservation = budget.reserve(0, "query")?; + let planned = vector::plan( + &Nearest { + table: table.id, + vector: vector.id, + column: vector, + mode, + scope: &keys, + query: &[], + k: usize::try_from(*k).unwrap_or(usize::MAX), + // Whether rows are turned away is known only when the search runs. + filtered: false, + }, + read, + &mut reservation, + )?; + + let indexed = mode == VectorSearchMode::Indexed; + let allowed: u64 = planned.iter().map(|partition| partition.live).sum(); + // AUTO over fewer vectors than the column's size is EXACT, as the search settles it. + let resolved = match mode { + VectorSearchMode::Auto if allowed < vector.index_at() => VectorSearchMode::Exact, + mode => mode, + }; + let mut lines = vec![ + format!("requested_mode: {}", mode_name(mode)), + format!("resolved_mode: {}", mode_name(resolved)), + format!("effective_auto_index_at: {}", vector.index_at()), + format!("aggregate_allowed_vectors: {allowed}"), + format!("partition_key_columns: {}", key_columns(vector)), + match every { + None => format!("partitions: {}", planned.len()), + Some(name) => format!("partitions: every partition (${name} is not bound)"), + }, + ]; + for partition in &planned { + let route = match partition.route { + Route::Exact => "exact", + Route::Indexed => "indexed", + }; + let grouping = |number: u32| match partition.groups { + 0 => "none".to_owned(), + _ => number.to_string(), + }; + lines.extend([ + format!("partition: {}", key_object(vector, &partition.tuple)), + format!(" route: {route}"), + format!(" lists: {}", grouping(partition.groups)), + format!(" probes: {}", grouping(partition.probes)), + format!(" live_rows: {}", partition.live), + format!(" outside_grouping: {}", partition.tail + partition.pending), + ]); + } + // A search under INDEXED with a partition that is not usable is refused whole. + lines.push( + if indexed && planned.iter().any(|partition| !partition.usable) { + format!( + "refusal: {}", + Error::VectorIndexedRouteUnavailable { index } + ) + } else { + "refusal: none".to_owned() + }, + ); + Ok(lines.join("\n")) +} diff --git a/v2/contextdb/src/query/mod.rs b/v2/contextdb/src/query/mod.rs index ad87e195..c05f6001 100644 --- a/v2/contextdb/src/query/mod.rs +++ b/v2/contextdb/src/query/mod.rs @@ -7,6 +7,7 @@ mod access; mod expression; +mod inspect; mod search; #[cfg(test)] mod tests; @@ -16,6 +17,7 @@ use std::collections::HashMap; use std::fmt; pub(crate) use expression::eval; +pub(crate) use inspect::explain; pub(crate) use search::vector_from_text; use crate::memory::{Budget, Reservation}; @@ -37,6 +39,7 @@ pub struct QueryResult { rows: Vec>, rows_affected: u64, rows_examined: u64, + distance_evaluations: u64, /// Never read: held so that the memory the rows take stays counted against the limit, and /// dropping it with the result gives that memory back. _reservation: Option, @@ -49,6 +52,7 @@ impl fmt::Debug for QueryResult { .field("rows", &self.rows) .field("rows_affected", &self.rows_affected) .field("rows_examined", &self.rows_examined) + .field("distance_evaluations", &self.distance_evaluations) .finish_non_exhaustive() } } @@ -76,6 +80,18 @@ impl QueryResult { self.rows_examined } + /// How many distance evaluations a similarity search made to produce this result: the + /// group centres it compared with the vector searched for, the sketches it estimated, the + /// stored vectors it screened through their one-byte copies, and the rows whose distance + /// it worked out. An exact count, the same on every run over the same store, and zero for + /// any other statement. A search that used the index usually made fewer than the same + /// search under `USE VECTOR EXACT`, and can make more: where it leaves few groups out, or + /// goes over its sketches again because rows were turned away, the sketches it estimates + /// come to more than the stored vectors EXACT screens. + pub fn distance_evaluations(&self) -> u64 { + self.distance_evaluations + } + /// The result of a write that changed `rows` rows. pub(crate) fn affected(rows: u64) -> QueryResult { QueryResult { @@ -92,6 +108,7 @@ impl QueryResult { rows, rows_affected: 0, rows_examined: 0, + distance_evaluations: 0, _reservation: None, } } @@ -146,6 +163,9 @@ pub(crate) fn run( if let Some(found) = search::run(plan, params, catalog, read, budget, now)? { return Ok(found); } + if let Some(shown) = inspect::show(plan, catalog, read, budget)? { + return Ok(shown); + } let select = Select::of(plan)?; let table = catalog .table(select.table) @@ -265,6 +285,7 @@ pub(crate) fn run( rows, rows_affected: 0, rows_examined: examined, + distance_evaluations: 0, _reservation: Some(reservation), }) } diff --git a/v2/contextdb/src/query/search.rs b/v2/contextdb/src/query/search.rs index 757c304c..5a9be8f8 100644 --- a/v2/contextdb/src/query/search.rs +++ b/v2/contextdb/src/query/search.rs @@ -3,6 +3,7 @@ //! which rows may be seen, and works out each row's distance from its stored vector; the vector //! part finds the rows. +use std::cell::Cell; use std::collections::HashMap; use super::access::{Source, as_key}; @@ -13,7 +14,7 @@ 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::types::{ColumnType, Error, Result, Value, VectorIndexRef}; use crate::vector::{self, Nearest}; /// The most key tuples a filter may name; past that the search reads every partition. @@ -68,9 +69,6 @@ pub(super) fn run( 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), @@ -124,8 +122,9 @@ pub(super) fn run( key, params, now, - } - .keys(filter)?; + unbound: Cell::default(), + }; + let (scope, on_rows) = (scope.keys(filter)?, scope.on_rows(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]; @@ -141,9 +140,12 @@ pub(super) fn run( let nearest = Nearest { table: table.id, vector: vector.id, + column: vector, + mode: mode.unwrap_or_else(|| vector.mode()), scope: &scope, query: &near, k, + filtered: on_rows, }; let mut distance_of = |key: &Key| -> Result> { examined += 1; @@ -157,7 +159,14 @@ pub(super) fn run( reservation.shrink(bytes); Ok(distance) }; - let hits = vector::search(&nearest, read, &mut searching, &mut distance_of)?; + // `None` is a search under INDEXED of a partition whose grouping is behind its writes, or + // one whose filter on rows passes fewer rows than it asks for. + let Some(found) = vector::search(&nearest, read, &mut searching, &mut distance_of)? else { + return Err(Error::VectorIndexedRouteUnavailable { index }); + }; + // Every row the vector part asked about had its distance worked out or was turned away. + let asked_for = examined; + let hits = found.hits; let moved = returned.moved(list)?; let mut rows = Vec::new(); @@ -173,6 +182,7 @@ pub(super) fn run( rows, rows_affected: 0, rows_examined: examined, + distance_evaluations: found.evaluations + asked_for, _reservation: Some(reservation), })) } @@ -182,21 +192,82 @@ pub(super) fn run( type Tuple = Vec>; /// Reads a filter for the partitions it selects. -struct Scope<'a> { - table: &'a Table, +pub(super) struct Scope<'a> { + pub(super) table: &'a Table, /// Where each PARTITION_KEY column is in a row, in declared order. - key: Vec, - params: &'a HashMap, - now: i64, + pub(super) key: Vec, + pub(super) params: &'a HashMap, + pub(super) now: i64, + /// The `$name` beside a key column that nothing is bound to, once one was met; its column + /// is left open. A search refuses such a name before it asks; `explain` binds nothing, + /// lists every partition and names it. + pub(super) unbound: Cell>, } impl Scope<'_> { + /// True when `filter` does more than select partitions, so that some row of the + /// partitions [`Scope::keys`] selects may not pass it: a filter on rows. False with no + /// filter, and for a filter made only of the selections of key tuples `keys` understood. + /// It fills `Nearest::filtered`: a search with a filter on rows uses the index as any + /// other does, and only the refusal of INDEXED asks whether there is one. + pub(super) fn on_rows(&self, filter: Option<&Expr>) -> Result { + let Some(filter) = filter else { + return Ok(false); + }; + // Each tuple selects whole partitions when the columns it fixes are the leading ones. + let leading = |tuple: &Tuple| { + let fixed = tuple.iter().take_while(|value| value.is_some()).count(); + fixed > 0 && tuple[fixed..].iter().all(Option::is_none) + }; + let tuples = self.said(filter)?; + let selects = tuples.len() <= MOST_TUPLES && tuples.iter().all(leading); + Ok(!selects || !self.only_selects(filter)?) + } + + /// True when every part of `filter` is one [`Scope::said`] reads whole: `AND` of two + /// sides that fix different key columns, `OR`, and `=` or `IN` whose every pair fixes a + /// key column. + fn only_selects(&self, filter: &Expr) -> Result { + let whole = |columns: &Expr, values: &Expr| -> Result { + let (columns, values) = (sides(columns), sides(values)); + let fixed = self.tuple(columns, values)?.iter().flatten().count(); + Ok(fixed == columns.len() && fixed == values.len()) + }; + Ok(match filter { + Expr::BinaryOp { left, op, right } if matches!(op, BinOp::And | BinOp::Or) => { + let (one, other) = (self.said(left)?, self.said(right)?); + let fixes = + |said: &[Tuple], at: usize| said.iter().any(|tuple| tuple[at].is_some()); + let clash = (0..self.key.len()).any(|at| fixes(&one, at) && fixes(&other, at)); + !(clash && *op == BinOp::And) + && self.only_selects(left)? + && self.only_selects(right)? + } + Expr::BinaryOp { left, op, right } if *op == BinOp::Eq => match **left { + Expr::Column(_) => whole(left, right)?, + _ => whole(right, left)?, + }, + Expr::InList { + expr, + list, + negated, + } => { + let mut every = !*negated; + for item in list { + every &= whole(expr, item)?; + } + every + } + _ => false, + }) + } + /// 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> { + pub(super) fn keys(&self, filter: Option<&Expr>) -> Result> { let (Some(filter), false) = (filter, self.key.is_empty()) else { return Ok(Vec::new()); }; @@ -272,9 +343,17 @@ impl Scope<'_> { 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); + let Some(place) = self.key.iter().position(|key| *key == at) else { + continue; + }; + match value { + Expr::Parameter(name) if !self.params.contains_key(name) => { + self.unbound.set(Some(name.clone())); + } + value => { + let value = eval(value, self.table, &[], self.params, self.now)?; + tuple[place] = as_key(self.table.meta.columns[at].column_type, value); + } } } Ok(tuple) diff --git a/v2/contextdb/src/read/mod.rs b/v2/contextdb/src/read/mod.rs index 6b751556..1cce806d 100644 --- a/v2/contextdb/src/read/mod.rs +++ b/v2/contextdb/src/read/mod.rs @@ -133,6 +133,19 @@ impl ReadSession { query::run(&plan, params, &catalog, &view, &self.budget, now) } + /// What a statement would do, without running it: as [`crate::Database::explain`], for + /// the newest commit this session can read. + /// + /// # Errors + /// + /// As [`ReadSession::execute`] for a statement that reads. + pub fn explain(&self, sql: &str) -> Result { + let plan = sql::plan(&sql::parse(sql)?)?; + let now = system::now_millis(&*self.clock); + let (view, catalog) = self.newest()?; + query::explain(&plan, &catalog, &view, &self.budget, now) + } + /// The names of the store's tables as of the newest commit, in name order. A store with no /// table gives an empty list. /// diff --git a/v2/contextdb/src/schema/check.rs b/v2/contextdb/src/schema/check.rs index b48abbac..29b79e1d 100644 --- a/v2/contextdb/src/schema/check.rs +++ b/v2/contextdb/src/schema/check.rs @@ -1,14 +1,14 @@ //! Checking a declaration: what CREATE TABLE and CREATE INDEX may say, and the table each gives. -use super::{Catalog, RESERVED_PREFIXES, Table, VectorColumn}; -use crate::sql::ast::{self, DataType, Expr, Literal}; +use super::{Catalog, Consolidation, Quantization, RESERVED_PREFIXES, Table, VectorColumn}; +use crate::sql::ast::{self, ConsolidationClause, DataType, Expr, Literal}; use crate::sql::not_supported; use crate::sql::plan::{CreateIndexPlan, CreateTablePlan}; use crate::store::{IndexId, TableId, VectorId}; use crate::types::{ ColumnDef, ColumnDefault, ColumnType, Error, IndexDecl, Result, SortDirection, TableMeta, Value, VectorIndexRef, VectorPartitionDeclarationIssue as PartitionIssue, - VectorPolicyDeclarationIssue, + VectorPolicyDeclarationIssue as PolicyIssue, }; fn invalid(reason: String) -> Error { @@ -109,18 +109,63 @@ fn vector_column( table: table.to_owned(), column: declared.name.clone(), }; + let policy = |issue| Error::InvalidVectorPolicyDeclaration { + index: index.clone(), + issue, + }; if declared.data_type == DataType::Vector(0) { - let issue = VectorPolicyDeclarationIssue::DimensionsNotPositive; - return Err(Error::InvalidVectorPolicyDeclaration { index, issue }); + return Err(policy(PolicyIssue::DimensionsNotPositive)); } - let max_partitions = partitions(declared, columns) - .map_err(|issue| Error::InvalidVectorPartitionDeclaration { index, issue })?; + let max_partitions = partitions(declared, columns).map_err(|issue| { + Error::InvalidVectorPartitionDeclaration { + index: index.clone(), + issue, + } + })?; + let written = declared.vector.quantization.as_deref(); + let quantization = match written.map(str::to_ascii_uppercase).as_deref() { + None | Some("F32") => Quantization::F32, + Some("SQ8") => Quantization::Sq8, + Some("SQ4") => return Err(not_supported("WITH (quantization = 'SQ4')")), + Some(_) => { + let member = "quantization"; + return Err(policy(PolicyIssue::InvalidMember { member })); + } + }; + let auto_index_at = match &declared.vector.auto_index_at { + None => None, + Some(written) => match written.parse::() { + Ok(at) if at > 0 => Some(at), + Err(_) if !written.starts_with('-') => { + return Err(policy(PolicyIssue::AutoIndexAtOutOfRange)); + } + _ => return Err(policy(PolicyIssue::AutoIndexAtNotPositive)), + }, + }; + let percent = |written: &Option| match written { + None => Ok(None), + Some(written) => match written.parse::() { + Ok(percent @ 1..=100) => Ok(Some(percent)), + _ => Err(policy(PolicyIssue::ConsolidationPercentOutOfRange)), + }, + }; + let consolidation = match &declared.vector.consolidation { + None => None, + Some(ConsolidationClause::None) => Some(Consolidation::None), + Some(ConsolidationClause::Percent { change, tombstone }) => Some(Consolidation::Percent { + change: percent(change)?, + tombstone: percent(tombstone)?, + }), + }; Ok(VectorColumn { id, column: declared.name.clone(), partition_key: declared.vector.partition_key.clone(), max_partitions, search_mode: declared.vector.search_mode, + auto_index_at, + consolidation, + quantization, }) } diff --git a/v2/contextdb/src/schema/mod.rs b/v2/contextdb/src/schema/mod.rs index 34260dda..312632f1 100644 --- a/v2/contextdb/src/schema/mod.rs +++ b/v2/contextdb/src/schema/mod.rs @@ -57,8 +57,48 @@ pub(crate) struct VectorColumn { pub max_partitions: Option, /// SEARCH_MODE as declared; `None` when the declaration names none. pub search_mode: Option, + /// AUTO_INDEX_AT as declared; `None` when the declaration names none. + pub auto_index_at: Option, + /// CONSOLIDATION as declared; `None` when the declaration names none. + pub consolidation: Option, + /// How the column keeps a vector: what `WITH (quantization = ...)` declared. + pub quantization: Quantization, } +/// What `CONSOLIDATION` declares: when the vectors written since a partition was grouped, and +/// the slots its rows have left behind, call for grouping it again. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Consolidation { + /// `CONSOLIDATION NONE`: never for those reasons. + None, + /// `CONSOLIDATION (CHANGE_PERCENT = n, TOMBSTONE_PERCENT = n)`; a member that was not + /// written is `None` and takes its default. + Percent { + change: Option, + tombstone: Option, + }, +} + +/// How a VECTOR column keeps its vectors. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub(crate) enum Quantization { + /// Every number as it was written, in 32 bits. + #[default] + F32, + /// One scale and one byte a number: a write keeps the nearest such vector. + Sq8, +} + +/// AUTO_INDEX_AT of a column that declares none: a comparison of every vector of this many +/// takes about a hundredth of a second on one core, and past it the groupings are worth using. +const DEFAULT_INDEX_AT: u64 = 100_000; +/// A partition is given its first grouping when it holds this share of AUTO_INDEX_AT: one in +/// a hundred, so that the groupings of a store spread over many partitions are in place before +/// AUTO begins to use them. +const GROUPED_SHARE: u64 = 100; +/// CHANGE_PERCENT and TOMBSTONE_PERCENT of a column that declares none. +const DEFAULT_CONSOLIDATION: (u8, u8) = (20, 10); + 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, @@ -72,6 +112,32 @@ impl VectorColumn { pub(crate) fn mode(&self) -> VectorSearchMode { self.search_mode.unwrap_or_default() } + + /// AUTO_INDEX_AT in force: what was declared, else 100,000. The count of vectors in the + /// selected partitions together from which AUTO uses their groupings. + pub(crate) fn index_at(&self) -> u64 { + self.auto_index_at.unwrap_or(DEFAULT_INDEX_AT) + } + + /// How many slots a partition's tail holds when its first grouping is due: a hundredth of + /// AUTO_INDEX_AT in force and one at least, so 1,000 where nothing is declared. + pub(crate) fn grouped_from(&self) -> u64 { + (self.index_at() / GROUPED_SHARE).max(1) + } + + /// CHANGE_PERCENT and TOMBSTONE_PERCENT in force, 20 and 10 where none was declared; + /// `None` under `CONSOLIDATION NONE`. + pub(crate) fn consolidation(&self) -> Option<(u8, u8)> { + let (change, tombstone) = DEFAULT_CONSOLIDATION; + match self.consolidation { + None => Some((change, tombstone)), + Some(Consolidation::None) => None, + Some(Consolidation::Percent { + change: c, + tombstone: t, + }) => Some((c.unwrap_or(change), t.unwrap_or(tombstone))), + } + } } impl Table { @@ -81,6 +147,16 @@ impl Table { self.vectors.iter().find(|vector| vector.column == column) } + /// How many numbers a vector of this VECTOR column has. + pub(crate) fn dimension(&self, column: &VectorColumn) -> usize { + let columns = self.meta.columns.iter(); + let mut sizes = columns.filter_map(|declared| match declared.column_type { + ColumnType::Vector(size) if declared.name == column.column => Some(size), + _ => None, + }); + sizes.next().unwrap_or_default() + } + /// The column as a refusal names it. pub(crate) fn vector_ref(&self, vector: &VectorColumn) -> VectorIndexRef { VectorIndexRef { diff --git a/v2/contextdb/src/schema/record.rs b/v2/contextdb/src/schema/record.rs index a74cf567..4a868acc 100644 --- a/v2/contextdb/src/schema/record.rs +++ b/v2/contextdb/src/schema/record.rs @@ -8,7 +8,7 @@ //! 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, VectorColumn}; +use super::{Consolidation, Quantization, Table, VectorColumn}; use crate::store::{IndexId, TableId, VectorId}; use crate::types::{ ColumnDef, ColumnDefault, ColumnType, IndexDecl, SortDirection, TableMeta, Value, @@ -94,6 +94,16 @@ pub(super) fn encode(table: &Table) -> Vec { Some(VectorSearchMode::Exact) => 2, Some(VectorSearchMode::Indexed) => 3, }); + // A declared AUTO_INDEX_AT and a declared percentage are never zero either. + out.extend(vector.auto_index_at.unwrap_or(0).to_le_bytes()); + out.extend(match vector.consolidation { + None => [0, 0, 0], + Some(Consolidation::None) => [1, 0, 0], + Some(Consolidation::Percent { change, tombstone }) => { + [2, change.unwrap_or(0), tombstone.unwrap_or(0)] + } + }); + out.push(u8::from(vector.quantization == Quantization::Sq8)); } out } @@ -162,6 +172,20 @@ pub(super) fn decode(record: &[u8]) -> Table { 2 => Some(VectorSearchMode::Exact), _ => Some(VectorSearchMode::Indexed), }, + auto_index_at: Some(u64::from_le_bytes(read.array())).filter(|at| *at != 0), + consolidation: match read.array::<3>() { + [0, ..] => None, + [1, ..] => Some(Consolidation::None), + [_, change, tombstone] => Some(Consolidation::Percent { + change: Some(change).filter(|percent| *percent != 0), + tombstone: Some(tombstone).filter(|percent| *percent != 0), + }), + }, + quantization: if read.flag() { + Quantization::Sq8 + } else { + Quantization::F32 + }, }); Table { id, diff --git a/v2/contextdb/src/schema/render.rs b/v2/contextdb/src/schema/render.rs index 5920090f..b4113a05 100644 --- a/v2/contextdb/src/schema/render.rs +++ b/v2/contextdb/src/schema/render.rs @@ -1,6 +1,6 @@ //! Printing a declaration back as SQL. -use super::{RESERVED_PREFIXES, Table}; +use super::{Consolidation, Quantization, RESERVED_PREFIXES, Table}; 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 @@ -58,8 +58,12 @@ fn column(table: &Table, column: &ColumnDef) -> String { ColumnType::Vector(size) => text.push_str(&format!("VECTOR({size})")), } if let Some(vector) = table.vector(&column.name) { + if vector.quantization == Quantization::Sq8 { + text.push_str(" WITH (quantization = 'SQ8')"); + } // 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. + // PARTITION_KEY, whether or not it was written. The clauses after it are printed only + // when they were written. if let Some(limit) = vector.limit() { let key: Vec = vector.partition_key.iter().map(|key| name(key)).collect(); text.push_str(&format!( @@ -73,6 +77,21 @@ fn column(table: &Table, column: &ColumnDef) -> String { Some(VectorSearchMode::Exact) => " SEARCH_MODE EXACT", Some(VectorSearchMode::Indexed) => " SEARCH_MODE INDEXED", }); + if let Some(at) = vector.auto_index_at { + text.push_str(&format!(" AUTO_INDEX_AT {at}")); + } + match vector.consolidation { + None => {} + Some(Consolidation::None) => text.push_str(" CONSOLIDATION NONE"), + Some(Consolidation::Percent { change, tombstone }) => { + let members = [("CHANGE_PERCENT", change), ("TOMBSTONE_PERCENT", tombstone)]; + let written = members.iter().filter_map(|(member, percent)| { + percent.map(|percent| format!("{member} = {percent}")) + }); + let written: Vec = written.collect(); + text.push_str(&format!(" CONSOLIDATION ({})", written.join(", "))); + } + } } 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 2166feb4..2f98af46 100644 --- a/v2/contextdb/src/schema/tests.rs +++ b/v2/contextdb/src/schema/tests.rs @@ -322,3 +322,60 @@ fn a_declared_table_gets_its_numbers_and_the_index_behind_each_rule() { let keyless = create("keyless", vec![declared("word", DataType::Text)], &[]); assert!(not_supported(&keyless, "a table without a PRIMARY KEY")); } + +#[test] +fn the_index_clauses_of_a_vector_column_are_stored_and_give_the_numbers_in_force() { + use super::{Consolidation, Quantization, VectorColumn}; + use crate::store::VectorId; + let column = |id, auto_index_at, consolidation, quantization| VectorColumn { + id: VectorId(id), + column: format!("v{id}"), + partition_key: Vec::new(), + max_partitions: None, + search_mode: None, + auto_index_at, + consolidation, + quantization, + }; + let percent = |change, tombstone| Some(Consolidation::Percent { change, tombstone }); + let vectors = vec![ + column(0, None, None, Quantization::F32), + column(1, Some(1000), Some(Consolidation::None), Quantization::Sq8), + column( + 2, + Some(u64::MAX), + percent(Some(50), None), + Quantization::F32, + ), + column(3, Some(1), percent(None, Some(100)), Quantization::Sq8), + column(4, Some(50), percent(Some(1), Some(2)), Quantization::F32), + ]; + // The size from which the column is grouped, the size from which a partition of it is, + // and the two shares at which a partition is grouped again. + let in_force: Vec<_> = vectors + .iter() + .map(|vector| { + ( + vector.index_at(), + vector.grouped_from(), + vector.consolidation(), + ) + }) + .collect(); + assert_eq!( + in_force, + [ + (100_000, 1000, Some((20, 10))), + (1000, 10, None), + (u64::MAX, u64::MAX / 100, Some((50, 10))), + (1, 1, Some((20, 100))), + (50, 1, Some((1, 2))), + ] + ); + // What was written is what the record gives back: a clause that was not stays unwritten. + let table = Table { + vectors: vectors.clone(), + ..notes() + }; + assert_eq!(record::decode(&record::encode(&table)).vectors, vectors); +} diff --git a/v2/contextdb/src/sql/ast.rs b/v2/contextdb/src/sql/ast.rs index 074ef877..31ad3d63 100644 --- a/v2/contextdb/src/sql/ast.rs +++ b/v2/contextdb/src/sql/ast.rs @@ -12,6 +12,14 @@ pub enum Statement { Select(Select), SetMemoryLimit(SetMemoryLimitValue), ShowMemoryLimit, + ShowVectorIndexes, + ShowVectorPartitions { + /// The table and column after `FOR`. + column: Option<(String, String)>, + limit: Option, + /// Zero when OFFSET was not written. + offset: u64, + }, } #[derive(Debug, Clone)] @@ -179,6 +187,25 @@ pub struct VectorClauses { pub max_partitions: Option, /// `SEARCH_MODE AUTO | EXACT | INDEXED`, when it was written. pub search_mode: Option, + /// The number after `AUTO_INDEX_AT`, kept as written for `schema` to judge. + pub auto_index_at: Option, + /// `CONSOLIDATION NONE` or `CONSOLIDATION (...)`, when it was written. + pub consolidation: Option, + /// The text of `WITH (quantization = '...')`, kept as written for `schema` to judge. + pub quantization: Option, +} + +/// `CONSOLIDATION` as it was written. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ConsolidationClause { + /// `CONSOLIDATION NONE`. + None, + /// `CONSOLIDATION (CHANGE_PERCENT = n, TOMBSTONE_PERCENT = n)`: each number as written, + /// `None` for a member that was not. + Percent { + change: Option, + tombstone: Option, + }, } #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/v2/contextdb/src/sql/grammar.pest b/v2/contextdb/src/sql/grammar.pest index 827cc715..c50013f8 100644 --- a/v2/contextdb/src/sql/grammar.pest +++ b/v2/contextdb/src/sql/grammar.pest @@ -12,6 +12,8 @@ stmt = _{ | update_stmt | set_memory_limit | show_memory_limit + | show_vector_indexes + | show_vector_partitions | select_stmt | later } @@ -25,8 +27,7 @@ stmt = _{ later = { later_name ~ ANY* } later_name = @{ ( - ^"SHOW" ~ SPACE+ ~ (^"VECTOR_INDEXES" | ^"VECTOR_PARTITIONS") - | ^"ALTER" | ^"DROP" | ^"WITH" | ^"BEGIN" | ^"COMMIT" | ^"ROLLBACK" | ^"PURGE" | ^"DISCARD" + ^"ALTER" | ^"DROP" | ^"WITH" | ^"BEGIN" | ^"COMMIT" | ^"ROLLBACK" | ^"PURGE" | ^"DISCARD" | ^"DECLARE" | ^"DISTINCT" | ^"GRAPH_TABLE" | ((^"LEFT" | ^"INNER" | ^"CROSS") ~ SPACE+)? ~ ^"JOIN" | ^"GROUP" ~ SPACE+ ~ ^"BY" | ^"HAVING" | ^"UNION" | ^"OFFSET" @@ -35,7 +36,7 @@ later_name = @{ | (^"NOT" ~ SPACE+)? ~ (^"LIKE" | ^"BETWEEN") | ^"FOREIGN" ~ SPACE+ ~ ^"KEY" | ^"REFERENCES" | ^"EXPIRES" | ^"IMMUTABLE" | ^"RANK_POLICY" | ^"CONTEXT_ID" | ^"SCOPE_LABEL_READ" | ^"SCOPE_LABEL" | ^"ACL" - | ^"AUTO_INDEX_AT" | ^"SEARCH" | ^"IVF" | ^"CONSOLIDATION" + | ^"SEARCH" | ^"IVF" | ^"RETAIN" | ^"HISTORY" | ^"SYNC" | ^"STATE_MACHINE" | ^"STATE" ~ SPACE+ ~ ^"MACHINE" | ^"DAG" | ^"PROPAGATE" | ^"DELIVERY" ~ SPACE+ ~ ^"MANIFEST" | ^"EDGE" ~ SPACE+ ~ ^"DISCARD" | ^"MAINTAIN" ~ SPACE+ ~ ^"COUNT" @@ -52,6 +53,13 @@ show_memory_limit = { ^"SHOW" ~ ^"MEMORY_LIMIT" } memory_limit_value = { ^"none" | size_with_unit } size_with_unit = @{ ASCII_DIGIT+ ~ ("G" | "M" | "K") } +show_vector_indexes = { ^"SHOW" ~ ^"VECTOR_INDEXES" } +show_vector_partitions = { + ^"SHOW" ~ ^"VECTOR_PARTITIONS" ~ (^"FOR" ~ identifier ~ "." ~ identifier)? + ~ (limit_clause ~ offset_clause?)? +} +offset_clause = { ^"OFFSET" ~ integer } + select_stmt = { ^"SELECT" ~ ( @@ -92,11 +100,28 @@ data_type = { } vector_type = { ^"VECTOR" ~ "(" ~ integer ~ ")" - ~ partition_key? ~ (^"MAX_PARTITIONS" ~ partition_count)? ~ (^"SEARCH_MODE" ~ search_mode)? + ~ quantization? ~ partition_key? ~ max_partitions? ~ (^"SEARCH_MODE" ~ search_mode)? + ~ auto_index_at? ~ consolidation? + // WITH stands first or nowhere: after another clause it is no clause of a later version. + ~ !(^"WITH" ~ !ident_char) } +quantization = { ^"WITH" ~ "(" ~ ^"quantization" ~ "=" ~ string ~ ")" } partition_key = { ^"PARTITION_KEY" ~ "(" ~ identifier ~ ("," ~ identifier)* ~ ")" } +max_partitions = { ^"MAX_PARTITIONS" ~ signed } +auto_index_at = { ^"AUTO_INDEX_AT" ~ signed } +consolidation = { + ^"CONSOLIDATION" + ~ ( + ^"NONE" + | "(" + ~ (change_percent ~ ("," ~ tombstone_percent)? | tombstone_percent ~ ("," ~ change_percent)?) + ~ ")" + ) +} +change_percent = { ^"CHANGE_PERCENT" ~ "=" ~ signed } +tombstone_percent = { ^"TOMBSTONE_PERCENT" ~ "=" ~ signed } // Read with its sign, so that `schema` refuses zero and a negative number by their kind. -partition_count = @{ "-"? ~ ASCII_DIGIT+ } +signed = @{ "-"? ~ ASCII_DIGIT+ } column_constraint = _{ not_null | nullable_marker | primary_key | unique | default_clause } not_null = { ^"NOT" ~ ^"NULL" } nullable_marker = { ^"NULL" } diff --git a/v2/contextdb/src/sql/parser/mod.rs b/v2/contextdb/src/sql/parser/mod.rs index 287970cd..b64b1911 100644 --- a/v2/contextdb/src/sql/parser/mod.rs +++ b/v2/contextdb/src/sql/parser/mod.rs @@ -76,6 +76,20 @@ pub(super) fn parse(text: &str) -> Result { })), Rule::select_stmt => select(pair).map(Statement::Select), Rule::set_memory_limit => memory_limit(pair).map(Statement::SetMemoryLimit), + Rule::show_vector_indexes => Ok(Statement::ShowVectorIndexes), + Rule::show_vector_partitions => { + let mut names = idents(pair.clone()).into_iter(); + let count = |rule, what| { + child(pair.clone(), rule) + .map(|clause| number(&clause.into_inner().next().expect("a count"), what)) + .transpose() + }; + Ok(Statement::ShowVectorPartitions { + column: names.next().zip(names.next()), + limit: count(Rule::limit_clause, "LIMIT")?, + offset: count(Rule::offset_clause, "OFFSET")?.unwrap_or(0), + }) + } // The one statement left: SHOW MEMORY_LIMIT. _ => Ok(Statement::ShowMemoryLimit), } @@ -290,6 +304,11 @@ fn where_clause(pair: Pair<'_>) -> Result> { .transpose() } +/// The number of a clause of a VECTOR column, as written. +fn signed(clause: Pair<'_>) -> Option { + child(clause, Rule::signed).map(|number| number.as_str().to_owned()) +} + /// 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() { @@ -430,12 +449,28 @@ fn column_def(pair: Pair<'_>) -> Result { let dimension = parts.next().expect("a dimension"); for clause in parts { match clause.as_rule() { + Rule::quantization => { + let text = clause.into_inner().next().expect("a quoted text"); + vector.quantization = Some(string(&text)); + } Rule::partition_key => vector.partition_key = idents(clause), - Rule::partition_count => { - vector.max_partitions = Some(clause.as_str().to_owned()); + Rule::max_partitions => vector.max_partitions = signed(clause), + Rule::search_mode => vector.search_mode = Some(search_mode(&clause)), + Rule::auto_index_at => vector.auto_index_at = signed(clause), + // The one clause left: CONSOLIDATION, which is NONE when it has no member. + _ => { + let member = |rule| child(clause.clone(), rule).and_then(signed); + let members = ( + member(Rule::change_percent), + member(Rule::tombstone_percent), + ); + vector.consolidation = Some(match members { + (None, None) => ConsolidationClause::None, + (change, tombstone) => { + ConsolidationClause::Percent { change, tombstone } + } + }); } - // The one clause left: SEARCH_MODE. - _ => vector.search_mode = Some(search_mode(&clause)), } } DataType::Vector(number(&dimension, "the VECTOR dimension")?) diff --git a/v2/contextdb/src/sql/parser/tests.rs b/v2/contextdb/src/sql/parser/tests.rs index 82aaa818..df385258 100644 --- a/v2/contextdb/src/sql/parser/tests.rs +++ b/v2/contextdb/src/sql/parser/tests.rs @@ -1,6 +1,7 @@ use super::parse; use crate::sql::ast::{ - BinOp, ColumnRef, DataType, Expr, Literal, Select, Statement, VectorClauses, + BinOp, ColumnRef, ConsolidationClause, DataType, Expr, Literal, Select, Statement, + VectorClauses, }; use crate::sql::plan; use crate::sql::plan::{Effect, PhysicalPlan}; @@ -237,10 +238,6 @@ fn what_a_later_version_reads_is_refused_by_its_name() { "CREATE TABLE t (id UUID PRIMARY KEY, o UUID REFERENCES p(id))", "REFERENCES", ), - ( - "CREATE TABLE t (id UUID PRIMARY KEY, v VECTOR(3) WITH (quantization = 'SQ8'))", - "WITH", - ), ( "CREATE TABLE t (a INT, b INT, FOREIGN KEY (a, b) REFERENCES o (a, b))", "FOREIGN KEY", @@ -494,20 +491,6 @@ fn a_row_value_and_any_right_side_of_a_distance_are_read_and_refused_when_planne 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] @@ -526,6 +509,7 @@ fn the_clauses_of_a_vector_column_and_in_lists_are_read() { partition_key: key.iter().map(|column| (*column).to_owned()).collect(), max_partitions: max.map(str::to_owned), search_mode: mode, + ..VectorClauses::default() }; let none = VectorClauses::default(); assert_eq!( @@ -571,3 +555,101 @@ fn the_clauses_of_a_vector_column_and_in_lists_are_read() { Expr::BinaryOp { op: BinOp::Or, .. } )); } + +#[test] +fn the_index_clauses_of_a_vector_column_are_read_as_written() { + let Statement::CreateTable(table) = parse( + "CREATE TABLE t (id INTEGER PRIMARY KEY, \ + a VECTOR(4) WITH (quantization = 'sq8') PARTITION_KEY (id) MAX_PARTITIONS 2 \ + SEARCH_MODE AUTO AUTO_INDEX_AT -5 \ + CONSOLIDATION (TOMBSTONE_PERCENT = 0, CHANGE_PERCENT = 101), \ + b VECTOR(4) auto_index_at 18446744073709551616 consolidation none, \ + c VECTOR(4) CONSOLIDATION (CHANGE_PERCENT = 7))", + ) + .unwrap() else { + panic!("not a CREATE TABLE"); + }; + // Every number and the quantization's text are kept as written: `schema` judges them. + let written = |text: &str| Some(text.to_owned()); + let wanted = [ + VectorClauses::default(), + VectorClauses { + partition_key: vec!["id".to_owned()], + max_partitions: written("2"), + search_mode: Some(VectorSearchMode::Auto), + auto_index_at: written("-5"), + consolidation: Some(ConsolidationClause::Percent { + change: written("101"), + tombstone: written("0"), + }), + quantization: written("sq8"), + }, + VectorClauses { + auto_index_at: written("18446744073709551616"), + consolidation: Some(ConsolidationClause::None), + ..VectorClauses::default() + }, + VectorClauses { + consolidation: Some(ConsolidationClause::Percent { + change: written("7"), + tombstone: None, + }), + ..VectorClauses::default() + }, + ]; + let read: Vec = table.columns.into_iter().map(|c| c.vector).collect(); + assert_eq!(read, wanted); + + // The clauses stand in one order, and a CONSOLIDATION list names one or both of its two + // members, each once. + for clauses in [ + "AUTO_INDEX_AT 5 SEARCH_MODE AUTO", + "CONSOLIDATION NONE AUTO_INDEX_AT 5", + "SEARCH_MODE AUTO WITH (quantization = 'SQ8')", + "CONSOLIDATION ()", + "CONSOLIDATION (CHANGE_PERCENT = 1, CHANGE_PERCENT = 2)", + "CONSOLIDATION (PROBES = 1)", + "AUTO_INDEX_AT", + ] { + refusal(&format!( + "CREATE TABLE t (id INTEGER PRIMARY KEY, v VECTOR(4) {clauses})" + )); + } +} + +#[test] +fn the_inspection_statements_are_planned_as_reads() { + assert!(matches!( + planned("show vector_indexes"), + PhysicalPlan::ShowVectorIndexes + )); + let partitions = |sql: &str| match planned(sql) { + PhysicalPlan::ShowVectorPartitions { + column, + limit, + offset, + } => (column, limit, offset), + other => panic!("{sql}\nwas planned as {other:?}"), + }; + let docs_v = Some(("docs".to_owned(), "v".to_owned())); + assert_eq!(partitions("SHOW VECTOR_PARTITIONS"), (None, None, 0)); + assert_eq!( + partitions("show vector_partitions for docs.v"), + (docs_v.clone(), None, 0) + ); + assert_eq!( + partitions("SHOW VECTOR_PARTITIONS LIMIT 3"), + (None, Some(3), 0) + ); + assert_eq!( + partitions("SHOW VECTOR_PARTITIONS FOR docs.v LIMIT 5 OFFSET 2"), + (docs_v, Some(5), 2) + ); + for sql in ["SHOW VECTOR_INDEXES", "SHOW VECTOR_PARTITIONS LIMIT 1"] { + assert_eq!(planned(sql).effect(), Effect::Reads, "{sql}"); + } + // OFFSET stands only after LIMIT, and FOR names a table and a column. + refusal("SHOW VECTOR_PARTITIONS OFFSET 2"); + refusal("SHOW VECTOR_PARTITIONS FOR docs"); + refusal("SHOW VECTOR_INDEXES FOR docs.v"); +} diff --git a/v2/contextdb/src/sql/plan.rs b/v2/contextdb/src/sql/plan.rs index 41b1a68f..b1bf826b 100644 --- a/v2/contextdb/src/sql/plan.rs +++ b/v2/contextdb/src/sql/plan.rs @@ -45,6 +45,16 @@ pub enum PhysicalPlan { }, SetMemoryLimit(SetMemoryLimitValue), ShowMemoryLimit, + /// `SHOW VECTOR_INDEXES`: one row for each VECTOR column of the store. + ShowVectorIndexes, + /// `SHOW VECTOR_PARTITIONS [FOR table.column] [LIMIT count [OFFSET start]]`: one row for + /// each partition, of the named column or of every VECTOR column. + ShowVectorPartitions { + /// The table and column after `FOR`. + column: Option<(String, String)>, + limit: Option, + offset: u64, + }, } /// What a planned statement does to the store. @@ -79,7 +89,9 @@ impl PhysicalPlan { | PhysicalPlan::Project { .. } | PhysicalPlan::Limit { .. } | PhysicalPlan::VectorSearch { .. } - | PhysicalPlan::ShowMemoryLimit => Effect::Reads, + | PhysicalPlan::ShowMemoryLimit + | PhysicalPlan::ShowVectorIndexes + | PhysicalPlan::ShowVectorPartitions { .. } => Effect::Reads, } } } diff --git a/v2/contextdb/src/sql/planner.rs b/v2/contextdb/src/sql/planner.rs index cdbbed2e..a561de2c 100644 --- a/v2/contextdb/src/sql/planner.rs +++ b/v2/contextdb/src/sql/planner.rs @@ -54,6 +54,16 @@ pub(super) fn plan(statement: &Statement) -> Result { Statement::Select(select) => plan_select(select)?, Statement::SetMemoryLimit(limit) => PhysicalPlan::SetMemoryLimit(limit.clone()), Statement::ShowMemoryLimit => PhysicalPlan::ShowMemoryLimit, + Statement::ShowVectorIndexes => PhysicalPlan::ShowVectorIndexes, + Statement::ShowVectorPartitions { + column, + limit, + offset, + } => PhysicalPlan::ShowVectorPartitions { + column: column.clone(), + limit: *limit, + offset: *offset, + }, }) } diff --git a/v2/contextdb/src/store/disk.rs b/v2/contextdb/src/store/disk.rs index b1e684f7..51150c91 100644 --- a/v2/contextdb/src/store/disk.rs +++ b/v2/contextdb/src/store/disk.rs @@ -324,9 +324,10 @@ impl Read for View { // 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 table = if keys::with_blocks(start[0]) { + &moment.blocks + } else { + &moment.vectors }; let records = table.range::<&[u8]>(start.as_slice()..end.as_slice()); let mut key = Key(Vec::new()); @@ -404,17 +405,18 @@ pub(super) fn write(data: &redb::Database, changes: &mut ChangeSet) -> Stored { - let row = record::encode_row(values, |large| -> Stored { - // 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 row = + record::encode_row(values, &row.compact, |large| -> Stored { + // 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())) } @@ -472,9 +474,10 @@ pub(super) fn write(data: &redb::Database, changes: &mut ChangeSet) -> Stored &mut blocks, - _ => &mut vectors, + let table = if keys::with_blocks(key.0[0]) { + &mut blocks + } else { + &mut vectors }; match value { Some(value) => drop(table.insert(key.0.as_slice(), value.as_slice())?), diff --git a/v2/contextdb/src/store/keys.rs b/v2/contextdb/src/store/keys.rs index f382547f..d27f955a 100644 --- a/v2/contextdb/src/store/keys.rs +++ b/v2/contextdb/src/store/keys.rs @@ -12,7 +12,7 @@ use std::ops::Bound; use super::{Key, VectorId, VectorRecord}; -use crate::types::{Error, Result, SortDirection, Value}; +use crate::types::{Error, Result, SortDirection, TxId, Uuid, Value}; /// No value's form begins with this byte, turned or not, so a key followed by it sorts after /// every key that begins with that key and before every other key above it. @@ -117,6 +117,20 @@ pub(super) fn stored(id: u32, key: &[u8]) -> Vec { pub(super) const BLOCK: u8 = 0; const PARTITION: u8 = 1; const PENDING: u8 = 2; +const GROUPING: u8 = 3; +/// The first byte of the stored key of a record of centres, which is kept with the blocks. +pub(super) const CENTRES: u8 = 4; +const BUILD: u8 = 5; +const ASSIGNED: u8 = 6; +/// The first byte of the stored key of a record of sketches, which is kept with the blocks. +pub(super) const SKETCH: u8 = 7; +const MEAN: u8 = 8; +const LEFT: u8 = 9; + +/// True for the kinds of record that are kept in the table of blocks. +pub(super) fn with_blocks(kind: u8) -> bool { + kind == BLOCK || kind == CENTRES || kind == SKETCH +} /// 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 @@ -136,6 +150,34 @@ pub(super) fn vector(record: VectorRecord<'_>) -> Vec { let numbers = [number(column.0), number(partition)]; [&[PENDING][..], numbers.as_flattened(), &row.0].concat() } + VectorRecord::Grouping(column, grouping) => { + let numbers = [number(column.0), number(grouping)]; + [&[GROUPING][..], numbers.as_flattened()].concat() + } + VectorRecord::Centres(column, grouping, place) => { + let numbers = [number(column.0), number(grouping), number(place)]; + [&[CENTRES][..], numbers.as_flattened()].concat() + } + VectorRecord::Build(column, partition) => { + let numbers = [number(column.0), number(partition)]; + [&[BUILD][..], numbers.as_flattened()].concat() + } + VectorRecord::Assigned(column, grouping, place) => { + let numbers = [number(column.0), number(grouping), number(place)]; + [&[ASSIGNED][..], numbers.as_flattened()].concat() + } + VectorRecord::Sketch(column, under, block) => { + let numbers = [number(column.0), number(under), number(block)]; + [&[SKETCH][..], numbers.as_flattened()].concat() + } + VectorRecord::Mean(column, partition) => { + let numbers = [number(column.0), number(partition)]; + [&[MEAN][..], numbers.as_flattened()].concat() + } + VectorRecord::Left(column, partition, row) => { + let numbers = [number(column.0), number(partition)]; + [&[LEFT][..], numbers.as_flattened(), &row.0].concat() + } } } @@ -151,10 +193,70 @@ pub(super) fn vector_of<'a>(stored: &[u8], key: &'a mut Key) -> Option { + PENDING => { let partition = number(5)?; key.0.extend_from_slice(&stored[9..]); VectorRecord::Pending(column, partition, key) } + GROUPING => VectorRecord::Grouping(column, number(5)?), + CENTRES => VectorRecord::Centres(column, number(5)?, number(9)?), + BUILD => VectorRecord::Build(column, number(5)?), + ASSIGNED => VectorRecord::Assigned(column, number(5)?, number(9)?), + SKETCH => VectorRecord::Sketch(column, number(5)?, number(9)?), + MEAN => VectorRecord::Mean(column, number(5)?), + LEFT => { + let partition = number(5)?; + key.0.extend_from_slice(&stored[9..]); + VectorRecord::Left(column, partition, key) + } + _ => return None, }) } + +/// The values a key of ascending columns was made from; `None` for bytes that are no such key. +pub(super) fn decode(mut key: &[u8]) -> Option> { + let whole = |bytes: [u8; 8]| i64::from_be_bytes(bytes) ^ i64::MIN; + let mut values = Vec::new(); + while let Some([tag]) = take(&mut key) { + values.push(match tag { + 0xFE => Value::Null, + 1 => Value::Bool(take::<1>(&mut key)? != [0]), + 2 => Value::Int64(whole(take(&mut key)?)), + 3 => { + let ordered = u64::from_be_bytes(take(&mut key)?); + let bits = if ordered >> 63 == 0 { + !ordered + } else { + ordered ^ (1 << 63) + }; + Value::Float64(f64::from_bits(bits)) + } + 4 => { + let mut text = Vec::new(); + loop { + match take(&mut key)? { + [0] => match take(&mut key)? { + [0] => break, + [0xFF] => text.push(0), + _ => return None, + }, + [byte] => text.push(byte), + } + } + Value::Text(String::from_utf8(text).ok()?) + } + 5 => Value::Uuid(Uuid::from_bytes(take(&mut key)?)), + 6 => Value::Timestamp(whole(take(&mut key)?)), + 7 => Value::TxId(TxId(u64::from_be_bytes(take(&mut key)?))), + _ => return None, + }); + } + Some(values) +} + +/// The next bytes of a key, `None` when it is cut short. +fn take(key: &mut &[u8]) -> Option<[u8; N]> { + let (taken, rest) = key.split_first_chunk()?; + *key = rest; + Some(*taken) +} diff --git a/v2/contextdb/src/store/mod.rs b/v2/contextdb/src/store/mod.rs index 7f9d6b6e..c30e87b5 100644 --- a/v2/contextdb/src/store/mod.rs +++ b/v2/contextdb/src/store/mod.rs @@ -91,8 +91,32 @@ pub(crate) enum VectorRecord<'a> { /// 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. + /// One block: the number its blocks are kept under, then the block's own. A partition's + /// tail blocks are kept under the partition's number and the blocks of a grouping under + /// the grouping's; both come from the column's one run of numbers, so they never meet. Block(VectorId, u32, u32), + /// A grouping of a partition's blocks, under the grouping's number: how far from its + /// nearest group a search scans, and where each group's slots begin in its blocks. + Grouping(VectorId, u32), + /// The centres of a grouping, some to a record: the grouping's number, then the record's + /// place among them. Kept with the blocks, as records of the same size. + Centres(VectorId, u32, u32), + /// A grouping being built for a partition, under the partition's number: how far it has + /// come. + Build(VectorId, u32), + /// The group each slot of one block goes to in a grouping being built: the number of that + /// grouping, then the block's place among the blocks it is built from. + Assigned(VectorId, u32, u32), + /// The sketches of one block, one for each of its slots in their order: the two numbers + /// of the block. Kept with the blocks, an eighth of a block each. + Sketch(VectorId, u32, u32), + /// The mean the sketches of a partition's tail blocks are taken from, under the + /// partition's number. + Mean(VectorId, u32), + /// A row that had a sealed slot in a partition has left it: the partition's number and the + /// row's key. The value is how many slots had been sealed into the partition's tail by + /// then, so every slot of that key sealed earlier is stale. + Left(VectorId, u32, &'a Key), } /// A number the store keeps for the vector part. @@ -106,6 +130,16 @@ pub(crate) enum VectorNumber { Pending, } +/// A vector cell that `write` asks the store to keep in one byte a number: where the cell is +/// in its row, and the scale and codes the vector part made of it. The row's value at that +/// place is each code times the scale, and that is what reading the row gives back. The store +/// works nothing out from a vector's numbers and recognises nothing in them. +pub(crate) struct Compact { + pub at: usize, + pub scale: f32, + pub codes: Vec, +} + /// 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); @@ -141,6 +175,12 @@ impl Key { self.0.starts_with(&prefix.0) } + /// The values the key was made from, ascending columns only: what [`Key::from_values`] + /// was given. Bytes that are no such key are `None`. + pub(crate) fn values(&self) -> Option> { + keys::decode(&self.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)?; @@ -303,6 +343,8 @@ type Staged = BTreeMap>; struct StagedRow { /// The new values, or `None` for a removal. values: Option>, + /// The vector cells to write as a scale and codes. + compact: Vec, /// What the transaction said when it first staged this key: it found no row there. new_key: bool, } @@ -379,6 +421,10 @@ impl ChangeSet { /// Stages a row, new or replacing the row with this key. Large text, JSON and vector values /// are kept beside the row when it is written; the caller does not split them. /// + /// `compact` names the vector cells to keep in one byte a number, each with its scale and + /// codes: such a cell is written as the tag `0x88`, its count of numbers, the scale and one + /// signed byte a number, and every other vector as its numbers. + /// /// `new_key` says the transaction found no row at this key. The answer given the first /// time a key is staged is the one kept, and [`ChangeSet::rows_changed`] hands it back. pub(crate) fn put_row( @@ -386,15 +432,16 @@ impl ChangeSet { table: TableId, key: Key, values: Vec, + compact: Vec, new_key: bool, ) -> Result<()> { - self.stage_row(table, key, Some(values), new_key) + self.stage_row(table, key, Some(values), compact, new_key) } /// Stages the removal of a row and of the values kept beside it. A key first staged here /// is one a row was at. pub(crate) fn delete_row(&mut self, table: TableId, key: &Key) -> Result<()> { - self.stage_row(table, key.clone(), None, false) + self.stage_row(table, key.clone(), None, Vec::new(), false) } fn stage_row( @@ -402,13 +449,20 @@ impl ChangeSet { table: TableId, key: Key, values: Option>, + compact: Vec, new_key: bool, ) -> Result<()> { let held: u64 = values.iter().flatten().map(value_bytes).sum(); - let bytes = self.counted(&key, held)?; + let codes: usize = compact.iter().map(|cell| cell.codes.len()).sum(); + let bytes = self.counted(&key, held + codes as u64)?; let rows = self.rows.entry(table.0).or_default(); let new_key = rows.get(&key).map_or(new_key, |before| before.new_key); - let before = rows.insert(key.clone(), StagedRow { values, new_key }); + let staged = StagedRow { + values, + compact, + new_key, + }; + let before = rows.insert(key.clone(), staged); self.calls.push((Undo::Row(table.0, key, before), bytes)); Ok(()) } diff --git a/v2/contextdb/src/store/record.rs b/v2/contextdb/src/store/record.rs index f3bc4f70..504a8e99 100644 --- a/v2/contextdb/src/store/record.rs +++ b/v2/contextdb/src/store/record.rs @@ -11,10 +11,14 @@ //! 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. +//! +//! A vector cell that comes with a scale and codes is written in a third form, in its row or +//! beside it: the tag `0x88`, the count of its numbers as a `u64`, the scale as an `f32` and +//! one signed byte a number. It reads back as each code times the scale. use lz4_flex::block::{compress_into, decompress_into, get_maximum_output_size}; -use super::{Beside, Cell}; +use super::{Beside, Cell, Compact}; use crate::memory::Reservation; use crate::types::{self, TxId, Uuid, Value}; @@ -43,25 +47,40 @@ const JSON: u8 = 7; const VECTOR: u8 = 8; const TXID: u8 = 9; const BESIDE: u8 = 10; +/// A vector as a scale and one signed byte a number. +const REDUCED: u8 = 0x88; + +/// The tag, the count and the scale in front of the codes of such a vector. +const REDUCED_FRONT: u64 = 13; /// 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. +/// holds the length of what was kept. A cell that `compact` names is written as its scale and +/// codes. pub(super) fn encode_row( values: &[Value], + compact: &[Compact], 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 { + for (at, value) in values.iter().enumerate() { if !matches!(value, Value::Text(_) | Value::Json(_) | Value::Vector(_)) { put_value(&mut row, value); continue; } record.clear(); - put_value(&mut record, value); + match compact.iter().find(|cell| cell.at == at) { + Some(cell) => { + let count = cell.codes.len() as u64; + put(&mut record, REDUCED, &count.to_le_bytes()); + record.extend(cell.scale.to_le_bytes()); + record.extend(cell.codes.iter().map(|code| code.cast_unsigned())); + } + None => put_value(&mut record, value), + } if record.len() > BESIDE_FROM { let place = beside(&mut record)?; row.push(BESIDE); @@ -84,7 +103,10 @@ pub(super) fn encode_row( 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() { + if matches!(record[0], VECTOR | REDUCED) + || len > COMPRESS_TO + || reservation.grow(room as u64).is_err() + { return 0; } let mut block = vec![0; room]; @@ -172,6 +194,17 @@ pub(super) fn load( let Some(front) = stored.first_chunk::() else { return Ok(Err("a large value is shorter than its tag and length")); }; + if front[0] == REDUCED { + let count = u64::from_le_bytes(take_array(&mut &front[1..])); + if (stored.len() as u64).checked_sub(REDUCED_FRONT) != Some(count) { + return Ok(Err("a reduced vector gives a length its bytes are not")); + } + // The numbers are four bytes each where the record has one. + reservation.grow(4 * count)?; + let value = decode_value(&stored); + reservation.shrink(REDUCED_FRONT - FRONT as u64 + count); + return Ok(Ok(value)); + } 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. @@ -241,6 +274,16 @@ fn take_value(tag: u8, bytes: &mut &[u8]) -> Value { .collect(), ) } + REDUCED => { + let count = u64::from_le_bytes(take_array(bytes)); + let scale = f32::from_le_bytes(take_array(bytes)); + let codes = take(bytes, count as usize).iter(); + Value::Vector( + codes + .map(|code| scale * f32::from(code.cast_signed())) + .collect(), + ) + } // TXID, the one tag of a value that is left. _ => Value::TxId(TxId(u64::from_le_bytes(take_array(bytes)))), } diff --git a/v2/contextdb/src/store/tests.rs b/v2/contextdb/src/store/tests.rs index a76a9f84..16df5e8d 100644 --- a/v2/contextdb/src/store/tests.rs +++ b/v2/contextdb/src/store/tests.rs @@ -12,7 +12,9 @@ use std::path::{Path, PathBuf}; use lz4_flex::block::{compress_into, get_maximum_output_size}; use redb::ReadableTable; -use super::{Cell, ChangeSet, Key, Read, Store, TableId, keys, record}; +use super::{ + Cell, ChangeSet, Compact, Key, Read, Store, TableId, VectorId, VectorRecord, keys, record, +}; use crate::memory::Budget; use crate::types::{Error, SortDirection, TxId, Uuid, Value}; @@ -149,7 +151,7 @@ fn a_row_reads_back_as_written_with_its_large_values_beside_it() { ]; let mut kept = Vec::new(); - let row = record::encode_row(&values, |large| { + let row = record::encode_row(&values, &[], |large| { kept.push(large.to_vec()); Ok::<_, ()>(kept.len() as u64 - 1) }) @@ -196,7 +198,7 @@ fn a_compressed_value_whose_recorded_length_is_wrong_is_damage_and_takes_no_memo 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) + .put_row(table, row_key.clone(), row, Vec::new(), true) .expect("stage"); store.commit(changes).expect("commit"); drop(store); @@ -340,7 +342,9 @@ fn store_one_row(folder: &Path, value: Value) { 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"); + changes + .put_row(table, row_key, row, Vec::new(), true) + .expect("stage"); store.commit(changes).expect("commit"); } @@ -497,3 +501,191 @@ fn a_compressed_json_that_decodes_to_bytes_that_are_not_json_is_damage() { let wrong = format!("{{{}}}", "a".repeat(LEN - 2)); a_block_of_other_bytes_is_damage("not-json", &json, wrong.as_bytes()); } + +#[test] +fn a_key_gives_its_values_back() { + let lists = [ + vec![], + vec![Value::Bool(false), Value::Bool(true)], + vec![ + Value::Int64(i64::MIN), + Value::Int64(-1), + Value::Int64(i64::MAX), + ], + vec![ + Value::Float64(-1.5), + Value::Float64(0.0), + Value::Float64(f64::INFINITY), + ], + vec![ + text(""), + text("a"), + text("a\0b\0"), + text("\u{65e5}\u{672c}"), + ], + vec![Value::Uuid(Uuid::from_u128(7)), Value::Timestamp(-1)], + vec![ + Value::TxId(TxId(u64::MAX)), + Value::Null, + text("after a NULL"), + ], + ]; + for values in lists { + let made = key(&values); + assert_eq!(made.values(), Some(values.clone()), "{values:?}"); + assert_eq!(key(&made.values().expect("the values")), made); + } + // Bytes that are no key of ascending columns: an unknown tag, and a value cut short. + assert_eq!(Key::from_bytes(&[0x7F]).values(), None); + for whole in [key(&[Value::Int64(5)]), key(&[text("cut")])] { + let bytes = whole.as_bytes(); + let short = Key::from_bytes(&bytes[..bytes.len() - 1]); + assert_eq!(short.values(), None, "{bytes:?} without its last byte"); + } +} + +/// `count` numbers that are each a whole code from -127 to 127 times one scale, with the +/// scale and the codes. +fn reduced(count: usize) -> (Vec, f32, Vec) { + let scale = 0.031_25_f32; + let codes: Vec = (0..count) + .map(|at| ((at % 255) as i16 - 127) as i8) + .collect(); + let numbers = codes.iter().map(|code| scale * f32::from(*code)).collect(); + (numbers, scale, codes) +} + +/// Writes a store of rows 1 to `rows`, each the number and then one vector of `count` reduced +/// numbers, handed to the store with its scale and codes when `compact`. Returns the bytes of +/// the data file once the store is closed. +fn store_vectors(folder: &Path, rows: i64, count: usize, compact: bool) -> u64 { + let (numbers, scale, codes) = reduced(count); + 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"); + for id in 1..=rows { + let row = vec![Value::Int64(id), Value::Vector(numbers.clone())]; + let kept = if compact { + vec![Compact { + at: 1, + scale, + codes: codes.clone(), + }] + } else { + Vec::new() + }; + changes + .put_row(TableId(1), key(&[Value::Int64(id)]), row, kept, true) + .expect("stage"); + } + store.commit(changes).expect("commit"); + drop(store); + std::fs::metadata(folder.join("data")) + .expect("the data file") + .len() +} + +#[test] +fn a_reduced_vector_reads_back_number_for_number_in_a_quarter_of_the_room() { + let folder = fresh_folder("reduced-vector"); + // 4,000 rows, so that both files are well past the size a data file starts at. + let compact = store_vectors(&folder, 4000, 768, true); + // 781 bytes stay in the row, and read back as exactly the numbers that were put. + let (numbers, ..) = reduced(768); + 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"); + for id in [1, 250, 500] { + let row = view.row(TableId(1), &key(&[Value::Int64(id)])); + let row = row.expect("read the row").expect("the row is there"); + match &row.cells[1] { + Cell::Here(Value::Vector(read)) => { + let bits = |numbers: &[f32]| -> Vec { + numbers.iter().map(|number| number.to_bits()).collect() + }; + assert_eq!(bits(read), bits(&numbers), "row {id}"); + } + other => panic!("the vector of row {id} was read as {other:?}"), + } + } + drop((view, store)); + let plain_folder = fresh_folder("reduced-vector-plain"); + let plain = store_vectors(&plain_folder, 4000, 768, false); + assert!( + 3 * compact < plain, + "4,000 vectors of 768 numbers take {compact} bytes of file reduced and {plain} as written" + ); + for folder in [folder, plain_folder] { + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); + } +} + +#[test] +fn a_reduced_record_of_another_length_is_damage() { + // 2,000 reduced numbers are a record of 2,013 bytes, kept beside the row. + const COUNT: usize = 2000; + let folder = fresh_folder("reduced-length"); + store_vectors(&folder, 1, COUNT, true); + let (loaded, beside, _, _) = load_the_value(&folder); + assert!(loaded.expect("load the value as it was written") == Value::Vector(reduced(COUNT).0)); + assert_eq!( + beside.bytes, + 13 + COUNT as u64, + "the record of a reduced vector" + ); + + // The record now says it holds one number more than its bytes do. + 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 mut first = first.expect("the first slice is there").value().to_vec(); + assert_eq!(first[0], 0x88, "the tag of a reduced vector"); + assert_eq!(first[1..9], (COUNT as u64).to_le_bytes()); + first[1..9].copy_from_slice(&(COUNT as u64 + 1).to_le_bytes()); + kept.insert(beside.place, first.as_slice()) + .expect("write the slice"); + } + write.commit().expect("commit"); + drop(data); + + loading_is_damage(&folder, "a reduced vector whose count is not its bytes'"); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +#[test] +fn the_records_of_a_grouping_keep_their_keys_apart() { + use VectorRecord::*; + let column = VectorId(3); + let (tuple, row) = (key(&[text("a")]), key(&[Value::Int64(1)])); + // A partition numbered 5 and its grouping numbered 6: every kind of record of the two. + let records = [ + Partition(column, &tuple), + Pending(column, 5, &row), + Block(column, 5, 0), + Block(column, 6, 0), + Grouping(column, 6), + Centres(column, 6, 0), + Build(column, 5), + Assigned(column, 6, 0), + Sketch(column, 5, 0), + Sketch(column, 6, 0), + Mean(column, 5), + Left(column, 5, &row), + ]; + let stored = records.map(keys::vector); + let apart: std::collections::BTreeSet<&Vec> = stored.iter().collect(); + assert_eq!(apart.len(), records.len(), "two records share a key"); + for (record, stored) in records.iter().zip(&stored) { + let mut held = key(&[]); + assert_eq!(keys::vector_of(stored, &mut held), Some(*record)); + // Blocks, centres and sketches are kept together; everything else is small. + let large = matches!(record, Block(..) | Centres(..) | Sketch(..)); + assert_eq!(keys::with_blocks(stored[0]), large, "{record:?}"); + } + // The sketches of a grouping's blocks come in the order of the blocks, before the next + // number's. + let sketch = |under, block| keys::vector(Sketch(column, under, block)); + assert!(sketch(6, 0) < sketch(6, 1) && sketch(6, 1) < sketch(7, 0)); +} diff --git a/v2/contextdb/src/types/mod.rs b/v2/contextdb/src/types/mod.rs index 2be9229e..2dc7cee3 100644 --- a/v2/contextdb/src/types/mod.rs +++ b/v2/contextdb/src/types/mod.rs @@ -252,8 +252,10 @@ pub enum Error { 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. + /// `SEARCH_MODE INDEXED`, and the index cannot answer it now: a partition it selects is + /// behind its writes, or its filter on rows leaves the index nothing to choose among. Such + /// a search is never answered with an exact scan under the indexed name; `EXACT` and + /// `AUTO` answer, and a maintenance cycle brings the index up to its writes. VectorIndexedRouteUnavailable { index: VectorIndexRef }, /// `USE VECTOR` was written in a SELECT that does not order by `<=>`. UseVectorRequiresVectorOrder, @@ -388,8 +390,9 @@ impl fmt::Display for Error { ), Error::VectorIndexedRouteUnavailable { index } => write!( f, - "no vector index ships in this version, so {index} cannot be searched INDEXED; \ - use EXACT or AUTO" + "{index} cannot be searched INDEXED for this statement now: its index is behind \ + its writes, or the statement's filter leaves the index nothing to choose among; \ + run a maintenance cycle, or use EXACT or AUTO" ), Error::UseVectorRequiresVectorOrder => { f.write_str("USE VECTOR requires ORDER BY embedding <=> $param in the same query") @@ -489,13 +492,31 @@ impl fmt::Display for VectorPartitionDeclarationIssue { pub enum VectorPolicyDeclarationIssue { /// `VECTOR(0)`. DimensionsNotPositive, + /// `WITH (...)` gives this member a value it does not take. + InvalidMember { member: &'static str }, + /// `AUTO_INDEX_AT` is zero or negative. + AutoIndexAtNotPositive, + /// `AUTO_INDEX_AT` is larger than the engine keeps. + AutoIndexAtOutOfRange, + /// A percentage of `CONSOLIDATION` is not a whole number from 1 to 100. + ConsolidationPercentOutOfRange, } 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", - }) + match self { + Self::DimensionsNotPositive => f.write_str("vector dimensions must be at least 1"), + Self::InvalidMember { member } => { + write!(f, "{member} has an invalid vector policy value") + } + Self::AutoIndexAtNotPositive => f.write_str("AUTO_INDEX_AT must be a positive integer"), + Self::AutoIndexAtOutOfRange => { + f.write_str("AUTO_INDEX_AT is outside the persisted numeric range") + } + Self::ConsolidationPercentOutOfRange => { + f.write_str("vector consolidation percentages must be integers from 1 through 100") + } + } } } @@ -503,12 +524,13 @@ impl fmt::Display for VectorPolicyDeclarationIssue { /// `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. + /// The engine chooses: through the index from the column's `AUTO_INDEX_AT` vectors on, + /// wherever the index is usable, and exactly otherwise. The answer is never short. #[default] Auto, /// Every allowed vector is compared: the true nearest rows. Exact, - /// Only through an index; refused when there is none. + /// Only through the index; refused when the index cannot answer the statement now. Indexed, } diff --git a/v2/contextdb/src/upkeep/mod.rs b/v2/contextdb/src/upkeep/mod.rs index a663cf59..b15ad756 100644 --- a/v2/contextdb/src/upkeep/mod.rs +++ b/v2/contextdb/src/upkeep/mod.rs @@ -1,8 +1,10 @@ -//! 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. +//! upkeep: background work in small steps, and the report of one cycle of it. A cycle seals, +//! which is moving the vectors that wait on a column's pending lists into blocks, and then +//! advances by one step every grouping the vector part has due or is building. The sealing is +//! staged here; `database` drives the steps, each through its one commit function. //! -//! 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 +//! It never blocks a write or a query: every commit of a cycle holds a bounded amount of work, +//! and what a cycle has not done waits for the next. Nothing here starts a thread or reads a //! clock; a cycle runs when it is called. use crate::memory::Budget; @@ -26,15 +28,25 @@ pub struct MaintenanceReport { #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] #[non_exhaustive] pub struct VectorMaintenanceReport { - /// How many vectors this cycle moved from pending lists into blocks. + /// How many vectors this cycle moved from pending lists into blocks, added to how many + /// it packed into the blocks of a grouping being built. 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. + /// How many vectors still wait after this cycle, over every column of the store: those on + /// a pending list, those a build under way has not packed, and those of partitions whose + /// grouping is due and not begun. Those of a stalled build are not counted: no further + /// cycle packs them until it has the memory. pub maintenance_vectors_remaining: u64, + /// How many partitions have a grouping due or being built after this cycle. + pub remaining_partitions: u64, + /// How many of those had their last step refused for want of memory. A further cycle does + /// nothing for them until the memory limit is raised or memory is let go, and with no + /// limit declared there are none. A host stops calling when `remaining_partitions` is no + /// more than this and `maintenance_vectors_remaining` is zero. + pub stalled_partitions: 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. +/// Stages the sealing of one cycle into `changes`, for the caller to commit, and says what it +/// staged: the vectors sealed and those still pending. The caller holds the commit lock from before `changes` was begun. pub(crate) fn cycle( changes: &mut ChangeSet, catalog: &Catalog, @@ -45,6 +57,7 @@ pub(crate) fn cycle( vector: VectorMaintenanceReport { maintenance_vectors_done: done, maintenance_vectors_remaining: remaining, + ..VectorMaintenanceReport::default() }, }) } diff --git a/v2/contextdb/src/vector/block.rs b/v2/contextdb/src/vector/block.rs index 584d6e31..7ac8c0f3 100644 --- a/v2/contextdb/src/vector/block.rs +++ b/v2/contextdb/src/vector/block.rs @@ -1,5 +1,5 @@ //! 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. +//! (328 at 768 numbers a vector with a UUID key), 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 @@ -80,6 +80,13 @@ impl BlockWriter { self.ends.push(self.keys.len() as u32); } + /// Adds a slot as it was written for another block, with the key of its row. + pub(super) fn push_slot(&mut self, key: &Key, slot: &[u8]) { + self.slots.extend_from_slice(slot); + 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() diff --git a/v2/contextdb/src/vector/build.rs b/v2/contextdb/src/vector/build.rs new file mode 100644 index 00000000..a5a8258a --- /dev/null +++ b/v2/contextdb/src/vector/build.rs @@ -0,0 +1,976 @@ +//! Building a grouping for a partition, in steps: when one is due, the record of how far a +//! build has come, and the steps that find its centres and say which group each slot goes to. +//! +//! A build reads the partition as it stood when the build began (the grouping in force and the +//! tail blocks below `closed`) and writes beside it, under a number of its own. Nothing a +//! search reads changes until the last step, which makes the new grouping the one in force in +//! one commit. A step is worked out from one committed moment with no lock held, and staged +//! under the commit lock only if the build has not moved since. + +use super::block::{BROKEN, BlockReader}; +use super::directory::{self, Grouped, Partition, kept}; +use super::grouping::{self, Place, centres_to_a_record}; +use super::pack; +use super::seal::room; +use super::sketch::{sketch_len, sketches_of}; +use super::train::{ + DRAWN, Mean, ROUNDS, Round, SAMPLE_TO_A_CENTRE, first_centres, groups_for, nearest, sampled, +}; +use crate::memory::{Budget, Reservation}; +use crate::schema::{Catalog, VectorColumn}; +use crate::sql::not_supported; +use crate::store::{ChangeSet, Key, Read, VectorId, VectorNumber, VectorRecord}; +use crate::types::{Result, VectorSearchMode}; + +/// The most slots one step gives a group or packs, unless one group is larger. +pub(super) const STEP_SLOTS: u64 = 4096; + +/// The most comparisons of a slot with a centre in one step that gives slots their groups or +/// measures: 4,194,304. A round over the sample is one step whatever it compares. +pub(super) const STEP_COMPARISONS: u64 = 1 << 22; + +/// How many questions a build asks of its new grouping to measure how far a search scans +/// and how wide it takes a sketch's estimate to be. Question `j` is the middle of two +/// stored vectors: the direction of the sum of the vectors that the slots at `j * slots / +/// MEASURED` and at half the grouping further on keep. +pub(super) const MEASURED: u32 = 128; + +/// Bytes of a build's record before its lists. +const HEAD: usize = 158; + +/// What a build does next. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum Phase { + /// The next round over the sample, counted from zero; the centres so far are stored. + Train { round: u32 }, + /// Giving slots their groups: the next block to read, by its place among those the build + /// reads. + Assign { next: u32 }, + /// Packing: the next group to write, and the block of the new grouping it starts in. + Pack { group: u32 }, + /// Measuring the gap: how many questions have been asked, and the largest gap any of them + /// needed for its ten nearest, in ten-thousandths. + Measure { done: u32, most: u32 }, + /// Measuring the width, with the gap settled: how many questions have been asked again, + /// the most groups any of them had inside the gap, and for each width from 2 to 16 + /// quarters how many of their neighbours the sketches of those groups kept, of `asked` + /// neighbours in all. + Width { + gap: u32, + probes: u32, + done: u32, + asked: u32, + kept: [u32; 15], + }, + /// Everything is written: the next step makes the grouping the one in force. + Switch { gap: u32, probes: u32, width: u8 }, + /// The grouping is in force, and the build was of slots some of which were stale. The + /// next steps remove the `Left` records the build applied, and the last removes the + /// build's record. + Sweep, +} + +/// How far the build of a partition's next grouping has come. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) struct Build { + /// The number the new grouping's records are kept under. + pub number: u32, + /// How many steps of this build have been staged. A step that was worked out for another + /// count is not staged. + pub steps: u64, + /// The grouping the build reads, the one in force when it began; `None` when there was + /// none. + pub from: Option, + /// The tail blocks the build reads: from this number up to, not including, the next. + pub tail: (u32, u32), + /// How many slots those blocks held when the build began. + pub slots: u64, + /// The partition's `sealed` when the build began: the tail slots it reads are the last + /// before that count. + pub upto: u64, + /// Whether any of them had been left behind by its row, so that every slot's key is + /// looked up among the partition's `Left` records. No row is read. + pub stale: bool, + /// How many groups the new grouping has: one in a column that is not at its size, and + /// then the build finds no centre, writes no sketch and asks no question. + pub groups: u32, + pub phase: Phase, + /// How many slots each group has been given so far, by group. + pub counts: Vec, + /// Where each group packed so far begins in the new blocks, and where the next begins. + pub starts: Vec, + /// The ten slots nearest each question asked so far, in the order asked: what measuring + /// the gap found and measuring the width reads. Empty until the first question. + pub tens: Vec, + /// The bytes the step that was last refused for memory needed; zero when the last step + /// was staged. + pub needing: u64, +} + +impl Build { + /// The record of the build, every number low byte first. Six numbers of eight bytes: + /// `steps`, `needing`, `slots`, `upto`, and the `slots` and `upto` of `from`. One byte of + /// flags (1: some slot it reads is stale; 2: it has a `from`) and one for the phase, + /// counted from zero in the order of [`Phase`]. Then numbers of four bytes: `number`, the + /// two of `tail`, `groups`, the `number`, `blocks` and `groups` of `from`, nineteen for the + /// fields of the phase in their order, the count of `starts`, `counts`, and the two + /// numbers of every place of `starts` and then of `tens`. A number no field fills is zero. + pub(super) fn encode(&self) -> Vec { + let from = self.from.map(|from| { + let short = [from.number, from.blocks, from.groups]; + ([from.slots, from.upto], short) + }); + let (long, short) = from.unwrap_or_default(); + let (phase, fields) = match self.phase { + Phase::Train { round } => (0, vec![round]), + Phase::Assign { next } => (1, vec![next]), + Phase::Pack { group } => (2, vec![group]), + Phase::Measure { done, most } => (3, vec![done, most]), + Phase::Width { + gap, + probes, + done, + asked, + kept, + } => (4, [&[gap, probes, done, asked], &kept[..]].concat()), + Phase::Switch { gap, probes, width } => (5, vec![gap, probes, width.into()]), + Phase::Sweep => (6, Vec::new()), + }; + let mut carried = [0; 19]; + carried[..fields.len()].copy_from_slice(&fields); + let flags = u8::from(self.stale) | (u8::from(self.from.is_some()) << 1); + let places = self.starts.iter().chain(&self.tens); + let places = places.flat_map(|place| [place.block, place.slot]); + let head = [self.number, self.tail.0, self.tail.1, self.groups]; + let short = head.into_iter().chain(short).chain(carried); + let short = short.chain([self.starts.len() as u32]); + let short = short.chain(self.counts.iter().copied()).chain(places); + let mut record = Vec::new(); + let first = [self.steps, self.needing, self.slots, self.upto]; + record.extend(first.into_iter().chain(long).flat_map(u64::to_le_bytes)); + record.extend([flags, phase]); + record.extend(short.flat_map(u32::to_le_bytes)); + record + } + + /// The build a record was written from, `None` for a record that disagrees with itself: + /// a length that does not follow from its `groups` and its count of `starts`, an unknown + /// flag or phase, no group or more groups than slots, a round, block or group past those it + /// has, places that do not begin at the first slot or are out of order, places for groups + /// it has not packed, or nearest slots that are not as many a question as its packed + /// slots give (ten, or all of fewer) for the questions its phase has asked. + fn decode(record: &[u8]) -> Option { + let (head, lists) = record.split_at_checked(HEAD)?; + let (lists, []) = lists.as_chunks::<4>() else { + return None; + }; + let long = |at: usize| u64::from_le_bytes(std::array::from_fn(|byte| head[8 * at + byte])); + let short = + |at: usize| u32::from_le_bytes(std::array::from_fn(|byte| head[50 + 4 * at + byte])); + let listed = |at: usize| u32::from_le_bytes(lists[at]); + let (flags, groups, begun) = (head[48], short(3) as usize, short(26) as usize); + if flags > 3 || lists.len() < groups + 2 * begun || (lists.len() - groups) % 2 != 0 { + return None; + } + let phase = match head[49] { + 0 => Phase::Train { round: short(7) }, + 1 => Phase::Assign { next: short(7) }, + 2 => Phase::Pack { group: short(7) }, + 3 => Phase::Measure { + done: short(7), + most: short(8), + }, + 4 => Phase::Width { + gap: short(7), + probes: short(8), + done: short(9), + asked: short(10), + kept: std::array::from_fn(|at| short(11 + at)), + }, + 5 => Phase::Switch { + gap: short(7), + probes: short(8), + width: u8::try_from(short(9)).ok()?, + }, + 6 => Phase::Sweep, + _ => return None, + }; + let mut places = (groups..lists.len()).step_by(2).map(|at| Place { + block: listed(at), + slot: listed(at + 1), + }); + let build = Build { + number: short(0), + steps: long(0), + from: (flags & 2 != 0).then(|| Grouped { + number: short(4), + blocks: short(5), + groups: short(6), + slots: long(4), + upto: long(5), + }), + tail: (short(1), short(2)), + slots: long(2), + upto: long(3), + stale: flags & 1 != 0, + groups: short(3), + phase, + counts: (0..groups).map(listed).collect(), + starts: places.by_ref().take(begun).collect(), + tens: places.collect(), + needing: long(1), + }; + let (inside, grouped) = build.from.map_or((0, 0), |from| (from.blocks, from.slots)); + let blocks = (build.tail.1.checked_sub(build.tail.0)?).checked_add(inside)?; + let tail = build.slots.checked_sub(grouped)?; + let packed: u64 = build.counts.iter().map(|count| u64::from(*count)).sum(); + let each = packed.min(10) as usize; + let questions = build.tens.len().checked_div(each).unwrap_or(0); + let whole = build.tens.len() == each * questions && questions <= MEASURED as usize; + let ended = begun == groups + 1; + let agrees = match phase { + Phase::Train { round } => round < ROUNDS && begun == 0 && questions == 0, + Phase::Assign { next } => next < blocks && begun == 0 && questions == 0, + Phase::Pack { group } => (group as usize) < groups && begun == group as usize + 1, + Phase::Measure { done, .. } => ended && done as usize == questions, + Phase::Width { + done, asked, kept, .. + } => ended && done as usize <= questions && kept.iter().all(|kept| *kept <= asked), + Phase::Switch { .. } => ended, + Phase::Sweep => ended && questions == 0, + }; + let placed = build.starts.is_sorted() + && build + .starts + .first() + .is_none_or(|first| *first == Place::default()); + let sound = 0 < groups && groups as u64 <= build.slots && build.upto >= tail && placed; + (sound && whole && agrees).then_some(build) + } + + /// The build of the partition of this number, `None` when it has none. A record that + /// disagrees with itself is damage. + pub(super) fn read(read: &dyn Read, vector: VectorId, partition: u32) -> Result> { + let found = kept(read, VectorRecord::Build(vector, partition), Build::decode)?; + let damage = "the record of a build of the vector part disagrees with itself"; + let build = found.map(|build| build.ok_or_else(|| read.damaged(damage))); + build.transpose() + } + + /// Every record the build has written so far, for the commit that removes them: its own + /// record, the groups it gave, its centres and the blocks it packed with their sketches. + /// Once its grouping is in force those are the grouping's and it has its record alone. + pub(super) fn written( + &self, + vector: VectorId, + partition: u32, + dimension: usize, + ) -> impl Iterator> + use<> { + let number = self.number; + let many = u32::from(self.groups > 1); + let centres = self.groups.div_ceil(centres_to_a_record(dimension) as u32); + let end = self.starts.last().copied().unwrap_or_default(); + let packed = end.block + u32::from(end.slot > 0); + let (given, centres, packed) = match self.phase { + Phase::Sweep => (0, 0, 0), + Phase::Train { round } => (0, centres * u32::from(round > 0), packed), + Phase::Assign { next } => (next, centres * many, packed), + _ => (self.blocks(), centres * many, packed), + }; + std::iter::once(VectorRecord::Build(vector, partition)) + .chain((0..given).map(move |place| VectorRecord::Assigned(vector, number, place))) + .chain((0..centres).map(move |place| VectorRecord::Centres(vector, number, place))) + .chain((0..packed).map(move |block| VectorRecord::Block(vector, number, block))) + .chain((0..packed * many).map(move |block| VectorRecord::Sketch(vector, number, block))) + } + + /// How many slots the build has still to pack: every slot it reads until each has its + /// group, then those of the groups it has not written. + pub(super) fn unpacked(&self) -> u64 { + let counts = match self.phase { + Phase::Train { .. } | Phase::Assign { .. } => return self.slots, + Phase::Pack { group } => &self.counts[group as usize..], + _ => &[], + }; + counts.iter().map(|count| u64::from(*count)).sum() + } + + /// How many blocks the build reads: those of `from`, then its tail blocks. + pub(super) fn blocks(&self) -> u32 { + self.from.map_or(0, |from| from.blocks) + (self.tail.1 - self.tail.0) + } + + /// The block the build reads at this place, counted from zero: `from`'s blocks in order, + /// then the tail blocks of the partition. + pub(super) fn block( + &self, + vector: VectorId, + partition: u32, + place: u32, + ) -> VectorRecord<'static> { + match self.from { + Some(from) if place < from.blocks => VectorRecord::Block(vector, from.number, place), + from => { + let after = place - from.map_or(0, |from| from.blocks); + VectorRecord::Block(vector, partition, self.tail.0 + after) + } + } + } +} + +/// One partition with work for a build: one that is due, or one under way. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Work { + pub vector: VectorId, + /// The partition's key tuple. + pub tuple: Key, +} + +/// What one step read and worked out, to be staged under the commit lock. +pub(crate) struct Step { + pub(super) work: Work, + /// The column as it was declared when the step was worked out: what the step that begins + /// a build asks [`due`] of again under the lock. + pub(super) column: VectorColumn, + /// The partition's number when the step was worked out, and the build's number and count + /// of steps then; `None` for the step that begins a build. + pub(super) at: (u32, Option<(u32, u64)>), + /// The records the step writes and removes. A build's records hold no key, so they + /// outlive what they were read from. + pub(super) puts: Vec<(VectorRecord<'static>, Vec)>, + pub(super) removes: Vec>, + /// The keys of the partition's `Left` records the step removes. + pub(super) left: Vec, + /// What the partition's record takes from the step: nothing, or the grouping the step + /// puts in force. + pub(super) switch: Option, + /// The centres a round made, in the order of the groups; empty for any other step. + pub(super) centres: Vec, + /// The partition's mean, from the step that stages it: its record says so from then on. + pub(super) mean: Option>, + /// How many vectors the step finished: those it packed. + pub(super) done: u64, + /// Never read: holds what the records above take until they are staged. + pub(super) held: Reservation, +} + +/// What a build is given to work on: one partition at one committed moment. +pub(super) struct Ground<'a> { + pub read: &'a dyn Read, + pub column: &'a VectorColumn, + pub dimension: usize, + pub partition: Partition, + pub build: Build, +} + +/// True when a grouping is due for a partition that is building none. `vectors` is how many +/// rows have a vector in the column, every partition together; the column is at its size when +/// that is its `index_at` or more, or when it declares `SEARCH_MODE INDEXED`. +/// +/// In a column at its size: the partition has no grouping and its tail holds at least the +/// column's `grouped_from` slots (a hundredth of `index_at`, and one at least); or it has a +/// grouping of one group, holds that many slots, and a build now would give it more than one +/// group, the square root of its slots rounding to two or more, so that a partition whose +/// grouping already has the groups a new build would give it is never due on its size alone; +/// or it has a grouping and the vectors outside it (in tail blocks and pending) together with +/// the slots left behind come to `CHANGE_PERCENT` of its slots, or the slots left behind alone +/// to `TOMBSTONE_PERCENT`. In any other column only the +/// slots left behind count, against every slot the partition has: the build then has one +/// group and writes no centre and no sketch. Under `CONSOLIDATION NONE` only a first grouping +/// of a column at its size is ever due. The slots left behind are the slots of the grouping +/// and the tail less the rows with a vector that are not pending. +pub(super) fn due( + partition: &Partition, + pending: u64, + column: &VectorColumn, + vectors: u64, +) -> bool { + let grouped = partition.grouping.map_or(0, |grouping| grouping.slots); + let slots = grouped + partition.slots; + let out = partition.slots + pending; + let stale = (slots + pending).saturating_sub(partition.live); + let behind = left_behind(partition, pending, column, vectors); + if !sized(column, vectors) { + return behind; + } + let Some(grouping) = partition.grouping else { + return partition.slots >= column.grouped_from(); + }; + let grows = grouping.groups == 1 && slots >= column.grouped_from() && groups_for(slots) > 1; + let reached = |(change, _): (u8, u8)| { + grows || behind || (out + stale) * 100 >= u64::from(change) * grouped + }; + // A partition with no slot has nothing to group: its pending rows are sealed first. + slots > 0 && column.consolidation().is_some_and(reached) +} + +/// True when the slots rows have left behind are alone enough for a grouping to be due: they +/// come to the column's `TOMBSTONE_PERCENT` of the slots of the grouping in force in a column +/// at its size, and of every slot of the partition in any other. [`due`] asks it, and +/// inspection gives it as the reason. +pub(super) fn left_behind( + partition: &Partition, + pending: u64, + column: &VectorColumn, + vectors: u64, +) -> bool { + let grouped = partition.grouping.map_or(0, |grouping| grouping.slots); + let slots = grouped + partition.slots; + let stale = (slots + pending).saturating_sub(partition.live); + let of = if sized(column, vectors) { + grouped + } else { + slots + }; + let reached = |(_, share): (u8, u8)| stale * 100 >= u64::from(share) * of; + stale > 0 && column.consolidation().is_some_and(reached) +} + +/// True when the column is at its size, `vectors` being its rows with a vector. +fn sized(column: &VectorColumn, vectors: u64) -> bool { + vectors >= column.index_at() || column.mode() == VectorSearchMode::Indexed +} + +/// True while a search of the nearest groups may be given this partition: the vectors +/// outside its grouping, in tail blocks and pending, are fewer than twice what makes a +/// grouping due. A partition with no grouping, or with one of one group, is searched whole, +/// and so while the vectors outside its grouping are fewer than twice the column's +/// `index_at`. Under `CONSOLIDATION NONE` a grouping of more groups is always usable. +pub(super) fn usable(partition: &Partition, pending: u64, column: &VectorColumn) -> bool { + let out = partition.slots + pending; + match partition.grouping.filter(|grouping| grouping.groups > 1) { + None => out < column.index_at().saturating_mul(2), + Some(grouping) => column.consolidation().is_none_or(|(change, _)| { + let due = (100 * column.grouped_from()).max(u64::from(change) * grouping.slots); + out.saturating_mul(100) < due.saturating_mul(2) + }), + } +} + +/// Calls `each` with every partition of a column and its key tuple. +fn entries( + read: &dyn Read, + vector: VectorId, + each: &mut dyn FnMut(&Key, Partition) -> Result<()>, +) -> Result<()> { + let first = Key::from_values(&[], &[])?; + let from = VectorRecord::Partition(vector, &first); + read.vectors(from, &mut |record, entry| match record { + VectorRecord::Partition(column, tuple) if column == vector => { + each(tuple, Partition::decode(read, entry)?)?; + Ok(true) + } + _ => Ok(false), + }) +} + +/// How many rows have a vector in a column, every partition together. +fn live(read: &dyn Read, vector: VectorId) -> Result { + let mut live = 0; + entries(read, vector, &mut |_, partition| { + live += partition.live; + Ok(()) + })?; + Ok(live) +} + +/// How many rows of the partition of this number are pending. +pub(super) fn pending(read: &dyn Read, vector: VectorId, number: u32) -> Result { + let first = Key::from_values(&[], &[])?; + let mut pending = 0; + let from = VectorRecord::Pending(vector, number, &first); + read.vectors(from, &mut |record, _| { + let more = matches!(record, VectorRecord::Pending(column, partition, _) + if (column, partition) == (vector, number)); + pending += u64::from(more); + Ok(more) + })?; + Ok(pending) +} + +/// Calls `each` with every partition that is building a grouping or has one due, and its +/// column and key tuple: the one walk [`waiting`] and [`remaining`] make. A column's rows +/// with a vector are added up before any of its partitions is asked. +fn working( + read: &dyn Read, + catalog: &Catalog, + each: &mut dyn FnMut(&VectorColumn, &Key, &Partition) -> Result<()>, +) -> Result<()> { + for table in catalog.tables() { + for column in &table.vectors { + let vectors = live(read, column.id)?; + entries(read, column.id, &mut |tuple, partition| { + let pending = pending(read, column.id, partition.number)?; + if partition.building.is_some() || due(&partition, pending, column, vectors) { + each(column, tuple, &partition)?; + } + Ok(()) + })?; + } + } + Ok(()) +} + +/// The partitions with work for a build at this moment: those with none in force first, then +/// the others, each in the order of column and partition number. The list and each key tuple +/// are taken from `held` before they are made. +pub(crate) fn waiting( + read: &dyn Read, + catalog: &Catalog, + held: &mut Reservation, +) -> Result> { + let mut found = Vec::new(); + working(read, catalog, &mut |column, tuple, partition| { + let grouped = partition.grouping.is_some(); + room(&mut found, usize::MAX, held)?; + held.grow(tuple.as_bytes().len() as u64)?; + found.push((grouped, column.id, partition.number, tuple.clone())); + Ok(()) + })?; + found.sort(); + held.grow((found.len() * size_of::()) as u64)?; + let work = |(_, vector, _, tuple)| Work { vector, tuple }; + Ok(found.into_iter().map(work).collect()) +} + +/// Works out the next step for one partition from this moment: the beginning of a build that +/// is due, or the next step of the one under way. `None` when the partition has no work left. +/// It reads and works out only, and holds what it made through a reservation named +/// `maintenance`; a refusal for memory is the step's answer, with nothing changed. +pub(crate) fn prepare( + work: &Work, + read: &dyn Read, + catalog: &Catalog, + budget: &Budget, +) -> Result> { + let vector = work.vector; + let found = directory::partition(read, vector, &work.tuple)?; + let Some(((table, column), partition)) = catalog.vector(vector).zip(found) else { + return Ok(None); + }; + let mut step = Step { + work: work.clone(), + column: column.clone(), + at: (partition.number, None), + puts: Vec::new(), + removes: Vec::new(), + left: Vec::new(), + switch: None, + centres: Vec::new(), + mean: None, + done: 0, + held: budget.reserve(0, "maintenance")?, + }; + // A build's record is read into lists as long as itself: its bytes are taken first. + let own = VectorRecord::Build(vector, partition.number); + step.held + .grow(kept(read, own, <[u8]>::len)?.unwrap_or(0) as u64)?; + let Some(build) = Build::read(read, vector, partition.number)? else { + let pending = pending(read, vector, partition.number)?; + let due = due(&partition, pending, column, live(read, vector)?); + return Ok(due.then_some(step)); + }; + step.at.1 = Some((build.number, build.steps)); + let dimension = table.dimension(column); + let ground = Ground { + read, + column, + dimension, + partition, + build, + }; + match ground.build.phase { + Phase::Train { .. } => train(&ground, &mut step)?, + Phase::Assign { .. } => assign(&ground, &mut step)?, + Phase::Pack { .. } => pack::pack(&ground, &mut step)?, + Phase::Measure { .. } => pack::measure(&ground, &mut step)?, + Phase::Width { .. } => pack::width(&ground, &mut step)?, + Phase::Switch { .. } => pack::switch(&ground, &mut step)?, + Phase::Sweep => pack::sweep(&ground, &mut step)?, + } + Ok(Some(step)) +} + +/// Stages a step, under the commit lock, over the newest moment. Stages nothing and answers +/// zero when the partition is gone or has another number, or is not building the grouping, +/// at the count of steps, the step was worked out for: the one place that is checked. A step +/// that begins a build stages nothing either when no grouping is due any more. Otherwise +/// answers how many vectors the step finished. +/// +/// The step that begins a build takes the build's number from the column's run of partition +/// numbers, writes the build's record, and closes the partition's tail blocks. Any other +/// stages its records, and writes the partition's record again, from what it says now, when +/// it stages the mean, puts a grouping in force or removes the build's record. +pub(crate) fn stage_step(step: Step, changes: &mut ChangeSet) -> Result { + let (vector, tuple) = (step.work.vector, &step.work.tuple); + let found = directory::partition(&*changes, vector, tuple)?; + let Some(mut partition) = found.filter(|partition| partition.number == step.at.0) else { + return Ok(0); + }; + let entry = VectorRecord::Partition(vector, tuple); + let own = VectorRecord::Build(vector, partition.number); + let build = Build::read(&*changes, vector, partition.number)?; + let mut held = step.held; + let Some((number, steps)) = step.at.1 else { + let vectors = live(&*changes, vector)?; + let pending = pending(&*changes, vector, partition.number)?; + if build.is_some() || !due(&partition, pending, &step.column, vectors) { + return Ok(0); + } + let next = VectorNumber::NextPartition(vector); + let number = changes.vector_number(next)?; + changes.put_vector_number(next, number + 1)?; + let number = u32::try_from(number).map_err(|_| { + not_supported("more than 4,294,967,296 partitions over the life of a VECTOR column") + })?; + let slots = partition.grouping.map_or(0, |grouping| grouping.slots) + partition.slots; + let grouped = sized(&step.column, vectors) && slots >= step.column.grouped_from(); + let groups = if grouped { groups_for(slots) } else { 1 }; + held.grow(HEAD as u64 + 8 * u64::from(groups))?; + let build = Build { + number, + steps: 0, + from: partition.grouping, + tail: (partition.first, partition.blocks), + slots, + upto: partition.sealed, + stale: slots + pending > partition.live, + groups, + phase: match groups { + 1 => Phase::Assign { next: 0 }, + _ => Phase::Train { round: 0 }, + }, + counts: vec![0; groups as usize], + starts: Vec::new(), + tens: Vec::new(), + needing: 0, + }; + changes.put_vector(own, build.encode())?; + (partition.building, partition.closed) = (Some(number), partition.blocks); + changes.put_vector(entry, partition.encode())?; + return Ok(0); + }; + let build = build.filter(|build| (build.number, build.steps) == (number, steps)); + let Some(build) = build.filter(|_| partition.building == Some(number)) else { + return Ok(0); + }; + let was = partition; + for (record, value) in step.puts { + held.shrink(value.len() as u64); + changes.put_vector(record, value)?; + } + for record in &step.removes { + changes.delete_vector(*record)?; + } + for key in &step.left { + changes.delete_vector(VectorRecord::Left(vector, partition.number, key))?; + } + if !step.centres.is_empty() { + let dimension = step.centres.len() / build.groups as usize; + grouping::put_centres(changes, vector, number, dimension, &step.centres)?; + } + if let Some(mean) = &step.mean { + grouping::put_mean(changes, vector, partition.number, mean)?; + partition.mean = true; + // The tail blocks sealed since the build began were sealed with no mean to take a + // sketch from: each is given its record of sketches here, from this mean. + for at in partition.closed..partition.blocks { + let block = VectorRecord::Block(vector, partition.number, at); + let mut sketches = None; + changes.vectors(block, &mut |found, bytes| { + let slots = BlockReader::new(bytes, mean.len()); + if found == block && slots.is_whole() { + held.grow((4 + slots.len() * sketch_len(mean.len())) as u64)?; + sketches = Some(sketches_of(bytes, mean.len(), &mut |_| mean)); + } + Ok(false) + })?; + let sketches = sketches.ok_or_else(|| changes.damaged(BROKEN))?; + held.shrink(sketches.len() as u64); + changes.put_vector(VectorRecord::Sketch(vector, partition.number, at), sketches)?; + } + } + if let Some(grouped) = step.switch { + // The tail blocks the build read go with the switch, and their slots with them. + let read = build.slots - build.from.map_or(0, |from| from.slots); + let slots = partition.slots.checked_sub(read); + partition.slots = slots.ok_or_else(|| changes.damaged(directory::SHORT))?; + (partition.grouping, partition.first) = (Some(grouped), build.tail.1); + } + if step.removes.contains(&own) { + partition.building = None; + } + if partition != was { + changes.put_vector(entry, partition.encode())?; + } + Ok(step.done) +} + +/// Records, in a commit of its own, that the next step of a partition's build was refused +/// for memory and how many bytes it needed, so that every reader's inspection can say so. +/// Stages nothing when the build's record says so already. +pub(crate) fn stage_refusal(work: &Work, needed: u64, changes: &mut ChangeSet) -> Result<()> { + let Some(partition) = directory::partition(&*changes, work.vector, &work.tuple)? else { + return Ok(()); + }; + if let Some(mut build) = Build::read(&*changes, work.vector, partition.number)? + && build.needing != needed + { + build.needing = needed; + let own = VectorRecord::Build(work.vector, partition.number); + changes.put_vector(own, build.encode())?; + } + Ok(()) +} + +/// What is left to do, over every VECTOR column, at one moment. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub(crate) struct Remaining { + /// Slots of builds under way not yet packed, and the slots of partitions whose grouping + /// is due and not begun. Those of a stalled build are left out: no further cycle packs + /// them until memory is found. + pub vectors: u64, + /// Partitions with a build due or under way. + pub partitions: u64, + /// Those of them whose last step was refused for memory. + pub stalled: u64, +} + +/// What is left to do at this moment. +pub(crate) fn remaining(read: &dyn Read, catalog: &Catalog) -> Result { + let mut left = Remaining::default(); + working(read, catalog, &mut |column, _, partition| { + let build = Build::read(read, column.id, partition.number)?; + let grouped = partition.grouping.map_or(0, |grouping| grouping.slots); + left.partitions += 1; + match build { + Some(build) if build.needing > 0 => left.stalled += 1, + Some(build) => left.vectors += build.unpacked(), + None => left.vectors += grouped + partition.slots, + } + Ok(()) + })?; + Ok(left) +} + +/// Gives `each` the block the build reads at this place, read in place, and answers what it +/// answers. A block that is missing or is not whole is damage. +fn input( + ground: &Ground<'_>, + place: u32, + each: &mut dyn FnMut(&BlockReader<'_>) -> Result, +) -> Result { + let record = (ground.build).block(ground.column.id, ground.partition.number, place); + let mut answer = None; + ground.read.vectors(record, &mut |found, bytes| { + let block = BlockReader::new(bytes, ground.dimension); + if found == record && block.is_whole() { + answer = Some(each(&block)?); + } + Ok(false) + })?; + answer.ok_or_else(|| ground.read.damaged(BROKEN)) +} + +/// Gives `each` the slots the build samples, each with its number among them, in one walk +/// over the blocks it reads. +fn sample(ground: &Ground<'_>, each: &mut dyn FnMut(u64, &[u8]) -> Result<()>) -> Result<()> { + let build = &ground.build; + let mut wanted = sampled(build.slots, build.groups).zip(0..).peekable(); + let mut before = 0; + for place in 0..build.blocks() { + input(ground, place, &mut |block| { + let end = before + block.len() as u64; + while let Some((slot, at)) = wanted.next_if(|(slot, _)| *slot < end) { + each(at, block.slot((slot - before) as usize))?; + } + before = end; + Ok(()) + })?; + } + Ok(()) +} + +/// Adds the build's record as a step leaves it to the step's records: one more step staged, +/// and none refused. +pub(super) fn moved(ground: &Ground<'_>, step: &mut Step, mut build: Build) { + (build.steps, build.needing) = (ground.build.steps + 1, 0); + let own = VectorRecord::Build(ground.column.id, ground.partition.number); + step.puts.push((own, build.encode())); +} + +/// The step that finds centres: one round over the sample. The first round of a build first +/// takes its centres from the sample, and that of a partition's first build also takes the +/// sample's mean, which the step stages as the partition's. After the last round the build +/// gives slots their groups. +/// +/// The step's reservation is grown first by what it takes: the build's record and its copy, +/// four bytes a number of every centre, the sums, and in the first round eight bytes a group +/// for the places of the first centres and eight a number for the mean and for one centre. +pub(super) fn train(ground: &Ground<'_>, step: &mut Step) -> Result<()> { + let Phase::Train { round } = ground.build.phase else { + unreachable!("a round is the step of a build that is finding its centres") + }; + let (build, dimension) = (&ground.build, ground.dimension); + let (groups, held) = (build.groups, &mut step.held); + held.grow(HEAD as u64 + 8 * u64::from(groups))?; + let mut centres = if round == 0 { + let sample_of = build.slots.min(SAMPLE_TO_A_CENTRE * u64::from(groups)); + held.grow((4 * dimension as u64 + 8 + DRAWN) * u64::from(groups))?; + let first = first_centres(sample_of, groups); + held.shrink(DRAWN * u64::from(groups)); + let mut first = first.into_iter().peekable(); + let mut centres = Vec::with_capacity(dimension * groups as usize); + let mut mean = None; + if !ground.partition.mean { + mean = Some(Mean::new(dimension, held)?); + } + sample(ground, &mut |at, slot| { + // The vector a slot keeps is the mean of that slot alone. + if first.next_if_eq(&at).is_some() { + let mut centre = Mean::new(dimension, held)?; + centre.take(slot); + held.grow(4 * dimension as u64)?; + centres.extend(centre.finish()); + held.shrink(12 * dimension as u64); + } + mean.iter_mut().for_each(|mean| mean.take(slot)); + Ok(()) + })?; + if let Some(mean) = mean { + // The mean's bytes are taken before it is made; its sum's are given back after. + held.grow(4 * dimension as u64)?; + step.mean = Some(mean.finish()); + held.shrink(8 * dimension as u64); + } + centres + } else { + let (read, vector) = (ground.read, ground.column.id); + grouping::centres(read, vector, build.number, groups, dimension, held)? + }; + let mut sums = Round::new(groups, dimension, held)?; + sample(ground, &mut |_, slot| { + sums.take(¢res, slot); + Ok(()) + })?; + sums.finish(&mut centres); + step.centres = centres; + let phase = match round + 1 { + ROUNDS => Phase::Assign { next: 0 }, + round => Phase::Train { round }, + }; + moved( + ground, + step, + Build { + phase, + ..build.clone() + }, + ); + Ok(()) +} + +/// When the row of this key last left a sealed slot of the partition: how many slots had +/// been sealed into the partition's tail by then, `None` when it never did. A record that is +/// not eight bytes, or says more than the partition has sealed, is damage. +fn left(ground: &Ground<'_>, key: &Key) -> Result> { + let record = VectorRecord::Left(ground.column.id, ground.partition.number, key); + kept(ground.read, record, |value| left_at(ground, value))?.transpose() +} + +/// The count a `Left` record of the partition holds. A value that is not eight bytes, or +/// says more than the partition has sealed, is damage. +pub(super) fn left_at(ground: &Ground<'_>, value: &[u8]) -> Result { + let sound = |left: &u64| *left <= ground.partition.sealed; + let left = value.try_into().ok().map(u64::from_le_bytes).filter(sound); + let damage = "a record of a row that left its slot is not what its partition says"; + left.ok_or_else(|| ground.read.damaged(damage)) +} + +/// The step that gives the slots of the next blocks their groups: whole blocks while their +/// slots times the build's groups come to at most `STEP_COMPARISONS`, and one block at least. +/// A block's groups are one record, two bytes a slot, `0xFFFF` for a stale slot. After the +/// last block the build packs, from the first place. +/// +/// A slot goes to the group whose centre is nearest; a build of one group reads no centre. +/// A build with stale slots first looks up when the row of the slot's key last left a slot, +/// and the slot is stale when it was sealed before then: a slot of the grouping the build +/// reads, all sealed before that grouping's `upto`, when the row left at that count or +/// later; a tail slot when its own place among the slots ever sealed is below the count. +/// +/// The step's reservation is grown first by the build's record and its copy, the centres, +/// and two bytes for each slot of a block. +pub(super) fn assign(ground: &Ground<'_>, step: &mut Step) -> Result<()> { + let Phase::Assign { next } = ground.build.phase else { + unreachable!("slots are given their groups by a build that has its centres") + }; + let (build, read, dimension) = (&ground.build, ground.read, ground.dimension); + let (vector, groups, held) = (ground.column.id, build.groups, &mut step.held); + held.grow(HEAD as u64 + 8 * u64::from(groups))?; + let mut after = build.clone(); + let mut centres = Vec::new(); + if groups > 1 { + centres = grouping::centres(read, vector, build.number, groups, dimension, held)?; + } + let from = build.from.map(|from| (from.blocks, from.slots, from.upto)); + let (inside, grouped, applied) = from.unwrap_or_default(); + // The place of the next tail slot among the slots ever sealed. The build's tail slots are + // the last before its `upto`, and the records of the tail blocks whose slots have their + // groups say how many came before this step's. + let mut sealed = build.upto - (build.slots - grouped); + if build.stale { + let given = VectorRecord::Assigned(vector, build.number, inside); + read.vectors(given, &mut |record, given| { + let earlier = matches!(record, VectorRecord::Assigned(column, number, place) + if (column, number) == (vector, build.number) && place < next); + if earlier && given.len() % 2 != 0 { + return Err(read.damaged(pack::ASSIGNED)); + } + sealed += u64::from(earlier) * given.len() as u64 / 2; + Ok(earlier) + })?; + } + let (mut place, mut taken) = (next, 0); + while place < build.blocks() { + let more = input(ground, place, &mut |block| { + let slots = block.len() as u64; + if taken > 0 && (taken + slots) * u64::from(groups) > STEP_COMPARISONS { + return Ok(false); + } + held.grow(2 * slots)?; + let mut given = Vec::with_capacity(2 * block.len()); + for at in 0..block.len() { + let left = build.stale.then(|| left(ground, &block.key(at))); + let stale = left.transpose()?.flatten().is_some_and(|left| { + let tail = place >= inside && sealed + (at as u64) < left; + tail || (place < inside && left >= applied) + }); + let group = match (stale, groups) { + (true, _) => u16::MAX, + (false, 1) => 0, + (false, _) => nearest(¢res, dimension, block.slot(at)).0 as u16, + }; + if !stale { + after.counts[usize::from(group)] += 1; + } + given.extend(group.to_le_bytes()); + } + let record = VectorRecord::Assigned(vector, build.number, place); + step.puts.push((record, given)); + sealed += if place < inside { 0 } else { slots }; + taken += slots; + Ok(true) + })?; + if !more { + break; + } + place += 1; + } + after.phase = Phase::Assign { next: place }; + if place == build.blocks() { + after.phase = Phase::Pack { group: 0 }; + after.starts = vec![Place::default()]; + } + moved(ground, step, after); + Ok(()) +} diff --git a/v2/contextdb/src/vector/directory.rs b/v2/contextdb/src/vector/directory.rs index af2c061b..1e475f04 100644 --- a/v2/contextdb/src/vector/directory.rs +++ b/v2/contextdb/src/vector/directory.rs @@ -2,7 +2,9 @@ //! vector. use std::collections::BTreeMap; +use std::collections::btree_map::Entry; +use crate::memory::Reservation; use crate::store::{Key, Read, VectorId, VectorRecord}; use crate::types::Result; @@ -13,42 +15,129 @@ pub(super) struct Partition { 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. + /// The number after the partition's last tail block. The tail blocks are the blocks + /// sealing adds to: they are kept under `number` and numbered from `first`. pub blocks: u32, + /// The number of the partition's first tail block. The tail blocks before it went into the + /// grouping in force and are gone. + pub first: u32, + /// Tail blocks numbered below this are never added to again: a grouping being built reads + /// them. Every tail block from here on but the last is full. + pub closed: u32, + /// How many slots the tail blocks hold, the slots rows have left behind among them. + pub slots: u64, + /// How many slots have ever been sealed into the partition's tail. It only rises: the tail + /// slots are the last `slots` of that count, in the order they were sealed. + pub sealed: u64, + /// The grouping in force, `None` until one has been built. + pub grouping: Option, + /// The number a grouping being built is kept under, `None` when none is. + pub building: Option, + /// Whether the partition has a mean: its first build keeps one from its first round on, + /// and from then every tail block at or after `closed` has its sketches. + pub mean: bool, +} + +/// Where a grouping's records are and how many it has: all a search of every slot, and a +/// commit that removes the grouping, need of it. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct Grouped { + /// The number its blocks, its centres and its record are kept under. + pub number: u32, + /// How many blocks it has, numbered from zero; every one but the last is full. + pub blocks: u32, + /// How many groups it has. + pub groups: u32, + /// How many slots its blocks hold, the slots rows have left behind since among them. + pub slots: u64, + /// The partition's `sealed` when the build that made it began. A row that left its slot + /// since, a `Left` record of that value or a greater, leaves a slot of this grouping stale. + pub upto: u64, } impl Partition { - /// The record of the partition: its three numbers, low byte first. + /// The record of the partition, 73 bytes, every number low byte first: `number`, `live`, + /// `blocks`, `first`, `closed`, `slots`, `sealed`, one byte of flags (1: it has a grouping; + /// 2: it is building one; 4: it has a mean), the five numbers of the grouping in force and + /// the number being built. A field its flag does not set is zero. pub(super) fn encode(&self) -> Vec { - let mut record = Vec::with_capacity(16); + let grouping = self.grouping.unwrap_or(Grouped { + number: 0, + blocks: 0, + groups: 0, + slots: 0, + upto: 0, + }); + let flags = u8::from(self.grouping.is_some()) + | (u8::from(self.building.is_some()) << 1) + | (u8::from(self.mean) << 2); + let mut record = Vec::with_capacity(ENTRY); record.extend(self.number.to_le_bytes()); record.extend(self.live.to_le_bytes()); record.extend(self.blocks.to_le_bytes()); + record.extend(self.first.to_le_bytes()); + record.extend(self.closed.to_le_bytes()); + record.extend(self.slots.to_le_bytes()); + record.extend(self.sealed.to_le_bytes()); + record.push(flags); + record.extend(grouping.number.to_le_bytes()); + record.extend(grouping.blocks.to_le_bytes()); + record.extend(grouping.groups.to_le_bytes()); + record.extend(grouping.slots.to_le_bytes()); + record.extend(grouping.upto.to_le_bytes()); + record.extend(self.building.unwrap_or(0).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])), + /// The partition a stored record was written from. A record of another length, with a + /// flag above the three, with more slots in its tail than were ever sealed, with its first + /// tail block after its closed one or its closed one after its last, with a grouping of no + /// group, or with a grouping of more than one group and no mean, is damage. + pub(super) fn decode(read: &dyn Read, record: &[u8]) -> Result { + if record.len() != ENTRY { + return Err(read.damaged("a partition record of the vector part is not 73 bytes")); + } + let short = |at: usize| u32::from_le_bytes(std::array::from_fn(|byte| record[at + byte])); + let long = |at: usize| u64::from_le_bytes(std::array::from_fn(|byte| record[at + byte])); + let flags = record[40]; + let partition = Partition { + number: short(0), + live: long(4), + blocks: short(12), + first: short(16), + closed: short(20), + slots: long(24), + sealed: long(32), + grouping: (flags & 1 != 0).then(|| Grouped { + number: short(41), + blocks: short(45), + groups: short(49), + slots: long(53), + upto: long(61), + }), + building: (flags & 2 != 0).then(|| short(69)), + mean: flags & 4 != 0, + }; + let groups = partition.grouping.map_or(1, |grouping| grouping.groups); + if flags > 7 + || partition.sealed < partition.slots + || partition.first > partition.closed + || partition.closed > partition.blocks + || groups == 0 + || (groups > 1 && !partition.mean) + { + return Err(read.damaged("a partition record of the vector part disagrees with itself")); } + Ok(partition) } } +/// Bytes of a partition's record. +const ENTRY: usize = 73; + /// 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( @@ -57,7 +146,7 @@ pub(super) fn partition( tuple: &Key, ) -> Result> { let record = VectorRecord::Partition(vector, tuple); - kept(read, record, |entry| checked(read, entry))?.transpose() + kept(read, record, |entry| Partition::decode(read, entry))?.transpose() } /// What `read` gives for the value of one record of the vector part at the committed moment, @@ -75,10 +164,21 @@ pub(super) fn kept( Ok(found) } +/// Bytes [`selected`] reserves for each partition it answers, beside its key tuple: its entry +/// while the partitions are gathered and its entry in the list they are answered in. +pub(super) const HELD: u64 = 2 * size_of::<(Key, Partition)>() as u64; + /// 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> { +/// partition of the column when `scope` is empty. Each comes once with its key tuple, in the +/// order of the numbers, whatever the order of `scope`. A tuple no row has a vector under +/// selects nothing. [`HELD`] bytes and those of its key tuple are taken from `reservation` for +/// each before it is kept; the caller gives them back when it lets the list go. +pub(super) fn selected( + read: &dyn Read, + vector: VectorId, + scope: &[Key], + reservation: &mut Reservation, +) -> 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 }; @@ -90,8 +190,11 @@ pub(super) fn selected(read: &dyn Read, vector: VectorId, scope: &[Key]) -> Resu VectorRecord::Partition(column, tuple) if column == vector && tuple.starts_with(prefix) => { - let partition = checked(read, entry)?; - found.insert(partition.number, partition); + let partition = Partition::decode(read, entry)?; + if let Entry::Vacant(place) = found.entry(partition.number) { + reservation.grow(HELD + tuple.as_bytes().len() as u64)?; + place.insert((tuple.clone(), partition)); + } Ok(true) } _ => Ok(false), diff --git a/v2/contextdb/src/vector/grouping.rs b/v2/contextdb/src/vector/grouping.rs new file mode 100644 index 00000000..0eb3174d --- /dev/null +++ b/v2/contextdb/src/vector/grouping.rs @@ -0,0 +1,242 @@ +//! The grouping of a partition: its slots sorted by the centre each is nearest, packed into full +//! blocks, with the centres, the place where each group's slots begin, and a sketch of every +//! slot taken from its group's centre. +//! +//! A grouping is written once, by the build that made it, and never changed: the partition's +//! record names the grouping in force, and one commit replaces that name. A partition's mean +//! is written once too, by the first build that finds centres, and stays while the partition does. + +use std::cmp::Ordering; + +use super::block::BLOCK_BYTES; +use super::directory::{Grouped, kept}; +use super::slot::summed; +use crate::memory::Reservation; +use crate::store::{ChangeSet, Read, VectorId, VectorRecord}; +use crate::types::Result; + +/// A slot of a grouping: the block it is in and its place in that block, both from zero. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)] +pub(super) struct Place { + pub block: u32, + pub slot: u32, +} + +/// What a search of the nearest groups reads of a grouping. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) struct Grouping { + /// How far a search scans, in ten-thousandths of a product, from 0 to 20,000: every group + /// whose centre's product with the vector of the search is within this of the best + /// centre's. The most any of the questions the build asked needed for its ten nearest. + pub gap: u32, + /// The most groups that gap took in for any of those questions, from one to the number of + /// groups: what inspection prints. No search reads it. + pub probes: u32, + /// How wide a search takes a sketch's estimate to be, in quarters of the sketch's factor: + /// the least of 2 to 16 that kept 99 in 100 of the neighbours the build measured, and 0 + /// when none did, in which case a search screens the slots of its groups directly. + pub width: u8, + /// Where each group's slots begin, in the order of the groups, and after the last of them + /// where the slots end: one more place than there are groups. Group `g` is the slots from + /// `starts[g]` up to `starts[g + 1]`. + pub starts: Vec, +} + +impl Grouping { + /// The record of the grouping: `gap`, `width` and `probes` as four bytes each, then every place + /// as two numbers of four bytes, all low byte first. + pub(super) fn encode(&self) -> Vec { + let head = [self.gap, u32::from(self.width), self.probes]; + let places = self.starts.iter(); + let places = places.flat_map(|place| [place.block, place.slot]); + let numbers = head.into_iter().chain(places); + numbers.flat_map(u32::to_le_bytes).collect() + } + + /// The grouping kept under `grouped.number`. A record that is missing, of a length other + /// than `12 + 8 * (grouped.groups + 1)`, with a `gap` above 20,000, with `probes` outside + /// `1..=grouped.groups`, with a + /// `width` of 1 or above 16, with a first place other than the first slot, with a place + /// before the one ahead of it or a last place past `grouped.blocks`, is damage. + pub(super) fn read(read: &dyn Read, vector: VectorId, grouped: Grouped) -> Result { + let end = Place { + block: grouped.blocks, + slot: 0, + }; + let record = VectorRecord::Grouping(vector, grouped.number); + let found = kept(read, record, |value| { + if value.len() as u64 != 20 + 8 * u64::from(grouped.groups) { + return None; + } + let (numbers, _) = value.as_chunks::<4>(); + let number = |at: usize| u32::from_le_bytes(numbers[at]); + let (gap, width, probes) = (number(0), number(1), number(2)); + let places = (3..numbers.len()).step_by(2); + let starts: Vec = places + .map(|at| Place { + block: number(at), + slot: number(at + 1), + }) + .collect(); + let agrees = gap <= 20_000 + && (1..=grouped.groups).contains(&probes) + && (width == 0 || (2..=16).contains(&width)) + && starts[0] == Place::default() + && starts.is_sorted() + && starts[starts.len() - 1] <= end; + agrees.then_some(Grouping { + gap, + probes, + width: width as u8, + starts, + }) + })?; + let damage = "a grouping of the vector part is not what its partition says"; + found.flatten().ok_or_else(|| read.damaged(damage)) + } +} + +/// The numbers of four bytes, low byte first, that these bytes hold. +fn numbers(bytes: &[u8]) -> impl Iterator { + let (numbers, _) = bytes.as_chunks::<4>(); + numbers.iter().map(|number| f32::from_le_bytes(*number)) +} + +/// How many centres of this dimension one record holds: as many as fit the bytes of a block. +pub(super) fn centres_to_a_record(dimension: usize) -> usize { + (BLOCK_BYTES / (4 * dimension)).max(1) +} + +/// Stages the centres of a grouping, `centres_to_a_record` to a record, in the order of the +/// groups. Each centre is `dimension` numbers of four bytes, low byte first, of length one. +pub(super) fn put_centres( + changes: &mut ChangeSet, + vector: VectorId, + number: u32, + dimension: usize, + centres: &[f32], +) -> Result<()> { + let records = centres.chunks(centres_to_a_record(dimension) * dimension); + for (place, record) in records.enumerate() { + let value = record.iter().flat_map(|number| number.to_le_bytes()); + let record = VectorRecord::Centres(vector, number, place as u32); + changes.put_vector(record, value.collect())?; + } + Ok(()) +} + +/// Hands `each` the bytes of every centre of a grouping, in the order of the groups, read in +/// place. Records that do not hold `groups` centres of this dimension are damage. +fn each_centre( + read: &dyn Read, + vector: VectorId, + number: u32, + groups: u32, + dimension: usize, + each: &mut dyn FnMut(&[u8]), +) -> Result<()> { + let (mut count, mut whole) = (0, true); + let from = VectorRecord::Centres(vector, number, 0); + read.vectors(from, &mut |record, value| match record { + VectorRecord::Centres(column, grouping, _) if column == vector && grouping == number => { + let more = value.len() / (4 * dimension); + whole = value.len() % (4 * dimension) == 0 && more <= (groups - count) as usize; + if whole { + value.chunks(4 * dimension).for_each(&mut *each); + count += more as u32; + } + Ok(whole) + } + _ => Ok(false), + })?; + if whole && count == groups { + Ok(()) + } else { + Err(read.damaged("the centres of a grouping are not as many as its groups")) + } +} + +/// The centres of a grouping, one after another, read into memory that `held` is grown by +/// first: four bytes a number. Records that do not hold `groups` centres of this dimension, or +/// a number that is not finite, are damage. +pub(super) fn centres( + read: &dyn Read, + vector: VectorId, + number: u32, + groups: u32, + dimension: usize, + held: &mut Reservation, +) -> Result> { + held.grow(4 * dimension as u64 * u64::from(groups))?; + let mut centres = Vec::with_capacity(dimension * groups as usize); + each_centre(read, vector, number, groups, dimension, &mut |centre| { + centres.extend(numbers(centre)); + })?; + if centres.iter().all(|number| number.is_finite()) { + Ok(centres) + } else { + Err(read.damaged("a centre of a grouping has a number that is not finite")) + } +} + +/// Every group of a grouping with the product of its centre with `unit`, nearest first: by that product, +/// greatest first, groups of one product in the order of their numbers. The centres are read +/// in place; `held` is grown by eight bytes a group for the list. A grouping of one group has +/// no centre: its group is answered with a product of zero and nothing is read. A product +/// that is not finite is damage. +pub(super) fn nearest( + read: &dyn Read, + vector: VectorId, + grouped: Grouped, + unit: &[f32], + held: &mut Reservation, +) -> Result> { + held.grow(8 * u64::from(grouped.groups))?; + if grouped.groups == 1 { + return Ok(vec![(0, 0.0)]); + } + let mut found = Vec::with_capacity(grouped.groups as usize); + let centre = &mut |centre: &[u8]| { + let (centre, _) = centre.as_chunks::<4>(); + found.push((found.len() as u32, summed(centre, unit, f32::from_le_bytes))); + }; + let Grouped { number, groups, .. } = grouped; + each_centre(read, vector, number, groups, unit.len(), centre)?; + if found.iter().any(|(_, product)| !product.is_finite()) { + return Err(read.damaged("a centre of a grouping has a number that is not finite")); + } + // The sort keeps groups of one product in the order they were read, that of their numbers. + found.sort_by(|one, other| other.1.partial_cmp(&one.1).unwrap_or(Ordering::Equal)); + Ok(found) +} + +/// Stages the mean of a partition: `dimension` numbers of four bytes, low byte first. +pub(super) fn put_mean( + changes: &mut ChangeSet, + vector: VectorId, + partition: u32, + mean: &[f32], +) -> Result<()> { + let value = mean.iter().flat_map(|number| number.to_le_bytes()); + changes.put_vector(VectorRecord::Mean(vector, partition), value.collect()) +} + +/// The mean of a partition, read into memory that `held` is grown by first. A record that is +/// missing for a partition whose entry says it has one, of another length than four bytes a +/// number, or with a number that is not finite, is damage. +pub(super) fn mean( + read: &dyn Read, + vector: VectorId, + partition: u32, + dimension: usize, + held: &mut Reservation, +) -> Result> { + held.grow(4 * dimension as u64)?; + let record = VectorRecord::Mean(vector, partition); + let found = kept(read, record, |value| { + (value.len() == 4 * dimension).then(|| numbers(value).collect::>()) + })?; + let found = found.flatten(); + let found = found.filter(|mean| mean.iter().all(|number| number.is_finite())); + found.ok_or_else(|| read.damaged("the mean of a partition is not what its entry says")) +} diff --git a/v2/contextdb/src/vector/inspect.rs b/v2/contextdb/src/vector/inspect.rs new file mode 100644 index 00000000..76c36a45 --- /dev/null +++ b/v2/contextdb/src/vector/inspect.rs @@ -0,0 +1,161 @@ +//! What inspection reads of a VECTOR column: its partitions as their records describe them. +//! It reads the vector part's small records and no block, no centre and no row. + +use super::build::{self, Build}; +use super::directory::Partition; +use super::grouping::Grouping; +use super::seal::room; +use crate::memory::Reservation; +use crate::schema::VectorColumn; +use crate::store::{Key, Read, VectorRecord}; +use crate::types::{Error, Result}; + +/// What a partition is doing about its grouping, in the order a column's state is taken: the +/// last of these that any of its partitions is in. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub(crate) enum State { + /// Nothing is due. + Idle, + /// A grouping is due and its build has not begun. + Pending, + /// A build is under way. + Building, + /// A build is under way and its last step was refused for memory. + Stalled, +} + +/// One partition as its records describe it. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Inspected { + /// The partition's key tuple. + pub tuple: Key, + /// Rows with a vector in the partition. + pub live: u64, + /// Those of them not yet in a block. + pub pending: u64, + /// Slots in its blocks that rows have left behind, and whether they alone are enough for + /// a grouping to be due. + pub stale: u64, + pub left_behind: bool, + /// Slots in tail blocks: sealed, and in no grouping yet. + pub tail: u64, + /// The bytes of its slots; and of its grouping's centres and record, its sketches and its + /// mean. + pub slot_bytes: u64, + pub index_bytes: u64, + /// The groups of the grouping in force and the most of them a search scanned for the + /// questions its build asked; zero and zero when there is none. + pub groups: u32, + pub probes: u32, + /// Whether a search of the nearest groups may be given this partition. + pub usable: bool, + pub state: State, + /// Of a build under way: the slots it reads and those it has packed; of a grouping that + /// is due and not begun: the partition's slots and zero. + pub total: u64, + pub done: u64, + /// The bytes a refused step needed; zero unless `state` is `Stalled`. + pub needing: u64, + /// The reason the grouping in force cannot be read, when it cannot. + pub damage: Option, +} + +/// The partitions of a column in the order of their key tuples, leaving out the first `skip` +/// and stopping after `most`. A column with no PARTITION_KEY has its one partition before its +/// first vector. +pub(crate) fn partitions( + read: &dyn Read, + column: &VectorColumn, + dimension: usize, + skip: u64, + most: Option, + held: &mut Reservation, +) -> Result> { + let first = Key::from_values(&[], &[])?; + let mut entries = Vec::new(); + let from = VectorRecord::Partition(column.id, &first); + read.vectors(from, &mut |record, entry| match record { + VectorRecord::Partition(vector, tuple) if vector == column.id => { + room(&mut entries, usize::MAX, held)?; + held.grow(tuple.as_bytes().len() as u64)?; + entries.push((tuple.clone(), Partition::decode(read, entry)?)); + Ok(true) + } + _ => Ok(false), + })?; + if entries.is_empty() && column.partition_key.is_empty() { + entries.push((first, Partition::default())); + } + // Whether a grouping is due depends on the rows with a vector in the whole column. + let vectors = entries.iter().map(|(_, partition)| partition.live).sum(); + let count = |number: u64| usize::try_from(number).unwrap_or(usize::MAX); + let asked = entries.into_iter().skip(count(skip)); + let asked = asked.take(most.map_or(usize::MAX, count)); + held.grow((asked.len() * size_of::()) as u64)?; + asked + .map(|(tuple, partition)| { + inspected(read, column, dimension as u64, vectors, tuple, &partition) + }) + .collect() +} + +/// One partition from its entry, its pending count, its grouping's record and its build's. +fn inspected( + read: &dyn Read, + column: &VectorColumn, + n: u64, + vectors: u64, + tuple: Key, + partition: &Partition, +) -> Result { + let pending = build::pending(read, column.id, partition.number)?; + let grouped = partition.grouping.map_or(0, |grouping| grouping.slots); + let slots = grouped + partition.slots; + let mut index_bytes = 4 * n * u64::from(partition.mean); + let (mut groups, mut probes, mut damage) = (0, 0, None); + if let Some(grouping) = partition.grouping { + groups = grouping.groups; + let (g, tail) = (u64::from(groups), partition.blocks - partition.first); + index_bytes += 12 + 8 * (g + 1); + if groups > 1 { + let sketches = slots * (n.div_ceil(8) + 2) + 4 * u64::from(grouping.blocks + tail); + index_bytes += 4 * n * g + sketches; + } + match Grouping::read(read, column.id, grouping) { + Ok(grouping) => probes = grouping.probes, + Err(Error::StoreCorrupted { reason, .. }) => damage = Some(reason), + Err(other) => return Err(other), + } + } + let (state, total, done, needing) = match Build::read(read, column.id, partition.number)? { + Some(build) => { + let state = if build.needing > 0 { + State::Stalled + } else { + State::Building + }; + let done = build.slots.saturating_sub(build.unpacked()); + (state, build.slots, done, build.needing) + } + None if build::due(partition, pending, column, vectors) => (State::Pending, slots, 0, 0), + None => (State::Idle, 0, 0, 0), + }; + Ok(Inspected { + tuple, + live: partition.live, + pending, + stale: (slots + pending).saturating_sub(partition.live), + left_behind: build::left_behind(partition, pending, column, vectors), + tail: partition.slots, + slot_bytes: slots * (8 + n), + index_bytes, + groups, + probes, + usable: build::usable(partition, pending, column), + state, + total, + done, + needing, + damage, + }) +} diff --git a/v2/contextdb/src/vector/mod.rs b/v2/contextdb/src/vector/mod.rs index 691fc8b5..4ce14cc4 100644 --- a/v2/contextdb/src/vector/mod.rs +++ b/v2/contextdb/src/vector/mod.rs @@ -1,31 +1,47 @@ -//! vector: the blocks of a VECTOR column, exact search over them, and its partitions. +//! vector: the blocks of a VECTOR column, their groupings and sketches, 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. +//! Every vector is stored once with its row, at full precision unless the column is declared +//! SQ8. 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. +//! An exact search, and any search of a partition searched whole, 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. +//! +//! A partition that has grown is given a grouping: its slots sorted around centres, and for +//! every slot a sketch of one bit a number. A search that may leave vectors out ranks the +//! centres, estimates the sketches of the nearest groups and of the tail, and screens the +//! one-byte copies of only the slots whose estimate can still reach the answer. mod block; +mod build; mod directory; +mod grouping; +mod inspect; +mod pack; mod seal; mod search; +mod sketch; mod slot; mod stage; mod tally; #[cfg(test)] mod tests; +mod train; +pub(crate) use build::{prepare, remaining, stage_refusal, stage_step, waiting}; +pub(crate) use inspect::{Inspected, State, partitions}; pub(crate) use seal::{Sealed, seal}; -pub(crate) use search::{Nearest, search}; -pub(crate) use slot::distance; +pub(crate) use search::{Nearest, Route, plan, search}; +pub(crate) use slot::{Reduced, distance}; pub(crate) use stage::stage; pub(crate) use tally::Tally; diff --git a/v2/contextdb/src/vector/pack.rs b/v2/contextdb/src/vector/pack.rs new file mode 100644 index 00000000..10d532a8 --- /dev/null +++ b/v2/contextdb/src/vector/pack.rs @@ -0,0 +1,534 @@ +//! The last steps of a build: packing the slots of each group into the blocks of the new +//! grouping with a sketch of each, measuring how far a search scans and how wide it +//! takes a sketch's estimate to be, and putting the grouping in force. + +use super::block::{BLOCK_BYTES, BROKEN, BlockReader, BlockWriter}; +use super::build::{ + Build, Ground, MEASURED, Phase, STEP_COMPARISONS, STEP_SLOTS, Step, left_at, moved, +}; +use super::directory::{Grouped, kept}; +use super::grouping::{Grouping, Place, centres, centres_to_a_record, mean, nearest}; +use super::seal::room; +use super::search::{gap_needed, kept_by_width}; +use super::sketch::{WIDEST, sketch_len, sketches_of}; +use super::slot::{SLOT_HEADER, slot_len}; +use crate::store::{Key, Read, VectorId, VectorRecord}; +use crate::types::Result; + +/// More bytes than the fields of a build's record take besides its lists. +const FIELDS: u64 = 256; + +/// The most `Left` records one step removes. +const SWEPT: usize = 4096; + +/// What a store is told when a build's record of the groups of a block's slots is not two +/// bytes for each slot of the block. +pub(super) const ASSIGNED: &str = + "the groups a build of the vector part gave a block are not one for each slot"; + +/// What a store is told when a build's `Assigned` records give the groups of a step more +/// slots than the build counted for them. +const COUNTED: &str = "a build of the vector part was given more slots than it counted"; + +/// A copy of the build for a step to move on. What its lists can come to (a count and a +/// place a group, ten places a question) is taken twice first: for the copy, and for the +/// record the step puts. +fn copy(ground: &Ground<'_>, step: &mut Step) -> Result { + let lists = 12 * u64::from(ground.build.groups) + 8 + 80 * u64::from(MEASURED); + step.held.grow(2 * lists + FIELDS)?; + Ok(ground.build.clone()) +} + +/// The grouping a build has packed, as the partition's record names it once it is in force. +fn packed(build: &Build) -> Grouped { + let end = build.starts.last().copied().unwrap_or_default(); + Grouped { + number: build.number, + blocks: end.block + u32::from(end.slot > 0), + groups: build.groups, + slots: build.counts.iter().map(|count| u64::from(*count)).sum(), + upto: build.upto, + } +} + +/// The grouping of the build's places with these three numbers, its places reserved first. +fn grouping_of( + build: &Build, + step: &mut Step, + (gap, probes, width): (u32, u32, u8), +) -> Result { + step.held.grow(8 * build.starts.len() as u64)?; + Ok(Grouping { + gap, + probes, + width, + starts: build.starts.clone(), + }) +} + +/// Hands `each` the block kept under `record`, read in place. A block that is missing or not +/// whole is damage. +fn with_block( + read: &dyn Read, + record: VectorRecord<'_>, + dimension: usize, + each: &mut dyn FnMut(&BlockReader<'_>) -> Result<()>, +) -> Result<()> { + let mut whole = false; + read.vectors(record, &mut |first, bytes| { + let block = BlockReader::new(bytes, dimension); + whole = first == record && block.is_whole(); + if whole { + each(&block)?; + } + Ok(false) + })?; + if whole { + Ok(()) + } else { + Err(read.damaged(BROKEN)) + } +} + +/// Finishes a block of the build and puts it, its bytes reserved first. With centres, the +/// record of its sketches is put with it, made from the whole block: the reference of a slot +/// is the centre of the last group that begins at or before the slot's place. +fn put_block( + step: &mut Step, + build: &Build, + block: u32, + writer: BlockWriter, + centres: &[f32], + dimension: usize, +) -> Result<()> { + let vector = step.work.vector; + step.held.grow(writer.bytes() as u64)?; + let value = writer.finish(); + if !centres.is_empty() { + let count = BlockReader::new(&value, dimension).len(); + step.held.grow((4 + count * sketch_len(dimension)) as u64)?; + let mut centre = |slot: usize| { + let place = Place { + block, + slot: slot as u32, + }; + let group = build.starts.partition_point(|start| *start <= place) - 1; + ¢res[group * dimension..][..dimension] + }; + let sketches = sketches_of(&value, dimension, &mut centre); + let record = VectorRecord::Sketch(vector, build.number, block); + step.puts.push((record, sketches)); + } + let record = VectorRecord::Block(vector, build.number, block); + step.puts.push((record, value)); + Ok(()) +} + +/// The step that packs the next groups: whole groups while they come to at most +/// `STEP_SLOTS` slots, and one group at least. Within a group the slots are in the order of +/// their rows' keys and a key comes once. Every block it writes comes with the record of its +/// sketches, each taken from the centre of its slot's group. +/// +/// Reserved before it is taken: for each slot of the step's groups its bytes, its key, and +/// the entries of the list the slots are sorted in and of the two that list is parted into; +/// the lists of the block being written, and every block and record of sketches the step +/// puts. +pub(super) fn pack(ground: &Ground<'_>, step: &mut Step) -> Result<()> { + let Phase::Pack { group } = ground.build.phase else { + unreachable!("a build that packs is in its packing phase") + }; + let (read, vector, dimension) = (ground.read, step.work.vector, ground.dimension); + let number = ground.partition.number; + let mut build = copy(ground, step)?; + let count = |group: u32| u64::from(ground.build.counts[group as usize]); + let (mut end, mut total) = (group + 1, count(group)); + while end < build.groups && total + count(end) <= STEP_SLOTS { + total += count(end); + end += 1; + } + let len = slot_len(dimension); + let each = len + size_of::<(u32, Key, u32)>() + size_of::<(Key, u32, u32)>(); + step.held.grow(total * each as u64 + BLOCK_BYTES as u64)?; + let centres = match build.groups { + 1 => Vec::new(), + groups => centres( + read, + vector, + build.number, + groups, + dimension, + &mut step.held, + )?, + }; + + // Every slot of the step's groups: its group, its key and where its bytes were copied. + let total = total as usize; + let mut found: Vec<(u32, Key, u32)> = Vec::with_capacity(total); + let mut slots: Vec = Vec::with_capacity(total * len); + let held = &mut step.held; + let from = VectorRecord::Assigned(vector, build.number, 0); + read.vectors(from, &mut |record, given| { + let VectorRecord::Assigned(column, grouping, place) = record else { + return Ok(false); + }; + if (column, grouping) != (vector, build.number) { + return Ok(false); + } + let (groups, []) = given.as_chunks::<2>() else { + return Err(read.damaged(ASSIGNED)); + }; + let to = |at: usize| u32::from(u16::from_le_bytes(groups[at])); + if !(0..groups.len()).any(|at| (group..end).contains(&to(at))) { + return Ok(true); + } + let block = build.block(vector, number, place); + with_block(read, block, dimension, &mut |block| { + if given.len() != 2 * block.len() { + return Err(read.damaged(ASSIGNED)); + } + for at in (0..groups.len()).filter(|at| (group..end).contains(&to(*at))) { + if found.len() == total { + return Err(read.damaged(COUNTED)); + } + let key = block.key(at); + held.grow(key.as_bytes().len() as u64)?; + slots.extend_from_slice(block.slot(at)); + found.push((to(at), key, found.len() as u32)); + } + Ok(()) + })?; + Ok(true) + })?; + found.sort_unstable(); + let parted = found.into_iter().map(|(to, key, at)| (key, (to, at))); + let (keys, order): (Vec, Vec<(u32, u32)>) = parted.unzip(); + + // The slots go on from the last block the build wrote while it has room. `place` is where + // the next slot goes if its block takes it. + let mut place = build.starts.pop().unwrap_or_default(); + let (mut writer, mut left) = if place.slot == 0 { + BlockWriter::new(dimension, None, &keys) + } else { + let last = VectorRecord::Block(vector, build.number, place.block); + let carried = kept(read, last, |last| { + let whole = BlockReader::new(last, dimension).is_whole(); + whole.then(|| BlockWriter::new(dimension, Some(last), &keys)) + })?; + carried.flatten().ok_or_else(|| read.damaged(BROKEN))? + }; + build.counts[group as usize..end as usize].fill(0); + for (done, (key, (to, at))) in keys.iter().zip(&order).enumerate() { + if left == 0 { + // A block that was full before this step is left as it is. + if !writer.is_empty() { + put_block(step, &build, place.block, writer, ¢res, dimension)?; + } + place = Place { + block: place.block + 1, + slot: 0, + }; + (writer, left) = BlockWriter::new(dimension, None, &keys[done..]); + } + // The slot's group begins here if this is its first, and so does every group before + // it that has no slot. + build.starts.resize(*to as usize + 1, place); + writer.push_slot(key, &slots[*at as usize * len..][..len]); + build.counts[*to as usize] += 1; + left -= 1; + place.slot += 1; + } + if !writer.is_empty() { + put_block(step, &build, place.block, writer, ¢res, dimension)?; + } + build.starts.resize(end as usize + 1, place); + + step.done = keys.len() as u64; + build.phase = if end < build.groups { + Phase::Pack { group: end } + } else if build.groups == 1 || packed(&build).slots < 2 { + let (gap, probes, width) = (0, 1, 0); + Phase::Switch { gap, probes, width } + } else { + Phase::Measure { done: 0, most: 0 } + }; + moved(ground, step, build); + Ok(()) +} + +/// Makes question `number` of a build over `slots` slots in `query`, from the vectors that +/// slots of the new grouping keep (scale times codes), brought to length one. An even +/// question is the middle of two: the slot at `number * slots / M`, with `M` the lesser of +/// `MEASURED` and `slots`, and the one half the grouping further on, added. An odd question is +/// the first of those alone, a stored vector, whose nearest slot is itself: a search asks +/// both kinds. A slot is found from the running sum of the groups' counts and the place +/// where its group begins; no row is read. +fn question( + ground: &Ground<'_>, + vector: VectorId, + slots: u64, + number: u32, + query: &mut Vec, +) -> Result<()> { + let build = &ground.build; + query.clear(); + query.resize(ground.dimension, 0.0); + let first = u64::from(number) * slots / slots.min(u64::from(MEASURED)); + let taken = 2 - number as usize % 2; + for mut ordinal in [first, (first + slots / 2) % slots].into_iter().take(taken) { + let mut group = 0; + while ordinal >= u64::from(build.counts[group]) { + ordinal -= u64::from(build.counts[group]); + group += 1; + } + let Place { mut block, slot } = build.starts[group]; + let mut at = Some(slot as usize + ordinal as usize); + while let Some(place) = at { + let record = VectorRecord::Block(vector, build.number, block); + with_block(ground.read, record, ground.dimension, &mut |block| { + at = place.checked_sub(block.len()); + if at.is_none() { + let slot = block.slot(place); + let scale = f32::from_le_bytes([slot[0], slot[1], slot[2], slot[3]]); + for (sum, code) in query.iter_mut().zip(&slot[SLOT_HEADER..]) { + *sum += scale * f32::from(code.cast_signed()); + } + } + Ok(()) + })?; + block += 1; + } + } + let length = query + .iter() + .map(|number| number * number) + .sum::() + .sqrt(); + if length > 0.0 { + query.iter_mut().for_each(|number| *number /= length); + } + Ok(()) +} + +/// The step that asks the new grouping the next of its questions (`build::MEASURED`) and +/// keeps the largest gap any of them needed for its ten nearest: how far the centre of the +/// furthest group holding one of them falls short of the best centre's product. +pub(super) fn measure(ground: &Ground<'_>, step: &mut Step) -> Result<()> { + let Phase::Measure { mut done, mut most } = ground.build.phase else { + unreachable!("a build that measures its gap is in that phase") + }; + let (read, vector, dimension) = (ground.read, step.work.vector, ground.dimension); + let mut build = copy(ground, step)?; + let grouped = packed(&build); + let grouping = grouping_of(&build, step, (0, 1, 0))?; + step.held.grow(4 * dimension as u64)?; + let mut query = Vec::with_capacity(dimension); + let questions = grouped.slots.min(u64::from(MEASURED)); + let fitting = (STEP_COMPARISONS / grouped.slots).max(1); + let until = (u64::from(done) + fitting).min(questions) as u32; + while done < until { + question(ground, vector, grouped.slots, done, &mut query)?; + let taken = step.held.bytes(); + step.held.grow(80)?; + let held = &mut step.held; + let (gap, places) = gap_needed(read, vector, grouped, &grouping, &query, held)?; + most = most.max(gap); + build.tens.extend(places); + step.held.shrink(step.held.bytes() - taken); + done += 1; + } + build.phase = if u64::from(done) < questions { + Phase::Measure { done, most } + } else { + Phase::Width { + gap: most, + probes: 1, + done: 0, + asked: 0, + kept: [0; 15], + } + }; + moved(ground, step, build); + Ok(()) +} + +/// The step that asks the next of the same questions again, of the groups inside the gap +/// alone, keeps the most groups any of them had inside, and counts for each width how many of +/// the question's ten the sketches kept: both the sketch a slot has from its group's centre +/// and the one a tail block would keep of it, from the partition's mean, so that the one +/// width serves a search of the grouping and of its tail. When the last is asked, the width +/// is the least that kept 99 in 100 of all the neighbours asked for, and 0 when none did. +pub(super) fn width(ground: &Ground<'_>, step: &mut Step) -> Result<()> { + let Phase::Width { + gap, + mut probes, + mut done, + mut asked, + kept: mut keeps, + } = ground.build.phase + else { + unreachable!("a build that measures its width is in that phase") + }; + let (read, vector, dimension) = (ground.read, step.work.vector, ground.dimension); + let build = copy(ground, step)?; + let grouped = packed(&build); + let grouping = grouping_of(&build, step, (gap, 1, 0))?; + step.held.grow(4 * dimension as u64)?; + let mut query = Vec::with_capacity(dimension); + let mean = mean( + read, + vector, + ground.partition.number, + dimension, + &mut step.held, + )?; + let questions = grouped.slots.min(u64::from(MEASURED)) as u32; + let each = build.tens.len() / questions as usize; + let (first, mut estimated) = (done, 0); + while done < questions { + question(ground, vector, grouped.slots, done, &mut query)?; + // The sketches the question has estimated: the slots of the groups inside the gap. + let taken = step.held.bytes(); + let inside: u64 = { + let groups = nearest(read, vector, grouped, &query, &mut step.held)?; + let short = |product: f32| f64::from(groups[0].1 - product) * 10_000.0; + let inside = groups + .iter() + .take_while(|(_, product)| short(*product) <= f64::from(gap)); + inside + .map(|(group, _)| u64::from(build.counts[*group as usize])) + .sum() + }; + step.held.shrink(step.held.bytes() - taken); + if done > first && estimated + inside > STEP_COMPARISONS { + break; + } + estimated += inside; + let wanted = &build.tens[done as usize * each..][..each]; + let held = &mut step.held; + let both = (query.as_slice(), mean.as_slice()); + let (more, inside) = kept_by_width(read, vector, grouped, &grouping, both, wanted, held)?; + step.held.shrink(step.held.bytes() - taken); + asked += each as u32; + keeps + .iter_mut() + .zip(more) + .for_each(|(kept, more)| *kept += more); + probes = probes.max(inside); + done += 1; + } + let mut build = build; + build.phase = if done < questions { + Phase::Width { + gap, + probes, + done, + asked, + kept: keeps, + } + } else { + let enough = + |width: &u8| u64::from(keeps[usize::from(*width) - 2]) * 100 >= u64::from(asked) * 99; + let width = (2..=WIDEST).find(enough).unwrap_or(0); + Phase::Switch { gap, probes, width } + }; + moved(ground, step, build); + Ok(()) +} + +/// The last step: writes the grouping's record and has the partition's record name it, and +/// removes what it replaces: the grouping the build read, the tail blocks it read with their +/// sketches, and the records of the build: all of them, or all but its own when some slot +/// it read was stale, which then stays for [`sweep`]. +pub(super) fn switch(ground: &Ground<'_>, step: &mut Step) -> Result<()> { + let Phase::Switch { gap, probes, width } = ground.build.phase else { + unreachable!("a build that switches is in that phase") + }; + let (vector, dimension) = (step.work.vector, ground.dimension); + let number = ground.partition.number; + let mut build = copy(ground, step)?; + let grouping = grouping_of(&build, step, (gap, probes, width))?; + // The grouping's record, and the list of what is removed: at most a record, the centres + // and two records a block of the grouping the build read, two a tail block, and the + // build's own with one for each block it read. + let removed = 3 * u64::from(build.blocks()) + build.from.map_or(0, |from| from.groups.into()); + let listed = (removed + 2) * size_of::>() as u64; + step.held + .grow(12 + 8 * build.starts.len() as u64 + listed)?; + let record = VectorRecord::Grouping(vector, build.number); + step.puts.push((record, grouping.encode())); + step.switch = Some(packed(&build)); + if let Some(from) = build.from { + step.removes.extend(records_of(vector, from, dimension)); + } + for block in build.tail.0..build.tail.1 { + let sketch = VectorRecord::Sketch(vector, number, block); + step.removes + .extend([VectorRecord::Block(vector, number, block), sketch]); + } + // The centres, blocks and sketches the build wrote are the grouping's now. + let written = build.written(vector, number, dimension); + step.removes.extend(written.filter(|record| match record { + VectorRecord::Assigned(..) => true, + VectorRecord::Build(..) => !build.stale, + _ => false, + })); + if build.stale { + build.tens.clear(); + build.phase = Phase::Sweep; + moved(ground, step, build); + } + Ok(()) +} + +/// A step after the switch: removes up to 4,096 of the partition's `Left` records whose value +/// is below the build's `upto`. The build applied each of them to every slot it read, and no +/// slot left can be stale by one; a record of `upto` itself may have been written after the +/// build read its slot, and stays for the next build. When none is left the step removes the +/// build's record. +pub(super) fn sweep(ground: &Ground<'_>, step: &mut Step) -> Result<()> { + let (vector, number) = (step.work.vector, ground.partition.number); + let upto = ground.build.upto; + let first = Key::from_values(&[], &[])?; + let (held, left) = (&mut step.held, &mut step.left); + let from = VectorRecord::Left(vector, number, &first); + ground.read.vectors(from, &mut |record, value| { + let VectorRecord::Left(column, partition, key) = record else { + return Ok(false); + }; + if (column, partition) != (vector, number) { + return Ok(false); + } + if left_at(ground, value)? < upto { + room(left, SWEPT, held)?; + held.grow(key.as_bytes().len() as u64)?; + left.push(key.clone()); + } + Ok(left.len() < SWEPT) + })?; + if step.left.is_empty() { + step.removes.push(VectorRecord::Build(vector, number)); + } else { + let build = copy(ground, step)?; + moved(ground, step, build); + } + Ok(()) +} + +/// The records of a grouping, for a commit that removes it: its record, its centres, its +/// blocks and their sketches. +pub(super) fn records_of( + vector: VectorId, + grouped: Grouped, + dimension: usize, +) -> impl Iterator> { + let Grouped { number, blocks, .. } = grouped; + let centres = grouped + .groups + .div_ceil(centres_to_a_record(dimension) as u32); + std::iter::once(VectorRecord::Grouping(vector, number)) + .chain((0..centres).map(move |place| VectorRecord::Centres(vector, number, place))) + .chain((0..blocks).flat_map(move |block| { + let sketch = VectorRecord::Sketch(vector, number, block); + [VectorRecord::Block(vector, number, block), sketch] + })) +} diff --git a/v2/contextdb/src/vector/seal.rs b/v2/contextdb/src/vector/seal.rs index 6ad7e03c..2ab94cb5 100644 --- a/v2/contextdb/src/vector/seal.rs +++ b/v2/contextdb/src/vector/seal.rs @@ -2,6 +2,8 @@ use super::block::{BLOCK_BYTES, BROKEN, BlockReader, BlockWriter}; use super::directory; +use super::grouping; +use super::sketch::{sketch_len, sketches_of}; use super::tally::{position, tuples}; use crate::memory::{Budget, Reservation}; use crate::schema::Catalog; @@ -24,17 +26,20 @@ const WAITING: &str = "a pending record of the vector part is of no row with a v /// 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. +/// The vectors of a partition are added to its last tail block while that block has room and +/// is not closed, and to new blocks after it, so every block from the first that is not closed +/// on but the last is full. Each is read from its row at full precision, one at a time, and +/// its pending record is removed. In a partition that has a mean every block that is staged +/// comes with the record of its sketches, each taken from the mean. 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. +/// and kept for the call), the partition record it is adding to, the partition's mean, the +/// record of one block's sketches until it is staged, 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. @@ -108,13 +113,7 @@ fn seal_partition( 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); - } + room(&mut keys, most, held)?; held.grow(key.as_bytes().len() as u64)?; keys.push(key.clone()); Ok(keys.len() < most) @@ -138,15 +137,28 @@ fn seal_partition( 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); + // The mean is read once, and held until the partition is done. + let mean = partition + .mean + .then(|| grouping::mean(&*changes, vector, number, dimension, held)) + .transpose()?; + let mean = mean.as_deref(); + + // Its last block is added to while it has room and is not closed. A writer is made, after + // the block before it is staged, with room for exactly the slots still to come, which + // `left` counts. + let open = partition.blocks > partition.closed; + let mut block = partition.blocks - u32::from(open); 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)? { + let last = if open { + directory::kept(&*changes, VectorRecord::Block(vector, number, block), begun)? + } else { + None + }; + let (mut writer, mut left) = match last { Some(begun) => begun.ok_or_else(|| changes.damaged(BROKEN))?, None => BlockWriter::new(dimension, None, &keys), }; @@ -169,8 +181,7 @@ fn seal_partition( 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)?; + put_block(changes, held, (vector, number, block), writer, mean)?; } block += 1; (writer, left) = BlockWriter::new(dimension, None, &keys[done..]); @@ -181,30 +192,59 @@ fn seal_partition( 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)?; + put_block(changes, held, (vector, number, block), writer, mean)?; partition.blocks = block + 1; - changes.put_vector(VectorRecord::Partition(vector, &tuple), partition.encode())?; let sealed = keys.len() as u64; + partition.slots += sealed; + partition.sealed += sealed; + changes.put_vector(VectorRecord::Partition(vector, &tuple), partition.encode())?; 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. +/// Stages a block of `bytes` bytes under its column, the number its blocks are kept under and +/// its own. `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. +/// +/// With the `mean` of a partition that has one, the record of the block's sketches is staged +/// under the same numbers, each sketch taken from the mean: its bytes are asked for before +/// they are made and given back when the change set counts them. fn put_block( changes: &mut ChangeSet, held: &mut Reservation, - record: VectorRecord<'_>, + (vector, under, number): (VectorId, u32, u32), writer: BlockWriter, + mean: Option<&[f32]>, ) -> Result<()> { let bytes = writer.bytes(); held.grow(bytes as u64)?; let block = writer.finish(); held.shrink(bytes as u64); - changes.put_vector(record, block) + if let Some(mean) = mean { + let slots = BlockReader::new(&block, mean.len()).len(); + let bytes = (4 + slots * sketch_len(mean.len())) as u64; + held.grow(bytes)?; + let sketches = sketches_of(&block, mean.len(), &mut |_| mean); + held.shrink(bytes); + changes.put_vector(VectorRecord::Sketch(vector, under, number), sketches)?; + } + changes.put_vector(VectorRecord::Block(vector, under, number), block) +} + +/// Makes room for one more entry in a list that has no place left: the list is made again +/// with as many more, and never past `most` entries, the bytes of the new one taken from +/// `held` before it is made and those of the old one given back once it is gone. What an +/// entry points at is the caller's to take. +pub(super) fn room(list: &mut Vec, most: usize, held: &mut Reservation) -> Result<()> { + if list.len() == list.capacity() { + let more = list.capacity().max(4).min(most - list.len()); + held.grow(((list.len() + more) * size_of::()) as u64)?; + list.reserve_exact(more); + held.shrink((list.len() * size_of::()) as u64); + } + Ok(()) } diff --git a/v2/contextdb/src/vector/search.rs b/v2/contextdb/src/vector/search.rs index 9aae8ebb..56603db8 100644 --- a/v2/contextdb/src/vector/search.rs +++ b/v2/contextdb/src/vector/search.rs @@ -1,22 +1,35 @@ -//! 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. +//! The one search: score the rows that are in no block exactly, screen the slots of the +//! selected partitions (every slot of a partition searched whole; of one searched by its +//! nearest groups, the slots its sketches keep), 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 std::collections::btree_map::Entry; +use std::collections::{BTreeMap, BTreeSet, BinaryHeap}; +use std::ops::Range; -use super::block::{BROKEN, BlockReader}; -use super::directory::{self, Partition}; -use super::slot::Screen; +use super::block::{BLOCK_BYTES, BROKEN, BlockReader}; +use super::build::{pending, usable}; +use super::directory::{self, Grouped, Partition, kept}; +use super::grouping::{self, Grouping, Place}; +use super::sketch::{Estimate, Sketches, WIDEST, sketch_len, write_sketch}; +use super::slot::{SLOT_HEADER, Screen, product, slot_len, summed}; +use super::train; use crate::memory::Reservation; +use crate::schema::VectorColumn; use crate::store::{Key, Read, TableId, VectorId, VectorRecord}; -use crate::types::Result; +use crate::types::{Error, Result, VectorSearchMode}; /// 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, + /// What the column declares: `AUTO_INDEX_AT` and `CONSOLIDATION` decide what AUTO runs + /// and when INDEXED may. + pub column: &'a VectorColumn, + /// The mode in force: the statement's, else the column's. + pub mode: VectorSearchMode, /// 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], @@ -24,6 +37,10 @@ pub(crate) struct Nearest<'a> { pub query: &'a [f32], /// How many rows to return at most. pub k: usize, + /// Whether the statement has a filter on rows: one that does more than select partitions. + /// The filter itself is asked through `distance_of`. This says one thing more: a search + /// under INDEXED that runs out of sketches to choose among is refused. + pub filtered: bool, } /// One row of an answer. @@ -34,6 +51,452 @@ pub(crate) struct Hit { pub distance: f64, } +/// What a search gives. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct Outcome { + /// The rows, nearest first. + pub hits: Vec, + /// How many centres it compared with the vector of the search, how many sketches it + /// estimated and how many slots it screened, together. The rows whose distance it asked + /// for are `query`'s to count. + pub evaluations: u64, +} + +/// How one partition of a search is searched. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Route { + /// Every slot of it. + Exact, + /// The sketches of its nearest groups and of its tail, and the slots of those whose + /// estimate can still reach the answer. + Indexed, +} + +/// What a search does with one of the partitions it selects, settled from the partition's +/// small records before a block is read. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Planned { + /// The partition's key tuple. + pub tuple: Key, + /// Rows with a vector in it, and those of them that are pending. + pub live: u64, + pub pending: u64, + /// Slots in its tail blocks, outside its grouping. + pub tail: u64, + /// The groups of its grouping and the most of them a search scanned for its build's own + /// questions; zero and zero when it has none. + pub groups: u32, + pub probes: u32, + pub route: Route, + /// Whether INDEXED may be given this partition. A search under INDEXED with a partition + /// that is not usable is refused whole. + pub usable: bool, +} + +/// One selected partition as a search goes over it. +struct Target { + partition: Partition, + /// What the plan says of it. + planned: Planned, + /// The record of its grouping, read when its nearest groups are searched and it has more + /// than one group. + grouping: Option, + /// How many of its groups the last pass took, nearest first, and how far in + /// ten-thousandths the next one lies past its gap: infinitely when it took them all. + taken: usize, + next: f64, +} + +impl Target { + /// How many sketches a search reads of it at most: one for each slot, and none of a + /// partition whose slots are screened with no sketch read. + fn sketches(&self) -> u64 { + let grouped = self.partition.grouping.map_or(0, |grouped| grouped.slots); + match &self.grouping { + Some(grouping) if grouping.width != 0 => self.partition.slots + grouped, + _ => 0, + } + } +} + +/// The selected partitions in the order of their numbers, each with its pending rows counted +/// and its route settled. [`directory::HELD`], its key tuple and the bytes of a `Target` are +/// taken from `reservation` for each. +fn targets( + nearest: &Nearest<'_>, + read: &dyn Read, + reservation: &mut Reservation, +) -> Result> { + let partitions = directory::selected(read, nearest.vector, nearest.scope, reservation)?; + reservation.grow((partitions.len() * size_of::()) as u64)?; + let live: u64 = partitions.iter().map(|(_, partition)| partition.live).sum(); + let mode = match nearest.mode { + VectorSearchMode::Auto if live < nearest.column.index_at() => VectorSearchMode::Exact, + mode => mode, + }; + let mut targets = Vec::with_capacity(partitions.len()); + // The places of a grouping's record, eight bytes a group and one more, are taken first. + let mut record_of = |grouped: Grouped| { + reservation.grow(8 * (u64::from(grouped.groups) + 1))?; + Grouping::read(read, nearest.vector, grouped) + }; + for (tuple, partition) in partitions { + let pending = pending(read, nearest.vector, partition.number)?; + let usable = usable(&partition, pending, nearest.column); + let grouped = partition.grouping.filter(|grouped| grouped.groups > 1); + // AUTO searches a partition whose grouping record is damaged exactly. + let grouping = match (mode, grouped) { + (VectorSearchMode::Indexed, Some(grouped)) => Some(record_of(grouped)?), + (VectorSearchMode::Auto, Some(grouped)) if usable => match record_of(grouped) { + Err(Error::StoreCorrupted { .. }) => None, + found => Some(found?), + }, + _ => None, + }; + let indexed = mode == VectorSearchMode::Indexed || grouping.is_some(); + let route = if indexed { + Route::Indexed + } else { + Route::Exact + }; + let planned = Planned { + tuple, + live: partition.live, + pending, + tail: partition.slots, + groups: partition.grouping.map_or(0, |grouped| grouped.groups), + probes: grouping.as_ref().map_or(0, |grouping| grouping.probes), + route, + usable, + }; + targets.push(Target { + partition, + planned, + grouping, + taken: 0, + next: f64::INFINITY, + }); + } + Ok(targets) +} + +/// What a search would do with each partition it selects, in the order of their numbers: +/// the one plan [`search`] runs and `.explain` prints. EXACT searches every slot of every +/// one. INDEXED searches the nearest groups of every one, which for a partition with no +/// grouping or a grouping of one group is every slot, through the slot screen alone. AUTO +/// searches every slot of every one while the rows with a vector in all of them together are +/// fewer than the column's `AUTO_INDEX_AT`, and from there on the nearest groups of each +/// partition that has a grouping and is usable, and every slot of the others. A filter on +/// rows changes no route: it is asked of each row as the search visits it. +/// +/// Only partition, pending and grouping records are read; EXACT reads no grouping record and +/// answers no `probes`. What the list takes is taken from `reservation` first and stays taken +/// until the caller lets it go. +pub(crate) fn plan( + nearest: &Nearest<'_>, + read: &dyn Read, + reservation: &mut Reservation, +) -> Result> { + let targets = targets(nearest, read, reservation)?; + Ok(targets.into_iter().map(|target| target.planned).collect()) +} + +/// What a search reads blocks and sketches with. +struct Asked<'a> { + read: &'a dyn Read, + vector: VectorId, + dimension: usize, + screen: &'a Screen, +} + +/// What a store is told when the sketches of a block are missing or are not one for each of +/// the block's slots. +const SKETCHES: &str = "the sketches of a block of the vector part are not one for each slot"; + +/// Hands `each` the blocks kept under `number` from block `from` on, each with its own +/// number, while it answers true. A block that is not whole is damage. +fn blocks( + asked: &Asked<'_>, + number: u32, + from: u32, + each: &mut dyn FnMut(u32, &BlockReader<'_>) -> Result, +) -> Result<()> { + let run = (asked.vector, number); + let from = VectorRecord::Block(run.0, run.1, from); + asked.read.vectors(from, &mut |record, bytes| match record { + VectorRecord::Block(vector, number, at) if (vector, number) == run => { + let block = BlockReader::new(bytes, asked.dimension); + if !block.is_whole() { + return Err(asked.read.damaged(BROKEN)); + } + each(at, &block) + } + _ => Ok(false), + }) +} + +/// Hands `each` the records of sketches of the blocks kept under `number` from block `from` +/// on, each with its block's number and how many sketches it holds, while it answers true. +/// A record it asks for that is not there, or is not whole, is damage. +fn sketches( + asked: &Asked<'_>, + number: u32, + from: u32, + each: &mut dyn FnMut(u32, &Sketches<'_>, usize) -> Result, +) -> Result<()> { + let (mut at, mut more) = (from, true); + let first = VectorRecord::Sketch(asked.vector, number, from); + asked.read.vectors(first, &mut |record, bytes| { + let sketches = Sketches::new(bytes, asked.dimension); + let count = bytes.len().saturating_sub(4) / sketch_len(asked.dimension); + if record != VectorRecord::Sketch(asked.vector, number, at) || !sketches.is_whole(count) { + return Err(asked.read.damaged(SKETCHES)); + } + more = each(at, &sketches, count)?; + at += 1; + Ok(more) + })?; + if more { + return Err(asked.read.damaged(SKETCHES)); + } + Ok(()) +} + +/// How far in ten-thousandths of a product a group lies from the nearest group, by the +/// products of their centres with the vector of a search; held to 20,000. +fn apart(nearest: f32, product: f32) -> f64 { + ((f64::from(nearest) - f64::from(product)) * 10_000.0).min(20_000.0) +} + +/// How many of the groups, nearest first, lie no farther from the nearest than a gap in +/// ten-thousandths widened by `fall`, a difference of products. +fn within(ranked: &[(u32, f32)], gap: f64, fall: f64) -> usize { + let past = |(_, product): &&(u32, f32)| apart(ranked[0].1, *product) - gap; + ranked + .iter() + .take_while(|group| past(group) <= fall * 10_000.0) + .count() +} + +/// The first place of a block. +fn first_of(block: u32) -> Place { + Place { block, slot: 0 } +} + +/// The slots of block `at`, which has `count` of them, that lie between two places. +fn part(places: &Range, at: u32, count: usize) -> Range { + let begin = if at == places.start.block { + places.start.slot as usize + } else { + 0 + }; + let end = if at == places.end.block { + places.end.slot as usize + } else { + count + }; + begin.min(end.min(count))..end.min(count) +} + +/// What sketches are estimated with: what the vector of the search gives once, and room for +/// the ranges of one block's sketches. +struct Wide { + estimate: Estimate, + ranges: Vec<(f32, f32)>, +} + +impl Wide { + /// Takes the planes of the vector, 32 bytes for every 64 numbers begun, and eight bytes + /// for each slot the largest block can have, before either is made. + fn new(query: &[f32], reservation: &mut Reservation) -> Result { + let slots = (BLOCK_BYTES / slot_len(query.len())).max(1); + reservation.grow((32 * query.len().div_ceil(64) + 8 * slots) as u64)?; + Ok(Wide { + estimate: Estimate::new(query), + ranges: Vec::with_capacity(slots), + }) + } + + /// Hands `each` the range of every sketch between two places of the blocks kept under + /// `number`, with its place: `width` quarters of the sketch's factor wide around `base`, + /// the product of the sketches' reference with the vector of the search. + fn span( + &mut self, + asked: &Asked<'_>, + number: u32, + places: Range, + (base, width): (f32, u8), + each: &mut dyn FnMut(Place, f32, f32) -> Result<()>, + ) -> Result<()> { + if places.start >= places.end { + return Ok(()); + } + let each = &mut |block, record: &Sketches<'_>, count| { + self.ranges.clear(); + let slots = part(&places, block, count); + let estimate = &self.estimate; + estimate.ranges(record, slots.clone(), base, width, &mut self.ranges); + for (slot, (low, high)) in slots.zip(&self.ranges) { + let slot = slot as u32; + each(Place { block, slot }, *low, *high)?; + } + Ok(first_of(block + 1) < places.end) + }; + sketches(asked, number, places.start.block, each) + } +} + +/// How far from its nearest group a grouping has to be scanned for the ten slots nearest +/// `query`, and where those slots are. Nearest is by the product of the vector a slot keeps, +/// its codes times its scale, with `query` at length one, and equal products go by place: no +/// row is read. Every group is scanned. The gap is the best centre's product with `query` +/// less the least such product among the groups those slots are in, in ten-thousandths, +/// rounded up and held to 20,000; zero for a grouping of no slot, which has no place to give. +pub(super) fn gap_needed( + read: &dyn Read, + vector: VectorId, + grouped: Grouped, + grouping: &Grouping, + query: &[f32], + reservation: &mut Reservation, +) -> Result<(u32, Vec)> { + let before = reservation.bytes(); + reservation.grow(4 * query.len() as u64)?; + let screen = Screen::new(query); + let asked = Asked { + read, + vector, + dimension: query.len(), + screen: &screen, + }; + let mut ten: Vec<(f32, Place)> = Vec::with_capacity(11); + blocks(&asked, grouped.number, 0, &mut |block, slots| { + for slot in 0..slots.len() { + let kept = slots.slot(slot); + let scale = f32::from_le_bytes([kept[0], kept[1], kept[2], kept[3]]); + let near = scale * product(&kept[SLOT_HEADER..], screen.unit()); + if ten.len() < 10 || near > ten[9].0 { + let slot = slot as u32; + ten.push((near, Place { block, slot })); + // The sort keeps slots of one product in the order of their places. + ten.sort_by(|one, other| other.0.total_cmp(&one.0)); + ten.truncate(10); + } + } + Ok(true) + })?; + let ranked = grouping::nearest(read, vector, grouped, screen.unit(), reservation)?; + let nearest = ranked[0].1; + let of = |place: &Place| { + let group = grouping.starts.partition_point(|start| start <= place) - 1; + let found = ranked.iter().find(|(number, _)| *number as usize == group); + found.map_or(nearest, |(_, product)| *product) + }; + let least = ten.iter().map(|(_, place)| of(place)); + let least = least.fold(nearest, f32::min); + reservation.shrink(reservation.bytes() - before); + let places = ten.into_iter().map(|(_, place)| place); + Ok((apart(nearest, least).ceil() as u32, places.collect())) +} + +/// For each width from 2 to 16 quarters, how many of `wanted` (slots of the grouping's blocks, +/// the places [`gap_needed`] gave) a search keeps by its sketches, wherever the slot is kept: +/// the sketches of the groups inside `grouping.gap` are estimated, and a slot is kept at a +/// width when the range of its sketch at that width reaches the tenth best low end of those +/// ranges at that width, and so does the range of the sketch a tail block would keep of it, +/// which is made here from the slot's vector and `mean`, the partition's. A search estimates +/// a tail's sketches after the groups', against that same low end. With the counts, how many +/// groups were inside the gap: one at least, the nearest. +pub(super) fn kept_by_width( + read: &dyn Read, + vector: VectorId, + grouped: Grouped, + grouping: &Grouping, + (query, mean): (&[f32], &[f32]), + wanted: &[Place], + reservation: &mut Reservation, +) -> Result<([u32; 15], u32)> { + let before = reservation.bytes(); + let dimension = query.len(); + let found = wanted.len() * size_of::<(Place, f32, f32)>(); + reservation.grow((4 * dimension + found + 4 + sketch_len(dimension)) as u64)?; + let screen = Screen::new(query); + let asked = Asked { + read, + vector, + dimension: query.len(), + screen: &screen, + }; + let mut wide = Wide::new(query, reservation)?; + let mut ranked = grouping::nearest(read, vector, grouped, screen.unit(), reservation)?; + let inside = within(&ranked, f64::from(grouping.gap), 0.0); + ranked.truncate(inside); + ranked.sort_unstable_by_key(|(group, _)| *group); + // Each sketch is estimated once, at the widest: its middle and a quarter of its factor + // give its range at every width. The ten best low ends at each width, the least first. + let mut lows = [[f32::NEG_INFINITY; 10]; 15]; + let halves = + |(low, high): (f32, f32)| ((low + high) / 2.0, (high - low) / f32::from(2 * WIDEST)); + let mut found: Vec<(Place, f32, f32)> = Vec::with_capacity(wanted.len()); + let each = &mut |place, low: f32, high: f32| { + let (middle, quarter) = halves((low, high)); + if wanted.contains(&place) { + found.push((place, middle, quarter)); + } + for (width, lows) in (2..=WIDEST).zip(&mut lows) { + let low = middle - f32::from(width) * quarter; + if low > lows[0] { + lows[0] = low; + lows.sort_unstable_by(f32::total_cmp); + } + } + Ok(()) + }; + for (group, base) in &ranked { + let places = grouping.starts[*group as usize]..grouping.starts[*group as usize + 1]; + wide.span(&asked, grouped.number, places, (*base, WIDEST), each)?; + } + // The slots found, block by block: each one's sketch from the mean, and whether both of + // its sketches reach the low end at each width. + let base = summed(mean, screen.unit(), |number| number); + let mut sketch = vec![0; 4 + sketch_len(dimension)]; + let mut kept = [0; 15]; + for places in found.chunk_by(|one, next| one.0.block == next.0.block) { + let mut whole = false; + blocks( + &asked, + grouped.number, + places[0].0.block, + &mut |at, block| { + whole = at == places[0].0.block + && places.iter().all(|slot| slot.0.slot < block.len() as u32); + for (place, middle, quarter) in places.iter().filter(|_| whole) { + let slot = block.slot(place.slot as usize); + let own = train::product(screen.unit(), slot); + write_sketch(slot, mean, &mut sketch[4..]); + wide.ranges.clear(); + let one = Sketches::new(&sketch, dimension); + wide.estimate + .ranges(&one, 0..1, base, WIDEST, &mut wide.ranges); + let from_mean = halves(wide.ranges[0]); + for (width, (kept, lows)) in (2..=WIDEST).zip(kept.iter_mut().zip(&lows)) { + let reaches = |(middle, quarter): (f32, f32)| { + middle + f32::from(width) * quarter >= lows[0].max(own) + }; + *kept += u32::from(reaches((*middle, *quarter)) && reaches(from_mean)); + } + } + Ok(false) + }, + )?; + if !whole { + return Err(read.damaged(BROKEN)); + } + } + reservation.shrink(reservation.bytes() - before); + Ok((kept, inside as u32)) +} + /// Bytes reserved for a row of the answer besides its key. const PAIR: u64 = 32; @@ -57,39 +520,107 @@ const LOW: u64 = 8; /// 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. +/// It runs [`plan`]. Over partitions searched whole the answer is exactly the `k` least of +/// `(distance, key)` over the rows `distance_of` gives a distance for. Of a partition searched +/// by its nearest groups it screens the slots that the sketches of the groups inside the +/// grouping's gap and of the tail keep, or all of them where the grouping's width is 0. +/// Blocks and sketches are read in place from the page cache; everything that grows is taken +/// from `reservation` first. When the answer is returned, `reservation` holds 32 bytes and +/// the key's for each of its rows and nothing else of the search. +/// +/// A row is turned away when `distance_of` gives no distance for it. While the answer is short +/// of `k` rows that were not turned away, the search takes four times as many sketches from +/// the groups it has, then twice as many groups. When it holds `k`, it widens each gap by how +/// far its last row lies below the `k`-th best slot screened, and keeps every sketch that +/// reaches that row. Having kept every sketch of every group it has done the work of EXACT and +/// returns what there is; under INDEXED with `nearest.filtered` it answers `None` at the step +/// that would keep them all. It answers `None` too, having read only small records, when the +/// mode is INDEXED and a selected partition is not usable. With `None` the caller lets go of +/// what `reservation` took. +/// +/// A record of the index that is damaged leaves the vectors whole. Under AUTO the search is +/// then run again as EXACT, which reads no such record and answers in full, or tells the +/// damage of a block; under INDEXED the damage is the answer. Nothing is rebuilt. pub(crate) fn search( nearest: &Nearest<'_>, read: &dyn Read, reservation: &mut Reservation, distance_of: &mut dyn FnMut(&Key) -> Result>, -) -> Result> { +) -> Result> { + let before = reservation.bytes(); + match answer(nearest, read, reservation, distance_of) { + Err(Error::StoreCorrupted { .. }) if nearest.mode == VectorSearchMode::Auto => { + reservation.shrink(reservation.bytes() - before); + let mode = VectorSearchMode::Exact; + answer( + &Nearest { mode, ..*nearest }, + read, + reservation, + distance_of, + ) + } + other => other, + } +} + +/// [`search`], in the mode it is given. +fn answer( + nearest: &Nearest<'_>, + read: &dyn Read, + reservation: &mut Reservation, + distance_of: &mut dyn FnMut(&Key) -> Result>, +) -> Result> { if nearest.k == 0 { - return Ok(Vec::new()); + return Ok(Some(Outcome { + hits: Vec::new(), + evaluations: 0, + })); + } + let before = reservation.bytes(); + let mut targets = targets(nearest, read, reservation)?; + let indexed = nearest.mode == VectorSearchMode::Indexed; + if indexed && targets.iter().any(|target| !target.planned.usable) { + return Ok(None); } - 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)?; + reservation.grow(4 * nearest.query.len() as u64)?; let screen = Screen::new(nearest.query); + let asked = Asked { + read, + vector: nearest.vector, + dimension: nearest.query.len(), + screen: &screen, + }; + let mut wide = None; + if targets.iter().any(|target| target.sketches() > 0) { + wide = Some(Wide::new(nearest.query, reservation)?); + } let mut found = Found { k: nearest.k, answer: BTreeSet::new(), - lows: BinaryHeap::new(), + lows: Best::new(nearest.k), reservation, distance_of, + threshold: f64::NEG_INFINITY, + rises: true, + kept: Kept::default(), + marks: BTreeMap::new(), + evaluations: 0, + refused: indexed && nearest.filtered, }; // 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))?; + // With fewer than `k` rows left to ask, the answer stays short: a search that is refused + // then is refused here, before any of them is asked. + let live: u64 = targets.iter().map(|target| target.partition.live).sum(); + if found.refused && (found.answer.len() as u64).saturating_add(live) < nearest.k as u64 { + return Ok(None); + } let first = Key::from_values(&[], &[])?; - for partition in &partitions { - let run = (nearest.vector, partition.number); + for target in &targets { + let run = (nearest.vector, target.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 => { @@ -99,25 +630,63 @@ pub(crate) fn search( })?; } - // 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))?; + // The slots. A pass leaves out every slot whose `high` is below the threshold its screen + // ended with, and every sketch whose high end is below the one its sketches ended with. + // So when the last row of a full answer is at least as similar as both, and no group left + // out lies within the gap widened by the fall, nothing left out could have entered. + // Otherwise rows were turned away or counted twice, and a further pass goes on from there. + let mut most = nearest.k; + let mut ended = found.pass(&asked, &mut targets, &mut wide, None, None, 0.0)?; + loop { + let (floor, sketches, fall) = match found.last() { + // Short of `k` rows: four times the sketches, and twice the groups once every + // sketch of the groups taken was kept. + None => { + most = most.max(ended.kept).saturating_mul(4); + let mut all = most as u64 >= targets.iter().map(Target::sketches).sum::(); + for target in &mut targets { + target.taken = target.taken.saturating_mul(if ended.all { 2 } else { 1 }); + let grouping = target.grouping.as_ref(); + all &= + grouping.is_none_or(|grouping| target.taken >= grouping.starts.len() - 1); + } + // A filter on rows left the index nothing to choose among: the pass in hand + // kept everything there is, or the next one would. + if found.refused && (ended.whole || all) { + return Ok(None); + } + if ended.whole { + break; + } + (f64::NEG_INFINITY, Some(most), 0.0) + } + Some(last) => { + // A row a fall below the nearest can lie in a group whose centre is four times + // as far below the best centre's: twice the angle from the vector of the + // search. The cover takes the groups that far out. + let fall = 4.0 * (found.lows.least() - last).max(0.0); + let nearer = |target: &Target| target.next <= fall * 10_000.0; + let proven = last >= ended.threshold && last >= ended.sketches; + if proven && !targets.iter().any(nearer) { + break; + } + (last, None, fall) + } + }; + ended = found.pass(&asked, &mut targets, &mut wide, Some(floor), sketches, fall)?; } - let lows = LOW * found.lows.len() as u64; - found.reservation.shrink(unit + kept.reserved() + lows); + let pairs = found.answer.iter(); + let pairs = pairs.map(|(_, key)| PAIR + key.as_bytes().len() as u64); + let held = before + pairs.sum::(); + found.reservation.shrink(found.reservation.bytes() - held); + let evaluations = found.evaluations; let hits = found.answer.into_iter(); - Ok(hits - .map(|(Order(distance), key)| Hit { key, distance }) - .collect()) + let hits = hits.map(|(Order(distance), key)| Hit { key, distance }); + Ok(Some(Outcome { + hits: hits.collect(), + evaluations, + })) } /// A distance or a similarity where numbers are put in order. None is NaN: a distance is @@ -145,17 +714,81 @@ impl PartialEq for Order { impl Eq for Order {} +/// The greatest `most` of the numbers offered, the least of them on top. +struct Best { + most: usize, + held: BinaryHeap>, +} + +impl Best { + fn new(most: usize) -> Best { + let held = BinaryHeap::new(); + Best { most, held } + } + + /// Offers a number. Eight bytes are taken for each place the list has, before it is made + /// again with as many more. + fn offer(&mut self, number: f64, reservation: &mut Reservation) -> Result<()> { + if self.held.len() < self.most { + if self.held.len() == self.held.capacity() { + let more = self.held.capacity().max(4).min(self.most - self.held.len()); + reservation.grow(LOW * (self.held.len() + more) as u64)?; + self.held.reserve_exact(more); + reservation.shrink(LOW * self.held.len() as u64); + } + self.held.push(Reverse(Order(number))); + } else if let Some(mut least) = self.held.peek_mut() + && Order(number) > least.0 + { + *least = Reverse(Order(number)); + } + Ok(()) + } + + /// The least of `most` numbers, and nothing while there are fewer. + fn least(&self) -> f64 { + let least = self.held.peek().filter(|_| self.held.len() == self.most); + least.map_or(f64::NEG_INFINITY, |Reverse(Order(number))| *number) + } +} + +/// How a pass ended: what the `high` of a slot had to reach when the screen ended and the +/// high end of a sketch when the sketches did; how many sketches were kept and whether all +/// it estimated were; and whether every sketch of every group and every slot screened was. +struct Ended { + threshold: f64, + sketches: f64, + kept: usize, + all: bool, + whole: bool, +} + /// 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>, + /// vector. With `k` of them, a slot whose `high` is below the least cannot be among the + /// `k` nearest of the vectors seen. + lows: Best, reservation: &'a mut Reservation, distance_of: &'a mut dyn FnMut(&Key) -> Result>, + /// What the `high` of a slot must reach in the pass in hand, and whether that rises with + /// the lower bounds, as it does in the first pass alone. + threshold: f64, + rises: bool, + /// The slots the pass in hand keeps. + kept: Kept, + /// One bit a slot of each block an earlier pass kept a slot of, by the number the block + /// is kept under and its own: set for the slots those passes kept, which are not screened + /// and whose rows are not asked for again. + marks: BTreeMap<(u32, u32), Vec>, + /// Centres compared, sketches estimated and slots screened. + evaluations: u64, + /// Whether the search is refused when it runs out of sketches to choose among: the mode + /// is INDEXED and the statement has a filter on rows. + refused: bool, } impl Found<'_> { @@ -165,7 +798,7 @@ impl Found<'_> { if let Some(distance) = (self.distance_of)(key)? && self.offer(distance, key)? { - self.offer_low(1.0 - distance)?; + self.lows.offer(1.0 - distance, self.reservation)?; } Ok(()) } @@ -176,15 +809,16 @@ impl Found<'_> { /// /// 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 { + // The pair's bytes are asked for before the key is copied. + let bytes = PAIR + key.as_bytes().len() as u64; + self.reservation.grow(bytes)?; 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) { + let worse = full && self.answer.last().is_some_and(|last| pair > *last); + if worse || self.answer.contains(&pair) { + self.reservation.shrink(bytes); 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); } @@ -192,94 +826,324 @@ impl Found<'_> { 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)); + /// 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 one slot of a block kept under `number`, unless an earlier pass kept it. The + /// slot is kept, with its `high`, its place and a copy of its key, unless its `high` is + /// below the threshold, and then offers its `low` to the lower bounds. + fn look( + &mut self, + asked: &Asked<'_>, + number: u32, + at: Place, + block: &BlockReader<'_>, + ) -> Result<()> { + let slot = at.slot as usize; + let bits = self.marks.get(&(number, at.block)); + let word = bits.and_then(|bits| bits.get(slot / 64)); + if word.is_some_and(|word| word >> (slot % 64) & 1 == 1) { + return Ok(()); + } + self.evaluations += 1; + let (low, high) = asked.screen.range(block.slot(slot)); + if f64::from(high) < self.threshold { + return Ok(()); + } + let key = block.key(slot); + self.kept + .keep(high, (number, at), &key, self.threshold, self.reservation)?; + self.lows.offer(f64::from(low), self.reservation)?; + if self.rises { + self.threshold = self.lows.least(); } 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) + /// Screens every slot from one place up to another of the blocks kept under `number`. + fn between(&mut self, asked: &Asked<'_>, number: u32, places: Range) -> Result<()> { + if places.start >= places.end { + return Ok(()); + } + blocks(asked, number, places.start.block, &mut |at, block| { + for slot in part(&places, at, block.len()) { + let slot = slot as u32; + self.look(asked, number, Place { block: at, slot }, block)?; + } + Ok(first_of(at + 1) < places.end) + }) } - /// 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 the slots of the sketches kept among the blocks under `number`, and those + /// only, and empties the list. A block is read here for the first time, so here its + /// sketches are found to be one for each of its slots, or damage. + fn screen_kept(&mut self, asked: &Asked<'_>, number: u32, chosen: &mut Chosen) -> Result<()> { + let threshold = chosen.threshold(); + chosen + .held + .retain(|(high, _)| f64::from(*high) >= threshold); + chosen.kept += chosen.held.len(); + for places in chosen + .held + .chunk_by(|one, next| one.1.block == next.1.block) + { + let at = places[0].1.block; + let record = VectorRecord::Sketch(asked.vector, number, at); + let sketches = kept(asked.read, record, <[u8]>::len)?; + let mut whole = false; + blocks(asked, number, at, &mut |found, block| { + let wanted = 4 + block.len() * sketch_len(asked.dimension); + whole = found == at && sketches == Some(wanted); + if whole { + for (_, place) in places { + self.look(asked, number, *place, block)?; + } + } + Ok(false) + })?; + if !whole { + return Err(asked.read.damaged(SKETCHES)); + } + } + chosen.held.clear(); + Ok(()) } - /// Screens every slot of the partitions and visits the ones kept; returns the threshold - /// the screen ended with. + /// One pass over the slots of the partitions, and the visit of the slots it keeps. /// - /// 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. + /// A slot is kept unless its `high` is below the threshold. With no `floor` the threshold + /// is the least of `k` lower bounds 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. + /// + /// Every slot of a partition searched whole is screened. Of one searched by its nearest + /// groups: the slots of its tail and of the groups within its gap and `fall` of the + /// nearest, and of as many groups as it was told to take. Where the grouping has a width + /// those slots are first chosen by their sketches: with no `floor`, every sketch whose + /// high end reaches the `k`-th best low end so far; with `most`, the `most` with the best + /// high ends; else every one whose high end reaches the floor. The slots the pass before + /// kept are marked first, and none of them is screened, or its row asked for, again. fn pass( &mut self, - nearest: &Nearest<'_>, - partitions: &[Partition], - read: &dyn Read, - screen: &Screen, - kept: &mut Kept, + asked: &Asked<'_>, + targets: &mut [Target], + wide: &mut Option, 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)); + most: Option, + fall: f64, + ) -> Result { + let words = (BLOCK_BYTES / slot_len(asked.dimension)) + .div_ceil(64) + .max(1); + for slot in &self.kept.slots { + let bits = match self.marks.entry((slot.number, slot.at.block)) { + Entry::Occupied(bits) => bits.into_mut(), + Entry::Vacant(none) => { + // The bits of the block, and 64 bytes for its place in the list. + self.reservation.grow(8 * words as u64 + 64)?; + none.insert(vec![0_u64; words]) + } + }; + if let Some(word) = bits.get_mut(slot.at.slot as usize / 64) { + *word |= 1 << (slot.at.slot % 64); + } + } + self.kept.clear(); + self.threshold = floor.unwrap_or_else(|| self.lows.least()); + self.rises = floor.is_none(); + let mut chosen = Chosen { + ends: Best::new(most.unwrap_or(self.k)), + by_high: most.is_some(), + fixed: floor.filter(|_| most.is_none()), + held: Vec::new(), + room: 0, + estimated: 0, + kept: 0, + }; + for target in targets.iter_mut() { + let partition = target.partition; + let tail = first_of(partition.first)..first_of(partition.blocks); + let (Some(grouping), Some(grouped)) = (&target.grouping, partition.grouping) else { + if let Some(grouped) = partition.grouping { + let every = Place::default()..first_of(grouped.blocks); + self.between(asked, grouped.number, every)?; + } + self.between(asked, partition.number, tail)?; + continue; + }; + let unit = asked.screen.unit(); + let mut ranked = + grouping::nearest(asked.read, asked.vector, grouped, unit, self.reservation)?; + self.evaluations += ranked.len() as u64; + let (nearest, gap) = (ranked[0].1, f64::from(grouping.gap)); + let take = within(&ranked, gap, fall).max(target.taken); + let take = take.min(ranked.len()); + let next = ranked + .get(take) + .map(|(_, product)| apart(nearest, *product) - gap); + (target.taken, target.next) = (take, next.unwrap_or(f64::INFINITY)); + ranked.truncate(take); + ranked.sort_unstable_by_key(|(group, _)| *group); + let spans = ranked.iter().map(|(group, base)| { + let group = *group as usize; + (grouping.starts[group]..grouping.starts[group + 1], *base) + }); + match (grouping.width, wide.as_mut()) { + (0, _) | (_, None) => { + for (places, _) in spans { + self.between(asked, grouped.number, places)?; + } + self.between(asked, partition.number, tail)?; + } + (width, Some(wide)) => { + let dimension = asked.dimension; + let mean = (asked.vector, partition.number); + let mean = { + let reservation = &mut *self.reservation; + let mean = + grouping::mean(asked.read, mean.0, mean.1, dimension, reservation)?; + summed(mean.as_slice(), unit, |number| number) + }; + self.reservation.shrink(4 * dimension as u64); + self.evaluations += 1; + let each = + &mut |place, low, high| chosen.offer(place, low, high, self.reservation); + for (places, base) in spans { + wide.span(asked, grouped.number, places, (base, width), each)?; } - 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(); - } + self.screen_kept(asked, grouped.number, &mut chosen)?; + // A tail's sketches are held against the low end the groups left, as the + // build measured them: none of them raises it for the next. + let rising = chosen.fixed.is_none() && !chosen.by_high; + if rising { + chosen.fixed = Some(chosen.threshold()); + } + let each = + &mut |place, low, high| chosen.offer(place, low, high, self.reservation); + wide.span(asked, partition.number, tail, (mean, width), each)?; + self.screen_kept(asked, partition.number, &mut chosen)?; + if rising { + chosen.fixed = None; } - Ok(true) } - _ => Ok(false), - })?; + } + self.reservation.shrink(8 * u64::from(grouped.groups)); } - kept.drop_below(threshold); + self.evaluations += chosen.estimated as u64; + let every = targets.iter().all(|target| target.next == f64::INFINITY); + let all = chosen.kept == chosen.estimated; + let ended = Ended { + threshold: self.threshold, + sketches: chosen.threshold(), + kept: chosen.kept, + all, + whole: all && every && self.threshold == f64::NEG_INFINITY, + }; + // The lists of the sketches are let go before their bytes are given back. + let places = chosen.ends.held.capacity() as u64; + let chosen_bytes = (chosen.room * size_of::<(f32, Place)>()) as u64; + drop(chosen); + self.reservation.shrink(LOW * places + chosen_bytes); + // The pass kept every slot there is and the answer cannot reach `k` rows: a search + // that is refused then is refused here, before the filter is asked of every row. + if ended.whole && self.refused { + return Ok(ended); + } + + self.kept.drop_below(self.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); + self.kept.slots.sort_unstable_by(by_high); + let kept = std::mem::take(&mut self.kept); 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))?; + // The slot's low end is among the lower bounds already: its row's similarity is + // not offered beside it, or one row would count twice and the fall never be zero. + let key = Key::from_bytes(&kept.keys[slot.start..][..slot.len as usize]); + if let Some(distance) = (self.distance_of)(&key)? { + self.offer(distance, &key)?; + } + } + self.kept = kept; + Ok(ended) + } +} + +/// The sketches one pass keeps: the high end of each and its place, in the order of the +/// places, twelve bytes each. +struct Chosen { + /// With no fixed threshold: the `k` best low ends so far, or the best high ends so far + /// of a pass that keeps the sketches with the best high ends. + ends: Best, + by_high: bool, + /// What a high end must reach in a pass that keeps every sketch reaching one similarity. + fixed: Option, + held: Vec<(f32, Place)>, + /// How many sketches the reservation holds room for. + room: usize, + /// How many sketches were offered, and how many of them were kept in the end. + estimated: usize, + kept: usize, +} + +impl Chosen { + /// What the high end of a sketch must reach to stay in the running. + fn threshold(&self) -> f64 { + self.fixed.unwrap_or_else(|| self.ends.least()) + } + + /// Offers the range of one sketch. It is kept unless its high end is below the + /// threshold; when the list is full, those now below the threshold are dropped first. + fn offer( + &mut self, + place: Place, + low: f32, + high: f32, + reservation: &mut Reservation, + ) -> Result<()> { + self.estimated += 1; + let threshold = self.threshold(); + if f64::from(high) < threshold { + return Ok(()); } - Ok(threshold) + if self.held.len() == self.room { + self.held.retain(|(high, _)| f64::from(*high) >= threshold); + widen(&mut self.held, &mut self.room, reservation)?; + } + self.held.push((high, place)); + if self.fixed.is_none() { + let end = if self.by_high { high } else { low }; + self.ends.offer(f64::from(end), reservation)?; + } + Ok(()) } } -/// A slot still in the running: the top of its range, and where its key is in [`Kept::keys`]. +/// Makes room in a list that is full after it dropped what it could: for 1,024 entries at +/// first, and for as many again only if more than half are left. +fn widen(list: &mut Vec, room: &mut usize, reservation: &mut Reservation) -> Result<()> { + let more = match *room { + 0 => 1024, + room if list.len() > room / 2 => room, + _ => 0, + }; + reservation.grow((more * size_of::()) as u64)?; + list.reserve_exact(*room + more - list.len()); + *room += more; + Ok(()) +} + +/// A slot still in the running: the top of its range, the number its block is kept under, +/// its place, and where its key is in [`Kept::keys`]. struct Slot { high: f32, + number: u32, + at: Place, len: u32, start: usize, } @@ -296,14 +1160,6 @@ struct Kept { } 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(); @@ -311,25 +1167,18 @@ impl Kept { } /// 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. + /// first. fn keep( &mut self, high: f32, + (number, at): (u32, Place), 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; + widen(&mut self.slots, &mut self.room, reservation)?; } let key = key.as_bytes(); let start = self.keys.len(); @@ -342,7 +1191,13 @@ impl Kept { 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 }); + self.slots.push(Slot { + high, + number, + at, + len, + start, + }); Ok(()) } diff --git a/v2/contextdb/src/vector/sketch.rs b/v2/contextdb/src/vector/sketch.rs new file mode 100644 index 00000000..46087aaa --- /dev/null +++ b/v2/contextdb/src/vector/sketch.rs @@ -0,0 +1,218 @@ +//! One vector as a sketch keeps it, and the range a search gets from a sketch. +//! +//! A sketch is one bit for each number of the vector, then two bytes. The vector is the one its +//! slot keeps (the slot's codes times its scale), less a reference: the centre of its group for +//! a slot of a grouping, the partition's mean for a slot of a tail block. Bit `i` is set when +//! number `i` of that difference is zero or more. The two bytes are a factor, low byte first: +//! the squared length of the difference over the sum of the sizes of its numbers, times 32,768 +//! and rounded, held to 65,535; zero for a difference of all zeros. +//! +//! The product of the slot's vector with the unit vector `w` of a search is estimated as the +//! product of the reference with `w`, plus the factor times the sum of `w`'s numbers, each +//! taken with the sign its bit says. The estimate is not proven: how far it may be off is +//! measured when a grouping is built, as a width in quarters of the factor, and a search keeps +//! every slot whose estimate, widened by it, can still reach the answer. + +use super::block::BlockReader; +use super::slot::{SLOT_HEADER, length}; + +/// Bytes of a sketch after its bits. +pub(super) const SKETCH_TAIL: usize = 2; + +/// Bytes the sketch of a vector of this dimension takes. +pub(super) fn sketch_len(dimension: usize) -> usize { + dimension.div_ceil(8) + SKETCH_TAIL +} + +/// The widest a sketch's estimate is ever taken to be, in quarters of its factor. +pub(super) const WIDEST: u8 = 16; + +/// Writes the sketch of the vector a slot keeps, less `reference`, into `out`, which is +/// `sketch_len` bytes long. `reference` is of the slot's dimension. +pub(super) fn write_sketch(slot: &[u8], reference: &[f32], out: &mut [u8]) { + let scale = f32::from_le_bytes([slot[0], slot[1], slot[2], slot[3]]); + let (bits, factor) = out.split_at_mut(reference.len().div_ceil(8)); + bits.fill(0); + let (mut squares, mut sizes) = (0.0_f64, 0.0_f64); + for (at, (code, from)) in slot[SLOT_HEADER..].iter().zip(reference).enumerate() { + let off = scale * f32::from(code.cast_signed()) - from; + bits[at / 8] |= u8::from(off >= 0.0) << (at % 8); + squares += f64::from(off) * f64::from(off); + sizes += f64::from(off).abs(); + } + // The cast holds the factor to 65,535, and gives zero for the nothing over nothing of a + // difference of all zeros. + let kept = (squares / sizes * 32_768.0).round() as u16; + factor.copy_from_slice(&kept.to_le_bytes()); +} + +/// The record of the sketches of a block: the block's count of slots as four bytes, low byte +/// first, then one sketch for each of its slots in their order. `reference` answers the +/// reference of the slot at a place of the block. The caller has reserved the record's bytes, +/// four and `sketch_len` a slot, before it calls. +pub(super) fn sketches_of<'a>( + block: &[u8], + dimension: usize, + reference: &mut dyn FnMut(usize) -> &'a [f32], +) -> Vec { + let block = BlockReader::new(block, dimension); + let mut record = vec![0; 4 + block.len() * sketch_len(dimension)]; + let (count, sketches) = record.split_at_mut(4); + count.copy_from_slice(&(block.len() as u32).to_le_bytes()); + let sketches = sketches.chunks_exact_mut(sketch_len(dimension)); + for (place, sketch) in sketches.enumerate() { + write_sketch(block.slot(place), reference(place), sketch); + } + record +} + +/// The sketch at a place of a record of sketches. A record whose length is not four bytes and +/// `count` sketches of this dimension is not whole: the store is told so by the caller. +pub(super) struct Sketches<'a> { + bytes: &'a [u8], + dimension: usize, +} + +impl<'a> Sketches<'a> { + pub(super) fn new(bytes: &'a [u8], dimension: usize) -> Sketches<'a> { + Sketches { bytes, dimension } + } + + /// True when the record holds `count` sketches and nothing else. + pub(super) fn is_whole(&self, count: usize) -> bool { + self.bytes.len() == 4 + count * sketch_len(self.dimension) + && self.bytes[..4] == (count as u32).to_le_bytes() + } + + /// The sketch at this place, counted from zero. + pub(super) fn at(&self, place: usize) -> &'a [u8] { + let len = sketch_len(self.dimension); + &self.bytes[4 + place * len..][..len] + } +} + +/// What a search works out once from its vector and then uses for every sketch it estimates: +/// the vector at length one with each number brought to one of sixteen whole steps from -8 to +/// 7, as four planes of bits, the size of a step, and the sum of the numbers. +pub(super) struct Estimate { + /// The four planes of each 64 numbers, one group of four after another. + planes: Vec, + step: f32, + sum: f32, +} + +impl Estimate { + /// `query` is of the column's dimension and every number of it is finite. + /// The caller has reserved the planes, 32 bytes for every 64 numbers begun, before it + /// calls. + pub(super) fn new(query: &[f32]) -> Estimate { + let length = length(query); + // A vector of zeros has no direction: its numbers at length one are zeros too. + let unit = |x: &f32| { + if length == 0.0 { + 0.0 + } else { + (f64::from(*x) / length) as f32 + } + }; + let largest = query.iter().map(|x| unit(x).abs()).fold(0.0, f32::max); + let step = (largest / 7.0).max(f32::MIN_POSITIVE); + let mut planes = vec![0_u64; 4 * query.len().div_ceil(64)]; + let mut sum = 0_i32; + for (at, x) in query.iter().enumerate() { + let steps = (unit(x) / step).round().clamp(-7.0, 7.0) as i32; + sum += steps; + for (plane, word) in planes[at / 64 * 4..][..4].iter_mut().enumerate() { + *word |= u64::from(steps >> plane & 1 == 1) << (at % 64); + } + } + Estimate { + planes, + step, + sum: sum as f32, + } + } + + /// The ranges the sketches at `places` of one record give for the similarity of their + /// slots' vectors, pushed onto `out` in order: each `(low, high)` around `base`, the + /// product of the sketches' reference with the search's unit vector, plus the factor + /// times the signed sum, and `width` quarters of the factor wide on each side. The sum is + /// taken by counting the bits a sketch shares with each plane: the planes count one, two + /// and four steps, and the fourth takes eight away. + /// + /// This is the loop a search spends its time in, and it has two arms over one body. On an + /// x86-64 processor that says it has the bit-count instruction + /// (`std::arch::is_x86_feature_detected!("popcnt")`, asked once a call) the body compiled + /// for that instruction runs; otherwise, and on every other kind of processor, the body + /// as the compiler builds it for the target. Calling the first arm is the one `unsafe` + /// call of the library, sound because the processor was asked first. Both arms give the + /// same ranges to the bit, and nothing outside this function knows which ran. `out` never + /// grows past the slots of one block: the search has reserved eight bytes for each slot + /// of the largest block before its first sketch. + pub(super) fn ranges( + &self, + sketches: &Sketches<'_>, + places: std::ops::Range, + base: f32, + width: u8, + out: &mut Vec<(f32, f32)>, + ) { + #[cfg(target_arch = "x86_64")] + { + if std::arch::is_x86_feature_detected!("popcnt") { + // SAFETY: the processor was asked on the line before and has the instruction + // this body is compiled for. + unsafe { self.counted(sketches, places, base, width, out) }; + return; + } + } + self.each(sketches, places, base, width, out); + } + + /// The body of [`Estimate::ranges`] compiled for the bit-count instruction. + #[cfg(target_arch = "x86_64")] + #[target_feature(enable = "popcnt")] + fn counted( + &self, + sketches: &Sketches<'_>, + places: std::ops::Range, + base: f32, + width: u8, + out: &mut Vec<(f32, f32)>, + ) { + self.each(sketches, places, base, width, out); + } + + /// The body of [`Estimate::ranges`]. A sketch's bits are read as words, low byte first, + /// and the last word is completed with zeros. + #[inline(always)] + fn each( + &self, + sketches: &Sketches<'_>, + places: std::ops::Range, + base: f32, + width: u8, + out: &mut Vec<(f32, f32)>, + ) { + let (planes, _) = self.planes.as_chunks::<4>(); + for place in places { + let (bits, factor) = sketches.at(place).split_at(sketches.dimension.div_ceil(8)); + let (words, rest) = bits.as_chunks::<8>(); + let mut last = [0_u8; 8]; + last[..rest.len()].copy_from_slice(rest); + let words = words.iter().chain((!rest.is_empty()).then_some(&last)); + let mut shared = [0_u32; 4]; + for (word, planes) in words.zip(planes) { + let word = u64::from_le_bytes(*word); + for (count, plane) in shared.iter_mut().zip(planes) { + *count += (word & plane).count_ones(); + } + } + let under = (shared[0] + 2 * shared[1] + 4 * shared[2]) as i32 - 8 * shared[3] as i32; + let factor = f32::from(u16::from_le_bytes([factor[0], factor[1]])) / 32_768.0; + let estimate = base + factor * self.step * ((2 * under) as f32 - self.sum); + let wide = f32::from(width) / 4.0 * factor; + out.push((estimate - wide, estimate + wide)); + } + } +} diff --git a/v2/contextdb/src/vector/slot.rs b/v2/contextdb/src/vector/slot.rs index ef576d84..9de89629 100644 --- a/v2/contextdb/src/vector/slot.rs +++ b/v2/contextdb/src/vector/slot.rs @@ -44,7 +44,7 @@ pub(crate) fn distance(stored: &[f32], query: &[f32]) -> f64 { } /// The length of a vector, in 64 bits: the square of the largest `f32` does not fit in 32. -fn length(vector: &[f32]) -> f64 { +pub(super) fn length(vector: &[f32]) -> f64 { let squares = vector.iter().map(|x| f64::from(*x) * f64::from(*x)); squares.sum::().sqrt() } @@ -76,24 +76,29 @@ pub(super) fn write_slot(vector: &[f32], out: &mut [u8]) { 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 { +/// The numbers `number` gives for `kept` 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. +pub(super) fn summed(kept: &[T], unit: &[f32], number: impl Fn(T) -> f32) -> f32 { let mut lanes = [0.0_f32; 16]; - let (code_runs, code_rest) = codes.as_chunks::<16>(); + let (kept_runs, kept_rest) = kept.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 (kept, unit) in kept_runs.iter().zip(unit_runs) { + for ((lane, one), other) in lanes.iter_mut().zip(kept).zip(unit) { + *lane += number(*one) * other; } } - for ((lane, code), number) in lanes.iter_mut().zip(code_rest).zip(unit_rest) { - *lane += f32::from(code.cast_signed()) * number; + for ((lane, one), other) in lanes.iter_mut().zip(kept_rest).zip(unit_rest) { + *lane += number(*one) * other; } lanes.iter().sum() } +/// The codes of a slot multiplied by the numbers of `unit` and added up, by [`summed`]. +pub(super) fn product(codes: &[u8], unit: &[f32]) -> f32 { + summed(codes, unit, |code| f32::from(code.cast_signed())) +} + /// 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. @@ -141,6 +146,11 @@ impl Screen { } } + /// The vector of the search at length one. + pub(super) fn unit(&self) -> &[f32] { + &self.unit + } + /// 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. @@ -190,3 +200,51 @@ impl Screen { ) } } + +/// A vector in one byte a number: what a column declared `WITH (quantization = 'SQ8')` keeps +/// in place of the vector a statement wrote. This is the one place a vector is brought to +/// whole steps for storing: `write` asks for it, and the store writes the scale and the codes +/// it is handed. +pub(crate) struct Reduced { + /// What every code is multiplied by. + pub scale: f32, + /// One whole number from -127 to 127 for each number of the vector. + pub codes: Vec, +} + +impl Reduced { + /// The reduced form nearest `vector`; `None` for a vector of zeros or of no number, which + /// is kept as it is. `scale` begins as the largest size among the numbers over 127 and is + /// moved until `(scale * 127.0) / 127.0` gives it back in 32-bit arithmetic; each code is + /// its number over the scale, rounded. `Reduced::of(&r.numbers())` is `r` to the bit. The + /// caller has reserved a byte for each number before it calls. + pub(crate) fn of(vector: &[f32]) -> Option { + let largest = vector.iter().fold(0.0_f32, |most, x| most.max(x.abs())); + let mut scale = largest / 127.0; + // Four moves to the value the arithmetic gives back, then one step up, and again. + let mut moves = 0; + while (scale * 127.0) / 127.0 != scale { + scale = if moves < 4 { + (scale * 127.0) / 127.0 + } else { + scale.next_up() + }; + moves = (moves + 1) % 5; + } + if scale == 0.0 { + return None; + } + let code = |x: &f32| (x / scale).round().clamp(-127.0, 127.0) as i8; + Some(Reduced { + scale, + codes: vector.iter().map(code).collect(), + }) + } + + /// The vector the column keeps: each code times the scale, in 32-bit arithmetic. The + /// caller has reserved four bytes for each number before it calls. + pub(crate) fn numbers(&self) -> Vec { + let codes = self.codes.iter(); + codes.map(|code| self.scale * f32::from(*code)).collect() + } +} diff --git a/v2/contextdb/src/vector/stage.rs b/v2/contextdb/src/vector/stage.rs index 67ea2d1f..8f16bf8e 100644 --- a/v2/contextdb/src/vector/stage.rs +++ b/v2/contextdb/src/vector/stage.rs @@ -4,7 +4,10 @@ use std::collections::BTreeMap; use std::collections::btree_map::Entry; +use super::build::Build; use super::directory::{self, Partition}; +use super::pack; +use super::seal::room; use super::tally::{position, tuples}; use crate::memory::{Budget, Reservation}; use crate::schema::Catalog; @@ -41,9 +44,10 @@ struct Count { /// 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. +/// was on one, lets it go, and a vector that was on none leaves a sealed slot, which a `Left` +/// record says from then on. A partition is made when its first vector arrives and removed, +/// with every record kept for it, 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. @@ -73,16 +77,26 @@ pub(crate) fn stage(changes: &mut ChangeSet, catalog: &Catalog, budget: &Budget) changes.delete_vector(waiting)?; let fewer = pending.checked_sub(1); pending = fewer.ok_or_else(|| changes.damaged(directory::SHORT))?; + } else { + // The row leaves a sealed slot: every slot of its key sealed so far is stale. + let sealed = left.sealed.to_le_bytes().to_vec(); + changes.put_vector( + VectorRecord::Left(moved.vector, left.number, &moved.row), + sealed, + )?; } } - // A partition whose last vector left goes, with its blocks. - for ((vector, _), Touched { now: emptied, .. }) in &mut touched { + // A partition whose last vector left goes, with every record kept for it. + for ((vector, tuple), 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))?; + empty(changes, catalog, *vector, emptied, &mut held)?; + let leaving = + |moved: &&Move| moved.vector == *vector && moved.leaves.as_ref() == Some(tuple); + for moved in moves.iter().filter(leaving) { + changes.delete_vector(VectorRecord::Left(*vector, emptied.number, &moved.row))?; } counts.entry(*vector).or_default().change -= 1; - emptied.blocks = 0; + *emptied = Partition::default(); } } for moved in &moves { @@ -142,6 +156,58 @@ pub(crate) fn stage(changes: &mut ChangeSet, catalog: &Catalog, budget: &Budget) Ok(()) } +/// Stages the removal of every record kept for a partition but its entry and the `Left` +/// records this commit staged: its tail blocks with their sketches, its mean, the records of +/// its grouping and of its build, and its `Left` records, whose keys are held until then. +fn empty( + changes: &mut ChangeSet, + catalog: &Catalog, + vector: VectorId, + emptied: &Partition, + held: &mut Reservation, +) -> Result<()> { + let number = emptied.number; + for block in emptied.first..emptied.blocks { + changes.delete_vector(VectorRecord::Block(vector, number, block))?; + changes.delete_vector(VectorRecord::Sketch(vector, number, block))?; + } + changes.delete_vector(VectorRecord::Mean(vector, number))?; + let (table, column) = catalog.vector(vector).expect("a staged row's column"); + let dimension = table.dimension(column); + if let Some(grouped) = emptied.grouping { + for record in pack::records_of(vector, grouped, dimension) { + changes.delete_vector(record)?; + } + } + if emptied.building.is_some() + && let Some(build) = Build::read(&*changes, vector, number)? + { + for record in build.written(vector, number, dimension) { + changes.delete_vector(record)?; + } + } + let first = Key::from_values(&[], &[])?; + let mut left: Vec = Vec::new(); + changes.vectors( + VectorRecord::Left(vector, number, &first), + &mut |record, _| match record { + VectorRecord::Left(column, partition, key) + if (column, partition) == (vector, number) => + { + room(&mut left, usize::MAX, held)?; + held.grow(key.as_bytes().len() as u64)?; + left.push(key.clone()); + Ok(true) + } + _ => Ok(false), + }, + )?; + for key in &left { + changes.delete_vector(VectorRecord::Left(vector, number, key))?; + } + 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. @@ -198,11 +264,17 @@ fn touch<'a>( tuple: &Key, held: &mut Reservation, ) -> Result<&'a mut Touched> { + // The tuple is copied to look it up: its bytes are taken first, and kept with a new entry. + let bytes = tuple.as_bytes().len() as u64 + 64; + held.grow(bytes)?; Ok(match touched.entry((vector, tuple.clone())) { - Entry::Occupied(known) => known.into_mut(), + Entry::Occupied(known) => { + let known = known.into_mut(); + held.shrink(bytes); + known + } 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/tests.rs b/v2/contextdb/src/vector/tests.rs index 33218864..fe285719 100644 --- a/v2/contextdb/src/vector/tests.rs +++ b/v2/contextdb/src/vector/tests.rs @@ -8,14 +8,21 @@ use std::collections::HashMap; use std::path::{Path, PathBuf}; use super::block::{BLOCK_BYTES, BlockReader, BlockWriter, takes}; -use super::directory::Partition; +use super::build::{Build, Phase}; +use super::directory::{Grouped, Partition, partition}; +use super::grouping::{ + Grouping, Place, centres, centres_to_a_record, mean, nearest as nearest_groups, put_centres, + put_mean, +}; use super::seal::seal; +use super::sketch::{Estimate, Sketches, WIDEST, sketch_len, sketches_of, write_sketch}; use super::slot::{Screen, distance, slot_len, write_slot}; use super::stage::stage; +use super::train::{Mean, Round, first_centres, groups_for, nearest, product, sampled}; use crate::Database; use crate::memory::Budget; use crate::schema::Catalog; -use crate::store::{ChangeSet, Key, Store, TableId}; +use crate::store::{ChangeSet, Key, Read, Store, TableId, VectorId, VectorRecord, View}; use crate::types::{Error, Value}; /// A fixed scatter of numbers: a xorshift that starts from a seed. @@ -221,41 +228,6 @@ fn a_block_reads_back_every_slot_and_key_and_stays_within_its_bytes() { 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 @@ -290,7 +262,9 @@ fn a_pending_record_of_a_row_with_no_vector_is_reported_as_damage() { 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"); + changes + .put_row(table, key, row, Vec::new(), 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 @@ -350,11 +324,14 @@ fn a_vector_entry_named(test: &str, named: u8) -> Misnamed { // 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] - ); + // The VECTOR entry is the last thing in the record, and in it the column's name follows + // the length of the name: the last place the record holds that length and `v`. + let entry = [1, 0, 0, 0, b'v']; + let at = record + .windows(entry.len()) + .rposition(|window| window == entry) + .expect("the VECTOR entry of the column v"); + let name = at + entry.len() - 1; record[name] = named; let mut changes = ChangeSet::new(view, &budget).expect("a change set"); changes.put_schema("docs", record).expect("stage"); @@ -381,7 +358,9 @@ fn a_write_is_reported_as_damage(test: &str, named: u8, what: &str) { 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"); + changes + .put_row(table, key, row, Vec::new(), true) + .expect("stage"); let staged = Catalog::load(&view).and_then(|catalog| stage(&mut changes, &catalog, &budget)); told_as_damage(staged, &folder, what); drop(changes); @@ -430,3 +409,1218 @@ fn a_cycle_over_a_vector_column_that_is_not_of_the_vector_type_is_reported_as_da 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); } + +// ---- the records of a grouping ---- + +/// The column every record of the tests below is kept for. +const COLUMN: VectorId = VectorId(3); + +/// An empty store in a folder of its own. +fn a_store(test: &str) -> (PathBuf, Store, Budget) { + 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 budget = Budget::new(None).expect("a budget"); + let store = Store::open_for_writing(&folder, &budget, None).expect("open a fresh store"); + (folder, store, budget) +} + +/// Commits one record of the vector part and returns the store as of that commit. +fn put(store: &Store, budget: &Budget, record: VectorRecord<'_>, value: Vec) -> View { + let mut changes = ChangeSet::new(store.view().expect("a view"), budget).expect("begin"); + changes.put_vector(record, value).expect("stage"); + store.commit(changes).expect("commit"); + store.view().expect("the newest commit") +} + +/// A record with four of its bytes, from `at` on, set to this number. +fn with(record: &[u8], at: usize, number: u32) -> Vec { + let mut record = record.to_vec(); + record[at..at + 4].copy_from_slice(&number.to_le_bytes()); + record +} + +fn unit(vector: &[f32]) -> Vec { + let length = vector + .iter() + .map(|x| f64::from(*x).powi(2)) + .sum::() + .sqrt(); + vector + .iter() + .map(|x| (f64::from(*x) / length) as f32) + .collect() +} + +fn slot_of(vector: &[f32]) -> Vec { + let mut slot = vec![0_u8; slot_len(vector.len())]; + write_slot(vector, &mut slot); + slot +} + +/// A slot written by hand: this scale, no rounding error, these codes. +fn hand_slot(scale: f32, codes: &[i8]) -> Vec { + let mut slot = scale.to_le_bytes().to_vec(); + slot.extend(0.0_f32.to_le_bytes()); + slot.extend(codes.iter().map(|code| code.cast_unsigned())); + slot +} + +/// The vector a slot keeps: its scale times each code. +fn kept_by(slot: &[u8]) -> Vec { + let scale = f32::from_le_bytes(std::array::from_fn(|at| slot[at])); + let codes = slot[8..].iter(); + codes + .map(|code| scale * f32::from(code.cast_signed())) + .collect() +} + +fn close_to(found: &[f32], wanted: &[f32], what: &str) { + assert_eq!(found.len(), wanted.len(), "{what}"); + for (at, (found, wanted)) in found.iter().zip(wanted).enumerate() { + assert!( + (found - wanted).abs() < 0.01, + "{what}, number {at}: {found:e}, not {wanted:e}" + ); + } +} + +/// A partition with every field set: three tail blocks from the third on, a grouping of 17 +/// groups in force, another being built, and a mean. +fn a_partition() -> Partition { + Partition { + number: 7, + live: 300, + blocks: 6, + first: 2, + closed: 4, + slots: 500, + sealed: 9000, + grouping: Some(Grouped { + number: 9, + blocks: 3, + groups: 17, + slots: 290, + upto: 8500, + }), + building: Some(11), + mean: true, + } +} + +#[test] +fn a_partition_record_reads_back() { + // The 73 bytes: the number, the live rows, the tail's last, first and closed blocks, its + // slots, the slots ever sealed, the flags, the grouping in force and the one being built. + let whole = a_partition(); + let mut wanted = Vec::new(); + wanted.extend(7_u32.to_le_bytes()); + wanted.extend(300_u64.to_le_bytes()); + for number in [6_u32, 2, 4] { + wanted.extend(number.to_le_bytes()); + } + wanted.extend(500_u64.to_le_bytes()); + wanted.extend(9000_u64.to_le_bytes()); + wanted.push(1 | 2 | 4); + for number in [9_u32, 3, 17] { + wanted.extend(number.to_le_bytes()); + } + wanted.extend(290_u64.to_le_bytes()); + wanted.extend(8500_u64.to_le_bytes()); + wanted.extend(11_u32.to_le_bytes()); + assert_eq!(wanted.len(), 73); + assert_eq!(whole.encode(), wanted); + + // A field its flag does not set is zero. + let bare = Partition { + number: 0x0102_0304, + live: 0x0506_0708_090A_0B0C, + blocks: 0x0D0E_0F10, + sealed: u64::MAX, + ..Partition::default() + }; + let mut wanted = vec![ + 4, 3, 2, 1, 0x0C, 0x0B, 0x0A, 9, 8, 7, 6, 5, 0x10, 0x0F, 0x0E, 0x0D, + ]; + wanted.extend([0; 16]); + wanted.extend([0xFF; 8]); + wanted.extend([0; 33]); + assert_eq!(bare.encode(), wanted); + + let (folder, store, budget) = a_store("partition-record"); + for (at, entry) in [whole, bare].into_iter().enumerate() { + let tuple = key(at as i64 + 1); + let record = VectorRecord::Partition(COLUMN, &tuple); + let view = put(&store, &budget, record, entry.encode()); + assert_eq!(partition(&view, COLUMN, &tuple).expect("read"), Some(entry)); + let other = key(50); + assert_eq!(partition(&view, COLUMN, &other).expect("read"), None); + } + drop(store); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +#[test] +fn a_partition_entry_that_disagrees_with_itself_is_damage() { + let (folder, store, budget) = a_store("partition-damage"); + let good = a_partition().encode(); + let flags = |flags: u8| { + let mut record = good.clone(); + record[40] = flags; + record + }; + let cases = [ + ("of sixteen bytes", good[..16].to_vec()), + ("one byte longer", [good.as_slice(), &[0]].concat()), + ("with a flag above the three", flags(15)), + ( + "whose tail holds more slots than were ever sealed", + with(&good, 32, 499), + ), + ("with a grouping of 17 groups and no mean", flags(3)), + ( + "whose first tail block is after its closed one", + with(&good, 16, 5), + ), + ("whose closed block is after its last", with(&good, 20, 7)), + ("with a grouping of no group", with(&good, 49, 0)), + ]; + let tuple = key(1); + for (what, record) in cases { + let view = put( + &store, + &budget, + VectorRecord::Partition(COLUMN, &tuple), + record, + ); + let what = format!("a partition entry {what}"); + told_as_damage(partition(&view, COLUMN, &tuple), &folder, &what); + } + drop(store); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +/// What a partition entry says of the grouping of the next tests, and that grouping: three +/// groups in two blocks. +const GROUPED: Grouped = Grouped { + number: 9, + blocks: 2, + groups: 3, + slots: 500, + upto: 500, +}; + +fn a_grouping() -> Grouping { + let place = |block, slot| Place { block, slot }; + Grouping { + gap: 780, + probes: 2, + width: 4, + starts: vec![place(0, 0), place(0, 200), place(1, 50), place(1, 300)], + } +} + +#[test] +fn a_grouping_record_reads_back() { + let grouping = a_grouping(); + // The gap, the width and the probes, then each place as its block and its slot. + let numbers = [780_u32, 4, 2, 0, 0, 0, 200, 1, 50, 1, 300]; + let wanted: Vec = numbers.iter().flat_map(|n| n.to_le_bytes()).collect(); + assert_eq!(wanted.len(), 12 + 8 * 4); + assert_eq!(grouping.encode(), wanted); + let (folder, store, budget) = a_store("grouping-record"); + let view = put(&store, &budget, VectorRecord::Grouping(COLUMN, 9), wanted); + assert_eq!( + Grouping::read(&view, COLUMN, GROUPED).expect("read"), + grouping + ); + drop((view, store)); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +#[test] +fn a_grouping_that_disagrees_with_its_partition_is_damage() { + let (folder, store, budget) = a_store("grouping-damage"); + let good = a_grouping().encode(); + let view = store.view().expect("a view"); + let read = Grouping::read(&view, COLUMN, GROUPED); + told_as_damage(read, &folder, "a grouping with no record"); + drop(view); + let cases = [ + ( + "of the places of two groups", + good[..good.len() - 8].to_vec(), + ), + ("with a gap above 20,000", with(&good, 0, 20_001)), + ("with a width of one quarter", with(&good, 4, 1)), + ("with a width above sixteen quarters", with(&good, 4, 17)), + ("that scans no group", with(&good, 8, 0)), + ("that scans more groups than it has", with(&good, 8, 4)), + ("that does not begin at the first slot", with(&good, 16, 1)), + ( + "with a place before the one ahead of it", + with(&with(&good, 28, 0), 32, 100), + ), + ("that ends past its blocks", with(&good, 36, 5)), + ]; + for (what, record) in cases { + let view = put(&store, &budget, VectorRecord::Grouping(COLUMN, 9), record); + let what = format!("a grouping {what}"); + told_as_damage(Grouping::read(&view, COLUMN, GROUPED), &folder, &what); + } + drop(store); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +#[test] +fn centres_read_back_in_records_of_a_block() { + const DIMENSION: usize = 16; + const GROUPS: u32 = 5000; + assert_eq!(centres_to_a_record(768), 85); + assert_eq!(centres_to_a_record(DIMENSION), 4095); + let all: Vec = (0..u64::from(GROUPS)) + .flat_map(|group| unit(&vector(group, DIMENSION))) + .collect(); + let (folder, store, budget) = a_store("centres"); + let mut changes = ChangeSet::new(store.view().expect("a view"), &budget).expect("begin"); + put_centres(&mut changes, COLUMN, 9, DIMENSION, &all).expect("stage the centres"); + store.commit(changes).expect("commit"); + + // One full record and the rest in the next. + let view = store.view().expect("the newest commit"); + let mut records = Vec::new(); + let from = VectorRecord::Centres(COLUMN, 9, 0); + view.vectors(from, &mut |record, value| match record { + VectorRecord::Centres(COLUMN, 9, place) => { + records.push((place, value.len())); + Ok(true) + } + _ => Ok(false), + }) + .expect("read the records"); + assert_eq!(records, [(0, 4095 * 64), (1, 905 * 64)]); + + let mut held = budget.reserve(0, "query").expect("a reservation"); + let read = centres(&view, COLUMN, 9, GROUPS, DIMENSION, &mut held).expect("read the centres"); + assert!(read == all, "the centres read back"); + assert_eq!(held.bytes(), 4 * DIMENSION as u64 * u64::from(GROUPS)); + + // Ranked for the vector that is centre 1,234: that group first, then by product, groups + // of one product in the order of their numbers. + let grouped = Grouped { + groups: GROUPS, + ..GROUPED + }; + let asked = &all[1234 * DIMENSION..1235 * DIMENSION]; + let ranked = nearest_groups(&view, COLUMN, grouped, asked, &mut held).expect("rank"); + assert_eq!(ranked.len(), GROUPS as usize); + assert_eq!(ranked[0].0, 1234); + for pair in ranked.windows(2) { + let ((first, high), (second, low)) = (pair[0], pair[1]); + assert!(high > low || high == low && first < second, "{pair:?}"); + } + drop(held); + drop((view, store)); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +#[test] +fn centres_that_do_not_come_to_their_groups_are_damage() { + const DIMENSION: usize = 16; + let (folder, store, budget) = a_store("centres-damage"); + let three: Vec = (0..3) + .flat_map(|group| unit(&vector(group, DIMENSION))) + .collect(); + let bytes: Vec = three + .iter() + .flat_map(|number| number.to_le_bytes()) + .collect(); + let mut not_finite = bytes.clone(); + not_finite[4..8].copy_from_slice(&f32::NAN.to_le_bytes()); + // Each under a number of its own, read as the centres of this many groups. + let cases = [ + ("three centres for a grouping of four", 20, 4, bytes.clone()), + ( + "a record that is no whole number of centres", + 21, + 3, + bytes[..bytes.len() - 1].to_vec(), + ), + ( + "a centre with a number that is not finite", + 22, + 3, + not_finite, + ), + ("no record", 23, 3, Vec::new()), + ]; + for (what, number, groups, record) in cases { + let view = if record.is_empty() { + store.view().expect("a view") + } else { + put( + &store, + &budget, + VectorRecord::Centres(COLUMN, number, 0), + record, + ) + }; + let mut held = budget.reserve(0, "query").expect("a reservation"); + let read = centres(&view, COLUMN, number, groups, DIMENSION, &mut held); + told_as_damage(read, &folder, what); + } + drop(store); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +#[test] +fn a_mean_reads_back_and_one_of_another_length_is_damage() { + const DIMENSION: usize = 16; + let (folder, store, budget) = a_store("mean"); + let kept = vector(77, DIMENSION); + let mut changes = ChangeSet::new(store.view().expect("a view"), &budget).expect("begin"); + put_mean(&mut changes, COLUMN, 7, &kept).expect("stage the mean"); + store.commit(changes).expect("commit"); + let view = store.view().expect("the newest commit"); + let mut held = budget.reserve(0, "query").expect("a reservation"); + let read = mean(&view, COLUMN, 7, DIMENSION, &mut held).expect("read the mean"); + assert!(read == kept, "the mean read back"); + drop(view); + + let bytes: Vec = kept + .iter() + .flat_map(|number| number.to_le_bytes()) + .collect(); + let mut not_finite = bytes.clone(); + not_finite[..4].copy_from_slice(&f32::INFINITY.to_le_bytes()); + let cases = [ + ( + "a mean of fifteen numbers", + 8, + bytes[..bytes.len() - 4].to_vec(), + ), + ("a mean with a number that is not finite", 9, not_finite), + ("a mean with no record", 10, Vec::new()), + ]; + for (what, number, record) in cases { + let view = if record.is_empty() { + store.view().expect("a view") + } else { + put(&store, &budget, VectorRecord::Mean(COLUMN, number), record) + }; + told_as_damage( + mean(&view, COLUMN, number, DIMENSION, &mut held), + &folder, + what, + ); + } + drop(held); + drop(store); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +// ---- finding centres ---- + +#[test] +fn the_groups_of_a_partition_follow_the_square_root() { + let wanted = [ + (0, 1), + (1, 1), + (2, 1), + (3, 2), + (2000, 45), + (100_008, 316), + (1_000_000, 1000), + (5_000_000, 2236), + (u64::MAX, 65_535), + ]; + for (slots, groups) in wanted { + assert_eq!(groups_for(slots), groups, "{slots} slots"); + } +} + +#[test] +fn the_first_centres_are_the_same_for_the_same_numbers() { + let first = first_centres(1000, 10); + assert_eq!(first, first_centres(1000, 10)); + assert_eq!(first.len(), 10); + assert!(first.windows(2).all(|pair| pair[0] < pair[1]), "{first:?}"); + assert!(first[9] < 1000, "{first:?}"); + assert_eq!(first_centres(10, 10), (0..10).collect::>()); + // The sample: every slot of a small partition, and 64 a group evenly spread of a large. + assert_eq!( + sampled(10, 1).collect::>(), + (0..10).collect::>() + ); + let spread: Vec = sampled(1000, 2).collect(); + assert_eq!(spread.len(), 128); + assert_eq!((&spread[..4], spread[127]), (&[0, 7, 15, 23][..], 992)); +} + +#[test] +fn a_round_moves_each_centre_to_its_slots() { + const DIMENSION: usize = 4; + let budget = Budget::new(None).expect("a budget"); + let mut held = budget.reserve(0, "maintenance").expect("a reservation"); + let mut found = vec![ + 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0_f32, + ]; + let near_first = [ + slot_of(&[1.0, 0.2, 0.0, 0.0]), + slot_of(&[1.0, -0.2, 0.0, 0.0]), + ]; + let near_second = slot_of(&[0.0, 1.0, 0.5, 0.0]); + // The nearest centre is the one with the greatest product, the lower number of two alike. + assert_eq!(nearest(&found, DIMENSION, &near_first[0]).0, 0); + let (group, greatest) = nearest(&found, DIMENSION, &near_second); + assert_eq!(group, 1); + assert!((greatest - 0.894).abs() < 0.01, "the product {greatest:e}"); + assert_eq!( + product(&found[4..8], &near_second).to_bits(), + greatest.to_bits() + ); + assert_eq!( + nearest(&found, DIMENSION, &slot_of(&[1.0, 1.0, 0.0, 0.0])).0, + 0 + ); + + let mut round = Round::new(3, DIMENSION, &mut held).expect("a round"); + assert_eq!(held.bytes(), 8 * 4 * 3 + 4 * 3); + for slot in near_first.iter().chain([&near_second]) { + round.take(&found, slot); + } + round.finish(&mut found); + // Each centre that drew a slot is the direction of its slots' sum; the third drew none. + close_to(&found[..4], &[1.0, 0.0, 0.0, 0.0], "the first centre"); + close_to(&found[4..8], &[0.0, 0.894, 0.447, 0.0], "the second centre"); + assert_eq!(found[8..], [0.0, 0.0, 1.0, 0.0]); + for centre in found.chunks(DIMENSION) { + let length: f32 = centre.iter().map(|x| x * x).sum(); + assert!( + (length - 1.0).abs() < 1e-5, + "a centre of squared length {length:e}" + ); + } + + // The mean is not brought to length one, and is zeros when nothing was taken. + let mut middle = Mean::new(DIMENSION, &mut held).expect("a mean"); + middle.take(&near_first[0]); + middle.take(&near_second); + close_to(&middle.finish(), &[0.490, 0.545, 0.224, 0.0], "the mean"); + let nothing = Mean::new(DIMENSION, &mut held).expect("a mean"); + assert_eq!(nothing.finish(), [0.0; DIMENSION]); +} + +// ---- the record of a build ---- + +/// A build of 45 groups over the one tail block of 2,000 slots, at this phase. +fn a_build(phase: Phase) -> Build { + Build { + number: 12, + steps: 5, + from: None, + tail: (0, 1), + slots: 2000, + upto: 2000, + stale: false, + groups: 45, + phase, + counts: vec![0; 45], + starts: Vec::new(), + tens: Vec::new(), + needing: 0, + } +} + +/// The slots each group was given, and where each group then begins in the one new block. +fn given() -> (Vec, Vec) { + let mut counts = vec![44_u32; 45]; + counts[0] = 64; + let mut starts = vec![Place::default()]; + for count in &counts { + let slot = starts[starts.len() - 1].slot + count; + starts.push(Place { block: 0, slot }); + } + (counts, starts) +} + +#[test] +fn a_build_record_reads_back() { + let (counts, starts) = given(); + let tens = |questions: usize| vec![Place { block: 0, slot: 7 }; 10 * questions]; + let packed = |phase: Phase, questions: usize| Build { + counts: counts.clone(), + starts: starts.clone(), + tens: tens(questions), + ..a_build(phase) + }; + let builds = [ + a_build(Phase::Train { round: 3 }), + Build { + from: Some(Grouped { + number: 9, + blocks: 1, + groups: 40, + slots: 1600, + upto: 1600, + }), + tail: (2, 3), + stale: true, + needing: 4096, + ..a_build(Phase::Assign { next: 1 }) + }, + Build { + counts: counts.clone(), + starts: starts[..3].to_vec(), + ..a_build(Phase::Pack { group: 2 }) + }, + packed(Phase::Measure { done: 2, most: 150 }, 2), + packed( + Phase::Width { + gap: 150, + probes: 3, + done: 1, + asked: 10, + kept: [9; 15], + }, + 128, + ), + packed( + Phase::Switch { + gap: 150, + probes: 3, + width: 4, + }, + 128, + ), + Build { + stale: true, + ..packed(Phase::Sweep, 0) + }, + ]; + let (folder, store, budget) = a_store("build-record"); + let view = store.view().expect("a view"); + assert_eq!(Build::read(&view, COLUMN, 7).expect("read"), None); + drop(view); + for build in builds { + let view = put( + &store, + &budget, + VectorRecord::Build(COLUMN, 7), + build.encode(), + ); + let read = Build::read(&view, COLUMN, 7).expect("read the build"); + assert_eq!(read, Some(build)); + } + drop(store); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +#[test] +fn a_build_whose_cursor_is_past_its_blocks_is_damage() { + let (counts, _) = given(); + let whole = a_build(Phase::Train { round: 3 }).encode(); + let cases = [ + ( + "a build that gives groups from a block it does not read", + a_build(Phase::Assign { next: 5 }).encode(), + ), + ( + "a build that packs a group it does not have", + Build { + counts, + ..a_build(Phase::Pack { group: 46 }) + } + .encode(), + ), + ( + "a build's record cut short", + whole[..whole.len() - 1].to_vec(), + ), + ]; + let (folder, store, budget) = a_store("build-damage"); + for (what, record) in cases { + let view = put(&store, &budget, VectorRecord::Build(COLUMN, 7), record); + told_as_damage(Build::read(&view, COLUMN, 7), &folder, what); + } + drop(store); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +// ---- sketches ---- + +#[test] +fn a_sketch_holds_the_signs_and_the_factor_of_its_slot() { + // Sixteen numbers worked out by hand. The slot keeps an eighth of each code; the + // reference is a half for the first eight numbers and zero for the rest. + let codes = [8, -8, 4, -4, 2, -2, 1, -1, 0, 0, 8, 8, -8, -8, 4, 4]; + let slot = hand_slot(0.125, &codes); + let reference = [[0.5_f32; 8], [0.0; 8]].concat(); + // The differences are 0.5, -1.5, 0, -1, -0.25, -0.75, -0.375, -0.625, then 0, 0, 1, 1, + // -1, -1, 0.5, 0.5: bits 0 and 2 of the first byte, and all but bits 4 and 5 of the + // second. Their squares come to 9.15625 and their sizes to 10: the factor is 30,003. + assert_eq!(sketch_len(16), 4); + let mut sketch = [0xAA_u8; 4]; + write_sketch(&slot, &reference, &mut sketch); + assert_eq!(sketch, [0b0000_0101, 0b1100_1111, 0x33, 0x75]); + + // No difference at all: every bit, and a factor of zero. + write_sketch(&slot, &kept_by(&slot), &mut sketch); + assert_eq!(sketch, [0xFF, 0xFF, 0, 0]); + + // Twelve numbers: the four bits past them stay zero. Differences of one give a factor + // of 32,768, and differences of two one that is held to 65,535. + assert_eq!(sketch_len(12), 4); + for (code, factor) in [(8_i8, 32_768_u16), (16, u16::MAX)] { + let mut sketch = [0xAA_u8; 4]; + write_sketch(&hand_slot(0.125, &[code; 12]), &[0.0; 12], &mut sketch); + let [low, high] = factor.to_le_bytes(); + assert_eq!(sketch, [0xFF, 0x0F, low, high]); + } +} + +/// The vector of a search at length one, as numbers of 32 bits. +fn asked_unit(query: &[f32]) -> Vec { + unit(query) +} + +/// The record of the sketches of these slots, each from this reference. +fn record_of(slots: &[Vec], reference: &[f32]) -> Vec { + let mut record = (slots.len() as u32).to_le_bytes().to_vec(); + for slot in slots { + let mut sketch = vec![0_u8; sketch_len(reference.len())]; + write_sketch(slot, reference, &mut sketch); + record.extend(sketch); + } + record +} + +#[test] +fn a_sketch_estimates_a_product_within_its_factor() { + const DIMENSION: usize = 16; + for pair in 0..20_u64 { + let slot = slot_of(&vector(3 * pair, DIMENSION)); + let reference: Vec = vector(3 * pair + 1, DIMENSION) + .iter() + .map(|number| 0.1 * number) + .collect(); + let query = vector(3 * pair + 2, DIMENSION); + let asked = asked_unit(&query); + let dot = |a: &[f32], b: &[f32]| -> f64 { + a.iter() + .zip(b) + .map(|(x, y)| f64::from(*x) * f64::from(*y)) + .sum() + }; + let truth = dot(&kept_by(&slot), &asked); + let base = dot(&reference, &asked) as f32; + + let record = record_of(std::slice::from_ref(&slot), &reference); + let factor = f64::from(u16::from_le_bytes([record[6], record[7]])) / 32_768.0; + let mut ranges = Vec::new(); + let sketches = Sketches::new(&record, DIMENSION); + Estimate::new(&query).ranges(&sketches, 0..1, base, WIDEST, &mut ranges); + let [(low, high)] = ranges[..] else { + panic!("pair {pair}: {} ranges for one sketch", ranges.len()) + }; + // At the widest, a range is four factors to each side of the estimate. + let (low, high) = (f64::from(low), f64::from(high)); + assert!( + ((high - low) - 8.0 * factor).abs() < 1e-4, + "pair {pair}: a range from {low:e} to {high:e} with the factor {factor:e}" + ); + assert!( + low - 1e-5 <= truth && truth <= high + 1e-5, + "pair {pair}: the product {truth:e} is outside the range from {low:e} to {high:e}" + ); + } +} + +#[test] +fn the_ranges_of_a_run_of_sketches_are_those_of_the_plain_arithmetic() { + // A hundred numbers: a sketch's second word is completed with zeros. + const DIMENSION: usize = 100; + const WIDTH: u8 = 6; + let reference: Vec = vector(500, DIMENSION).iter().map(|x| 0.1 * x).collect(); + let slots: Vec> = (0..40) + .map(|n| slot_of(&vector(600 + n, DIMENSION))) + .collect(); + let record = record_of(&slots, &reference); + let query = vector(700, DIMENSION); + let base = 0.125_f32; + let mut ranges = Vec::new(); + let sketches = Sketches::new(&record, DIMENSION); + Estimate::new(&query).ranges(&sketches, 3..37, base, WIDTH, &mut ranges); + assert_eq!(ranges.len(), 34); + + // The same worked out bit by bit: the search's vector at length one in whole steps of a + // seventh of its largest number, and for each sketch the steps under its set bits. + let asked = asked_unit(&query); + let step = asked.iter().fold(0.0_f32, |most, x| most.max(x.abs())) / 7.0; + let steps: Vec = asked + .iter() + .map(|x| ((x / step).round() as i32).clamp(-7, 7)) + .collect(); + let sum: i32 = steps.iter().sum(); + for (place, (low, high)) in (3..37).zip(ranges) { + let sketch = sketches.at(place); + let set = |at: usize| sketch[at / 8] >> (at % 8) & 1 == 1; + let under: i32 = (0..DIMENSION) + .filter(|at| set(*at)) + .map(|at| steps[at]) + .sum(); + let factor = f32::from(u16::from_le_bytes([sketch[13], sketch[14]])) / 32_768.0; + let estimate = base + factor * step * (2 * under - sum) as f32; + let wide = f32::from(WIDTH) / 4.0 * factor; + for (found, wanted) in [(low, estimate - wide), (high, estimate + wide)] { + assert!( + (found - wanted).abs() <= 1e-5 * wanted.abs().max(1.0), + "sketch {place}: {found:e} where the plain arithmetic gives {wanted:e}" + ); + } + } +} + +#[test] +fn the_sketches_of_a_block_read_back() { + const DIMENSION: usize = 16; + let keys: Vec = (0..5).map(key).collect(); + let (mut writer, taken) = BlockWriter::new(DIMENSION, None, &keys); + assert_eq!(taken, 5); + for (n, key) in keys.iter().enumerate() { + writer.push(key, &vector(800 + n as u64, DIMENSION)); + } + let block = writer.finish(); + // Slot `place` is sketched from reference `place % 2`. + let references = [vector(900, DIMENSION), vector(901, DIMENSION)]; + let record = sketches_of(&block, DIMENSION, &mut |place| &references[place % 2]); + assert_eq!(record.len(), 4 + 5 * sketch_len(DIMENSION)); + assert_eq!(record[..4], 5_u32.to_le_bytes()); + let sketches = Sketches::new(&record, DIMENSION); + assert!(sketches.is_whole(5) && !sketches.is_whole(4) && !sketches.is_whole(6)); + assert!(!Sketches::new(&record[..record.len() - 1], DIMENSION).is_whole(5)); + let reader = BlockReader::new(&block, DIMENSION); + for place in 0..5 { + let mut wanted = vec![0_u8; sketch_len(DIMENSION)]; + write_sketch(reader.slot(place), &references[place % 2], &mut wanted); + assert_eq!(sketches.at(place), wanted.as_slice(), "sketch {place}"); + } +} + +// ---- damage in a built store, through the public calls ---- + +const BUILT: usize = 16; + +/// Row `n` of a built store: one of eight directions at right angles, moved a little. +fn gathered(n: u64) -> Vec { + let moved = vector(n, BUILT); + let direction = (n % 8) as usize; + let lifted = |at: usize| f32::from(u8::from(at / 2 == direction)); + (0..BUILT).map(|at| lifted(at) + 0.05 * moved[at]).collect() +} + +fn fresh(test: &str) -> PathBuf { + 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"); + } + folder +} + +fn insert_gathered(db: &Database, ids: std::ops::RangeInclusive) { + let tx = db.transaction().expect("begin"); + for id in ids { + tx.execute( + "INSERT INTO docs (id, v) VALUES ($id, $v)", + &crate::params! { "id" => id as i64, "v" => gathered(id) }, + ) + .expect("insert a row"); + } + tx.commit().expect("commit"); +} + +/// An open store whose `docs.v` holds 2,000 vectors in a column that is at its size, after +/// this many maintenance cycles: sixteen build its grouping of 45 groups. +fn a_store_with_vectors(test: &str, cycles: usize) -> (PathBuf, Database) { + let folder = fresh(test); + let db = Database::open(&folder).expect("open a fresh store"); + let table = "CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(16) AUTO_INDEX_AT 1000)"; + db.execute(table, &HashMap::new()).expect(table); + insert_gathered(&db, 1..=2000); + for _ in 0..cycles { + db.run_maintenance_cycle().expect("a maintenance cycle"); + } + (folder, db) +} + +/// The first record of the kind of `from` that the column of `from` keeps, with its value. +fn first_record(view: &View, from: VectorRecord<'static>) -> (VectorRecord<'static>, Vec) { + let mut found = None; + view.vectors(from, &mut |record, value| { + found = Some(match record { + VectorRecord::Grouping(column, number) => VectorRecord::Grouping(column, number), + VectorRecord::Sketch(column, under, block) => { + VectorRecord::Sketch(column, under, block) + } + VectorRecord::Assigned(column, number, place) => { + VectorRecord::Assigned(column, number, place) + } + other => panic!("looked for {from:?} and the next record is {other:?}"), + }) + .map(|record| (record, value.to_vec())); + Ok(false) + }) + .expect("read the records"); + let (record, value) = found.unwrap_or_else(|| panic!("no record from {from:?} on")); + assert_eq!( + std::mem::discriminant(&record), + std::mem::discriminant(&from), + "the first record from {from:?} on" + ); + (record, value) +} + +/// Closes the store, changes the first record from `from` on with `damage`, and opens it. +fn damaged( + folder: &Path, + db: Database, + from: fn(VectorId) -> VectorRecord<'static>, + damage: fn(&mut Vec), +) -> Database { + db.close().expect("close"); + drop(db); + 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 column = catalog.table("docs").expect("the table").vectors[0].id; + let (record, mut value) = first_record(&view, from(column)); + damage(&mut value); + drop(view); + put(&store, &budget, record, value); + drop(store); + Database::open(folder).expect("open the damaged store") +} + +/// The ids a search of ten returns under this mode. +fn ten(db: &Database, mode: &str) -> crate::types::Result> { + let sql = format!( + "SELECT id, v <=> $q AS distance FROM docs ORDER BY distance USE VECTOR {mode} LIMIT 10" + ); + let result = db.execute(&sql, &crate::params! { "q" => gathered(40_003) })?; + Ok(result.rows().iter().map(|row| row[0].clone()).collect()) +} + +/// Damages the index of a built store and asserts what its searches then give: INDEXED the +/// store's damage error, AUTO and EXACT the same ten rows. +fn a_damaged_index_refuses_indexed_and_auto_answers_exactly( + test: &str, + from: fn(VectorId) -> VectorRecord<'static>, + damage: fn(&mut Vec), + what: &str, +) { + let (folder, db) = a_store_with_vectors(test, 16); + assert_eq!(ten(&db, "INDEXED").expect("the built store").len(), 10); + let db = damaged(&folder, db, from, damage); + told_as_damage(ten(&db, "INDEXED"), &folder, what); + let exact = ten(&db, "EXACT").expect("EXACT reads no record of the index"); + assert_eq!(exact.len(), 10); + match ten(&db, "AUTO") { + Ok(auto) => assert_eq!(auto, exact, "AUTO over {what}"), + Err(error) => panic!("AUTO over {what} gave {error:?}"), + } + db.close().expect("close"); + drop(db); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +#[test] +fn a_sketch_record_that_does_not_match_its_block_is_damage() { + // The record of the one block's sketches, holding one sketch fewer than the block has + // slots, and saying so. + a_damaged_index_refuses_indexed_and_auto_answers_exactly( + "sketch-count", + |column| VectorRecord::Sketch(column, 0, 0), + |record| { + let count = u32::from_le_bytes(std::array::from_fn(|at| record[at])) - 1; + record[..4].copy_from_slice(&count.to_le_bytes()); + record.truncate(record.len() - sketch_len(BUILT)); + }, + "a record of sketches whose count is not its block's", + ); +} + +#[test] +fn auto_answers_exactly_where_a_grouping_is_damaged() { + a_damaged_index_refuses_indexed_and_auto_answers_exactly( + "grouping-gap", + |column| VectorRecord::Grouping(column, 0), + |record| record[..4].copy_from_slice(&20_001_u32.to_le_bytes()), + "a grouping whose gap is above 20,000", + ); +} + +#[test] +fn assigned_groups_that_do_not_match_their_block_are_damage() { + // Twelve cycles: sealed and begun, ten rounds, and every slot given its group. + let (folder, db) = a_store_with_vectors("assigned-count", 12); + let db = damaged( + &folder, + db, + |column| VectorRecord::Assigned(column, 0, 0), + |record| record.extend([0, 0]), + ); + told_as_damage( + db.run_maintenance_cycle(), + &folder, + "a cycle that packs from the groups of one slot more than the block has", + ); + db.close().expect("close"); + drop(db); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +/// Closes the store and opens it as these tests write to it: the store, its budget, the +/// column of `docs.v` and the entry of its one partition. +fn beneath(folder: &Path, db: Database) -> (Store, Budget, VectorId, Partition) { + db.close().expect("close"); + drop(db); + 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 column = catalog.table("docs").expect("the table").vectors[0].id; + let none = Key::from_values(&[], &[]).expect("the key of no values"); + let entry = partition(&view, column, &none).expect("read the partition"); + drop(view); + (store, budget, column, entry.expect("the one partition")) +} + +/// The build the one partition of a store is making, read as [`beneath`] opens the store. +fn build_under_way(store: &Store, column: VectorId, entry: &Partition) -> Build { + let view = store.view().expect("the newest commit"); + let build = Build::read(&view, column, entry.number).expect("read the build"); + build.expect("a build under way") +} + +/// Puts this record of a build in place of the partition's, opens the store, and asserts that +/// the next cycle tells the damage. +fn a_cycle_over_this_build_is_damage(test: &str, cycles: usize, damage: fn(Build) -> Build) { + let (folder, db) = a_store_with_vectors(test, cycles); + let (store, budget, column, entry) = beneath(&folder, db); + let build = damage(build_under_way(&store, column, &entry)); + let record = VectorRecord::Build(column, entry.number); + put(&store, &budget, record, build.encode()); + drop(store); + let db = Database::open(&folder).expect("open the damaged store"); + told_as_damage(db.run_maintenance_cycle(), &folder, test); + db.close().expect("close"); + drop(db); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +#[test] +fn a_build_of_more_groups_than_slots_is_damage() { + // One cycle: sealed and begun. The record then says 2,001 groups for its 2,000 slots. + a_cycle_over_this_build_is_damage("build-more-groups-than-slots", 1, |build| { + assert_eq!((build.slots, build.groups), (2000, 45)); + Build { + groups: 2001, + counts: vec![0; 2001], + ..build + } + }); +} + +#[test] +fn a_packed_build_whose_groups_begin_out_of_order_is_damage() { + // Thirteen cycles: every group is packed and its beginning recorded. The second group + // then begins after the third. + a_cycle_over_this_build_is_damage("build-starts-out-of-order", 13, |mut build| { + assert_eq!(build.starts.len(), 46); + assert!(build.starts[1] < build.starts[2], "{:?}", build.starts); + build.starts.swap(1, 2); + build + }); +} + +#[test] +fn assigned_groups_of_an_odd_length_are_damage() { + // Twelve cycles: sealed and begun, ten rounds, and every slot given its group in two + // bytes. The record then ends half way through a slot's. + let (folder, db) = a_store_with_vectors("assigned-odd", 12); + let db = damaged( + &folder, + db, + |column| VectorRecord::Assigned(column, 0, 0), + |record| record.push(0), + ); + told_as_damage( + db.run_maintenance_cycle(), + &folder, + "a cycle that packs from groups recorded in an odd number of bytes", + ); + db.close().expect("close"); + drop(db); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +/// A store whose grouping is being built again from slots one of which its row has left, so +/// that the partition keeps one `Left` record. Once the build has reached `sweeping` (or at +/// once, for `false`) that record's value is made `bytes` long, and the cycles that follow +/// must tell the damage: the build reads the record for the slot's key before it gives the +/// slot a group, and the steps after the switch read it to remove it. +fn a_left_record_of_this_length_is_damage(test: &str, sweeping: bool, bytes: usize) { + let (folder, mut db) = a_store_with_vectors(test, 16); + insert_gathered(&db, 2001..=2400); + let changed = crate::params! { "v" => gathered(50_000) }; + db.execute("UPDATE docs SET v = $v WHERE id = 5", &changed) + .expect("change a vector"); + db.run_maintenance_cycle().expect("the cycle that begins"); + let none = Key::from_values(&[], &[]).expect("the key of no values"); + for cycle in 0.. { + let (store, budget, column, entry) = beneath(&folder, db); + let phase = build_under_way(&store, column, &entry).phase; + if sweeping && phase != Phase::Sweep { + assert!(cycle < 40, "40 cycles and the build is at {phase:?}"); + drop(store); + db = Database::open(&folder).expect("open"); + db.run_maintenance_cycle().expect("a cycle of the build"); + continue; + } + let view = store.view().expect("the newest commit"); + let mut left = None; + view.vectors( + VectorRecord::Left(column, entry.number, &none), + &mut |record, value| { + if let VectorRecord::Left(_, _, key) = record { + left = Some((key.clone(), value.to_vec())); + } + Ok(false) + }, + ) + .expect("read the records"); + drop(view); + let (key, mut value) = left.expect("a record of the row that left its slot"); + assert_eq!(value.len(), 8); + value.resize(bytes, 0); + put( + &store, + &budget, + VectorRecord::Left(column, entry.number, &key), + value, + ); + drop(store); + db = Database::open(&folder).expect("open the damaged store"); + break; + } + let told = (0..40).find_map(|_| db.run_maintenance_cycle().err()); + told_as_damage(told.map_or(Ok(()), Err), &folder, test); + db.close().expect("close"); + drop(db); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} + +#[test] +fn a_left_record_that_is_not_eight_bytes_is_damage() { + a_left_record_of_this_length_is_damage("left-short-while-building", false, 7); + a_left_record_of_this_length_is_damage("left-short-while-sweeping", true, 7); + a_left_record_of_this_length_is_damage("left-long-while-sweeping", true, 9); +} + +/// How many records of the kind of `from` the column of `from` keeps, from `from` on. +fn records_from(view: &View, from: VectorRecord<'_>) -> usize { + let column = |record: VectorRecord<'_>| match record { + VectorRecord::Partition(column, _) + | VectorRecord::Pending(column, ..) + | VectorRecord::Block(column, ..) + | VectorRecord::Grouping(column, _) + | VectorRecord::Centres(column, ..) + | VectorRecord::Build(column, _) + | VectorRecord::Assigned(column, ..) + | VectorRecord::Sketch(column, ..) + | VectorRecord::Mean(column, _) + | VectorRecord::Left(column, ..) => column, + }; + let mut count = 0; + view.vectors(from, &mut |record, _| { + let same = std::mem::discriminant(&record) == std::mem::discriminant(&from) + && column(record) == column(from); + count += usize::from(same); + Ok(same) + }) + .expect("read the records"); + count +} + +#[test] +fn a_partition_that_empties_leaves_no_record_of_its_grouping_or_its_build() { + // A grouping in force; then a fifth more vectors sealed into its tail, a row that left + // its slot, and the regrouping those make due begun. + let (folder, db) = a_store_with_vectors("emptied-partition", 16); + insert_gathered(&db, 2001..=2400); + let changed = crate::params! { "v" => gathered(50_000) }; + db.execute("UPDATE docs SET v = $v WHERE id = 5", &changed) + .expect("change a vector"); + db.run_maintenance_cycle().expect("a maintenance cycle"); + // Every kind of record the column may keep, and how many of each it keeps at a moment. + let counts = |folder: &Path| -> Vec { + 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 column = catalog.table("docs").expect("the table").vectors[0].id; + let none = Key::from_values(&[], &[]).expect("the key of no values"); + let kinds = [ + VectorRecord::Partition(column, &none), + VectorRecord::Pending(column, 0, &none), + VectorRecord::Block(column, 0, 0), + VectorRecord::Grouping(column, 0), + VectorRecord::Centres(column, 0, 0), + VectorRecord::Build(column, 0), + VectorRecord::Assigned(column, 0, 0), + VectorRecord::Sketch(column, 0, 0), + VectorRecord::Mean(column, 0), + VectorRecord::Left(column, 0, &none), + ]; + kinds.map(|from| records_from(&view, from)).to_vec() + }; + db.close().expect("close"); + drop(db); + let before = counts(&folder); + for (kind, count) in [ + "partition", + "pending", + "block", + "grouping", + "centres", + "build", + ] + .iter() + .zip(&before) + { + // Nothing is pending after the cycle; of every other kind there is a record to lose. + assert_eq!( + *count == 0, + *kind == "pending", + "{kind} records before: {before:?}" + ); + } + assert!( + before[7] > 0 && before[8] == 1 && before[9] == 1, + "{before:?}" + ); + + let db = Database::open(&folder).expect("open"); + db.execute("DELETE FROM docs WHERE id > 0", &HashMap::new()) + .expect("delete every row"); + db.close().expect("close"); + drop(db); + assert_eq!( + counts(&folder), + [0; 10], + "records left by the emptied partition" + ); + std::fs::remove_dir_all(&folder).expect("remove the test's folder"); +} diff --git a/v2/contextdb/src/vector/train.rs b/v2/contextdb/src/vector/train.rs new file mode 100644 index 00000000..54934f7b --- /dev/null +++ b/v2/contextdb/src/vector/train.rs @@ -0,0 +1,161 @@ +//! Finding centres: the arithmetic of a build, over slots alone. Nothing here reads a store. +//! +//! The vector a slot keeps is its codes times its scale, a vector of length one within the +//! slot's rounding error. Centres are found from those, so a build reads blocks and no row. + +use std::collections::BTreeMap; + +use super::slot::{self, SLOT_HEADER}; +use crate::memory::Reservation; +use crate::types::Result; + +/// How many times a build goes over its sample. +pub(super) const ROUNDS: u32 = 10; + +/// How many slots a build samples for each centre it looks for. +pub(super) const SAMPLE_TO_A_CENTRE: u64 = 64; + +/// Bytes a group that [`first_centres`] takes while it draws: an entry of its list of moved +/// places and that entry's share of the list's nodes. +pub(super) const DRAWN: u64 = 48; + +/// How many groups a partition of this many slots is given: the square root, rounded, and +/// from 1 to 65,535. +pub(super) fn groups_for(slots: u64) -> u32 { + ((slots as f64).sqrt().round() as u32).clamp(1, 65_535) +} + +/// Which of `slots` slots, counted from zero in the order the build reads them, a round +/// samples: `min(slots, 64 * groups)` of them, evenly spread, ascending. +pub(super) fn sampled(slots: u64, groups: u32) -> impl Iterator { + let sample = slots.min(SAMPLE_TO_A_CENTRE * u64::from(groups)); + (0..sample).map(move |at| (u128::from(at) * u128::from(slots) / u128::from(sample)) as u64) +} + +/// Which of the sampled slots, counted from zero, the first centres are taken from: `groups` +/// different ones, ascending, the same for the same two numbers. +/// +/// The choice is the start of a shuffle of the sampled slots by one seeded sequence: for each +/// centre in turn, the slot at its place changes places with one at or after it. Only the +/// places that have changed are kept, never the whole sample. +pub(super) fn first_centres(sampled: u64, groups: u32) -> Vec { + let mut seeded: u64 = 0x5EED_C1A5_7E25 | 1; + let mut moved = BTreeMap::new(); + let mut first = Vec::with_capacity(groups as usize); + for place in 0..u64::from(groups) { + seeded ^= seeded << 13; + seeded ^= seeded >> 7; + seeded ^= seeded << 17; + let other = place + seeded % (sampled - place); + let here = moved.get(&place).copied().unwrap_or(place); + first.push(moved.insert(other, here).unwrap_or(other)); + } + first.sort_unstable(); + first +} + +/// The scale of a slot. +fn scale(slot: &[u8]) -> f32 { + f32::from_le_bytes([slot[0], slot[1], slot[2], slot[3]]) +} + +/// The product of a centre with the vector a slot keeps. +pub(super) fn product(centre: &[f32], slot: &[u8]) -> f32 { + scale(slot) * slot::product(&slot[SLOT_HEADER..], centre) +} + +/// The centre nearest the vector a slot keeps, by [`product`], and that product: the greatest, +/// and of equal ones the centre with the lowest number. +pub(super) fn nearest(centres: &[f32], dimension: usize, slot: &[u8]) -> (u32, f32) { + let mut best = (0, f32::NEG_INFINITY); + for (number, centre) in centres.chunks_exact(dimension).enumerate() { + let product = product(centre, slot); + if product > best.1 { + best = (number as u32, product); + } + } + best +} + +/// Adds the vector a slot keeps, its scale times each code, to `sum`. +fn add(sum: &mut [f64], slot: &[u8]) { + let scale = f64::from(scale(slot)); + for (sum, code) in sum.iter_mut().zip(&slot[SLOT_HEADER..]) { + *sum += scale * f64::from(code.cast_signed()); + } +} + +/// One round over the sample: the sum of the vectors that are nearest each centre. +pub(super) struct Round { + sums: Vec, + counts: Vec, + dimension: usize, +} + +impl Round { + /// A round for this many centres. `held` is grown by what the sums take: eight bytes a + /// number and four a centre. + pub(super) fn new(groups: u32, dimension: usize, held: &mut Reservation) -> Result { + let groups = groups as usize; + held.grow((8 * dimension * groups + 4 * groups) as u64)?; + Ok(Round { + sums: vec![0.0; dimension * groups], + counts: vec![0; groups], + dimension, + }) + } + + /// Adds the vector a slot keeps to the sum of the centre it is nearest. + pub(super) fn take(&mut self, centres: &[f32], slot: &[u8]) { + let (group, _) = nearest(centres, self.dimension, slot); + let at = group as usize * self.dimension; + add(&mut self.sums[at..at + self.dimension], slot); + self.counts[group as usize] += 1; + } + + /// Moves every centre to the direction of its sum, at length one. A centre that no + /// sampled slot was nearest, or whose sum has no direction, stays where it is. + pub(super) fn finish(self, centres: &mut [f32]) { + let sums = self.sums.chunks_exact(self.dimension); + let centres = centres.chunks_exact_mut(self.dimension); + for ((sum, centre), count) in sums.zip(centres).zip(self.counts) { + let length = sum.iter().map(|x| x * x).sum::().sqrt(); + if count > 0 && length > 0.0 { + for (number, sum) in centre.iter_mut().zip(sum) { + *number = (sum / length) as f32; + } + } + } + } +} + +/// The mean of the vectors the sampled slots keep, which a partition's first build takes in +/// its first round. +pub(super) struct Mean { + sum: Vec, + count: u64, +} + +impl Mean { + /// A mean of vectors of this dimension. `held` is grown by eight bytes a number. + pub(super) fn new(dimension: usize, held: &mut Reservation) -> Result { + held.grow(8 * dimension as u64)?; + Ok(Mean { + sum: vec![0.0; dimension], + count: 0, + }) + } + + /// Adds the vector a slot keeps. + pub(super) fn take(&mut self, slot: &[u8]) { + add(&mut self.sum, slot); + self.count += 1; + } + + /// The sum divided by the count, number by number; zeros when nothing was taken. The + /// caller has taken its four bytes a number before it calls. + pub(super) fn finish(self) -> Vec { + let count = self.count.max(1) as f64; + self.sum.iter().map(|sum| (sum / count) as f32).collect() + } +} diff --git a/v2/contextdb/src/write/mod.rs b/v2/contextdb/src/write/mod.rs index c1e12b74..7f0d947b 100644 --- a/v2/contextdb/src/write/mod.rs +++ b/v2/contextdb/src/write/mod.rs @@ -17,14 +17,14 @@ mod values; use std::collections::HashMap; use std::ops::Bound; -use crate::memory::Budget; +use crate::memory::{Budget, Reservation}; use crate::query; -use crate::schema::{Catalog, Table}; +use crate::schema::{Catalog, Quantization, Table}; use crate::sql::Plan; use crate::sql::plan::{DeletePlan, InsertPlan, UpdatePlan}; -use crate::store::{Cell, ChangeSet, IndexId, Key, Read, View}; +use crate::store::{Cell, ChangeSet, Compact, IndexId, Key, Read, View}; use crate::types::{ColumnDefault, Error, Result, SortDirection, Value}; -use crate::vector::Tally; +use crate::vector::{Reduced, 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. @@ -55,7 +55,7 @@ pub(crate) fn apply( now: i64, ) -> Result { match plan { - Plan::Insert(plan) => insert(plan, params, catalog, changes, partitions, now), + Plan::Insert(plan) => insert(plan, params, catalog, changes, partitions, budget, 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"), @@ -155,10 +155,12 @@ fn insert( catalog: &Catalog, changes: &mut ChangeSet, partitions: &mut Tally, + budget: &Budget, now: i64, ) -> Result { let table = find(catalog, &plan.table)?; let rules = Rules::of(table)?; + let mut held = budget.reserve(0, "INSERT")?; let columns = &table.meta.columns; let named: Vec = if plan.columns.is_empty() { (0..columns.len()).collect() @@ -194,7 +196,8 @@ fn insert( } let row = row.into_iter().enumerate(); let row = row.map(|(at, value)| rules.settle(at, value)); - rules.put(changes, partitions, None, row.collect::>()?)?; + let row = row.collect::>()?; + rules.put(changes, partitions, &mut held, None, row)?; } changes.add_count(table.id, plan.values.len() as i64)?; rules.settle_partitions(partitions)?; @@ -216,7 +219,7 @@ fn update( let set: Vec<_> = set .map(|(name, value)| Ok((column(table, name)?, value))) .collect::>()?; - let (found, _held) = query::matching( + let (found, mut held) = query::matching( table, plan.where_clause.as_ref(), params, @@ -233,7 +236,7 @@ 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, partitions, Some((key, old)), row)?; + rules.put(changes, partitions, &mut held, Some((key, old)), row)?; } } rules.settle_partitions(partitions)?; @@ -303,6 +306,8 @@ struct Rules<'a> { /// 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, + /// Where the VECTOR columns that keep a vector in one byte a number are in a row. + reduced: Vec, } /// One index of a table as a write uses it. @@ -338,10 +343,15 @@ impl<'a> Rules<'a> { }) }); let mut vectors = table.vectors.iter(); + let reduced = vectors + .clone() + .filter(|vector| vector.quantization == Quantization::Sq8) + .map(|vector| column(table, &vector.column)); Ok(Rules { table, key, indexes: indexes.collect::>()?, + reduced: reduced.collect::>()?, partitioned: vectors.any(|vector| !vector.partition_key.is_empty()), }) } @@ -395,13 +405,36 @@ 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. The row is noted in `partitions` first. + /// + /// Before anything else, a vector in a column that keeps one byte a number is replaced by + /// its reduced form, whose scale and codes go to the store with the row. `held` is grown + /// by a byte a number for the codes and four for the numbers until the row is staged. fn put( &self, changes: &mut ChangeSet, partitions: &mut Tally, + held: &mut Reservation, was: Option<(&Key, &[Value])>, - row: Vec, + mut row: Vec, ) -> Result<()> { + let mut compact = Vec::new(); + let mut grown = 0; + for &at in &self.reduced { + let Value::Vector(vector) = &row[at] else { + continue; + }; + let bytes = 5 * vector.len() as u64; + held.grow(bytes)?; + grown += bytes; + if let Some(reduced) = Reduced::of(vector) { + row[at] = Value::Vector(reduced.numbers()); + compact.push(Compact { + at, + scale: reduced.scale, + codes: reduced.codes, + }); + } + } 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), &[])?; @@ -434,7 +467,9 @@ impl<'a> Rules<'a> { { changes.delete_row(self.table.id, old_key)?; } - changes.put_row(self.table.id, key, row, new_key) + changes.put_row(self.table.id, key, row, compact, new_key)?; + held.shrink(grown); + Ok(()) } /// The columns of the primary key, as a refusal names them. diff --git a/v2/contextdb/tests/crash/main.rs b/v2/contextdb/tests/crash/main.rs index 37b6395e..0466f524 100644 --- a/v2/contextdb/tests/crash/main.rs +++ b/v2/contextdb/tests/crash/main.rs @@ -612,3 +612,133 @@ fn stopping_before_call(stop_at: u64) -> (WrapStorage, Arc) { }); (storage, calls) } + +// ---- a kill inside a build ---- + +/// A closed store whose `docs` has `rows` pending vectors in a column that is at its size from +/// a third of them on: the cycles seal them and build one grouping of `round(sqrt(rows))` +/// groups, 17 for 300 rows and 3 for 12. +fn store_to_build(name: &str, rows: i64) -> PathBuf { + let folder = support::folder(name); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, &format!("AUTO_INDEX_AT {}", rows / 3)); + support::load_docs(&db, 1..=rows, "a", |n| support::gathered(n, DIMENSION)); + support::close(db); + folder +} + +/// Asserts that an exact search returns every one of the `rows` rows once, and that INDEXED +/// answers with true distances or is refused by its own kind. +fn every_vector_is_found(db: &Database, rows: usize, state: &str) { + let stored = support::stored(db, support::DOCS); + assert_eq!(stored.rows().len(), rows, "{state}: the rows of docs"); + for seed in [98_101, 98_102] { + let query = support::gathered(seed, DIMENSION); + let wanted = stored.nearest(|_| true, &query, 1000); + let sql = support::by_mode(None, "EXACT", 1000); + support::assert_nearest(db, &sql, &query, &wanted, state); + match support::search(db, &support::by_mode(None, "INDEXED", 10), &query) { + Ok(result) => { + support::assert_sound(&result, &stored, |_| true, &query, 10, state); + } + Err(Error::VectorIndexedRouteUnavailable { .. }) => {} + Err(other) => panic!("{state}: INDEXED gave {other:?}"), + } + } +} + +fn groups(db: &Database) -> Value { + support::field(&support::the_partition(db), "lists").clone() +} + +#[test] +fn a_kill_at_any_storage_call_of_a_build_leaves_every_vector_found() { + if let Some(role) = support::role() { + return one_build(role); + } + let test = "a_kill_at_any_storage_call_of_a_build_leaves_every_vector_found"; + // Twelve rows: a build finds centres only for more than one group, which takes + // `round(sqrt(rows))` of 2 or more, and the INDEXED search asks for ten rows. The test's + // length is one store built for every storage call of the sixteen cycles. + const ROWS: i64 = 12; + // Run once to the end to learn which storage calls belong to the cycles of the build. + let folder = store_to_build("kill-build-count", ROWS); + let lines = Child::start(test, "builder stop_at=0", &folder).end(); + let begins = support::last_number(&lines, "building").expect("the builder announces it"); + let calls = support::last_number(&lines, "calls").expect("the builder announces its calls"); + assert_eq!(support::last_number(&lines, "cycles"), Some(16)); + assert_eq!(groups(&support::open(&folder)), Value::Int64(3)); + + for k in begins + 1..=calls { + let folder = store_to_build(&format!("kill-build-at-{k}"), ROWS); + let mut builder = Child::start(test, &format!("builder stop_at={k}"), &folder); + builder.until("stopped"); + builder.kill(); + // The store opens with nothing to mend and every vector is found, whichever step's + // commit the kill fell in; the next cycles finish the build. + let db = support::open(&folder); + let state = format!("killed at storage call {k} of a build"); + every_vector_is_found(&db, ROWS as usize, &state); + support::upkeep(&db); + assert_eq!(groups(&db), Value::Int64(3), "{state}"); + every_vector_is_found(&db, ROWS as usize, &format!("{state}, then built")); + } +} + +#[test] +fn a_killed_build_goes_on_from_its_record() { + if let Some(role) = support::role() { + return one_build(role); + } + let test = "a_killed_build_goes_on_from_its_record"; + let folder = store_to_build("killed-build-count", 300); + let lines = Child::start(test, "builder stop_at=0", &folder).end(); + let begins = support::last_number(&lines, "building").expect("the builder announces it"); + let calls = support::last_number(&lines, "calls").expect("the builder announces its calls"); + let whole = support::last_number(&lines, "cycles").expect("the builder counts its cycles"); + assert_eq!(whole, 16); + + // Killed at the middle storage call of the build's cycles. + let folder = store_to_build("killed-build-goes-on", 300); + let stop_at = begins + (calls - begins) / 2; + let mut builder = Child::start(test, &format!("builder stop_at={stop_at}"), &folder); + let told = builder.until("stopped"); + builder.kill(); + let returned = support::last_number(&told, "cycles").unwrap_or(0); + assert!( + 0 < returned && returned + 1 < whole, + "{returned} of {whole} cycles had returned at storage call {stop_at}" + ); + // The cycles after reopening are those the build had left: the killed cycle's step is + // there or it is not, and no step that was committed is worked out again. + let db = support::open(&folder); + let (ran, _) = support::upkeep(&db); + assert!( + ran == whole - returned || ran == whole - returned - 1, + "{returned} of {whole} cycles had returned; the build then took {ran} more" + ); + assert_eq!(groups(&db), Value::Int64(17)); + every_vector_is_found(&db, 300, "built after the kill"); +} + +/// Runs maintenance cycles until nothing is left. Announces `building ` with the storage +/// calls made so far before the first, `cycles ` as each returns, and `calls ` once +/// the store is closed. +fn one_build(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!("building {}", calls.load(Ordering::SeqCst))); + for cycles in 1_u64.. { + let report = support::cycle(&db); + support::say(format!("cycles {cycles}")); + if report.maintenance_vectors_remaining == 0 + && report.remaining_partitions <= report.stalled_partitions + { + break; + } + } + support::close(db); + support::say(format!("calls {}", calls.load(Ordering::SeqCst))); +} diff --git a/v2/contextdb/tests/memory_ceiling.rs b/v2/contextdb/tests/memory_ceiling.rs index 0ea84d8c..898eeca1 100644 --- a/v2/contextdb/tests/memory_ceiling.rs +++ b/v2/contextdb/tests/memory_ceiling.rs @@ -450,3 +450,91 @@ fn one_cycle_under_the_limit(role: &support::Role) { support::say(format!("before {before}")); support::say(format!("peak {}", ALLOCATED.peak.load(Ordering::Relaxed))); } + +// ---- a build and an indexed search under the limit ---- + +/// The limit of the build test: the cache's 8 MiB, the open's ninth, and room for the steps. +const BUILD_LIMIT: usize = 16 << 20; + +#[test] +fn an_indexed_search_and_a_build_stay_under_the_limit() { + if let Some(role) = support::role() { + return a_build_under_the_limit(&role); + } + let folder = support::folder("build-under-the-limit"); + let db = support::open(&folder); + write_the_filler(&db); + support::declare_docs(&db, 64, "AUTO_INDEX_AT 1000"); + support::load_docs(&db, 1..=2000, "a", |n| support::gathered(n, 64)); + support::close(db); + + let lines = support::Child::start( + "an_indexed_search_and_a_build_stay_under_the_limit", + "builder", + &folder, + ) + .end(); + let refused: Vec<&String> = lines + .iter() + .filter(|line| line.starts_with("refused ")) + .collect(); + assert!(refused.is_empty(), "under {BUILD_LIMIT} bytes: {refused:?}"); + let told = |key: &str| support::last_number(&lines, key); + assert_eq!( + (told("stalled"), told("lists"), told("rows")), + (Some(0), Some(45), Some(10)), + "the build's stalled partitions and groups, and the rows of the indexed search" + ); + // What the process took for the store, from before the open to the end, every step of + // the build and the search included, is within the limit. Each began with the cache full. + let before = told("before").expect("the builder announces it"); + let peak = told("peak").expect("the builder announces it"); + assert!( + peak <= before + BUILD_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 {BUILD_LIMIT}", + peak - before + ); +} + +/// Opens the store under [`BUILD_LIMIT`] and, with the cache filled before each, runs every +/// maintenance cycle and one indexed search. Announces `refused ` for a cycle or +/// a search that was refused, `stalled `, `lists `, `rows `, +/// and, once the store is closed, `before ` and `peak ` as the sealing test does. +fn a_build_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(BUILD_LIMIT); + let db = support::open_with(&role.folder, options).expect("open the store under the limit"); + for _ in 0..100 { + fill_the_cache(&db); + match db.run_maintenance_cycle() { + Ok(report) => { + let report = report.vector; + support::say(format!("stalled {}", report.stalled_partitions)); + if report.maintenance_vectors_remaining == 0 + && report.remaining_partitions <= report.stalled_partitions + { + break; + } + } + Err(error) => { + support::say(format!("refused {error:?}")); + break; + } + } + } + if let Value::Int64(lists) = support::field(&support::the_partition(&db), "lists") { + support::say(format!("lists {lists}")); + } + fill_the_cache(&db); + let query = support::gathered(93_000, 64); + match support::search(&db, &support::by_mode(None, "INDEXED", 10), &query) { + Ok(result) => support::say(format!("rows {}", result.rows().len())), + Err(error) => support::say(format!("refused {error:?}")), + } + support::close(db); + support::say(format!("before {before}")); + support::say(format!("peak {}", ALLOCATED.peak.load(Ordering::Relaxed))); +} diff --git a/v2/contextdb/tests/readers.rs b/v2/contextdb/tests/readers.rs index 2b043f4b..aac309d2 100644 --- a/v2/contextdb/tests/readers.rs +++ b/v2/contextdb/tests/readers.rs @@ -357,3 +357,125 @@ fn scanner_stopped_inside_its_scan(test: &str, folder: &Path, rows: i64) -> Chil scanner.expect("reading"); scanner } + +// ---- a reader beside a build ---- + +const GROUPED: usize = 16; + +/// The search of the two tests below, and the vector it asks for. +fn indexed_ten() -> String { + support::by_mode(None, "INDEXED", 10) +} + +fn asked() -> Vec { + support::gathered(97_000, GROUPED) +} + +/// A search's counted work and answer as one line: its distance evaluations, then each key +/// with the bits of its distance. Two searches of one commit give the same line. +fn told(result: &contextdb::QueryResult) -> String { + let rows: Vec = support::found(result, "the indexed search") + .iter() + .map(|near| format!("{:?}:{:016x}", near.key, near.distance.to_bits())) + .collect(); + format!("{} {}", result.distance_evaluations(), rows.join(",")) +} + +#[test] +fn a_reader_sees_one_grouping_whole_while_the_writer_builds_and_switches() { + if let Some(role) = support::role() { + return searcher(&role); + } + let test = "a_reader_sees_one_grouping_whole_while_the_writer_builds_and_switches"; + let (folder, db) = support::built("reader-beside-build", GROUPED, "AUTO_INDEX_AT 1000", 2000); + let line = |db: &contextdb::Database| told(&support::answer(db, &indexed_ten(), &asked())); + let first = line(&db); + + // One reader is stopped halfway through the storage reads of a search of the first + // grouping; where halfway is comes from a run of the same search to its end. + let lines = Child::start(test, "searcher pause_at=0", &folder).end(); + assert_eq!(lines.first(), Some(&first), "a search nothing stopped"); + let reads = support::last_number(&lines, "reads").expect("the reader announces its reads"); + assert!(reads > 1, "an indexed search made {reads} storage reads"); + let role = format!("searcher pause_at={}", reads.div_ceil(2)); + let mut held = Child::start(test, &role, &folder); + held.expect("reading"); + + // The writer adds a fifth, and builds and switches to the grouping that makes due. The + // switch removes every record of the grouping the reader is inside. + support::load_docs(&db, 2001..=2400, "a", |n| support::gathered(n, GROUPED)); + support::upkeep(&db); + assert_eq!( + support::field(&support::the_partition(&db), "lists"), + &Value::Int64(49) + ); + let second = line(&db); + assert!( + second != first, + "the new grouping changed nothing a search tells" + ); + + // The stopped reader answers from the grouping it began in, whole; its next search, and + // a reader that starts now, see the new one. + held.tell("go"); + assert_eq!( + held.line().as_ref(), + Some(&first), + "the search that was stopped" + ); + assert!(held.line().is_some_and(|line| line.starts_with("reads "))); + assert_eq!( + held.line().as_ref(), + Some(&second), + "the stopped reader's next search" + ); + held.end(); + let lines = Child::start(test, "searcher pause_at=0", &folder).end(); + assert_eq!( + lines.first(), + Some(&second), + "a reader that starts after the switch" + ); +} + +/// 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 `, searches once more and announces that line. +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"); + searching.store(true, Ordering::SeqCst); + let result = support::search(&reader, &indexed_ten(), &asked()) + .expect("a reader is not refused because a build ran"); + searching.store(false, Ordering::SeqCst); + support::say(told(&result)); + support::say(format!("reads {}", reads.load(Ordering::SeqCst))); + let next = support::search(&reader, &indexed_ten(), &asked()).expect("the next search"); + support::say(told(&next)); +} + +#[test] +fn a_new_process_searches_a_built_grouping_at_once() { + if let Some(role) = support::role() { + // Opens the store for writing and searches twice: each answer's line, then how far + // the two searches moved the committed position. + let db = support::open(&role.folder); + let before = db.committed_watermark(); + for _ in 0..2 { + support::say(told(&support::answer(&db, &indexed_ten(), &asked()))); + } + support::say(format!("moved {}", db.committed_watermark().0 - before.0)); + return; + } + let test = "a_new_process_searches_a_built_grouping_at_once"; + let (folder, db) = support::built("first-search", GROUPED, "AUTO_INDEX_AT 1000", 2000); + let wanted = told(&support::answer(&db, &indexed_ten(), &asked())); + support::close(db); + // The first search of a fresh process does the work of its second and of the process + // that built the grouping: nothing is built or mended at open. + let lines = Child::start(test, "first", &folder).end(); + assert_eq!(lines, [wanted.clone(), wanted, "moved 0".to_owned()]); +} diff --git a/v2/contextdb/tests/sql/vector_declarations.sql b/v2/contextdb/tests/sql/vector_declarations.sql index e7767bfc..83a945d5 100644 --- a/v2/contextdb/tests/sql/vector_declarations.sql +++ b/v2/contextdb/tests/sql/vector_declarations.sql @@ -67,14 +67,6 @@ 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 @@ -83,14 +75,10 @@ CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) IVF (LISTS = 64, PROBES = ---- 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 + v VECTOR(3) PARTITION_KEY (tenant) MAX_PARTITIONS 4 SEARCH_MODE AUTO IVF (LISTS = 64, PROBES = 8) ); ---- error: PlanError @@ -100,7 +88,7 @@ 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. +# as it was declared, with room for one partition and searched INDEXED, which answers. CREATE TABLE docs ( id INTEGER PRIMARY KEY, tenant TEXT NOT NULL, @@ -121,7 +109,8 @@ error: VectorPartitionLimitExceeded SELECT id FROM docs ORDER BY v <=> [1.0, 0.0, 0.0] LIMIT 1; ---- -error: VectorIndexedRouteUnavailable +id +1 SELECT id FROM docs ORDER BY v <=> [1.0, 0.0, 0.0] USE VECTOR EXACT LIMIT 1; ---- diff --git a/v2/contextdb/tests/sql/vector_index_small.sql b/v2/contextdb/tests/sql/vector_index_small.sql new file mode 100644 index 00000000..1871dad6 --- /dev/null +++ b/v2/contextdb/tests/sql/vector_index_small.sql @@ -0,0 +1,70 @@ +# INDEXED over rows that no maintenance cycle has grouped: it searches every one of them, so +# INDEXED, AUTO and EXACT give the same rows at the same distances. 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. + +CREATE TABLE passages (id INTEGER PRIMARY KEY, a TEXT NOT NULL, v VECTOR(2) PARTITION_KEY (a)); +INSERT INTO passages (id, a, v) +VALUES (6, 'y', [0.0, 2.0]), + (3, 'y', [1.0, 0.0]), + (1, 'x', [1.0, 0.0]), + (5, 'y', [1.0, 1.0]), + (2, 'x', [0.0, 1.0]), + (7, 'x', NULL), + (4, 'x', [-1.0, 0.0]); + +SELECT id, v <=> [1.0, 0.0] AS distance FROM passages ORDER BY distance USE VECTOR INDEXED LIMIT 5; +---- +id | distance +1 | 0.0 +3 | 0.0 +5 | * +2 | 1.0 +6 | 1.0 + +SELECT id, v <=> [1.0, 0.0] AS distance FROM passages ORDER BY distance USE VECTOR AUTO LIMIT 5; +---- +id | distance +1 | 0.0 +3 | 0.0 +5 | * +2 | 1.0 +6 | 1.0 + +SELECT id, v <=> [1.0, 0.0] AS distance FROM passages ORDER BY distance USE VECTOR EXACT LIMIT 5; +---- +id | distance +1 | 0.0 +3 | 0.0 +5 | * +2 | 1.0 +6 | 1.0 + +# One partition named, and a LIMIT above what it holds: every row of it that has a vector. +SELECT id, v <=> [1.0, 0.0] AS distance FROM passages WHERE a = 'x' ORDER BY distance USE VECTOR INDEXED LIMIT 5; +---- +id | distance +1 | 0.0 +2 | 1.0 +4 | 2.0 + +# Key values no row has select nothing and refuse nothing. +SELECT id FROM passages WHERE a = 'z' ORDER BY v <=> [1.0, 0.0] USE VECTOR INDEXED LIMIT 5; +---- +id + +# A column declared SEARCH_MODE INDEXED answers with no USE VECTOR, and as EXACT does. +CREATE TABLE declared (id INTEGER PRIMARY KEY, v VECTOR(2) SEARCH_MODE INDEXED); +INSERT INTO declared (id, v) VALUES (1, [0.0, 1.0]), (2, [1.0, 0.0]), (3, [-1.0, 0.0]); +SELECT id, v <=> [1.0, 0.0] AS distance FROM declared ORDER BY distance LIMIT 3; +---- +id | distance +2 | 0.0 +1 | 1.0 +3 | 2.0 + +SELECT id, v <=> [1.0, 0.0] AS distance FROM declared ORDER BY distance USE VECTOR EXACT LIMIT 3; +---- +id | distance +2 | 0.0 +1 | 1.0 +3 | 2.0 diff --git a/v2/contextdb/tests/sql/vector_inspection.sql b/v2/contextdb/tests/sql/vector_inspection.sql new file mode 100644 index 00000000..fe1622b9 --- /dev/null +++ b/v2/contextdb/tests/sql/vector_inspection.sql @@ -0,0 +1,63 @@ +# SHOW VECTOR_INDEXES gives one row for each VECTOR column, in the order of table and column; +# SHOW VECTOR_PARTITIONS one row for each partition, in the order of table, column and key +# values. No maintenance cycle runs here, so every vector is pending, no block or grouping holds +# a byte, nothing is due, and INDEXED may be given every partition. Three columns: one with no +# key and no vector, one with a key and three partitions (row 5 has no vector), and one that +# keeps its vectors reduced to a byte a number. + +CREATE TABLE a_plain (id INTEGER PRIMARY KEY, v VECTOR(3)); +CREATE TABLE b_parts (id INTEGER PRIMARY KEY, a TEXT NOT NULL, b INTEGER NOT NULL, v VECTOR(2) PARTITION_KEY (a, b)); +INSERT INTO b_parts (id, a, b, v) +VALUES (4, 'y', 1, [0.0, 2.0]), + (3, 'x', 2, [1.0, 0.0]), + (1, 'x', 1, [1.0, 0.0]), + (2, 'x', 1, [0.0, 1.0]), + (5, 'x', 1, NULL); +CREATE TABLE c_small (id INTEGER PRIMARY KEY, v VECTOR(4) WITH (quantization = 'SQ8')); +INSERT INTO c_small (id, v) VALUES (1, [1.0, 0.0, 0.0, 0.0]), (2, [0.0, 1.0, 0.0, 0.0]); + +SHOW VECTOR_INDEXES; +---- +table | column | dimension | quantization | vector_count | bytes | partition_key_columns | max_partitions | live_partitions | search_mode | declared_auto_index_at | effective_auto_index_at | declared_consolidation_mode | declared_consolidation_change_percent | declared_consolidation_tombstone_percent | effective_consolidation_mode | effective_consolidation_change_percent | effective_consolidation_tombstone_percent | pending_inserts | tombstones | durable_vector_bytes | durable_index_bytes | indexed_available | maintenance_state | maintenance_vectors_total | maintenance_vectors_done | maintenance_vectors_remaining | unavailable_partitions | stalled_partitions | broad_route | recovery_action +a_plain | v | 3 | F32 | 0 | 0 | [] | NULL | 0 | AUTO | NULL | 100000 | NULL | NULL | NULL | PERCENT | 20 | 10 | 0 | 0 | 0 | 0 | true | idle | 0 | 0 | 0 | 0 | 0 | fan_out | none +b_parts | v | 2 | F32 | 4 | 0 | ["a","b"] | 256 | 3 | AUTO | NULL | 100000 | NULL | NULL | NULL | PERCENT | 20 | 10 | 4 | 0 | 0 | 0 | true | idle | 0 | 0 | 0 | 0 | 0 | fan_out | none +c_small | v | 4 | SQ8 | 2 | 0 | [] | NULL | 1 | AUTO | NULL | 100000 | NULL | NULL | NULL | PERCENT | 20 | 10 | 2 | 0 | 0 | 0 | true | idle | 0 | 0 | 0 | 0 | 0 | fan_out | none + +# A column with no key has one partition, there before its first vector. +SHOW VECTOR_PARTITIONS; +---- +table | column | partition_key | live_rows | pending_inserts | tombstones | durable_vector_bytes | durable_index_bytes | indexed_available | availability_reason | maintenance_state | maintenance_reason | maintenance_vectors_total | maintenance_vectors_done | maintenance_vectors_remaining | lists | probes | recovery_action +a_plain | v | {} | 0 | 0 | 0 | 0 | 0 | true | none | idle | none | 0 | 0 | 0 | NULL | NULL | none +b_parts | v | {"a":"x","b":1} | 2 | 2 | 0 | 0 | 0 | true | none | idle | none | 0 | 0 | 0 | NULL | NULL | none +b_parts | v | {"a":"x","b":2} | 1 | 1 | 0 | 0 | 0 | true | none | idle | none | 0 | 0 | 0 | NULL | NULL | none +b_parts | v | {"a":"y","b":1} | 1 | 1 | 0 | 0 | 0 | true | none | idle | none | 0 | 0 | 0 | NULL | NULL | none +c_small | v | {} | 2 | 2 | 0 | 0 | 0 | true | none | idle | none | 0 | 0 | 0 | NULL | NULL | none + +SHOW VECTOR_PARTITIONS FOR b_parts.v; +---- +table | column | partition_key | live_rows | pending_inserts | tombstones | durable_vector_bytes | durable_index_bytes | indexed_available | availability_reason | maintenance_state | maintenance_reason | maintenance_vectors_total | maintenance_vectors_done | maintenance_vectors_remaining | lists | probes | recovery_action +b_parts | v | {"a":"x","b":1} | 2 | 2 | 0 | 0 | 0 | true | none | idle | none | 0 | 0 | 0 | NULL | NULL | none +b_parts | v | {"a":"x","b":2} | 1 | 1 | 0 | 0 | 0 | true | none | idle | none | 0 | 0 | 0 | NULL | NULL | none +b_parts | v | {"a":"y","b":1} | 1 | 1 | 0 | 0 | 0 | true | none | idle | none | 0 | 0 | 0 | NULL | NULL | none + +# LIMIT cuts the list with FOR or without; OFFSET, after LIMIT only, skips rows first. +SHOW VECTOR_PARTITIONS LIMIT 2; +---- +table | column | partition_key | live_rows | pending_inserts | tombstones | durable_vector_bytes | durable_index_bytes | indexed_available | availability_reason | maintenance_state | maintenance_reason | maintenance_vectors_total | maintenance_vectors_done | maintenance_vectors_remaining | lists | probes | recovery_action +a_plain | v | {} | 0 | 0 | 0 | 0 | 0 | true | none | idle | none | 0 | 0 | 0 | NULL | NULL | none +b_parts | v | {"a":"x","b":1} | 2 | 2 | 0 | 0 | 0 | true | none | idle | none | 0 | 0 | 0 | NULL | NULL | none + +SHOW VECTOR_PARTITIONS FOR b_parts.v LIMIT 2 OFFSET 1; +---- +table | column | partition_key | live_rows | pending_inserts | tombstones | durable_vector_bytes | durable_index_bytes | indexed_available | availability_reason | maintenance_state | maintenance_reason | maintenance_vectors_total | maintenance_vectors_done | maintenance_vectors_remaining | lists | probes | recovery_action +b_parts | v | {"a":"x","b":2} | 1 | 1 | 0 | 0 | 0 | true | none | idle | none | 0 | 0 | 0 | NULL | NULL | none +b_parts | v | {"a":"y","b":1} | 1 | 1 | 0 | 0 | 0 | true | none | idle | none | 0 | 0 | 0 | NULL | NULL | none + +# Refused by name: a table that is not there, and a column that is not a VECTOR column. +SHOW VECTOR_PARTITIONS FOR nope.v; +---- +error: TableNotFound + +SHOW VECTOR_PARTITIONS FOR b_parts.a; +---- +error: UnknownVectorIndex diff --git a/v2/contextdb/tests/sql/vector_policy.sql b/v2/contextdb/tests/sql/vector_policy.sql new file mode 100644 index 00000000..66157f1d --- /dev/null +++ b/v2/contextdb/tests/sql/vector_policy.sql @@ -0,0 +1,100 @@ +# The clauses of a VECTOR column that say how its vectors are kept and when they are grouped for +# search: WITH (quantization), AUTO_INDEX_AT and CONSOLIDATION. Each is declared alone and all +# together in the one order they are written in, after which SHOW VECTOR_INDEXES gives back what +# was declared and what is in force. A clause not written is NULL as declared and its default +# in force: 100000 vectors, and regrouping at 20 percent changed or 10 percent removed. + +CREATE TABLE a_reduced (id INTEGER PRIMARY KEY, v VECTOR(3) WITH (quantization = 'SQ8')); +CREATE TABLE b_sized (id INTEGER PRIMARY KEY, v VECTOR(3) AUTO_INDEX_AT 5000); +CREATE TABLE c_both (id INTEGER PRIMARY KEY, v VECTOR(3) CONSOLIDATION (CHANGE_PERCENT = 25, TOMBSTONE_PERCENT = 5)); +CREATE TABLE d_change (id INTEGER PRIMARY KEY, v VECTOR(3) CONSOLIDATION (CHANGE_PERCENT = 25)); +CREATE TABLE e_never (id INTEGER PRIMARY KEY, v VECTOR(3) CONSOLIDATION NONE); +CREATE TABLE f_whole ( + id INTEGER PRIMARY KEY, + tenant TEXT NOT NULL, + v VECTOR(3) WITH (quantization = 'SQ8') PARTITION_KEY (tenant) MAX_PARTITIONS 4 SEARCH_MODE EXACT + AUTO_INDEX_AT 5000 CONSOLIDATION (CHANGE_PERCENT = 25, TOMBSTONE_PERCENT = 5) +); + +SHOW VECTOR_INDEXES; +---- +table | column | dimension | quantization | vector_count | bytes | partition_key_columns | max_partitions | live_partitions | search_mode | declared_auto_index_at | effective_auto_index_at | declared_consolidation_mode | declared_consolidation_change_percent | declared_consolidation_tombstone_percent | effective_consolidation_mode | effective_consolidation_change_percent | effective_consolidation_tombstone_percent | pending_inserts | tombstones | durable_vector_bytes | durable_index_bytes | indexed_available | maintenance_state | maintenance_vectors_total | maintenance_vectors_done | maintenance_vectors_remaining | unavailable_partitions | stalled_partitions | broad_route | recovery_action +a_reduced | v | 3 | SQ8 | 0 | 0 | [] | NULL | 0 | AUTO | NULL | 100000 | NULL | NULL | NULL | PERCENT | 20 | 10 | 0 | 0 | 0 | 0 | true | idle | 0 | 0 | 0 | 0 | 0 | fan_out | none +b_sized | v | 3 | F32 | 0 | 0 | [] | NULL | 0 | AUTO | 5000 | 5000 | NULL | NULL | NULL | PERCENT | 20 | 10 | 0 | 0 | 0 | 0 | true | idle | 0 | 0 | 0 | 0 | 0 | fan_out | none +c_both | v | 3 | F32 | 0 | 0 | [] | NULL | 0 | AUTO | NULL | 100000 | PERCENT | 25 | 5 | PERCENT | 25 | 5 | 0 | 0 | 0 | 0 | true | idle | 0 | 0 | 0 | 0 | 0 | fan_out | none +d_change | v | 3 | F32 | 0 | 0 | [] | NULL | 0 | AUTO | NULL | 100000 | PERCENT | 25 | NULL | PERCENT | 25 | 10 | 0 | 0 | 0 | 0 | true | idle | 0 | 0 | 0 | 0 | 0 | fan_out | none +e_never | v | 3 | F32 | 0 | 0 | [] | NULL | 0 | AUTO | NULL | 100000 | NONE | NULL | NULL | NONE | NULL | NULL | 0 | 0 | 0 | 0 | true | idle | 0 | 0 | 0 | 0 | 0 | fan_out | none +f_whole | v | 3 | SQ8 | 0 | 0 | ["tenant"] | 4 | 0 | EXACT | 5000 | 5000 | PERCENT | 25 | 5 | PERCENT | 25 | 5 | 0 | 0 | 0 | 0 | true | idle | 0 | 0 | 0 | 0 | 0 | fan_out | none + +# What a declaration refuses, each by the name of its kind. A way of keeping a vector that is +# neither F32 nor SQ8, +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) WITH (quantization = 'PQ')); +---- +error: InvalidVectorPolicyDeclaration + +# an AUTO_INDEX_AT that is not a positive number, or is past the largest the engine keeps, +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) AUTO_INDEX_AT 0); +---- +error: InvalidVectorPolicyDeclaration + +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) AUTO_INDEX_AT -1); +---- +error: InvalidVectorPolicyDeclaration + +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) AUTO_INDEX_AT 99999999999999999999); +---- +error: InvalidVectorPolicyDeclaration + +# and a percentage that is not a whole number from 1 to 100. +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) CONSOLIDATION (CHANGE_PERCENT = 0)); +---- +error: InvalidVectorPolicyDeclaration + +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) CONSOLIDATION (TOMBSTONE_PERCENT = 101)); +---- +error: InvalidVectorPolicyDeclaration + +# SQ4, IVF (...) and SEARCH (...) are read by a later version: refused by name, not as a syntax +# error. +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) WITH (quantization = 'SQ4')); +---- +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) SEARCH (TARGET_ACCURACY = 95, EXPECTED_VECTORS = 1000000)); +---- +error: PlanError + +# A syntax error: a clause out of its order, and a CONSOLIDATION list that is empty, names a +# member twice or names another member. +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) AUTO_INDEX_AT 5000 SEARCH_MODE EXACT); +---- +error: ParseError + +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) CONSOLIDATION NONE AUTO_INDEX_AT 5000); +---- +error: ParseError + +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) AUTO_INDEX_AT 5000 WITH (quantization = 'SQ8')); +---- +error: ParseError + +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) CONSOLIDATION ()); +---- +error: ParseError + +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) CONSOLIDATION (CHANGE_PERCENT = 25, CHANGE_PERCENT = 30)); +---- +error: ParseError + +CREATE TABLE docs (id INTEGER PRIMARY KEY, v VECTOR(3) CONSOLIDATION (LISTS = 4)); +---- +error: ParseError + +# A refused declaration declares nothing. +SELECT id FROM docs; +---- +error: TableNotFound diff --git a/v2/contextdb/tests/sql/vector_search_refusals.sql b/v2/contextdb/tests/sql/vector_search_refusals.sql index add273dd..f1497f95 100644 --- a/v2/contextdb/tests/sql/vector_search_refusals.sql +++ b/v2/contextdb/tests/sql/vector_search_refusals.sql @@ -16,16 +16,6 @@ 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; ---- @@ -44,11 +34,13 @@ 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. +# INDEXED is not refused for rows that no maintenance cycle has grouped yet: asked for in the +# statement or declared on the column, it searches every one of them, as EXACT and AUTO do. SELECT id FROM docs ORDER BY v <=> [1.0, 0.0, 0.0] USE VECTOR INDEXED LIMIT 2; ---- -error: VectorIndexedRouteUnavailable +id +1 +2 SELECT id FROM docs ORDER BY v <=> [1.0, 0.0, 0.0] USE VECTOR EXACT LIMIT 2; ---- @@ -60,7 +52,9 @@ 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 +id +2 +1 SELECT id FROM indexed ORDER BY v <=> [1.0, 0.0, 0.0] USE VECTOR EXACT LIMIT 2; ---- diff --git a/v2/contextdb/tests/store.rs b/v2/contextdb/tests/store.rs index d8d02c8b..a8af7d68 100644 --- a/v2/contextdb/tests/store.rs +++ b/v2/contextdb/tests/store.rs @@ -204,3 +204,29 @@ fn a_store_closed_here_opens_again_while_a_child_process_holds_a_copy_of_its_loc heir.end(); support::assert_rows_are(&support::open(&folder), 1..=4); } + +#[test] +fn a_grouping_is_there_after_close_and_reopen() { + let (folder, db) = support::built("grouping-after-reopen", 16, "AUTO_INDEX_AT 1000", 2000); + let query = support::gathered(99_000, 16); + let sql = support::by_mode(None, "INDEXED", 10); + // The partitions as inspection lists them, and an indexed search's rows and work. + let state = |db: &contextdb::Database| { + let result = support::answer(db, &sql, &query); + ( + support::shown(db, "SHOW VECTOR_PARTITIONS"), + result.rows().to_vec(), + result.distance_evaluations(), + db.committed_watermark(), + ) + }; + let before = state(&db); + assert_eq!(support::field(&before.0[0], "lists"), &Value::Int64(45)); + support::close(db); + + let db = support::open(&folder); + assert!(state(&db) == before, "the grouping after reopen"); + // Nothing is left to build: a cycle makes no commit. + assert_eq!(support::upkeep(&db), (1, 0)); + assert!(state(&db) == before, "after one more cycle"); +} diff --git a/v2/contextdb/tests/support/mod.rs b/v2/contextdb/tests/support/mod.rs index febad0e8..95b29211 100644 --- a/v2/contextdb/tests/support/mod.rs +++ b/v2/contextdb/tests/support/mod.rs @@ -8,7 +8,7 @@ // Each test program uses a part of this module; the rest would be reported as never used. #![allow(dead_code)] -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::fmt; use std::fs::{self, File}; use std::io::{self, BufRead, BufReader, Write}; @@ -702,25 +702,246 @@ 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. +/// Runs maintenance cycles until nothing is left for one to do, and returns how many vectors +/// the cycles sealed and packed between them. Nothing else in the tests seals or groups. 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; + upkeep(db).1 +} + +/// One maintenance cycle, which must not be refused. +pub fn cycle(db: &Database) -> contextdb::VectorMaintenanceReport { + match db.run_maintenance_cycle() { + Ok(report) => report.vector, + Err(error) => panic!("a maintenance cycle was refused: {error:?}"), + } +} + +/// Runs maintenance cycles as a host does: while a vector is left, or a partition that is +/// not stalled. Returns how many cycles it ran, the last included, and the vectors they +/// reported done. A store that never comes to rest fails the test. +pub fn upkeep(db: &Database) -> (u64, u64) { + let mut done = 0; + for ran in 1..=100_000_u64 { + let report = cycle(db); + done += report.maintenance_vectors_done; + if report.maintenance_vectors_remaining == 0 + && report.remaining_partitions <= report.stalled_partitions + { + return (ran, done); } + } + panic!("100,000 maintenance cycles left work to do"); +} + +// ---- the index ---- + +/// The vector of row `n` in a scatter around eight directions at right angles: direction +/// `n % 8` has its numbers `2d` and `2d + 1` at one, and every number is moved by a twentieth +/// of what [`vector`] gives. Rows of one direction lie close together and far from the rest. +/// The dimension is 16 or more. +pub fn gathered(n: i64, dimension: usize) -> Vec { + let direction = n.rem_euclid(8) as usize; + let moved = vector(n, dimension); + (0..dimension) + .map(|at| f32::from(u8::from(at / 2 == direction)) + 0.05 * moved[at]) + .collect() +} + +/// The vector of row `n` of `of` spread evenly round one circle that lies across every number +/// of the dimension, each number moved by a hundredth of what [`vector`] gives: no direction +/// gathers rows, and the rows nearest any of them are its neighbours on the circle. +pub fn ringed(n: i64, of: i64, dimension: usize) -> Vec { + let turn = std::f64::consts::TAU * n as f64 / of as f64; + let moved = vector(n, dimension); + (0..dimension) + .map(|at| { + let phase = std::f64::consts::TAU * at as f64 / dimension as f64; + (turn - phase).cos() as f32 + 0.01 * moved[at] + }) + .collect() +} + +/// Inserts these rows of `docs` under one context in one commit, row `n` with `vector(n)`. +pub fn load_docs( + db: &Database, + numbers: RangeInclusive, + ctx: &str, + vector: impl Fn(i64) -> Vec, +) { + let tx = db.transaction().expect("begin"); + for n in numbers { + if let Err(error) = insert_doc(&tx, n, ctx, Some(vector(n))) { + panic!("inserting row {n} of docs was refused: {error:?}"); + } + } + tx.commit().expect("commit the rows of docs"); +} + +/// A store whose `docs` has these clauses after `VECTOR(dimension)` and rows 1 to `rows` of +/// [`gathered`] under the context `a`, with every maintenance cycle run. +pub fn built(name: &str, dimension: usize, clauses: &str, rows: i64) -> (PathBuf, Database) { + let folder = folder(name); + let db = open(&folder); + declare_docs(&db, dimension, clauses); + load_docs(&db, 1..=rows, "a", |n| gathered(n, dimension)); + upkeep(&db); + (folder, db) +} + +/// A search of `docs` for the `k` rows nearest `$q` under this mode: `EXACT`, `AUTO`, +/// `INDEXED`, or nothing for the column's own. +pub fn by_mode(filter: Option<&str>, mode: &str, k: usize) -> String { + let filter = filter.map_or_else(String::new, |filter| format!(" WHERE {filter}")); + let clause = if mode.is_empty() { + String::new() + } else { + format!("USE VECTOR {mode} ") + }; + format!("SELECT id, v <=> $q AS distance FROM docs{filter} ORDER BY distance {clause}LIMIT {k}") +} + +/// A search that must answer. +pub fn answer(on: &impl Run, sql: &str, query: &[f32]) -> QueryResult { + match search(on, sql, query) { + Ok(result) => result, + Err(error) => panic!("{sql}\nwas refused: {error:?}"), + } +} + +/// Asserts that an answer is one a search that may leave rows out can give: `count` rows, +/// nearest first and none twice, each a row `passes` lets through, with the distance of the +/// vector stored for it, to the bit. Returns the answer. +pub fn assert_sound( + result: &QueryResult, + stored: &Stored, + passes: impl Fn(&[Value]) -> bool, + query: &[f32], + count: usize, + what: &str, +) -> Vec { + let answer = found(result, what); + assert_eq!(answer.len(), count, "{what}: rows in the answer"); + let rows: BTreeMap> = stored + .rows() + .iter() + .map(|row| match row[0] { + Value::Int64(id) => (id, row), + ref other => panic!("the key {other:?} is no INTEGER"), + }) + .collect(); + for near in &answer { + let Value::Int64(id) = near.key else { + panic!("{what}: the key {:?}", near.key) + }; + let row = rows + .get(&id) + .unwrap_or_else(|| panic!("{what}: row {id} is not stored")); + assert!(passes(row), "{what}: row {id} does not pass the filter"); + let Some(Value::Vector(vector)) = row.last() else { + panic!("{what}: row {id} has no vector") + }; + assert_eq!( + near.distance.to_bits(), + distance(vector, query).to_bits(), + "{what}: the distance of row {id}" + ); + } + for pair in answer.windows(2) { + let order = pair[0].distance.total_cmp(&pair[1].distance); assert!( - report.maintenance_vectors_done > 0, - "a cycle sealed nothing and left {} vectors waiting", - report.maintenance_vectors_remaining + order + .then_with(|| key_order(&pair[0].key, &pair[1].key)) + .is_lt(), + "{what}: {:?} came before {:?}", + pair[0], + pair[1] ); } + answer +} + +/// How the refusal of an INDEXED search goes on after the column's name. +pub const INDEXED_REFUSED: &str = " cannot be searched INDEXED for this statement now: its index \ + is behind its writes, or the statement's filter leaves the index nothing to choose among; \ + run a maintenance cycle, or use EXACT or AUTO"; + +/// Asserts that a search of `docs.v` was refused as one that asked for the index, in the +/// sentence that names both causes. +pub fn assert_indexed_refused(outcome: contextdb::Result, what: &str) { + match outcome { + Err(error @ Error::VectorIndexedRouteUnavailable { .. }) => { + assert_eq!( + error.to_string(), + format!("docs.v{INDEXED_REFUSED}"), + "{what}" + ); + } + other => panic!("{what}: asked for the index and got {other:?}"), + } +} + +/// One row of an inspection statement: each column's name with its value. +pub type Shown = BTreeMap; + +/// The rows of `SHOW VECTOR_INDEXES` or `SHOW VECTOR_PARTITIONS`. +pub fn shown(on: &impl Run, sql: &str) -> Vec { + let result = run(on, sql); + let names = result.columns(); + let named = |row: &Vec| names.iter().cloned().zip(row.iter().cloned()).collect(); + result.rows().iter().map(named).collect() +} + +/// One field of an inspection row. +pub fn field<'a>(row: &'a Shown, name: &str) -> &'a Value { + row.get(name) + .unwrap_or_else(|| panic!("inspection has no field {name}: {row:?}")) +} + +/// A field that is a whole number. +pub fn whole(row: &Shown, name: &str) -> i64 { + match field(row, name) { + Value::Int64(number) => *number, + other => panic!("the field {name} is {other:?}, not a whole number"), + } +} + +/// A field that is text. +pub fn word(row: &Shown, name: &str) -> String { + match field(row, name) { + Value::Text(text) => text.clone(), + other => panic!("the field {name} is {other:?}, not text"), + } +} + +/// A field that holds JSON, whether the engine hands it over as JSON or as its text. +pub fn json(row: &Shown, name: &str) -> contextdb::serde_json::Value { + match field(row, name) { + Value::Json(value) => value.clone(), + Value::Text(text) => contextdb::serde_json::from_str(text) + .unwrap_or_else(|error| panic!("the field {name} is no JSON: {error}")), + other => panic!("the field {name} is {other:?}, not JSON"), + } +} + +/// The one partition of a store whose `docs.v` has one. +pub fn the_partition(on: &impl Run) -> Shown { + let mut rows = shown(on, "SHOW VECTOR_PARTITIONS FOR docs.v"); + assert_eq!(rows.len(), 1, "partitions of docs.v"); + rows.remove(0) +} + +/// A storage door that adds up the bytes read from the data file. +pub fn counting_bytes_read() -> (WrapStorage, Arc) { + let bytes = Arc::new(AtomicU64::new(0)); + let door = storage({ + let bytes = bytes.clone(); + move |call| { + if let Call::Read(_, read) = call { + bytes.fetch_add(read.len() as u64, Ordering::SeqCst); + } + } + }); + (door, bytes) } // ---- the storage door ---- diff --git a/v2/contextdb/tests/transactions.rs b/v2/contextdb/tests/transactions.rs index 96e396b3..be422b4c 100644 --- a/v2/contextdb/tests/transactions.rs +++ b/v2/contextdb/tests/transactions.rs @@ -466,3 +466,73 @@ fn a_transaction_that_declares_a_table_may_not_also_write_rows() { Error::TableNotFound(_) )); } + +#[test] +fn a_statement_commits_while_a_cycle_works_out_a_step() { + use std::sync::Mutex; + use std::sync::atomic::AtomicBool; + use std::sync::mpsc; + use std::time::Duration; + + // 2,000 sealed vectors of 256 numbers in three blocks, and their build begun. + const DIMENSION: usize = 256; + let folder = support::folder("commit-beside-a-step"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, "AUTO_INDEX_AT 1000"); + support::load_docs(&db, 1..=2000, "a", |n| support::gathered(n, DIMENSION)); + let report = support::cycle(&db); + assert_eq!( + (report.maintenance_vectors_done, report.remaining_partitions), + (2000, 1) + ); + support::close(db); + + // The storage door holds the first read of 64 KiB or more made once it is armed: a block + // read by the step the cycle works out. It says so, and waits to be let go. + let armed = Arc::new(AtomicBool::new(false)); + let (holding, held) = mpsc::channel::<()>(); + let (let_go, go) = mpsc::channel::<()>(); + let go = Mutex::new(go); + let mut options = DatabaseOpenOptions::default(); + options.storage = Some(support::storage({ + let armed = armed.clone(); + move |call| { + let large = matches!(&call, Call::Read(_, bytes) if bytes.len() >= 64 << 10); + if large && armed.swap(false, Ordering::SeqCst) { + holding.send(()).expect("the test listens"); + let go = go.lock().expect("one read is held"); + go.recv().expect("the test lets the read go"); + } + } + })); + let db = support::open_with(&folder, options).expect("open"); + let before = db.committed_watermark(); + armed.store(true, Ordering::SeqCst); + + let db = &db; + let (inserted, in_time, cycled) = std::thread::scope(|scope| { + let cycle = scope.spawn(|| db.run_maintenance_cycle()); + held.recv_timeout(Duration::from_secs(20)) + .expect("a step of the build reads a block"); + // The cycle is inside its step. An INSERT on another thread returns while it is. + let (done, returned) = mpsc::channel::<()>(); + let insert = scope.spawn(move || { + let row = support::gathered(9001, DIMENSION); + let outcome = support::insert_doc(db, 9001, "a", Some(row)); + done.send(()).expect("the test listens"); + outcome.map(|_| db.committed_watermark()) + }); + let in_time = returned.recv_timeout(Duration::from_secs(20)).is_ok(); + let_go.send(()).expect("the held read listens"); + let inserted = insert.join().expect("the thread of the INSERT"); + (inserted, in_time, cycle.join().expect("the cycle's thread")) + }); + assert!( + in_time, + "the INSERT waited for the step the cycle was working out" + ); + let after = inserted.expect("an INSERT beside a cycle"); + assert!(after > before, "the INSERT committed"); + cycled.expect("the cycle whose read was held"); + assert_eq!(support::int(db, "SELECT COUNT(*) FROM docs"), 2001); +} diff --git a/v2/contextdb/tests/vectors_index.rs b/v2/contextdb/tests/vectors_index.rs new file mode 100644 index 00000000..f8b85f48 --- /dev/null +++ b/v2/contextdb/tests/vectors_index.rs @@ -0,0 +1,828 @@ +//! The index of a VECTOR column: a partition of a column that is large enough is cut into +//! groups, and a search under INDEXED, or under AUTO from `AUTO_INDEX_AT` vectors on, looks at +//! the nearest groups first. Every row it returns has its true distance; a column below its +//! size is searched whole and pays nothing; and INDEXED is refused by name where the index +//! cannot serve the statement. + +mod support; + +use std::collections::BTreeSet; +use std::io; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use contextdb::{ + DatabaseOpenOptions, Error, QueryResult, ReadSession, ReadSessionOptions, Value, params, +}; + +const DIMENSION: usize = 16; + +/// The clause under which a partition of 2,000 rows is cut into groups. +const GROUPED: &str = "AUTO_INDEX_AT 1000"; + +/// Question `j`: a vector near one of the eight directions the rows are scattered around. +fn question(j: i64) -> Vec { + support::gathered(50_000 + j, DIMENSION) +} + +fn all(_: &[Value]) -> bool { + true +} + +fn gathered(n: i64) -> Vec { + support::gathered(n, DIMENSION) +} + +fn lists(partition: &support::Shown) -> Value { + support::field(partition, "lists").clone() +} + +#[test] +fn an_indexed_search_does_fewer_distance_evaluations_than_exact() { + let (_, db) = support::built("indexed-fewer-evaluations", DIMENSION, GROUPED, 2000); + assert_eq!(lists(&support::the_partition(&db)), Value::Int64(45)); + let stored = support::stored(&db, support::DOCS); + for j in 0..20 { + let query = question(j); + let wanted = stored.nearest(all, &query, 10); + let sql = support::by_mode(None, "EXACT", 10); + let exact = support::assert_nearest(&db, &sql, &query, &wanted, "EXACT"); + let indexed = support::answer(&db, &support::by_mode(None, "INDEXED", 10), &query); + let what = format!("INDEXED, question {j}"); + support::assert_sound(&indexed, &stored, all, &query, 10, &what); + assert!( + indexed.distance_evaluations() < exact.distance_evaluations(), + "{what}: {} distance evaluations, EXACT {}", + indexed.distance_evaluations(), + exact.distance_evaluations() + ); + } +} + +#[test] +fn an_indexed_search_reads_sketches_before_blocks() { + // 256 rows spread evenly round a circle, in vectors wide enough for five blocks of 63 + // slots. The size is fixed by the blocks: INDEXED reads every centre and the sketches + // first, so it reads less than EXACT only when the blocks it then leaves unread outweigh + // them, which takes several blocks; few rows in wide vectors make those with a short build. + const WIDE: usize = 4096; + const ROWS: i64 = 256; + let folder = support::folder("indexed-reads-sketches"); + let db = support::open(&folder); + support::declare_docs(&db, WIDE, "AUTO_INDEX_AT 256"); + support::load_docs(&db, 1..=ROWS, "a", |n| support::ringed(n, ROWS, WIDE)); + support::upkeep(&db); + assert_eq!(lists(&support::the_partition(&db)), Value::Int64(16)); + support::close(db); + + // Each search runs first in a session of its own, so every page it needs is read once. + let query = support::ringed(157, 2 * ROWS, WIDE); + let read_by = |mode: &str| -> u64 { + let (door, bytes) = support::counting_bytes_read(); + let mut options = ReadSessionOptions::default(); + options.storage = Some(door); + let reader = ReadSession::open_with_options(&folder, options).expect("open for reading"); + let opened = bytes.load(Ordering::SeqCst); + let result = support::answer(&reader, &support::by_mode(None, mode, 10), &query); + assert_eq!(result.rows().len(), 10, "rows under {mode}"); + bytes.load(Ordering::SeqCst) - opened + }; + // EXACT reads every block. INDEXED reads what EXACT does not, the centres and the + // sketches, and still reads less: it left blocks unread. + let (exact, indexed) = (read_by("EXACT"), read_by("INDEXED")); + assert!( + indexed < exact, + "INDEXED read {indexed} bytes of the data file and EXACT {exact}" + ); +} + +#[test] +fn a_column_below_its_size_pays_nothing_and_is_searched_whole() { + let (_, db) = support::built("below-its-size", DIMENSION, "", 2000); + // The three modes return the reference answer, and INDEXED does EXACT's work. + let agree = |k: usize, rows: usize, state: &str| { + let stored = support::stored(&db, support::DOCS); + for j in 0..3 { + let query = question(j); + let wanted = stored.nearest(all, &query, k); + assert_eq!(wanted.len(), rows, "{state}"); + let work = ["EXACT", "AUTO", "INDEXED"].map(|mode| { + let sql = support::by_mode(None, mode, k); + support::assert_nearest(&db, &sql, &query, &wanted, state).distance_evaluations() + }); + assert_eq!( + work[2], work[0], + "{state}: INDEXED's evaluations and EXACT's" + ); + } + }; + let index_bytes = || { + let index = support::shown(&db, "SHOW VECTOR_INDEXES").remove(0); + support::whole(&index, "durable_index_bytes") + }; + assert_eq!(lists(&support::the_partition(&db)), Value::Null); + assert_eq!(index_bytes(), 0); + agree(10, 10, "2,000 rows sealed"); + + // A fifth of the rows leave their slots: the one build such a column gets drops them. + assert_eq!( + support::run(&db, "DELETE FROM docs WHERE tag < 10").rows_affected(), + 400 + ); + support::upkeep(&db); + let partition = support::the_partition(&db); + assert_eq!(support::whole(&partition, "tombstones"), 0); + assert_eq!(lists(&partition), Value::Int64(1)); + assert_eq!(index_bytes(), 28, "the record of a grouping of one group"); + agree(10, 10, "a fifth deleted"); + agree(5000, 1600, "a LIMIT above the rows that are left"); +} + +#[test] +fn indexed_is_served_at_every_size() { + let folder = support::folder("indexed-every-size"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, GROUPED); + let query = question(0); + let sql = support::by_mode(None, "INDEXED", 10); + let mut next = 1; + for upto in [1, 10, 1999] { + support::load_docs(&db, next..=upto, "a", gathered); + next = upto + 1; + // With no grouping the answer is the exact one. + let wanted = support::stored(&db, support::DOCS).nearest(all, &query, 10); + let state = format!("{upto} rows and no cycle"); + support::assert_nearest(&db, &sql, &query, &wanted, &state); + } + support::upkeep(&db); + assert_eq!(lists(&support::the_partition(&db)), Value::Int64(45)); + let stored = support::stored(&db, support::DOCS); + let result = support::answer(&db, &sql, &query); + support::assert_sound(&result, &stored, all, &query, 10, "1,999 rows grouped"); +} + +/// One partition built and one given 2,000 rows since the last cycle, which is twice the +/// 1,000 a column of these clauses may have outside a grouping. `indexed` is how the search +/// asks for the index: `INDEXED`, or nothing on a column declared so. +fn refused_whole_while_one_partition_is_behind(name: &str, clauses: &str, indexed: &str) { + let (_, db) = support::built(name, DIMENSION, clauses, 2000); + support::load_docs(&db, 2001..=4000, "b", gathered); + let query = question(1); + let both = support::by_mode(None, indexed, 10); + let before = db.committed_watermark(); + support::assert_indexed_refused(support::search(&db, &both, &query), name); + assert_eq!(db.committed_watermark(), before, "{name}: a refused search"); + + let stored = support::stored(&db, support::DOCS); + let in_a = |row: &[Value]| row[1] == support::text("a"); + let alone = support::by_mode(Some("ctx = 'a'"), indexed, 10); + let result = support::answer(&db, &alone, &query); + support::assert_sound( + &result, + &stored, + in_a, + &query, + 10, + "the built partition alone", + ); + let wanted = stored.nearest(all, &query, 10); + let sql = support::by_mode(None, "EXACT", 10); + support::assert_nearest(&db, &sql, &query, &wanted, "EXACT"); + // AUTO searches the built partition as INDEXED does and the other exactly: ten rows with + // their true distances, nearest first, and of the partition with no grouping every row + // nearer than the last of them. + let result = support::answer(&db, &support::by_mode(None, "AUTO", 10), &query); + let answer = support::assert_sound(&result, &stored, all, &query, 10, "AUTO over both"); + let in_b = |row: &[Value]| row[1] == support::text("b"); + for near in stored.nearest(in_b, &query, 10) { + assert!( + near.distance >= answer[9].distance || answer.contains(&near), + "{name}: AUTO left out {near:?} of the partition it searches exactly; it answered \ + {answer:?}" + ); + } + + support::upkeep(&db); + let result = support::answer(&db, &both, &query); + support::assert_sound(&result, &stored, all, &query, 10, "after the cycles"); +} + +#[test] +fn indexed_is_refused_whole_when_one_partition_is_behind() { + refused_whole_while_one_partition_is_behind( + "behind-use-vector", + "PARTITION_KEY (ctx) AUTO_INDEX_AT 1000", + "INDEXED", + ); + refused_whole_while_one_partition_is_behind( + "behind-declared", + "PARTITION_KEY (ctx) SEARCH_MODE INDEXED AUTO_INDEX_AT 1000", + "", + ); +} + +/// The answers and the counted work of one statement under AUTO and under EXACT. +fn auto_and_exact( + db: &contextdb::Database, + filter: Option<&str>, + query: &[f32], +) -> [QueryResult; 2] { + ["AUTO", "EXACT"].map(|mode| support::answer(db, &support::by_mode(filter, mode, 10), query)) +} + +#[test] +fn auto_uses_groupings_from_auto_index_at_on() { + // Two built partitions of 1,200 under AUTO_INDEX_AT 2000: the count is over the selected + // partitions together. + let clauses = "PARTITION_KEY (ctx) AUTO_INDEX_AT 2000"; + let (_, db) = support::built("auto-threshold", DIMENSION, clauses, 1200); + support::load_docs(&db, 1201..=2400, "b", gathered); + support::upkeep(&db); + let grouped: Vec = support::shown(&db, "SHOW VECTOR_PARTITIONS") + .iter() + .map(lists) + .collect(); + assert_eq!(grouped, [Value::Int64(35), Value::Int64(35)]); + let stored = support::stored(&db, support::DOCS); + let query = question(2); + let [auto, exact] = auto_and_exact(&db, Some("ctx = 'a'"), &query); + assert_eq!( + support::found(&auto, "AUTO"), + support::found(&exact, "EXACT") + ); + assert_eq!(auto.distance_evaluations(), exact.distance_evaluations()); + let [auto, exact] = auto_and_exact(&db, None, &query); + support::assert_sound(&auto, &stored, all, &query, 10, "AUTO over 2,400"); + assert!( + auto.distance_evaluations() < exact.distance_evaluations(), + "AUTO over 2,400 vectors did {} distance evaluations, EXACT {}", + auto.distance_evaluations(), + exact.distance_evaluations() + ); + + // A column below its declared size stays ungrouped however many cycles run: 1,000 rows + // under AUTO_INDEX_AT 50000 are not 50,000, though one partition holds a hundredth of it. + let index_bytes = |db: &contextdb::Database| { + let index = support::shown(db, "SHOW VECTOR_INDEXES").remove(0); + support::whole(&index, "durable_index_bytes") + }; + let clauses = "PARTITION_KEY (ctx) AUTO_INDEX_AT 50000"; + let (_, db) = support::built("auto-below-declared", DIMENSION, clauses, 400); + support::load_docs(&db, 401..=1000, "b", gathered); + for _ in 0..3 { + support::upkeep(&db); + } + let grouped: Vec = support::shown(&db, "SHOW VECTOR_PARTITIONS") + .iter() + .map(lists) + .collect(); + assert_eq!(grouped, [Value::Null, Value::Null]); + assert_eq!(index_bytes(&db), 0); + + // With no AUTO_INDEX_AT the number in force is 100,000: a partition of 1,200 stays + // ungrouped, and AUTO over it is EXACT. + let (_, db) = support::built("auto-undeclared", DIMENSION, "", 1200); + let index = support::shown(&db, "SHOW VECTOR_INDEXES").remove(0); + assert_eq!( + support::field(&index, "declared_auto_index_at"), + &Value::Null + ); + assert_eq!(support::whole(&index, "effective_auto_index_at"), 100_000); + assert_eq!(lists(&support::the_partition(&db)), Value::Null); + assert_eq!(index_bytes(&db), 0); + let [auto, exact] = auto_and_exact(&db, None, &query); + assert_eq!( + support::found(&auto, "AUTO"), + support::found(&exact, "EXACT") + ); + assert_eq!(auto.distance_evaluations(), exact.distance_evaluations()); +} + +#[test] +fn auto_searches_each_partition_by_what_its_grouping_saves() { + let clauses = "PARTITION_KEY (ctx) AUTO_INDEX_AT 1000"; + let (_, db) = support::built("auto-each-partition", DIMENSION, clauses, 2000); + // Five pending rows under another context, nearer the question than any row of the first. + let query = question(3); + support::load_docs(&db, 9001..=9005, "b", |n| { + let moved = support::vector(n, DIMENSION); + query + .iter() + .zip(moved) + .map(|(q, by)| q + 0.001 * by) + .collect() + }); + let stored = support::stored(&db, support::DOCS); + let [auto, exact] = auto_and_exact(&db, None, &query); + let answer = support::assert_sound(&auto, &stored, all, &query, 10, "AUTO over both"); + let pending = |near: &&support::Near| matches!(near.key, Value::Int64(id) if id > 9000); + assert_eq!(answer.iter().filter(pending).count(), 5, "{answer:?}"); + assert!( + auto.distance_evaluations() < exact.distance_evaluations(), + "AUTO did {} distance evaluations, EXACT {}", + auto.distance_evaluations(), + exact.distance_evaluations() + ); +} + +#[test] +fn auto_answers_exactly_where_a_grouping_is_behind_or_damaged() { + let (_, db) = support::built("auto-behind", DIMENSION, GROUPED, 2000); + // 800 vectors outside the grouping: twice the 400 at which it is grouped again. + support::load_docs(&db, 2001..=2800, "a", gathered); + let query = question(4); + let wanted = support::stored(&db, support::DOCS).nearest(all, &query, 10); + let sql = support::by_mode(None, "INDEXED", 10); + support::assert_indexed_refused(support::search(&db, &sql, &query), "a grouping behind"); + let partition = support::the_partition(&db); + assert_eq!( + support::word(&partition, "availability_reason"), + "behind_writes" + ); + let work = ["EXACT", "AUTO"].map(|mode| { + let sql = support::by_mode(None, mode, 10); + support::assert_nearest(&db, &sql, &query, &wanted, mode).distance_evaluations() + }); + assert_eq!(work[1], work[0], "AUTO's evaluations and EXACT's"); +} + +#[test] +fn writes_after_a_grouping_are_found_at_once() { + let (_, db) = support::built("writes-after-grouping", DIMENSION, GROUPED, 2000); + let gone = gathered(7); + let inserted = support::vector(70_001, DIMENSION); + let changed = support::vector(70_002, DIMENSION); + support::insert_doc(&db, 3001, "a", Some(inserted.clone())).expect("insert a row"); + db.execute( + "UPDATE docs SET v = $v WHERE id = 5", + ¶ms! { "v" => changed.clone() }, + ) + .expect("change a vector"); + support::run(&db, "DELETE FROM docs WHERE id = 7"); + let check = |state: &str| { + let stored = support::stored(&db, support::DOCS); + for (query, id) in [(&inserted, 3001), (&changed, 5)] { + let result = support::answer(&db, &support::by_mode(None, "INDEXED", 1), query); + let answer = support::assert_sound(&result, &stored, all, query, 1, state); + assert_eq!(answer[0].key, Value::Int64(id), "{state}"); + } + let result = support::answer(&db, &support::by_mode(None, "INDEXED", 10), &gone); + let answer = support::assert_sound(&result, &stored, all, &gone, 10, state); + assert!( + answer.iter().all(|near| near.key != Value::Int64(7)), + "{state}: the deleted row is in {answer:?}" + ); + }; + check("before any cycle"); + assert_eq!(support::cycle(&db).maintenance_vectors_done, 2); + check("after sealing"); +} + +#[test] +fn a_search_with_a_filter_on_rows_widens_until_it_is_whole() { + let (_, db) = support::built("filter-widens", DIMENSION, GROUPED, 2000); + let stored = support::stored(&db, support::DOCS); + let query = question(5); + let search = |db: &contextdb::Database, filter: Option<&str>, mode: &str, k: usize| { + support::search(db, &support::by_mode(filter, mode, k), &query) + }; + + // A filter that passes every row costs what no filter costs. + let plain = search(&db, None, "INDEXED", 10).expect("no filter"); + let passing = search(&db, Some("tag >= 0"), "INDEXED", 10).expect("every row passes"); + assert_eq!( + support::found(&passing, "every row passes"), + support::found(&plain, "no filter") + ); + assert_eq!(passing.distance_evaluations(), plain.distance_evaluations()); + + // Nine rows in ten pass. INDEXED answers ten rows the filter allows, nearest first, each + // with the distance the exact route is held to for that row, and does less than the exact + // route's work. Which ten is the index's to say: it may leave out a row the exact ten hold. + let most = |row: &[Value]| matches!(row[2], Value::Int64(tag) if tag < 45); + let wanted = stored.nearest(most, &query, 10); + let sql = support::by_mode(Some("tag < 45"), "EXACT", 10); + let exact = support::assert_nearest(&db, &sql, &query, &wanted, "EXACT"); + let indexed = search(&db, Some("tag < 45"), "INDEXED", 10).expect("nine in ten pass"); + let what = "INDEXED, nine rows in ten pass"; + support::assert_sound(&indexed, &stored, most, &query, 10, what); + assert!( + indexed.distance_evaluations() < exact.distance_evaluations(), + "{what}: {} distance evaluations, EXACT {}", + indexed.distance_evaluations(), + exact.distance_evaluations() + ); + + // Twelve rows pass, spread over the groups, and ten are asked for. AUTO widens until it + // has done the exact work and returns the ten nearest of them; INDEXED holds ten allowed + // rows before the step that would take every group, and answers them. + let tagged = |upto: i64| move |row: &[Value]| matches!((&row[0], &row[2]), (Value::Int64(id), Value::Int64(7)) if *id <= upto); + let twelve = Some("tag = 7 AND id <= 600"); + assert_eq!(stored.nearest(tagged(600), &query, 100).len(), 12); + let wanted = stored.nearest(tagged(600), &query, 10); + let sql = support::by_mode(twelve, "AUTO", 10); + support::assert_nearest(&db, &sql, &query, &wanted, "AUTO, twelve rows pass"); + let indexed = search(&db, twelve, "INDEXED", 10).expect("twelve rows pass"); + let what = "INDEXED, twelve rows pass"; + support::assert_sound(&indexed, &stored, tagged(600), &query, 10, what); + + // Six rows pass and ten are asked for: AUTO returns the six, and INDEXED, short of its + // limit at the step that would take every group, is refused by name. + let six = Some("tag = 7 AND id <= 300"); + let wanted = stored.nearest(tagged(300), &query, 10); + assert_eq!(wanted.len(), 6); + let sql = support::by_mode(six, "AUTO", 10); + support::assert_nearest(&db, &sql, &query, &wanted, "AUTO, six rows pass"); + support::assert_indexed_refused(search(&db, six, "INDEXED", 10), "six rows pass"); + + // Three rows pass: AUTO returns the three, and INDEXED is refused by name. + let few = Some("tag = 7 AND id <= 150"); + let wanted = stored.nearest(tagged(150), &query, 10); + assert_eq!(wanted.len(), 3); + let sql = support::by_mode(few, "AUTO", 10); + support::assert_nearest(&db, &sql, &query, &wanted, "AUTO, three rows pass"); + support::assert_indexed_refused(search(&db, few, "INDEXED", 10), "USE VECTOR INDEXED"); + + // With no filter and a LIMIT above the rows, INDEXED returns every row. + let every = search(&db, None, "INDEXED", 5000).expect("every row"); + support::assert_sound(&every, &stored, all, &query, 2000, "a LIMIT above the rows"); + + let clauses = "SEARCH_MODE INDEXED AUTO_INDEX_AT 1000"; + let (_, declared) = support::built("filter-widens-declared", DIMENSION, clauses, 2000); + support::assert_indexed_refused(search(&declared, few, "", 10), "SEARCH_MODE INDEXED"); +} + +/// The keys of an answer, which must be `count` different rows. +fn different_keys(result: &QueryResult, count: usize, what: &str) { + let keys: BTreeSet = result + .rows() + .iter() + .map(|row| match row[0] { + Value::Int64(id) => id, + ref other => panic!("{what}: the key {other:?} is no INTEGER"), + }) + .collect(); + assert_eq!(result.rows().len(), count, "{what}: rows in the answer"); + assert_eq!(keys.len(), count, "{what}: different rows in the answer"); +} + +#[test] +fn a_search_that_widens_returns_no_row_twice() { + let (_, db) = support::built("widens-no-row-twice", DIMENSION, GROUPED, 2000); + let stored = support::stored(&db, support::DOCS); + // Filters that pass a half, a tenth and a fiftieth of the rows: the fewer pass, the more + // often the search goes back to groups and sketches it has already been through. + type Passing = (&'static str, fn(&[Value]) -> bool); + let filters: [Passing; 3] = [ + ( + "tag < 25", + |row| matches!(row[2], Value::Int64(tag) if tag < 25), + ), + ( + "tag < 5", + |row| matches!(row[2], Value::Int64(tag) if tag < 5), + ), + ("tag = 7", |row| row[2] == Value::Int64(7)), + ]; + for j in 0..8 { + let query = question(j); + for (filter, passes) in filters { + for k in [10, 30] { + let what = format!("AUTO WHERE {filter} LIMIT {k}, question {j}"); + let sql = support::by_mode(Some(filter), "AUTO", k); + let result = support::answer(&db, &sql, &query); + different_keys(&result, k, &what); + support::assert_sound(&result, &stored, passes, &query, k, &what); + } + } + } +} + +#[test] +fn an_indexed_search_that_takes_every_group_does_the_exact_evaluations() { + let (_, db) = support::built("every-group-evaluations", DIMENSION, GROUPED, 2000); + let stored = support::stored(&db, support::DOCS); + let query = question(6); + let evaluations = |mode: &str, k: usize, rows: usize| { + let result = support::answer(&db, &support::by_mode(None, mode, k), &query); + let what = format!("{mode} LIMIT {k}"); + support::assert_sound(&result, &stored, all, &query, rows, &what); + result.distance_evaluations() + }; + // Ten rows: the nearest groups, and fewer evaluations than the whole partition takes. + let (indexed, exact) = (evaluations("INDEXED", 10, 10), evaluations("EXACT", 10, 10)); + assert!( + indexed < exact, + "INDEXED did {indexed} distance evaluations for ten rows, EXACT {exact}" + ); + // Every row: every group is taken. An evaluation is a centre compared, a sketch + // estimated, a slot screened or a row asked for its distance. + let (indexed, exact) = ( + evaluations("INDEXED", 5000, 2000), + evaluations("EXACT", 5000, 2000), + ); + assert_eq!( + exact, + 2 * 2000, + "EXACT's distance evaluations for every row: each of the 2,000 slots is screened, \ + which is one, and each of the 2,000 rows is asked for its distance, which is one more" + ); + assert!( + indexed >= exact + 2000 + 45, + "INDEXED did {indexed} distance evaluations for every row: it screens every slot and \ + asks every row as EXACT does ({exact}), after comparing its 45 centres and estimating \ + the sketch of each of the 2,000 slots at least once" + ); +} + +#[test] +fn an_indexed_search_short_of_its_limit_is_refused_before_the_whole_table_is_read() { + // 256 rows spread evenly round a circle, in vectors wide enough for five blocks of 63 + // slots. The size is fixed by the blocks: a refusal reads less than the whole table only + // when there are blocks for it to stop short of, and the filter needs its three rows + // among the first 150; few rows in wide vectors make those with a short build. + const WIDE: usize = 4096; + const ROWS: i64 = 256; + let folder = support::folder("indexed-refused-before-the-table"); + let db = support::open(&folder); + support::declare_docs(&db, WIDE, "AUTO_INDEX_AT 256"); + support::load_docs(&db, 1..=ROWS, "a", |n| support::ringed(n, ROWS, WIDE)); + support::upkeep(&db); + assert_eq!(lists(&support::the_partition(&db)), Value::Int64(16)); + let stored = support::stored(&db, support::DOCS); + support::close(db); + + // Three rows pass the filter and the statement asks for ten. + let few = Some("tag = 7 AND id <= 150"); + let passes = |row: &[Value]| matches!((&row[0], &row[2]), (Value::Int64(id), Value::Int64(7)) if *id <= 150); + let query = support::ringed(157, 2 * ROWS, WIDE); + let wanted = stored.nearest(passes, &query, 10); + assert_eq!(wanted.len(), 3); + // Each search runs first in a session of its own, so every page it needs is read once. + let read_by = |mode: &str| -> (u64, contextdb::Result) { + let (door, bytes) = support::counting_bytes_read(); + let mut options = ReadSessionOptions::default(); + options.storage = Some(door); + let reader = ReadSession::open_with_options(&folder, options).expect("open for reading"); + let opened = bytes.load(Ordering::SeqCst); + let outcome = support::search(&reader, &support::by_mode(few, mode, 10), &query); + (bytes.load(Ordering::SeqCst) - opened, outcome) + }; + // AUTO asks the filter of every row, which is the whole table's work, and answers. + let (whole, auto) = read_by("AUTO"); + let auto = auto.expect("AUTO, three rows pass"); + support::assert_same_rows( + &support::found(&auto, "AUTO"), + &wanted, + "AUTO, three rows pass", + ); + // INDEXED is refused at the step that would take every group: before that work. + let (refused, indexed) = read_by("INDEXED"); + support::assert_indexed_refused(indexed, "INDEXED, three rows pass"); + assert!( + refused < whole, + "INDEXED was refused after reading {refused} bytes of the data file; AUTO, which asks \ + the filter of all {ROWS} rows, read {whole}" + ); +} + +/// One store of [`a_nearer_allowed_row_in_a_group_taken_later_is_in_the_answer`], with row +/// 507 at these two numbers. +fn the_nearer_allowed_row_is_first(name: &str, nearer: (f32, f32)) { + // The question is the first direction. Rows 7, 57, .. 457 (ten rows with tag 7) lie + // half way from it, each towards a direction of its own, and row 507 (tag 7 too) lies + // nearer the question but nearer still to the second direction, among whose 500 rows + // it is grouped. Every other row is gathered round one of four directions; those round + // the first are the nearest rows of all, and none of them passes the filter. + let fixed = |n: i64| -> Vec { + let mut at = [0.0_f32; DIMENSION]; + match n { + 507 => (at[0], at[1]) = nearer, + n if n % 50 == 7 && n < 507 => { + at[0] = 0.5; + at[2 + (n / 50) as usize] = 0.866; + } + n => at[[0, 1, 12, 13][n.rem_euclid(4) as usize]] = 1.0, + } + let moved = support::vector(n, DIMENSION); + (0..DIMENSION).map(|i| at[i] + 0.02 * moved[i]).collect() + }; + let folder = support::folder(name); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, GROUPED); + support::load_docs(&db, 1..=2000, "a", fixed); + support::upkeep(&db); + let groups = lists(&support::the_partition(&db)); + assert!( + matches!(groups, Value::Int64(groups) if groups > 1), + "the partition has {groups:?} groups" + ); + let stored = support::stored(&db, support::DOCS); + let mut query = vec![0.0_f32; DIMENSION]; + query[0] = 1.0; + + let filter = Some("tag = 7 AND id <= 550"); + let passes = |row: &[Value]| matches!((&row[0], &row[2]), (Value::Int64(id), Value::Int64(7)) if *id <= 550); + let wanted = stored.nearest(passes, &query, 10); + assert_eq!(stored.nearest(passes, &query, 100).len(), 11); + assert_eq!(wanted[0].key, Value::Int64(507), "the nearest allowed row"); + // AUTO answers ten of the eleven allowed rows with their true distances, and the + // nearest of them first: it is nearer than every row the answer could hold in its place. + let result = support::answer(&db, &support::by_mode(filter, "AUTO", 10), &query); + let what = format!("{name}: AUTO, eleven rows pass"); + let answer = support::assert_sound(&result, &stored, passes, &query, 10, &what); + assert_eq!(answer, wanted, "{what}"); +} + +#[test] +fn a_nearer_allowed_row_in_a_group_taken_later_is_in_the_answer() { + the_nearer_allowed_row_is_first("nearer-row-in-a-later-group", (0.7, 0.714)); + // Nearer the second direction still, on the line half way between the two: the group it + // is in lies four times as far behind the first as the row itself does. + the_nearer_allowed_row_is_first("nearer-row-half-way-to-a-later-group", (0.55, 0.835)); +} + +#[test] +fn a_grouping_whose_sketches_cannot_tell_its_vectors_apart_answers_as_exact() { + // Every row has one of eight vectors, so each vector of a group is the group's centre: + // the sketch of every one is the same and the width measured for it is nothing. + let same = |n: i64| -> Vec { + let direction = n.rem_euclid(8) as usize; + (0..DIMENSION) + .map(|at| f32::from(u8::from(at / 2 == direction)) + 0.25) + .collect() + }; + let folder = support::folder("sketches-all-the-same"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, GROUPED); + support::load_docs(&db, 1..=2000, "a", same); + support::upkeep(&db); + let groups = lists(&support::the_partition(&db)); + assert!( + matches!(groups, Value::Int64(groups) if groups > 1), + "the partition has {groups:?} groups" + ); + let stored = support::stored(&db, support::DOCS); + let half = |row: &[Value]| matches!(row[2], Value::Int64(tag) if tag < 25); + for direction in [0, 3, 7] { + let query = same(direction); + // Rows at one distance come in key order, so the ten are the same under every mode. + let wanted = stored.nearest(all, &query, 10); + let allowed = stored.nearest(half, &query, 10); + for mode in ["EXACT", "AUTO", "INDEXED"] { + let what = format!("{mode}, direction {direction}"); + let sql = support::by_mode(None, mode, 10); + support::assert_nearest(&db, &sql, &query, &wanted, &what); + let sql = support::by_mode(Some("tag < 25"), mode, 10); + support::assert_nearest(&db, &sql, &query, &allowed, &format!("{what}, filtered")); + } + } +} + +#[test] +fn a_search_with_no_room_for_its_sketches_is_refused_by_the_memory_limit() { + const MIB: usize = 1 << 20; + let (folder, db) = support::built("no-room-for-sketches", DIMENSION, GROUPED, 2000); + assert_eq!(lists(&support::the_partition(&db)), Value::Int64(45)); + let stored = support::stored(&db, support::DOCS); + support::close(db); + + // 9 MiB is the smallest limit a store opens under, and leaves a search nothing to keep + // the sketches of its groups in. + let mut options = DatabaseOpenOptions::default(); + options.memory_limit = Some(9 * MIB); + let db = support::open_with(&folder, options).expect("open under 9 MiB"); + let query = question(7); + let sql = support::by_mode(None, "INDEXED", 10); + let before = (db.memory_used(), db.committed_watermark()); + match support::search(&db, &sql, &query) { + Err(Error::MemoryBudgetExceeded { + budget_limit_bytes, .. + }) => assert_eq!(budget_limit_bytes, 9 * MIB as u64), + other => panic!("{sql}\nunder 9 MiB got {other:?}"), + } + assert_eq!( + (db.memory_used(), db.committed_watermark()), + before, + "after the refused search" + ); + + db.set_memory_limit(Some(64 * MIB)).expect("a larger limit"); + let result = support::answer(&db, &sql, &query); + support::assert_sound(&result, &stored, all, &query, 10, "under 64 MiB"); +} + +#[test] +fn a_filter_of_partition_keys_only_is_no_filter_on_rows() { + let clauses = "PARTITION_KEY (ctx) AUTO_INDEX_AT 1000"; + let (_, db) = support::built("keys-only-filter", DIMENSION, clauses, 2000); + support::load_docs(&db, 2001..=2005, "b", gathered); + support::upkeep(&db); + let stored = support::stored(&db, support::DOCS); + let query = question(8); + let in_b = |row: &[Value]| row[1] == support::text("b"); + + // Five rows are selected and ten asked for: every one of them, and no refusal. + let wanted = stored.nearest(in_b, &query, 10); + assert_eq!(wanted.len(), 5); + let sql = support::by_mode(Some("ctx IN ('b')"), "INDEXED", 10); + support::assert_nearest(&db, &sql, &query, &wanted, "one small partition selected"); + // A tuple no row is under adds nothing and refuses nothing. + let sql = support::by_mode(Some("ctx IN ('b', 'zz')"), "INDEXED", 10); + support::assert_nearest(&db, &sql, &query, &wanted, "a tuple with no row beside it"); + // Both partitions and a LIMIT above their 2,005 rows: every row. + let sql = support::by_mode(Some("ctx IN ('a', 'b')"), "INDEXED", 5000); + let result = support::answer(&db, &sql, &query); + different_keys(&result, 2005, "both partitions selected"); + support::assert_sound( + &result, + &stored, + all, + &query, + 2005, + "both partitions selected", + ); +} + +#[test] +fn a_read_that_fails_under_auto_is_returned_as_that_failure() { + let (folder, db) = support::built("auto-read-fails", DIMENSION, GROUPED, 2000); + support::close(db); + let query = question(9); + let sql = support::by_mode(None, "AUTO", 10); + // A search in a session of its own, whose storage door fails read number `fail` of the + // data file; 0 fails none. Returns the reads made before the search and by its end. + let searched = |fail: u64| -> (u64, u64, contextdb::Result) { + let reads = Arc::new(AtomicU64::new(0)); + let mut options = ReadSessionOptions::default(); + options.storage = Some(support::failing_storage({ + let reads = reads.clone(); + move |call| match call { + support::Call::Read(..) if reads.fetch_add(1, Ordering::SeqCst) + 1 == fail => { + Err(io::Error::other("a read this test makes fail")) + } + _ => Ok(()), + } + })); + let reader = ReadSession::open_with_options(&folder, options).expect("open for reading"); + let opened = reads.load(Ordering::SeqCst); + let outcome = support::search(&reader, &sql, &query); + (opened, reads.load(Ordering::SeqCst), outcome) + }; + let (opened, ended, whole) = searched(0); + assert_eq!(whole.expect("AUTO, no read failing").rows().len(), 10); + assert!(ended > opened, "the search read the data file"); + // Whichever read of the search fails, of a grouping's record as of any other, the search + // returns the failure: it never answers as though the record were only unreadable. + for fail in opened + 1..=ended { + match searched(fail).2 { + Err(Error::StoreIo { .. }) => {} + other => panic!( + "read {} of the {} of an AUTO search failed, and the search gave {other:?}", + fail - opened, + ended - opened + ), + } + } +} + +#[test] +fn an_indexed_search_whose_first_pass_takes_every_slot_is_refused_before_a_row_is_asked() { + // One row, and no grouping. Its vector is wide enough that asking the row for its + // distance shows in the bytes read: 65,536 bytes of numbers. + const WIDE: usize = 16384; + let folder = support::folder("indexed-refused-before-a-row"); + let db = support::open(&folder); + support::declare_docs(&db, WIDE, ""); + support::insert_doc(&db, 1, "a", Some(support::vector(1, WIDE))).expect("insert the row"); + support::close(db); + + let query = support::vector(2, WIDE); + // Each search runs first in a session of its own, so every page it needs is read once. + let read_by = |mode: &str| -> (u64, contextdb::Result) { + let (door, bytes) = support::counting_bytes_read(); + let mut options = ReadSessionOptions::default(); + options.storage = Some(door); + let reader = ReadSession::open_with_options(&folder, options).expect("open for reading"); + let opened = bytes.load(Ordering::SeqCst); + let sql = support::by_mode(Some("tag = 1"), mode, 2); + let outcome = support::search(&reader, &sql, &query); + (bytes.load(Ordering::SeqCst) - opened, outcome) + }; + let vector_bytes = 4 * WIDE as u64; + let (asked, exact) = read_by("EXACT"); + assert_eq!(exact.expect("EXACT, one row passes").rows().len(), 1); + assert!( + asked >= vector_bytes, + "EXACT asked the row for its distance and read {asked} bytes of the data file" + ); + // The filter passes one row and two are asked for: the pass that would keep every slot is + // the first, and INDEXED is refused there, with no row asked for its distance. + let (refused, indexed) = read_by("INDEXED"); + support::assert_indexed_refused(indexed, "one row, a filter and LIMIT 2"); + assert!( + refused < vector_bytes, + "INDEXED was refused after reading {refused} bytes of the data file; the row's vector \ + is {vector_bytes}" + ); +} diff --git a/v2/contextdb/tests/vectors_inspection.rs b/v2/contextdb/tests/vectors_inspection.rs new file mode 100644 index 00000000..d06b1a0c --- /dev/null +++ b/v2/contextdb/tests/vectors_inspection.rs @@ -0,0 +1,388 @@ +//! Inspection of VECTOR columns: `SHOW VECTOR_INDEXES`, `SHOW VECTOR_PARTITIONS` and `explain` +//! say what the vector part keeps and how a search would go, from its small records alone: +//! they read no block and change nothing. + +mod support; + +use std::sync::atomic::Ordering; + +use contextdb::{DatabaseOpenOptions, Error, ReadSession, Value, VectorIndexRef}; + +use support::{Run, Shown}; + +fn int(number: i64) -> Value { + Value::Int64(number) +} + +/// Asserts the fields of an inspection row that hold text, numbers, flags and NULL. The two +/// JSON fields are compared apart. +fn assert_fields(row: &Shown, wanted: &[(&str, Value)], what: &str) { + for (name, value) in wanted { + assert_eq!(support::field(row, name), value, "{what}: {name}"); + } +} + +/// Two columns: `docs.v`, keyed by context, with 600 sealed vectors under `a` and 400 pending +/// under `b`, which together bring the column to its declared size; and `plain.e`, with no +/// key and no vector. +fn two_columns(name: &str) -> (std::path::PathBuf, contextdb::Database) { + let clauses = "PARTITION_KEY (ctx) AUTO_INDEX_AT 1000 CONSOLIDATION (CHANGE_PERCENT = 50)"; + let (folder, db) = support::built(name, 8, clauses, 0); + support::run( + &db, + "CREATE TABLE plain (id INTEGER PRIMARY KEY, e VECTOR(4))", + ); + support::load_docs(&db, 1..=600, "a", |n| support::vector(n, 8)); + assert_eq!(support::upkeep(&db), (1, 600)); + support::load_docs(&db, 601..=1000, "b", |n| support::vector(n, 8)); + (folder, db) +} + +#[test] +fn show_vector_indexes_describes_every_vector_column() { + let (folder, db) = two_columns("show-vector-indexes"); + // The names of the fields are held by `tests/sql/vector_inspection.sql`, through the + // writer: every other door gives the same. + let fields = support::run(&db, "SHOW VECTOR_INDEXES").columns().to_vec(); + let through = |on: &dyn Run, door: &str| { + let result = support::run(&on, "SHOW VECTOR_INDEXES"); + assert_eq!(result.columns(), fields, "through {door}"); + let rows = support::shown(&on, "SHOW VECTOR_INDEXES"); + assert_eq!(rows.len(), 2, "through {door}"); + let text = support::text; + assert_fields( + &rows[0], + &[ + ("table", text("docs")), + ("column", text("v")), + ("dimension", int(8)), + ("quantization", text("F32")), + ("vector_count", int(1000)), + ("bytes", int(600 * 16)), + ("max_partitions", int(256)), + ("live_partitions", int(2)), + ("search_mode", text("AUTO")), + ("declared_auto_index_at", int(1000)), + ("effective_auto_index_at", int(1000)), + ("declared_consolidation_mode", text("PERCENT")), + ("declared_consolidation_change_percent", int(50)), + ("declared_consolidation_tombstone_percent", Value::Null), + ("effective_consolidation_mode", text("PERCENT")), + ("effective_consolidation_change_percent", int(50)), + ("effective_consolidation_tombstone_percent", int(10)), + ("pending_inserts", int(400)), + ("tombstones", int(0)), + ("durable_vector_bytes", int(600 * 16)), + ("durable_index_bytes", int(0)), + ("indexed_available", Value::Bool(true)), + ("maintenance_state", text("pending")), + ("maintenance_vectors_total", int(600)), + ("maintenance_vectors_done", int(0)), + ("maintenance_vectors_remaining", int(600)), + ("unavailable_partitions", int(0)), + ("stalled_partitions", int(0)), + ("broad_route", text("fan_out")), + ("recovery_action", text("run a maintenance cycle")), + ], + door, + ); + assert_eq!( + support::json(&rows[0], "partition_key_columns"), + contextdb::serde_json::json!(["ctx"]) + ); + assert_fields( + &rows[1], + &[ + ("table", text("plain")), + ("column", text("e")), + ("dimension", int(4)), + ("quantization", text("F32")), + ("vector_count", int(0)), + ("bytes", int(0)), + ("max_partitions", Value::Null), + ("live_partitions", int(0)), + ("search_mode", text("AUTO")), + ("declared_auto_index_at", Value::Null), + ("effective_auto_index_at", int(100_000)), + ("declared_consolidation_mode", Value::Null), + ("declared_consolidation_change_percent", Value::Null), + ("declared_consolidation_tombstone_percent", Value::Null), + ("effective_consolidation_mode", text("PERCENT")), + ("effective_consolidation_change_percent", int(20)), + ("effective_consolidation_tombstone_percent", int(10)), + ("pending_inserts", int(0)), + ("indexed_available", Value::Bool(true)), + ("maintenance_state", text("idle")), + ("maintenance_vectors_total", int(0)), + ("broad_route", text("fan_out")), + ("recovery_action", text("none")), + ], + door, + ); + assert_eq!( + support::json(&rows[1], "partition_key_columns"), + contextdb::serde_json::json!([]) + ); + }; + through(&db, "the writer"); + let tx = db.transaction().expect("begin"); + through(&tx, "a transaction"); + tx.rollback(); + through( + &ReadSession::open(&folder).expect("open for reading"), + "a read session", + ); +} + +#[test] +fn show_vector_partitions_lists_partitions_by_key_with_limit_and_offset() { + let (folder, db) = two_columns("show-vector-partitions"); + let reader = ReadSession::open(&folder).expect("open for reading"); + // The names of the fields are held by `tests/sql/vector_inspection.sql`, through the + // writer: a reader gives the same. + let result = support::run(&reader, "SHOW VECTOR_PARTITIONS"); + let by_writer = support::run(&db, "SHOW VECTOR_PARTITIONS"); + assert_eq!(result.columns(), by_writer.columns()); + // The key of each row a statement returns, with its column. + let keys = |sql: &str| -> Vec { + support::shown(&reader, sql) + .iter() + .map(|row| { + format!( + "{}.{} {}", + support::word(row, "table"), + support::word(row, "column"), + support::json(row, "partition_key") + ) + }) + .collect() + }; + let (a, b, plain) = ( + r#"docs.v {"ctx":"a"}"#, + r#"docs.v {"ctx":"b"}"#, + "plain.e {}", + ); + assert_eq!(keys("SHOW VECTOR_PARTITIONS"), [a, b, plain]); + assert_eq!(keys("SHOW VECTOR_PARTITIONS FOR docs.v"), [a, b]); + assert_eq!(keys("SHOW VECTOR_PARTITIONS FOR docs.v LIMIT 1"), [a]); + assert_eq!( + keys("SHOW VECTOR_PARTITIONS FOR docs.v LIMIT 1 OFFSET 1"), + [b] + ); + assert_eq!(keys("SHOW VECTOR_PARTITIONS LIMIT 5 OFFSET 2"), [plain]); + assert_eq!(keys("SHOW VECTOR_PARTITIONS LIMIT 2"), [a, b]); + assert_eq!(keys("SHOW VECTOR_PARTITIONS FOR plain.e"), [plain]); + + let rows = support::shown(&db, "SHOW VECTOR_PARTITIONS"); + let text = support::text; + // Sealed, in a column that has just come to its size: its first grouping is due. + assert_fields( + &rows[0], + &[ + ("live_rows", int(600)), + ("pending_inserts", int(0)), + ("tombstones", int(0)), + ("durable_vector_bytes", int(600 * 16)), + ("durable_index_bytes", int(0)), + ("indexed_available", Value::Bool(true)), + ("availability_reason", text("none")), + ("maintenance_state", text("pending")), + ("maintenance_reason", text("first_grouping")), + ("maintenance_vectors_total", int(600)), + ("maintenance_vectors_done", int(0)), + ("maintenance_vectors_remaining", int(600)), + ("lists", Value::Null), + ("probes", Value::Null), + ("recovery_action", text("run a maintenance cycle")), + ], + "the sealed partition", + ); + assert_fields( + &rows[1], + &[ + ("live_rows", int(400)), + ("pending_inserts", int(400)), + ("durable_vector_bytes", int(0)), + ("lists", Value::Null), + ], + "the pending partition", + ); + // A column with no key has one row before its first vector. + assert_fields( + &rows[2], + &[ + ("live_rows", int(0)), + ("pending_inserts", int(0)), + ("tombstones", int(0)), + ("durable_vector_bytes", int(0)), + ("durable_index_bytes", int(0)), + ("maintenance_state", text("idle")), + ("maintenance_reason", text("none")), + ("lists", Value::Null), + ("recovery_action", text("none")), + ], + "the column with no vector", + ); + + // The two refusals. + match support::refusal(&db, "SHOW VECTOR_PARTITIONS FOR nowhere.v") { + Error::TableNotFound(table) => assert_eq!(table, "nowhere"), + other => panic!("a table that is not there gave {other:?}"), + } + match support::refusal(&db, "SHOW VECTOR_PARTITIONS FOR docs.tag") { + Error::UnknownVectorIndex { index } => { + let wanted = VectorIndexRef { + table: "docs".to_owned(), + column: "tag".to_owned(), + }; + assert_eq!(index, wanted); + } + other => panic!("a column that is no VECTOR column gave {other:?}"), + } +} + +/// A search of `docs` with this filter and mode, with nothing bound. +fn searched(filter: Option<&str>, mode: &str) -> String { + support::by_mode(filter, mode, 10) +} + +#[test] +fn inspection_and_explain_read_no_block_and_change_nothing() { + // One block of 2,000 slots of 72 bytes with keys of 9: 170,004 bytes. + const BLOCK: u64 = 4 + 2000 * (72 + 4 + 9); + let (folder, db) = support::built("inspection-reads-no-block", 64, "AUTO_INDEX_AT 1000", 2000); + support::close(db); + let (door, bytes) = support::counting_bytes_read(); + let mut options = DatabaseOpenOptions::default(); + options.storage = Some(door); + let db = support::open_with(&folder, options).expect("open"); + let (opened, before) = (bytes.load(Ordering::SeqCst), db.committed_watermark()); + assert_eq!(support::shown(&db, "SHOW VECTOR_INDEXES").len(), 1); + assert_eq!( + support::field(&support::the_partition(&db), "lists"), + &int(45) + ); + let lines = db.explain(&searched(None, "INDEXED")).expect("explain"); + assert!(lines.contains("route: indexed"), "{lines}"); + let read = bytes.load(Ordering::SeqCst) - opened; + assert!( + read < BLOCK, + "inspection and explain read {read} bytes; the block is {BLOCK}" + ); + assert_eq!(db.committed_watermark(), before); +} + +/// The lines of an explanation, each trimmed of the spaces it begins with. +fn lines_of(explained: &str) -> Vec<&str> { + explained.lines().map(str::trim_start).collect() +} + +#[test] +fn explain_names_how_each_partition_is_searched() { + // Partition a is built; b has 2,000 vectors outside any grouping, twice the declared size. + let clauses = "PARTITION_KEY (ctx) AUTO_INDEX_AT 1000"; + let (folder, db) = support::built("explain-routes", 16, clauses, 2000); + support::load_docs(&db, 2001..=4000, "b", |n| support::gathered(n, 16)); + let reader = ReadSession::open(&folder).expect("open for reading"); + let explain = |sql: &str| -> String { + let from_writer = db.explain(sql).expect("the writer explains"); + assert_eq!(reader.explain(sql).expect("a reader explains"), from_writer); + from_writer + }; + let routes = |lines: &[&str]| -> Vec { + let routes = lines.iter().filter(|line| line.starts_with("route: ")); + routes + .map(|line| line["route: ".len()..].to_owned()) + .collect() + }; + + let explained = explain(&searched(None, "EXACT")); + let lines = lines_of(&explained); + for wanted in [ + "requested_mode: EXACT", + "resolved_mode: EXACT", + "effective_auto_index_at: 1000", + "aggregate_allowed_vectors: 4000", + r#"partition_key_columns: ["ctx"]"#, + "partitions: 2", + r#"partition: {"ctx":"a"}"#, + r#"partition: {"ctx":"b"}"#, + "refusal: none", + ] { + assert!( + lines.contains(&wanted), + "no line {wanted:?} in\n{explained}" + ); + } + assert_eq!(routes(&lines), ["exact", "exact"]); + assert_eq!(lines.last(), Some(&"refusal: none")); + + // The built partition alone, under INDEXED: its own six lines. + let explained = explain(&searched(Some("ctx = 'a'"), "INDEXED")); + let lines = lines_of(&explained); + let at = lines + .iter() + .position(|line| *line == r#"partition: {"ctx":"a"}"#) + .unwrap_or_else(|| panic!("no partition line in\n{explained}")); + assert_eq!(lines[at + 1..at + 3], ["route: indexed", "lists: 45"]); + assert!(lines[at + 3].starts_with("probes: "), "{explained}"); + assert_eq!( + lines[at + 4..at + 6], + ["live_rows: 2000", "outside_grouping: 0"] + ); + assert!(explained.contains("\n route: indexed\n"), "{explained}"); + assert!(lines.contains(&"partitions: 1") && lines.contains(&"refusal: none")); + + // Both under INDEXED: the partition that is behind refuses the search, in its sentence. + let explained = explain(&searched(None, "INDEXED")); + let refusal = format!("refusal: docs.v{}", support::INDEXED_REFUSED); + assert_eq!(lines_of(&explained).last(), Some(&refusal.as_str())); + assert!(explained.contains("outside_grouping: 2000"), "{explained}"); + + // Under AUTO each partition goes its own way and nothing is refused. + let explained = explain(&searched(None, "AUTO")); + let lines = lines_of(&explained); + assert!(lines.contains(&"requested_mode: AUTO") && lines.contains(&"resolved_mode: AUTO")); + assert_eq!(routes(&lines), ["indexed", "exact"]); + assert_eq!(lines.last(), Some(&"refusal: none")); + + // A partition named through a parameter is not known: every partition is listed. + let explained = explain(&searched(Some("ctx = $c"), "EXACT")); + let lines = lines_of(&explained); + let unbound = |line: &&str| { + line.starts_with("partitions: every partition (") && line.ends_with("is not bound)") + }; + assert!(lines.iter().any(unbound), "{explained}"); + assert_eq!(routes(&lines).len(), 2); + + // Any other statement is a later version's to explain. + match db.explain("SELECT id FROM docs") { + Err(error) => assert!( + error.to_string().ends_with( + "explaining this statement is not supported in this version of the library" + ), + "{error}" + ), + Ok(explained) => panic!("a plain SELECT was explained:\n{explained}"), + } +} + +#[test] +fn explain_says_that_auto_below_the_columns_size_resolves_to_exact() { + // Fifty vectors in a column whose size is the 100,000 of no AUTO_INDEX_AT. + let (_, db) = support::built("explain-auto-below-its-size", 16, "", 50); + let explained = db.explain(&searched(None, "AUTO")).expect("explain"); + let lines = lines_of(&explained); + for wanted in [ + "requested_mode: AUTO", + "resolved_mode: EXACT", + "effective_auto_index_at: 100000", + "aggregate_allowed_vectors: 50", + "refusal: none", + ] { + assert!( + lines.contains(&wanted), + "no line {wanted:?} in\n{explained}" + ); + } +} diff --git a/v2/contextdb/tests/vectors_memory.rs b/v2/contextdb/tests/vectors_memory.rs index 22c139fc..1db16408 100644 --- a/v2/contextdb/tests/vectors_memory.rs +++ b/v2/contextdb/tests/vectors_memory.rs @@ -792,3 +792,81 @@ fn the_memory_of_a_search_is_given_back_when_its_result_is_dropped() { tx.commit().expect("commit"); assert_eq!(support::int(&db, "SELECT COUNT(*) FROM docs"), 420); } + +#[test] +fn a_commit_with_no_room_for_its_vector_records_changes_nothing() { + const DIMENSION: usize = 16; + let (folder, db) = support::built("commit-no-room", DIMENSION, "AUTO_INDEX_AT 1000", 2000); + let rows = "SELECT id, ctx, tag, v FROM docs ORDER BY id"; + let untouched = (support::rows(&db, rows), db.committed_watermark()); + support::close(db); + let query = support::gathered(96_500, DIMENSION); + // Every search agrees with the rows that are there: EXACT returns the reference, and + // INDEXED rows with their true distances. + let searches_agree = |db: &Database, state: &str| { + let stored = support::stored(db, support::DOCS); + let wanted = stored.nearest(|_| true, &query, 10); + let sql = support::by_mode(None, "EXACT", 10); + support::assert_nearest(db, &sql, &query, &wanted, state); + let result = support::answer(db, &support::by_mode(None, "INDEXED", 10), &query); + support::assert_sound(&result, &stored, |_| true, &query, 10, state); + }; + + // One transaction under each limit, from one with no room at all upward: 300 sealed rows + // get another vector, which leaves a record for each slot they left and one for each + // vector that waits, and 300 rows are new. + let attempt = |db: &Database| -> contextdb::Result<()> { + let tx = db.transaction()?; + for id in 1..=300_i64 { + tx.execute( + "UPDATE docs SET v = $v WHERE id = $id", + ¶ms! { "id" => id, "v" => support::vector(96_600 + id, DIMENSION) }, + )?; + support::insert_doc( + &tx, + 5000 + id, + "a", + Some(support::gathered(5000 + id, DIMENSION)), + )?; + } + tx.commit() + }; + let mut refused = 0; + let mut committed = false; + let spares = (0..=32).map(|step| step * 64 * KIB).chain([55 * MIB]); + for spare in spares { + let limit = 9 * MIB + spare; + let db = open_limited(&folder, limit); + let before = used(&db); + let outcome = attempt(&db); + assert_eq!( + used(&db), + before, + "memory in use after a commit under {limit}" + ); + let db = reopen(db, &folder); + match outcome { + Ok(()) => { + assert_eq!(support::int(&db, "SELECT COUNT(*) FROM docs"), 2300); + searches_agree(&db, &format!("committed under {limit} bytes")); + committed = true; + break; + } + Err(Error::MemoryBudgetExceeded { .. }) => { + let state = format!("refused under {limit} bytes"); + assert!( + (support::rows(&db, rows), db.committed_watermark()) == untouched, + "{state}: the rows or the committed position changed" + ); + searches_agree(&db, &state); + refused += 1; + } + Err(other) => panic!("a commit under {limit} bytes gave {other:?}"), + } + support::close(db); + } + assert!( + refused > 0 && committed, + "{refused} limits refused the commit; committed at the last: {committed}" + ); +} diff --git a/v2/contextdb/tests/vectors_reduced.rs b/v2/contextdb/tests/vectors_reduced.rs new file mode 100644 index 00000000..2b2df785 --- /dev/null +++ b/v2/contextdb/tests/vectors_reduced.rs @@ -0,0 +1,129 @@ +//! A column declared `WITH (quantization = 'SQ8')` keeps each vector as one scale and one +//! whole number from -127 to 127 for each of its numbers. What a SELECT returns is that kept +//! vector, every search ranks by it, and such a column is grouped and searched as any other. + +mod support; + +use std::fs; + +use contextdb::Value; + +const DIMENSION: usize = 16; +const REDUCED: &str = "WITH (quantization = 'SQ8')"; + +/// What row `n` is written with: numbers up to three, so that the scale is not that of a +/// vector at length one. +fn written(n: i64, dimension: usize) -> Vec { + let scatter = support::vector(n, dimension); + scatter.iter().map(|number| 3.0 * number).collect() +} + +/// The vector the store returns for row `id` of `docs`. +fn kept(db: &contextdb::Database, id: i64) -> Vec { + let rows = support::rows(db, &format!("SELECT v FROM docs WHERE id = {id}")); + match rows.as_slice() { + [row] => match &row[0] { + Value::Vector(numbers) => numbers.clone(), + other => panic!("row {id} holds {other:?}"), + }, + rows => panic!("{} rows with the id {id}", rows.len()), + } +} + +/// The bytes of the data file of a closed store of 2,000 rows of 384 numbers each. +fn file_of(name: &str, clauses: &str) -> u64 { + let folder = support::folder(name); + let db = support::open(&folder); + support::declare_docs(&db, 384, clauses); + support::load_docs(&db, 1..=2000, "a", |n| written(n, 384)); + support::close(db); + fs::metadata(folder.join("data")) + .expect("the store's data file") + .len() +} + +#[test] +fn a_reduced_column_keeps_the_nearest_one_byte_vector() { + let folder = support::folder("reduced-keeps"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, REDUCED); + support::load_docs(&db, 1..=20, "a", |n| written(n, DIMENSION)); + let largest = |numbers: &[f32]| numbers.iter().fold(0.0_f32, |most, x| most.max(x.abs())); + for id in 1..=20 { + let (put, got) = (written(id, DIMENSION), kept(&db, id)); + assert_eq!(got.len(), DIMENSION); + // Every number is the scale times a whole code within 127, and the largest is 127. + let scale = largest(&got) / 127.0; + for (at, number) in got.iter().enumerate() { + let code = (number / scale).round(); + assert!(code.abs() <= 127.0, "row {id}, number {at}: code {code}"); + assert_eq!( + (scale * code).to_bits(), + number.to_bits(), + "row {id}, number {at} is {number:e}: the scale {scale:e} times {code}" + ); + } + let off = (largest(&got) - largest(&put)).abs(); + assert!( + off <= largest(&put) / 127.0, + "row {id}: the largest number was {:e} and is kept as {:e}", + largest(&put), + largest(&got) + ); + assert!(got != put, "row {id} came back as it was written"); + } + support::close(db); + + let (reduced, full) = ( + file_of("reduced-size", REDUCED), + file_of("reduced-size-f32", ""), + ); + assert!( + 2 * reduced < full, + "2,000 reduced rows take {reduced} bytes and the same rows at full precision {full}" + ); +} + +#[test] +fn exact_on_a_reduced_column_ranks_by_the_kept_vectors() { + let folder = support::folder("reduced-exact"); + let db = support::open(&folder); + support::declare_docs(&db, DIMENSION, REDUCED); + support::load_docs(&db, 1..=500, "a", |n| written(n, DIMENSION)); + // The reference is worked out from the vectors a SELECT returns, which are the kept ones. + let stored = support::stored(&db, support::DOCS); + assert!(kept(&db, 1) != written(1, DIMENSION)); + let check = |state: &str| { + for seed in [80_001, 80_002, 80_003] { + let query = support::vector(seed, DIMENSION); + let wanted = stored.nearest(|_| true, &query, 25); + let sql = support::by_mode(None, "EXACT", 25); + support::assert_nearest(&db, &sql, &query, &wanted, state); + } + }; + check("pending"); + assert_eq!(support::seal_all(&db), 500); + check("sealed"); +} + +#[test] +fn a_reduced_column_is_grouped_and_searched_as_any_other() { + let clauses = format!("{REDUCED} SEARCH_MODE INDEXED"); + let (_, db) = support::built("reduced-grouped", DIMENSION, &clauses, 0); + support::load_docs(&db, 1..=2000, "a", |n| support::gathered(n, DIMENSION)); + assert_eq!(support::upkeep(&db).0, 16, "the cycles of a build of 2,000"); + let index = support::shown(&db, "SHOW VECTOR_INDEXES").remove(0); + assert_eq!(support::word(&index, "quantization"), "SQ8"); + assert_eq!(support::whole(&index, "effective_auto_index_at"), 100_000); + assert_eq!( + support::field(&support::the_partition(&db), "lists"), + &Value::Int64(45) + ); + let stored = support::stored(&db, support::DOCS); + for j in 0..5 { + let query = support::gathered(81_000 + j, DIMENSION); + let result = support::answer(&db, &support::by_mode(None, "", 10), &query); + let state = format!("question {j} on the column's own mode"); + support::assert_sound(&result, &stored, |_| true, &query, 10, &state); + } +} diff --git a/v2/contextdb/tests/vectors_statements.rs b/v2/contextdb/tests/vectors_statements.rs index 8c78d0ea..1a0abe60 100644 --- a/v2/contextdb/tests/vectors_statements.rs +++ b/v2/contextdb/tests/vectors_statements.rs @@ -4,8 +4,7 @@ mod support; -use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; use contextdb::{ Database, Error, ReadSession, ReadSessionOptions, Value, VectorIndexRef, @@ -125,14 +124,11 @@ fn vector_search_mode_routes_each_request() { 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. + // USE VECTOR overrides the column's mode: each of the three answers on every + // table. No column here holds the 100,000 vectors from which AUTO uses a + // grouping, so AUTO and EXACT are one search: the same rows for the same + // counted work. And no partition has a grouping, so INDEXED searches it whole + // and returns the same rows. let exact = answer("USE VECTOR EXACT"); assert_eq!( answer("USE VECTOR AUTO"), @@ -140,12 +136,11 @@ fn vector_search_mode_routes_each_request() { "AUTO and EXACT on {table}" ); assert_eq!(answer("use vector exact"), exact, "{table}"); - refused("USE VECTOR INDEXED"); + answer("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}"); + let own = answer(""); + if table != "indexed" { + assert_eq!(own, exact, "no clause on {table}"); } } } @@ -176,34 +171,43 @@ fn vector_search_mode_routes_each_request() { assert_eq!(support::found(&found, "the table of one row"), [wanted]); } +/// The most pages holding rows that a refused INDEXED search may read. It reads small records +/// only, and a small record can lie in a page that also holds rows where the two kinds meet: +/// one page a level in a tree of at most four levels. +const SMALL_RECORD_READS: usize = 4; + #[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, "plain", "PARTITION_KEY (ctx) AUTO_INDEX_AT 1000"); + declare( + &db, + "indexed", + "PARTITION_KEY (ctx) SEARCH_MODE INDEXED AUTO_INDEX_AT 1000", + ); declare(&db, "empty", "SEARCH_MODE INDEXED"); - // Two thousand pending rows in each: a search that read rows would read them all. + declare(&db, "small", "PARTITION_KEY (ctx) AUTO_INDEX_AT 1000"); + let tx = db.transaction().expect("begin"); + fill(&tx, "small", 1..=300); + tx.commit().expect("commit"); + // Two thousand pending rows in each of three partitions: twice the vectors a column of + // these clauses may have outside a grouping, so the index is behind its writes. for table in ["plain", "indexed"] { let tx = db.transaction().expect("begin"); - fill(&tx, table, 1..=2000); + fill(&tx, table, 1..=6000); 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. + // Every way a search asks for an index that is behind: in the statement, by the column's + // declaration, and both. 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), @@ -212,8 +216,13 @@ fn indexed_search_is_refused_at_every_door() { "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)), + ]; + // A partition no row is in, and a table with no row, select nothing and refuse nothing. + let nothing = [ + search_sql("plain", Some("ctx = 'zz'"), "USE VECTOR INDEXED", 5), + search_sql("indexed", Some("ctx = 'zz'"), "", 5), + search_sql("empty", None, "", 5), + search_sql("empty", None, "USE VECTOR INDEXED", 5), ]; let before = ( db.committed_watermark(), @@ -227,8 +236,43 @@ fn indexed_search_is_refused_at_every_door() { refused_as_indexed(support::search(&tx, sql, &query), table, sql); refused_as_indexed(support::search(&reader, sql, &query), table, sql); } + // A partition with no grouping and fewer vectors outside one than twice AUTO_INDEX_AT is + // the index's one-group case: INDEXED answers it, and its answer and its work are EXACT's. + let small = support::stored(&db, "SELECT id, ctx, tag, v FROM small"); + type Scope = (Option<&'static str>, fn(&[Value]) -> bool); + let scopes: [Scope; 2] = [ + (None, |_| true), + (Some("ctx = 'c1'"), |row| row[1] == support::text("c1")), + ]; + for (filter, passes) in scopes { + let wanted = small.nearest(passes, &query, 5); + assert_eq!(wanted.len(), 5); + let exact = search_sql("small", filter, "USE VECTOR EXACT", 5); + let exact = support::assert_nearest(&db, &exact, &query, &wanted, "EXACT"); + let sql = search_sql("small", filter, "USE VECTOR INDEXED", 5); + let by_writer = support::assert_nearest(&db, &sql, &query, &wanted, "the writer"); + assert_eq!( + by_writer.distance_evaluations(), + exact.distance_evaluations(), + "{sql}\nthe one-group case does EXACT's distance evaluations" + ); + support::assert_nearest(&tx, &sql, &query, &wanted, "a transaction"); + support::assert_nearest(&reader, &sql, &query, &wanted, "a reader"); + } + for sql in ¬hing { + for outcome in [ + support::search(&db, sql, &query), + support::search(&tx, sql, &query), + support::search(&reader, sql, &query), + ] { + match outcome { + Ok(result) => assert!(result.rows().is_empty(), "{sql}"), + Err(error) => panic!("{sql}\nselects nothing and was refused: {error:?}"), + } + } + } // The refused transaction is as usable as before, and nothing changed. - fill(&tx, "indexed", 2001..=2001); + fill(&tx, "indexed", 6001..=6001); tx.rollback(); assert_eq!( ( @@ -245,47 +289,100 @@ fn indexed_search_is_refused_at_every_door() { 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)); + // storage door notes where every read of the data file falls. The pages that hold the + // table's rows are those a read of every row makes and a count of the table, which reads + // no row, does not. A refusal reads the small records of the partitions it selects (the + // record of each, and one for each of its pending rows, which is how far it is behind) + // and no row: of the pages of rows, at most `SMALL_RECORD_READS`. A search that answers + // reads the rows, so the comparison can tell. + type Reads = Vec<(u64, u64)>; + let reads_of = |sql: &str| -> (Reads, contextdb::Result) { + let reads = Arc::new(Mutex::new(Reads::new())); 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); + if let Call::Read(at, bytes) = call { + let read = (at, at + bytes.len() as u64); + reads.lock().expect("the reads").push(read); } } })); let session = ReadSession::open_with_options(&folder, options).expect("open for reading"); - let opened = reads.load(Ordering::SeqCst); + let opened = reads.lock().expect("the reads").len(); let outcome = support::search(&session, sql, &query); - (reads.load(Ordering::SeqCst) - opened, outcome) + let made = reads.lock().expect("the reads")[opened..].to_vec(); + (made, outcome) + }; + let touching = |reads: &Reads, pages: &Reads| -> usize { + let touches = + |read: &&(u64, u64)| pages.iter().any(|page| read.0 < page.1 && page.0 < read.1); + reads.iter().filter(touches).count() }; 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(), [[Value::Int64(6000)]]); assert_eq!(count.rows_examined(), 0); + let (every_read, every_row) = reads_of("SELECT id, ctx, tag, v FROM indexed"); + assert_eq!(every_row.expect("read every row").rows().len(), 6000); + let of_rows: Reads = every_read + .into_iter() + .filter(|read| touching(&vec![*read], &count_reads) == 0) + .collect(); for (table, sql) in asking.iter().filter(|(table, _)| *table == "indexed") { let (reads, outcome) = reads_of(sql); refused_as_indexed(outcome, table, sql); + let rows_read = touching(&reads, &of_rows); assert!( - reads <= count_reads, - "{sql}\nwas refused after {reads} storage reads; a count that reads no row made \ - {count_reads}" + rows_read <= SMALL_RECORD_READS, + "{sql}\nwas refused after {} storage reads, {rows_read} of them of the {} pages \ + that hold the table's rows", + reads.len(), + of_rows.len() ); } let (exact_reads, exact) = reads_of(&sql); assert_eq!(exact.expect("the exact search").rows().len(), 5); + let rows_read = touching(&exact_reads, &of_rows); assert!( - exact_reads > count_reads, - "an exact search of 2,000 pending rows made {exact_reads} storage reads, a count \ - {count_reads}" + rows_read > SMALL_RECORD_READS, + "an exact search of 6,000 pending rows made {} storage reads, {rows_read} of them of \ + the {} pages that hold the table's rows", + exact_reads.len(), + of_rows.len() ); } +#[test] +fn the_largest_auto_index_at_is_declared_and_served() { + let folder = support::folder("largest-auto-index-at"); + let db = support::open(&folder); + declare(&db, "largest", "AUTO_INDEX_AT 18446744073709551615"); + fill(&db, "largest", 1..=20); + let query = support::vector(91_500, DIMENSION); + // The column is below its size: every mode answers exactly, before a cycle and after. + let served = |state: &str| { + let stored = support::stored(&db, "SELECT id, ctx, tag, v FROM largest"); + let wanted = stored.nearest(|_| true, &query, 5); + assert_eq!(wanted.len(), 5); + for clause in [ + "USE VECTOR EXACT", + "", + "USE VECTOR AUTO", + "USE VECTOR INDEXED", + ] { + let sql = search_sql("largest", None, clause, 5); + support::assert_nearest(&db, &sql, &query, &wanted, state); + } + let partitions = support::shown(&db, "SHOW VECTOR_PARTITIONS FOR largest.v"); + assert_eq!(partitions.len(), 1, "{state}: partitions shown"); + }; + served("twenty pending rows"); + let report = db.run_maintenance_cycle().expect("a maintenance cycle"); + assert_eq!(report.vector.maintenance_vectors_done, 20); + served("twenty sealed rows"); +} + #[test] fn vector_query_errors_are_typed() { let folder = support::folder("vector-query-errors"); @@ -609,22 +706,14 @@ fn what_the_parser_reads_and_this_version_does_not_run_is_refused_when_planned() 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}" - ); - } + // A search with no LIMIT is refused as every statement of a later version is, by its + // name. + let sql = "SELECT id FROM pairs ORDER BY v <=> [1.0, 0.0, 0.0]"; + let sentence = refused_when_planned(sql); + assert!( + sentence.contains("vector search without LIMIT") && 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); @@ -891,13 +980,12 @@ fn vector_column_declarations_are_checked_stored_and_printed_back() { } other => panic!("{state}: a third partition of `capped` got {other:?}"), } - // And the mode that was printed is the mode in force. + // A column declared SEARCH_MODE INDEXED that no row has a vector in selects nothing + // and refuses nothing: its answer is complete and empty. 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:?}"), + Ok(result) => assert!(result.rows().is_empty(), "{state}: {sql}"), + Err(error) => panic!("{state}: {sql}\nselects nothing and was refused: {error:?}"), } }; holds(&db, "as declared"); @@ -1174,3 +1262,259 @@ fn a_distance_in_the_where_of_a_write_is_refused_on_a_table_with_no_row() { [[Value::Int64(1), Value::Float64(0.25)]] ); } + +/// A table `t` whose one VECTOR column `v` is declared with this text, as a statement. +fn with_vector(declared: &str) -> String { + format!("CREATE TABLE t (id INTEGER PRIMARY KEY, ctx TEXT NOT NULL, v {declared})") +} + +#[test] +fn the_index_clauses_are_read_checked_and_printed_back() { + let folder = support::folder("index-clauses"); + let db = support::open(&folder); + // Each column as it is declared, and the words its declaration prints back: a clause + // that was not written is not printed, F32 is the silent quantization, and + // MAX_PARTITIONS is printed for every column with a PARTITION_KEY. + let columns: [(&str, &str, &[&str]); 5] = [ + ( + "a", + "VECTOR(4) WITH (quantization = 'sq8')", + &["a", "VECTOR", "4", "WITH", "quantization", "=", "'SQ8'"], + ), + ( + "b", + "VECTOR(4) WITH (quantization = 'F32') AUTO_INDEX_AT 5000", + &["b", "VECTOR", "4", "AUTO_INDEX_AT", "5000"], + ), + ( + "c", + "VECTOR(4) PARTITION_KEY (ctx) SEARCH_MODE INDEXED \ + AUTO_INDEX_AT 18446744073709551615 CONSOLIDATION NONE", + &[ + "c", + "VECTOR", + "4", + "PARTITION_KEY", + "ctx", + "MAX_PARTITIONS", + "256", + "SEARCH_MODE", + "INDEXED", + "AUTO_INDEX_AT", + "18446744073709551615", + "CONSOLIDATION", + "NONE", + ], + ), + ( + "d", + "VECTOR(4) consolidation (tombstone_percent = 100)", + &[ + "d", + "VECTOR", + "4", + "CONSOLIDATION", + "TOMBSTONE_PERCENT", + "=", + "100", + ], + ), + ( + "e", + "VECTOR(4) CONSOLIDATION (CHANGE_PERCENT = 1, TOMBSTONE_PERCENT = 100)", + &[ + "e", + "VECTOR", + "4", + "CONSOLIDATION", + "CHANGE_PERCENT", + "=", + "1", + "TOMBSTONE_PERCENT", + "=", + "100", + ], + ), + ]; + let declared: Vec = columns + .iter() + .map(|(name, clauses, _)| format!("{name} {clauses}")) + .collect(); + support::run( + &db, + &format!( + "CREATE TABLE t (id INTEGER PRIMARY KEY, ctx TEXT NOT NULL, {})", + declared.join(", ") + ), + ); + let check = |printed: &str, state: &str| { + for (name, _, wanted) in columns { + assert_eq!(words(printed, name), wanted, "{state}: {name}"); + } + }; + let printed = db.declaration("t").expect("the declaration"); + check(&printed, "as declared"); + + // After reopen the same text, through the writer and through a reader; and the printed + // text declares the same table in a new store. + support::close(db); + let db = support::open(&folder); + assert_eq!(db.declaration("t").expect("the declaration"), printed); + let reader = ReadSession::open(&folder).expect("open the store for reading"); + assert_eq!(reader.declaration("t").expect("a reader's"), printed); + let again = support::open(&support::folder("index-clauses-again")); + for statement in printed.split_inclusive(";\n") { + support::run(&again, statement); + } + assert_eq!(again.declaration("t").expect("the declaration"), printed); +} + +#[test] +fn the_index_clauses_are_refused_by_kind() { + let folder = support::folder("index-clauses-refused"); + let db = support::open(&folder); + let fresh = db.committed_watermark(); + use VectorPolicyDeclarationIssue as Issue; + let quantization = Issue::InvalidMember { + member: "quantization", + }; + let refused: [(&str, Issue, &str); 9] = [ + ( + "VECTOR(4) WITH (quantization = 'SQ16')", + quantization, + "quantization has an invalid vector policy value", + ), + ( + "VECTOR(4) WITH (quantization = '')", + quantization, + "quantization has an invalid vector policy value", + ), + ( + "VECTOR(4) AUTO_INDEX_AT 0", + Issue::AutoIndexAtNotPositive, + "AUTO_INDEX_AT must be a positive integer", + ), + ( + "VECTOR(4) AUTO_INDEX_AT -1", + Issue::AutoIndexAtNotPositive, + "AUTO_INDEX_AT must be a positive integer", + ), + ( + "VECTOR(4) AUTO_INDEX_AT 18446744073709551616", + Issue::AutoIndexAtOutOfRange, + "AUTO_INDEX_AT is outside the persisted numeric range", + ), + ( + "VECTOR(4) CONSOLIDATION (CHANGE_PERCENT = 0)", + Issue::ConsolidationPercentOutOfRange, + "vector consolidation percentages must be integers from 1 through 100", + ), + ( + "VECTOR(4) CONSOLIDATION (CHANGE_PERCENT = 101)", + Issue::ConsolidationPercentOutOfRange, + "vector consolidation percentages must be integers from 1 through 100", + ), + ( + "VECTOR(4) CONSOLIDATION (TOMBSTONE_PERCENT = -5)", + Issue::ConsolidationPercentOutOfRange, + "vector consolidation percentages must be integers from 1 through 100", + ), + ( + "VECTOR(4) CONSOLIDATION (CHANGE_PERCENT = 20, TOMBSTONE_PERCENT = 0)", + Issue::ConsolidationPercentOutOfRange, + "vector consolidation percentages must be integers from 1 through 100", + ), + ]; + for (declared, wanted, sentence) in refused { + let sql = with_vector(declared); + let error = support::refusal(&db, &sql); + assert!(error.to_string().ends_with(sentence), "{sql}: {error}"); + match error { + Error::InvalidVectorPolicyDeclaration { index, issue } => { + assert_eq!(index, column("t", "v"), "{sql}"); + assert_eq!(issue, wanted, "{sql}"); + } + other => panic!("{sql}\nwas refused as {other:?}"), + } + } + // What a later version reads is refused by its name. + for (declared, name) in [ + ( + "VECTOR(4) WITH (quantization = 'SQ4')", + "WITH (quantization = 'SQ4')", + ), + ("VECTOR(4) IVF (LISTS = 4, PROBES = 2)", "IVF"), + ( + "VECTOR(4) SEARCH (TARGET_ACCURACY = 0.9, EXPECTED_VECTORS = 100)", + "SEARCH", + ), + ] { + let sql = with_vector(declared); + let sentence = support::refusal(&db, &sql).to_string(); + assert!( + sentence.ends_with(&format!("{name} {LATER_VERSION}")), + "{sql}\nwas refused in the sentence: {sentence}" + ); + } + // A CONSOLIDATION list that is empty, names a member twice or names another member, and + // a clause out of its place, are syntax errors. + for declared in [ + "VECTOR(4) CONSOLIDATION ()", + "VECTOR(4) CONSOLIDATION (CHANGE_PERCENT = 20, CHANGE_PERCENT = 30)", + "VECTOR(4) CONSOLIDATION (LISTS = 20)", + "VECTOR(4) AUTO_INDEX_AT 1000 SEARCH_MODE EXACT", + "VECTOR(4) AUTO_INDEX_AT 1000 WITH (quantization = 'SQ8')", + ] { + let sql = with_vector(declared); + match support::refusal(&db, &sql) { + Error::ParseError(_) => {} + other => panic!("{sql}\nwas refused as {other:?}"), + } + } + assert!(db.table_names().is_empty(), "a refused statement declared"); + assert_eq!(db.committed_watermark(), fresh); +} + +#[test] +fn the_inspection_statements_are_read() { + let folder = support::folder("inspection-statements"); + let db = support::open(&folder); + declare(&db, "docs", ""); + let reader = ReadSession::open(&folder).expect("open the store for reading"); + // Every form, in any letter case, with the rows it returns from a store of one column + // with no key and no vector: one index and its one partition. + let forms = [ + ("SHOW VECTOR_INDEXES", 1), + ("show vector_indexes", 1), + ("SHOW VECTOR_PARTITIONS", 1), + ("SHOW VECTOR_PARTITIONS FOR docs.v", 1), + ("show vector_partitions for docs.v limit 5", 1), + ("SHOW VECTOR_PARTITIONS LIMIT 0", 0), + ("SHOW VECTOR_PARTITIONS LIMIT 1 OFFSET 0", 1), + ("SHOW VECTOR_PARTITIONS FOR docs.v LIMIT 1 OFFSET 1", 0), + ]; + let before = db.committed_watermark(); + let tx = db.transaction().expect("begin"); + for (sql, rows) in forms { + assert_eq!(support::rows(&db, sql).len(), rows, "{sql}"); + assert_eq!( + support::rows(&tx, sql).len(), + rows, + "{sql}, in a transaction" + ); + assert_eq!(support::rows(&reader, sql).len(), rows, "{sql}, a reader"); + } + tx.rollback(); + assert_eq!(db.committed_watermark(), before); + // OFFSET stands only after LIMIT, and FOR names a table and a column. + for sql in [ + "SHOW VECTOR_PARTITIONS OFFSET 1", + "SHOW VECTOR_PARTITIONS FOR docs", + "SHOW VECTOR_INDEXES FOR docs.v", + ] { + match support::refusal(&db, sql) { + Error::ParseError(_) => {} + other => panic!("{sql}\nwas refused as {other:?}"), + } + } +} diff --git a/v2/contextdb/tests/vectors_upkeep.rs b/v2/contextdb/tests/vectors_upkeep.rs index 17a1d4fc..638714eb 100644 --- a/v2/contextdb/tests/vectors_upkeep.rs +++ b/v2/contextdb/tests/vectors_upkeep.rs @@ -9,7 +9,8 @@ use std::fs; use std::sync::atomic::Ordering; use contextdb::{ - Database, MaintenanceReport, QueryResult, ReadSession, ReadSessionOptions, params, + Database, DatabaseOpenOptions, MaintenanceReport, QueryResult, ReadSession, ReadSessionOptions, + Value, params, }; use support::{Child, Run}; @@ -31,7 +32,12 @@ 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)"); + // The column declares a size it never reaches, so no partition here is ever grouped. + support::declare_docs( + &db, + DIMENSION, + "PARTITION_KEY (ctx) AUTO_INDEX_AT 1000000000", + ); // 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) { @@ -387,7 +393,7 @@ fn a_small_fixed_store_stays_within_its_recorded_size() { let db = support::open(&folder); support::run( &db, - "CREATE TABLE notes (id INTEGER PRIMARY KEY, body TEXT NOT NULL, embedding VECTOR(384))", + "CREATE TABLE notes (id INTEGER PRIMARY KEY, body TEXT NOT NULL, embedding VECTOR(384) AUTO_INDEX_AT 1000000000)", ); // 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. @@ -421,3 +427,337 @@ fn a_small_fixed_store_stays_within_its_recorded_size() { bytes, and {PUT_IN} bytes were put in" ); } + +// ---- groupings ---- + +/// The dimension of the tests of a grouping, and the clause under which a partition of 2,000 +/// rows is cut into groups. +const GROUPED: usize = 16; +const AT_1000: &str = "AUTO_INDEX_AT 1000"; + +fn gathered(n: i64) -> Vec { + support::gathered(n, GROUPED) +} + +/// A partition's `maintenance_state` and `maintenance_reason`. +fn state_of(partition: &support::Shown) -> (String, String) { + ( + support::word(partition, "maintenance_state"), + support::word(partition, "maintenance_reason"), + ) +} + +/// A partition's `lists`: the groups of the grouping in force, `None` with none. +fn lists_of(partition: &support::Shown) -> Option { + match support::field(partition, "lists") { + Value::Null => None, + _ => Some(support::whole(partition, "lists")), + } +} + +/// What a build reports of itself in a partition's row: the slots it reads, has packed and +/// has left. +fn progress(partition: &support::Shown) -> [i64; 3] { + ["total", "done", "remaining"] + .map(|part| support::whole(partition, &format!("maintenance_vectors_{part}"))) +} + +#[test] +fn a_grouping_is_built_in_bounded_cycles_and_switched_in_one() { + let folder = support::folder("grouping-in-cycles"); + let db = support::open(&folder); + // 2,000 rows are this column's size, and half of the 4,000 vectors a partition of it may + // have outside a grouping and still be searched INDEXED: every search below is served. + support::declare_docs(&db, GROUPED, "AUTO_INDEX_AT 2000"); + support::load_docs(&db, 1..=2000, "a", gathered); + let query = gathered(60_000); + let late = support::vector(60_001, GROUPED); + let late_sql = support::by_mode(None, "INDEXED", 1); + let found_late = |state: &str| { + let stored = support::stored(&db, support::DOCS); + let result = support::answer(&db, &late_sql, &late); + let answer = support::assert_sound(&result, &stored, |_| true, &late, 1, state); + assert_eq!(answer[0].key, Value::Int64(2001), "{state}"); + }; + + for number in 1..=17_u64 { + if number == 6 { + // A write between two cycles commits and is found. + support::insert_doc(&db, 2001, "a", Some(late.clone())).expect("a write in a build"); + found_late("written between two cycles"); + } + let before = db.committed_watermark().0; + let report = support::cycle(&db); + let commits = db.committed_watermark().0 - before; + // The vectors the cycle reports done and left, the partitions left, and its commits: + // the first cycle seals and begins, ten find centres, one gives every slot its group, + // one packs, two measure, one switches, and the seventeenth has nothing to do. The + // sixth also seals the row written before it. + let wanted = match number { + 1 => (2000, 2000, 1, 2), + 6 => (1, 2000, 1, 2), + 13 => (2000, 0, 1, 1), + 14 | 15 => (0, 0, 1, 1), + 16 => (0, 0, 0, 1), + 17 => (0, 0, 0, 0), + _ => (0, 2000, 1, 1), + }; + assert_eq!( + ( + report.maintenance_vectors_done, + report.maintenance_vectors_remaining, + report.remaining_partitions, + commits + ), + wanted, + "cycle {number}" + ); + let partition = support::the_partition(&db); + let probes = support::field(&partition, "probes").clone(); + if number < 16 { + assert_eq!(state_of(&partition).0, "building", "after cycle {number}"); + assert_eq!((lists_of(&partition), probes), (None, Value::Null)); + let packed = if number < 13 { 0 } else { 2000 }; + assert_eq!(progress(&partition), [2000, packed, 2000 - packed]); + // Until the switch every mode returns what EXACT returns. + let wanted = support::stored(&db, support::DOCS).nearest(|_| true, &query, 10); + for mode in ["EXACT", "AUTO", "INDEXED"] { + let sql = support::by_mode(None, mode, 10); + let state = format!("{mode} after cycle {number}"); + support::assert_nearest(&db, &sql, &query, &wanted, &state); + } + } else { + assert_eq!(state_of(&partition).0, "idle", "after cycle {number}"); + assert_eq!(lists_of(&partition), Some(45), "round(sqrt(2000))"); + assert!( + matches!(probes, Value::Int64(1..=45)), + "probes after the switch: {probes:?}" + ); + assert_eq!(progress(&partition), [0, 0, 0]); + } + } + found_late("after the switch"); +} + +/// The bytes inspection reports for the index of the one partition. +fn index_bytes(db: &Database) -> i64 { + support::whole(&support::the_partition(db), "durable_index_bytes") +} + +#[test] +fn a_grouping_keeps_a_sketch_of_every_slot_and_measures_its_width() { + let (_, db) = support::built("grouping-index-bytes", GROUPED, AT_1000, 2000); + let partition = support::the_partition(&db); + assert_eq!(lists_of(&partition), Some(45)); + // 2,000 slots of 24 bytes with keys of 9 are one block. The index is 45 centres of 16 + // numbers, the record of 45 groups, the block's 2,000 sketches of 4 bytes behind their + // count, and the mean. + assert_eq!( + support::whole(&partition, "durable_vector_bytes"), + 2000 * 24 + ); + assert_eq!( + index_bytes(&db), + 45 * 4 * 16 + (12 + 8 * 46) + (4 + 2000 * 4) + 4 * 16 + ); +} + +#[test] +fn sealing_after_a_grouping_keeps_a_sketch_of_every_slot() { + let (_, db) = support::built("sealing-after-grouping", GROUPED, AT_1000, 2000); + let before = index_bytes(&db); + support::load_docs(&db, 2001..=2300, "a", gathered); + assert_eq!(support::cycle(&db).maintenance_vectors_done, 300); + // 300 sketches, and the count of the one tail block they are in. + assert_eq!(index_bytes(&db) - before, 300 * 4 + 4); + let stored = support::stored(&db, support::DOCS); + let sql = support::by_mode(None, "INDEXED", 1); + for id in 2001..=2300 { + let query = gathered(id); + let result = support::answer(&db, &sql, &query); + let state = format!("row {id}, sealed after the grouping"); + let answer = support::assert_sound(&result, &stored, |_| true, &query, 1, &state); + assert_eq!(answer[0].key, Value::Int64(id), "{state}"); + } +} + +#[test] +fn vectors_sealed_while_a_first_build_starts_are_found_after_the_switch() { + // Vectors of 64 numbers, of which a block takes 3,083: fewer than one cycle seals. + const WIDE: usize = 64; + let wide = |n: i64| support::gathered(n, WIDE); + let folder = support::folder("sealed-while-first-build-starts"); + let db = support::open(&folder); + // The grouping is never built again, so the vectors sealed after it stay in its tail. + support::declare_docs(&db, WIDE, "AUTO_INDEX_AT 1000 CONSOLIDATION NONE"); + support::load_docs(&db, 1..=2000, "a", wide); + // The first cycle seals and begins the partition's first build. The 3,500 rows written + // then are sealed by the second, before the build's first round: a whole block of them + // and the start of the next. + assert_eq!(support::cycle(&db).maintenance_vectors_done, 2000); + assert_eq!(state_of(&support::the_partition(&db)).0, "building"); + support::load_docs(&db, 2001..=5500, "a", wide); + assert_eq!(support::cycle(&db).maintenance_vectors_done, 3500); + support::upkeep(&db); + let partition = support::the_partition(&db); + assert_eq!(state_of(&partition).0, "idle"); + assert_eq!(lists_of(&partition), Some(45)); + + // Every search under INDEXED answers, and a row sealed in those cycles is the first row + // of a search for its own vector, from the block that was filled and from the one begun. + let stored = support::stored(&db, support::DOCS); + let sql = support::by_mode(None, "INDEXED", 1); + for id in (2001..=5500).step_by(233) { + let query = wide(id); + let result = support::answer(&db, &sql, &query); + let state = format!("row {id}, sealed while the first build was starting"); + let answer = support::assert_sound(&result, &stored, |_| true, &query, 1, &state); + assert_eq!(answer[0].key, Value::Int64(id), "{state}"); + } + let query = wide(60_002); + let result = support::answer(&db, &support::by_mode(None, "INDEXED", 10), &query); + support::assert_sound(&result, &stored, |_| true, &query, 10, "ten rows of 5,500"); +} + +#[test] +fn a_grouping_holds_every_live_vector_once_and_no_stale_slot() { + let clauses = "PARTITION_KEY (ctx) AUTO_INDEX_AT 1000"; + let (_, db) = support::built("grouping-no-stale", GROUPED, clauses, 1000); + support::load_docs(&db, 1001..=2000, "b", gathered); + support::upkeep(&db); + // 150 rows of a get another vector and 150 rows of b are deleted: each leaves a slot. + let tx = db.transaction().expect("begin"); + for id in 1..=150_i64 { + tx.execute( + "UPDATE docs SET v = $v WHERE id = $id", + ¶ms! { "id" => id, "v" => support::vector(61_000 + id, GROUPED) }, + ) + .expect("change a vector"); + } + tx.commit().expect("commit"); + support::run(&db, "DELETE FROM docs WHERE id > 1850"); + let stale = |db: &Database| -> Vec { + support::shown(db, "SHOW VECTOR_PARTITIONS") + .iter() + .map(|partition| support::whole(partition, "tombstones")) + .collect() + }; + assert_eq!(stale(&db), [150, 150]); + + support::upkeep(&db); + assert_eq!(stale(&db), [0, 0], "after the regrouping"); + let stored = support::stored(&db, support::DOCS); + let query = gathered(60_002); + for (ctx, rows) in [("a", 1000), ("b", 850)] { + let wanted = stored.nearest(|row| row[1] == support::text(ctx), &query, 5000); + assert_eq!(wanted.len(), rows); + let sql = support::by_mode(Some(&format!("ctx = '{ctx}'")), "EXACT", 5000); + support::assert_nearest(&db, &sql, &query, &wanted, ctx); + } +} + +#[test] +fn a_partition_is_grouped_again_only_when_its_changes_reach_the_declared_share() { + let state = |db: &Database| state_of(&support::the_partition(db)); + let lists = |db: &Database| lists_of(&support::the_partition(db)); + let named = |state: &str, reason: &str| (state.to_owned(), reason.to_owned()); + + // New vectors: 20 in 100 of the grouping's 2,000 slots when nothing is declared, and the + // declared 50. + for (name, clause, share) in [ + ("regroup-default", "", 400_i64), + ( + "regroup-declared", + " CONSOLIDATION (CHANGE_PERCENT = 50)", + 1000, + ), + ] { + let clauses = format!("{AT_1000}{clause}"); + let (_, db) = support::built(name, GROUPED, &clauses, 2000); + support::load_docs(&db, 2001..=1999 + share, "a", gathered); + support::upkeep(&db); + assert_eq!(state(&db), named("idle", "none"), "{name}: one below"); + assert_eq!(lists(&db), Some(45), "{name}: one below"); + support::load_docs(&db, 2000 + share..=2000 + share, "a", gathered); + support::cycle(&db); + assert_eq!(state(&db), named("building", "changes"), "{name}"); + support::upkeep(&db); + let groups = ((2000 + share) as f64).sqrt().round() as i64; + assert_eq!(lists(&db), Some(groups), "{name}: grouped again"); + } + + // Slots rows have left: 10 in 100. + let (_, db) = support::built("regroup-tombstones", GROUPED, AT_1000, 2000); + support::run(&db, "DELETE FROM docs WHERE id < 200"); + support::upkeep(&db); + assert_eq!(state(&db), named("idle", "none"), "199 slots left behind"); + support::run(&db, "DELETE FROM docs WHERE id = 200"); + support::cycle(&db); + assert_eq!(state(&db), named("building", "tombstones")); + + // Under CONSOLIDATION NONE the first grouping stays, whatever changes. + let clauses = format!("{AT_1000} CONSOLIDATION NONE"); + let (_, db) = support::built("regroup-none", GROUPED, &clauses, 2000); + support::load_docs(&db, 2001..=3000, "a", gathered); + support::run(&db, "DELETE FROM docs WHERE id <= 500"); + support::upkeep(&db); + assert_eq!(state(&db), named("idle", "none")); + assert_eq!(lists(&db), Some(45)); + let partition = support::the_partition(&db); + assert_eq!(support::whole(&partition, "tombstones"), 500); +} + +#[test] +fn a_cycle_under_a_tight_limit_records_the_stall_and_the_old_grouping_serves() { + // A round over 49 centres of 256 numbers takes 150,528 bytes for its sums and centres; + // the limit leaves 98,304 over what an open sets aside. + const WIDE: usize = 256; + const LIMIT: usize = (9 << 20) + (96 << 10); + let (folder, db) = support::built("stalled-build", WIDE, AT_1000, 2000); + support::load_docs(&db, 2001..=2400, "a", |n| support::gathered(n, WIDE)); + // This cycle seals the 400 and begins the regrouping they make due. + assert_eq!(support::cycle(&db).maintenance_vectors_done, 400); + let stored = support::stored(&db, support::DOCS); + support::close(db); + + let mut options = DatabaseOpenOptions::default(); + options.memory_limit = Some(LIMIT); + let db = support::open_with(&folder, options).expect("open under the limit"); + let report = match db.run_maintenance_cycle() { + Ok(report) => report.vector, + Err(error) => panic!("a cycle whose step does not fit was refused: {error:?}"), + }; + assert_eq!( + (report.stalled_partitions, report.remaining_partitions), + (1, 1) + ); + let partition = support::the_partition(&db); + assert_eq!( + state_of(&partition), + ("stalled".to_owned(), "memory_limit".to_owned()) + ); + assert_eq!( + support::word(&partition, "recovery_action"), + "raise the memory limit" + ); + let index = support::shown(&db, "SHOW VECTOR_INDEXES").remove(0); + assert_eq!(support::word(&index, "maintenance_state"), "stalled"); + assert_eq!(support::whole(&index, "stalled_partitions"), 1); + assert_eq!( + support::word(&index, "recovery_action"), + "raise the memory limit" + ); + + // The grouping in force still serves. + assert_eq!(lists_of(&partition), Some(45)); + let query = support::gathered(60_003, WIDE); + let result = support::answer(&db, &support::by_mode(None, "INDEXED", 10), &query); + support::assert_sound(&result, &stored, |_| true, &query, 10, "while stalled"); + + support::run(&db, "SET MEMORY_LIMIT 64M"); + support::upkeep(&db); + let partition = support::the_partition(&db); + assert_eq!(state_of(&partition), ("idle".to_owned(), "none".to_owned())); + assert_eq!(lists_of(&partition), Some(49), "round(sqrt(2400))"); + assert_eq!(support::cycle(&db).stalled_partitions, 0); +} diff --git a/v2/numbers.sh b/v2/numbers.sh index ce54d13c..6360be0d 100755 --- a/v2/numbers.sh +++ b/v2/numbers.sh @@ -10,17 +10,24 @@ # 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 +# Four 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 +# 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; +# PASSAGES_LARGER=: and a third, larger still # # 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. +# +# Every time of a similarity search is taken by one thread, in a read session opened under the +# memory limit that `index_reader_limit` prints unless the line's name says another, and each +# group of such lines opens with a line ending in `_session` that says both. NUMBERS_LARGE=1 +# adds five million generated sightings, which takes the better part of an hour and is run +# once for a landing. set -u cd "$(dirname "$0")/.." || exit 1 scratch=$(mktemp -d) || exit 1 @@ -56,6 +63,15 @@ fi # 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' '__') -" +# The first pass of an indexed search counts bits, with the processor's bit-count instruction +# where the processor says it has one: this line says which arm of that loop the times are of. +if grep -m 1 '^flags' /proc/cpuinfo 2>/dev/null | grep -qw popcnt; then + echo "find_counting popcnt -" +else + echo "find_counting plain -" +fi +echo "index_reader_limit $((1 << 30)) bytes" +echo "index_threads 1 threads" example() { cargo run --quiet --release --locked --manifest-path v2/Cargo.toml -p contextdb --example "$@" } @@ -77,16 +93,39 @@ 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 -" +# Similarity search under each mode on generated vectors: a million of the index pilot's, then +# a million sightings; and five million sightings when asked. +example measure -- index "$scratch/index" || echo "index_measurements unavailable -" +rm -rf -- "$scratch/index" +if [ "${NUMBERS_LARGE:-}" = 1 ]; then + example measure -- sightings 5000000 "$scratch/index" || + echo "sightings_5000000_measurements unavailable -" + rm -rf -- "$scratch/index" +fi + +# Search passages: load, seal, build the index, search through SQL under each mode, and the +# answers beside the recorded ones. First 10,000 passages of one partition of the second +# dump; then each dump whole, the second and the third also with a tail written after the +# build, the first and the second also kept as reduced vectors. +if [ -n "${PASSAGES_LARGE:-}" ]; then + example passages -- single "${PASSAGES_LARGE%%:*}" "$scratch/passages" || + echo "single_measurements unavailable -" rm -rf -- "$scratch/passages" -done -[ "$given" = 1 ] || echo "passages_measurements unavailable -" +else + echo "single_measurements unavailable -" +fi +passages() { + if [ -n "$2" ]; then + example passages -- "${2%%:*}" "${2#*:}" "$scratch/passages" "$3" || + echo "$1 unavailable -" + rm -rf -- "$scratch/passages" + else + echo "$1 unavailable -" + fi +} +passages passages_measurements "${PASSAGES:-}" sq8 +passages passages_large_measurements "${PASSAGES_LARGE:-}" sq8,tail +passages passages_larger_measurements "${PASSAGES_LARGER:-}" tail # From a fresh clone to a green `cargo test`: nothing built before. start=$SECONDS From 0f254add699224e4fcb9566200c4a0e1f3a701ea Mon Sep 17 00:00:00 2001 From: Lucid Programmer Date: Thu, 12 Nov 2026 00:22:06 +0530 Subject: [PATCH 2/2] contextdb: keep crash test child input open until the child is killed and reaped --- v2/contextdb/tests/support/mod.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/v2/contextdb/tests/support/mod.rs b/v2/contextdb/tests/support/mod.rs index 95b29211..9d6f6f38 100644 --- a/v2/contextdb/tests/support/mod.rs +++ b/v2/contextdb/tests/support/mod.rs @@ -1175,9 +1175,16 @@ impl Child { /// Kills the child where it stands. It must still be running: a child that is to be killed /// has announced a point and is waiting there. + /// + /// The child's standard input stays open until the child is gone. Waiting for a process + /// closes the input this side holds for it, and a kill has only been sent when it returns: + /// a child that read the end of its input before the kill reached it would leave the + /// point by itself. pub fn kill(&mut self) { + let input = self.process.stdin.take(); self.process.kill().expect("kill the child"); let status = self.process.wait().expect("wait for the child"); + drop(input); if status.signal().is_none() { self.fail("had already ended by itself when it was to be killed"); }