perf(knn): reduce memory for batch flat vector search by LeoReeYang · Pull Request #6950 · lance-format/lance

LeoReeYang · 2026-05-26T22:09:55Z

Summary

Closes #6940.

Batch flat KNN previously kept a full scan RecordBatch (including the vector column) for each top-k candidate, which could use on the order of m × k × d × 4 bytes when many query vectors and large k/d are used.

This PR changes batch flat KNN so that:

If the vector column is not in the final projection, the top-k heap stores only row identifiers and distances (not whole batches or vector data).
If the vector column is projected, each candidate keeps a single-row vector plus a shared scan batch without the vector column (for any other columns needed on output).
If the flat scan already reads non-vector columns (for example columns referenced by a filter), only the vector column is dropped from what we retain; other columns can still be assembled in the KNN node without keeping vectors in the heap.

Behavior

Scanner::flat_knn sets retain_vector from whether the user projection includes the vector field.
KNNVectorDistanceExec batch mode maintains a per-query top-k heap across all scan batches and emits one result batch with query_index and _distance.
Indexed batch KNN is unchanged; this targets flat batch vector search only.

Memory (qualitative)

Projection / scan input	What top-k heap retains
No vector in projection; scan is `vec + _rowid`	Row id + distance only
Vector in projection	Shared batch without vector column + one row of vector per candidate
No vector in projection; scan includes other columns	Shared batch without vector column (no vector rows stored)

Test plan

cargo test -p lance --lib test_batch_knn_flat
cargo test -p lance --lib test_batch_knn_flat_omits_vector_without_projection
cargo test -p lance --lib test_batch_knn_flat_filter_keeps_non_vector_columns
cargo test -p lance --lib test_batch_knn_flat_respects_distance_range
cargo fmt --all and cargo clippy -p lance --lib -- -D warnings

Avoid retaining whole scan RecordBatches (and vectors) in batch flat KNN heaps when the final projection does not include the vector column. When vectors are requested, retain only a per-row copy.

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

LeoReeYang · 2026-05-26T22:31:58Z

@BubbleCal Could you review this PR when you have time? It addresses #6940

codecov · 2026-05-26T22:43:04Z

Codecov Report

❌ Patch coverage is 86.61088% with 96 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
rust/lance/src/io/exec/knn.rs	79.55%	68 Missing and 24 partials ⚠️
rust/lance/src/dataset/scanner.rs	98.50%	1 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!

BubbleCal · 2026-05-27T03:52:39Z

            input
        };
+        let retain_vector = if self.is_batch_nearest {
+            let vector_field_id = self.dataset.schema().field_id(q.column.as_str())?;


does this work with nested vector field? can we add a test for this?

BubbleCal · 2026-05-27T03:54:59Z

-        let input_schema = Arc::new(input_schema);
+
+        let stored_schema = if is_batch && !retain_vector {
+            Arc::new(Schema::new(vec![ROW_ID_FIELD.clone()]))


something like intput_schema.drop_column(vector_column) would be better, we may have more columns other than row ID here

only vector column will be dropped now

BubbleCal · 2026-05-27T03:55:25Z

+        column: &str,
+        row_index: u32,
+    ) -> DataFusionResult<ArrayRef> {
+        let vectors = batch.column_by_name(column).ok_or_else(|| {


just do batch[column] is fine

BubbleCal · 2026-05-27T03:56:29Z

+            ))
+        })?;
+        let indices = UInt32Array::from(vec![row_index]);
+        arrow::compute::take(vectors, &indices, None)


wonder whether there is an easier way to avoid constructing this indices, probably just directly copy the data out?

would introducing Array::slice be a better choice？

BubbleCal · 2026-05-27T03:57:17Z

                .clone();

+            let slim_batch = if retain_vector {
+                Some(Arc::new(


Arc seems unnecessary

slim batch will be pinned when necessary

BubbleCal · 2026-05-27T04:02:26Z

+}
+
+#[cfg(test)]
+mod batch_knn_memory_tests {


I guess we don't need these memory tests
probably only test that if the query doesn't select vector column, then the batch shouldn't contain it

BubbleCal · 2026-05-27T04:08:05Z

-                        row_index,
+                    let candidate = if retain_vector {
+                        let row_index = row_index as u32;
+                        let vector_row = Self::take_vector_row(&batch, &column, row_index)?;


Don't copy the vector unless it will be pushed into the candidate heap

there will be a would_enter_heap check ahead and a slice for vector when necessary

BubbleCal · 2026-05-27T04:34:05Z

-    row_id: u64,
-    batch: RecordBatch,
-    row_index: u32,
+enum BatchKnnCandidate {


keep it a struct and make only the diff fields into enum

minimize the diff and make a nested enum inside

Refactor heap candidates to struct + extra enum, defer vector and slim batch copies until rows enter top-k, and use slice/batched take for assembly. Add correctness tests asserting vector column is omitted when not projected.

For batch flat KNN, keep stored schema as input without the vector column instead of row-id-only, and carry slim batches when needed so non-vector columns can be assembled without retaining vectors. Add a filter-based batch KNN regression test to verify non-vector outputs remain correct.

Add regression tests for payload.vec batch KNN omit/retain behavior and use qualified batch column access when retaining nested vectors.

BubbleCal · 2026-05-28T09:42:55Z

+
+#[derive(Clone)]
+enum BatchKnnExtra {
+    RowIdOnly,


Can you add a test which selects row_id and row_addr but no vectors?
I think it can't work with the current implementation

Rebuild batch KNN partition statistics from output schema, remove vector columns via nested path-aware helpers, and switch projected vector row capture from slice to take/copy to avoid retaining full buffers.

Add a batch flat KNN regression test for projecting row_id and row_addr without vectors to ensure system columns are materialized correctly.

Use parse_field_path-based vector resolution for batch retained vectors and remove the nested projected fallback that cloned full batches, then add regressions for escaped nested vector projection and nested slim-batch vector removal.

perf(knn): reduce memory for batch flat vector search

1226b70

Avoid retaining whole scan RecordBatches (and vectors) in batch flat KNN heaps when the final projection does not include the vector column. When vectors are requested, retain only a per-row copy.

claude Bot reviewed May 26, 2026

View reviewed changes

LeoReeYang mentioned this pull request May 26, 2026

perf(knn): reduce memory for batch flat vector search #6948

Closed

3 tasks

github-actions Bot added the performance label May 26, 2026

BubbleCal requested changes May 27, 2026

View reviewed changes

LeoReeYang added 3 commits May 27, 2026 13:22

fix(knn): address batch flat KNN review feedback

2699f9b

Refactor heap candidates to struct + extra enum, defer vector and slim batch copies until rows enter top-k, and use slice/batched take for assembly. Add correctness tests asserting vector column is omitted when not projected.

test(knn): cover nested vector batch flat KNN projection

d69c68d

Add regression tests for payload.vec batch KNN omit/retain behavior and use qualified batch column access when retaining nested vectors.

LeoReeYang requested a review from BubbleCal May 27, 2026 22:53

BubbleCal reviewed May 28, 2026

View reviewed changes

LeoReeYang added 2 commits May 28, 2026 19:21

fix(knn): align batch stats and trim nested vector retention

806fe98

Rebuild batch KNN partition statistics from output schema, remove vector columns via nested path-aware helpers, and switch projected vector row capture from slice to take/copy to avoid retaining full buffers.

test(knn): cover row id and row addr batch output

de2e9f6

Add a batch flat KNN regression test for projecting row_id and row_addr without vectors to ensure system columns are materialized correctly.

LeoReeYang requested a review from BubbleCal May 28, 2026 12:13

Conversation

LeoReeYang commented May 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Behavior

Memory (qualitative)

Test plan

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

LeoReeYang commented May 26, 2026

Uh oh!

codecov Bot commented May 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

LeoReeYang commented May 26, 2026 •

edited

Loading

codecov Bot commented May 26, 2026 •

edited

Loading