Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Build an In-Memory Vector Table and Match Its Index

In progress: This Rust course material is awaiting a deeper review from the author.

Chapter 1

Start from the two *-starter crates. Finish with ordinary Arrow tables, one explicitly attached vector index, and a conservative DataFusion optimizer rule.

Your first query uses the course’s small three-column table:

SELECT id, payload
FROM points
ORDER BY cosine_distance(embedding, [1.0, 0.0, 0.0])
LIMIT 3;

Without an index match, DataFusion scans the MemTable, computes every distance, and keeps the nearest three with a bounded sort:

SortExec: TopK(fetch=3), ...
  DataSourceExec: partitions=1, ...

That plan is exact for every valid query. A vector index can select candidates only when the SQL ordering refers to the same metric, literal, dimension, direction, and configured vector column. A match changes the leaf while leaving DataFusion’s final sort in place by default:

SortExec: TopK(fetch=3), ...
  VectorIndexScanExec: index=flat, metric=Cosine, query_dim=3, fetch=Some(3), ordered=false

Chapter 2 will put index=ivf_flat behind the same boundary.

Data Model

A vector index belongs to one field of an ordinary table. It does not own a special (id, payload, vector) row format.

The small VectorRow and vector_mem_table helper remain the first example because they make Arrow construction easy to inspect:

id         UInt64
payload    Utf8
embedding  FixedSizeList<Float32, dimension>

The public indexing surface is more general. Register any MemTable, then construct a VectorIndexAttachment with its table reference and selected vector-column name:

let attachment = VectorIndexAttachment::try_new(
    &context,
    "documents",
    &table,
    "text_embedding",
    Metric::Euclidean,
    IndexConfig::Flat,
)
.await?;
let context = with_vector_indexes(&context, vec![attachment]);

The rich Chapter 1 test table deliberately puts ordinary scalar fields around two vector fields:

doc_key         Utf8
tenant_id       UInt32
price           Float64
inventory       Int32
text_embedding  FixedSizeList<Float32, 3>  <- selected
image_embedding FixedSizeList<Float32, 3>
active          Boolean

Both vector columns have the same type and width, but their nearest-neighbor orders differ. A query ordered by text_embedding may use the attached index. The same query shape over image_embedding must remain on DataFusion’s exact scan and return the image-vector ranking. No field name or ordinal is inherently special; only the field selected by the attachment may use its index.

The attachment snapshots the registered MemTable batches. It copies only the selected vectors into the core Dataset and records a checked row location for each dataset ordinal:

index dataset ordinal -> snapshot RowId -> checked batch/row -> projected output

The source Arrow buffers remain shared with the ordinary MemTable. Scalar columns and the unselected vector column stay normal table data. User columns are never row identity.

DataFusion has no generic stable point-lookup API for arbitrary TableProvider implementations. This adapter is therefore intentionally limited to registered in-memory MemTable instances. A disk or distributed provider would need its own stable row locator and lookup implementation.

Build the Exact Path and Matcher in Rust

You will modify:

rust/vector-starter/core/src/dataset.rs
rust/vector-starter/datafusion/src/lib.rs

The starter exposes the same public API as the reference but leaves the Chapter 1 implementation points as TODOs. Metric math, index implementations, snapshot lookup scaffolding, examples, and tests are ready. Do not modify public APIs or tests while completing the exercises.

Correctness Requirements

A dataset must be nonempty, have a fixed nonzero dimension, and contain only finite f32 values. Reject invalid data before an index is built.

An attachment must resolve the exact registered MemTable instance and the configured field. The selected field must exist, be FixedSizeList<Float32>, have a positive width, and contain no null list or null element. Each source row must contribute exactly one dataset vector and one checked snapshot row location.

A different positive list width is a valid schema choice; the core dataset takes its dimension from the selected field. The SQL matcher later rejects a literal whose width differs from that dataset. A zero-width selected field is invalid at construction.

The optimizer may replace a scan only for one supported distance expression over the configured vector field, a literal query vector, a compatible metric and direction, a positive LIMIT, and a live source snapshot. Filters, multiple sort keys, non-literal vectors, another vector field, wrong metrics or directions, and invalid literals remain on DataFusion’s exact scan and sort.

Unless ordered output is explicitly enabled for the session, DataFusion retains the final bounded sort after the index selects candidates. Candidate order is not automatically SQL order.

Checkpoint 1: Validate the In-Memory Dataset

Implement the three TODOs in vector-starter/core/src/dataset.rs.

Dataset::try_new reads the first row to establish the dimension, rejects an empty dataset or zero-dimensional vector, then checks every row for equal length and finite components. Store the vectors as Arc<[Vec<f32>]>.

validate_for_metric rejects zero-norm stored rows for cosine distance. validate_query checks dimension, finiteness, and the same cosine boundary for a query. Use the existing VectorError variants.

cd rust
cargo test -p vector-core-starter --test indexes flat_search_is_deterministic_and_validates_queries
cargo test -p vector-core-starter --test indexes cosine_rejects_zero_norm_vectors

Checkpoint 2: Build the Introductory MemTable

Implement vector_mem_table in vector-starter/datafusion/src/lib.rs.

Build a Dataset from the VectorRow embeddings to validate their shared dimension. Create the three Arrow arrays in the same input order, assemble one RecordBatch, then return an ordinary MemTable.

FixedSizeListArray stores vector components in one flat Float32Array. For two three-dimensional rows, its child values are:

[x0, y0, z0, x1, y1, z1]
 `---row 0--' `---row 1--'

Use i32::try_from(dataset.dimension()) for Arrow’s list width.

Prediction: What breaks if the payload array is reordered while the embedding array keeps insertion order?

Checkpoint 3: Attach One Selected Vector Column

Implement VectorIndexAttachment::try_new.

  1. Resolve the table reference and prove the supplied Arc<MemTable> is the registered provider.
  2. Snapshot every partition and batch, requiring one shared schema.
  3. Resolve only the configured vector-column name.
  4. Validate its Arrow type, positive width, and non-null values.
  5. Copy those selected vectors into Dataset in batch/row order.
  6. Build the requested core index and record the corresponding checked row locations.

The rich-schema tests make the ownership rule observable: text and image vectors have identical shapes but different rankings.

cargo test -p vector-datafusion-starter --test sql rich_schema_matches_only_the_configured_vector_column
cargo test -p vector-datafusion-starter --test sql rich_schema_rejects_a_missing_selected_column
cargo test -p vector-datafusion-starter --test sql rich_schema_rejects_a_scalar_selected_column
cargo test -p vector-datafusion-starter --test sql rich_schema_rejects_a_zero_width_selected_column
cargo test -p vector-datafusion-starter --test sql rich_schema_rejects_a_null_selected_value

Checkpoint 4: Match and Rewrite One Safe Top-k

Implement match_vector_order and VectorIndexOptimizer::rewrite_sort.

The matcher accepts only:

  1. one physical sort expression;
  2. Euclidean array_distance/list_distance, cosine_distance, or dot inner_product/dot_product;
  3. ascending Euclidean/cosine or descending dot-product order;
  4. one vector Column and one literal;
  5. the exact configured vector-column name after projection;
  6. a finite literal with the index dataset’s dimension; and
  7. a nonzero cosine literal.

DataFusion widens the fixed-size Float32 list to List<Float64> for its distance functions. match_vector_column accepts exactly that planner-added cast, while scalar_vector admits only values that preserve their exact f32 representation.

The optimizer must also prove the physical MemorySourceConfig still matches the attached table, snapshot, schema, projection, and unambiguous live provider. On a match, construct VectorIndexScanExec; otherwise leave the plan unchanged.

cargo test -p vector-datafusion-starter --test sql compatible_top_k_uses_vector_index_scan_and_keeps_sort
cargo test -p vector-datafusion-starter --test sql unsafe_sort_shapes_are_not_lowered
cargo test -p vector-datafusion-starter --test sql filter_keeps_datafusion_exact_fallback
cargo test -p vector-datafusion-starter --test sql dot_product_requires_descending_order

Checkpoint 5: Search, Fetch, and Preserve ORDER BY

Implement VectorIndexScanExec::selected_rows and ExecutionPlan::with_fetch.

Search the selected index for at most fetch rows. Reject an index result that does not resolve to the snapshot. The supplied lookup scaffolding reconstructs the requested projection in index-result order.

For ordered=true, return the scan with its accepted ordering property. For the default ordered=false path, clear that property and wrap the scan in SortExec::new(ordering, scan).with_fetch(Some(k)). The index chooses candidates; DataFusion still owns SQL’s nearest-first result.

cargo test -p vector-datafusion-starter --test sql ordered_session_mode_allows_sort_elision
cargo test -p vector-datafusion-starter --test sqllogictest day1_table_and_optimizer_sql

The SQLLogicTest checks the simple table plus both rich-schema paths: selected text_embedding reaches VectorIndexScanExec, while image_embedding stays on DataSourceExec and returns its different ranking.

Chapter 1 Review

After the core tests, sql.rs, and the Chapter 1 SQLLogicTest pass, explain:

  • how the simple VectorRow helper becomes an ordinary MemTable;
  • why an attachment owns exactly one configured vector field;
  • how an index dataset ordinal resolves to a projected source row;
  • why the same-shaped image-vector query cannot use the text-vector index;
  • where DataFusion performs exact fallback and final ordering; and
  • how later approximate indexes reuse this boundary without weakening it.

IVFFlat, filtered pushdown, joins, DDL, persistence, and disk row lookup remain outside this chapter.

Your feedback is greatly appreciated. Join our Discord community.
Found an issue? Open an issue or pull request at github.com/skyzh/write-you-a-vector-db.
write-you-a-vector-db-book © 2024-2026 by Alex Chi Z is licensed under CC BY-NC-SA 4.0.