Skip to content

feat(knn): let a caller rebuild a batch KNN node with a new k - #9334

Open
hamersaw wants to merge 1 commit into
lance-format:mainfrom
hamersaw:feature/knn-batch-with-k
Open

hamersaw wants to merge 1 commit into
lance-format:mainfrom
hamersaw:feature/knn-batch-with-k

Conversation

@hamersaw

Copy link
Copy Markdown
Contributor

What

Adds KNNVectorDistanceExec::with_k(k), returning the node rebuilt with a different per-query k.

Why

A batch KNNVectorDistanceExec bounds every query's candidates by its own k and carries no enclosing top-k SortExec — the plan is just KNNVectorDistance(queries=N, k=K) over a scan. A downstream rewriter that wants to widen the candidate set has nothing to move, and it can't rebuild the node either: try_new_batch is pub(crate) and retain_vector is a private field.

k feeds only the execute-time cut (execute_batch's BatchKnnConfig), never the schema or plan properties, so with_k carries the rest of the node over unchanged. Zero k on a batch node is rejected, matching try_new_batch.

This is the brute-force counterpart to the ANNIvfBatchExec accessors added in #9262. Sophon's WAL union over-fetches the base arm so a block-list can drop superseded rows and still leave ~k live ones; on the indexed batch shape it rebuilds ANNIvfBatchExec, and on the unindexed batch shape it currently can't reach the cut at all.

Testing

test_batch_with_k_rebuilds_the_cut_and_keeps_the_schema covers the widened k, the preserved schema/query_count/is_batch, that the original node is left alone, and the zero-k rejection.

🤖 Generated with Claude Code

A batch `KNNVectorDistanceExec` bounds every query's candidates by its own
`k` and carries no enclosing top-k `SortExec`, so a plan rewriter that wants
to widen the candidate set has nothing to move — `try_new_batch` is
`pub(crate)` and `retain_vector` is private, so it can't rebuild the node
either.

`with_k` covers that: `k` feeds only the execute-time cut, never the schema
or plan properties, so everything else carries over unchanged.

Sophon's WAL union needs it to over-fetch the base arm on the unindexed
batch shape, the way it already does for `ANNIvfBatchExec`.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@github-actions github-actions Bot added the enhancement New feature or request label Sep 17, 2026

@lance-gatekeeper lance-gatekeeper Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Gate recommendation: approve.

This focused builder changes only the batch node per-query candidate bound while preserving its child, schema, distance bounds, vector-retention mode, and plan properties. It retains the constructor zero-k validation and resets execution metrics for the rebuilt node.

@lance-gatekeeper lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Sep 17, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request K-approved Latest Gatekeeper recommendation permits acceptance.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant