Skip to content

Cover reused search-query batch rows with tests - #91

Open
jnadeau207-collab wants to merge 1 commit into
xai-org:mainfrom
jnadeau207-collab:fix/stale-search-query-embeddings
Open

Cover reused search-query batch rows with tests#91
jnadeau207-collab wants to merge 1 commit into
xai-org:mainfrom
jnadeau207-collab:fix/stale-search-query-embeddings

Conversation

@jnadeau207-collab

@jnadeau207-collab jnadeau207-collab commented Sep 2, 2026

Copy link
Copy Markdown

Context

The original version of this PR fixed a stale-search-query-embedding bug in
InputBuffer::new_with_candidates by returning a zero vector instead of an
empty one. That bug was real at the base commit this branch was cut from
(7ba7768), where RankingBatchPrep wrote the reused search-query array with
no prior zeroing.

Upstream 85ac72a landed independently and fixed it directly, by adding
sq_slice.fill(0.0) before the per-item write. That fix is strictly more
complete than the one proposed here: it also clears rows the per-item write
never reaches. The util.rs change is therefore redundant on main and has
been dropped.

What is left

sq_slice.fill(0.0) is currently unguarded by any test. Nothing fails if it is
deleted, and the rows it protects are exactly the ones that are easiest to
overlook:

  • items whose input_buffer is None, which the if let Some(..) skips
  • rows past .take(length_of_input) when a batch follows a larger one
  • items carrying an empty or wrong-width query, which the
    src.len() == search_query_embedding_dim guard skips

Change

Extract the search-query fill from RankingBatchPrep::prepare into
util::fill_candidate_search_query_embeddings so it is reachable from a unit
test without a Python interpreter. The move is mechanical and production
behavior is unchanged.

Five tests cover the rows above, including a query-then-queryless reuse of one
array. Each of them fails if sq_slice.fill(0.0) is removed.

InputBuffer gains #[derive(Default)] so a test can build one with only the
two fields the fill path reads.

@jnadeau207-collab

Copy link
Copy Markdown
Author

Contribution by Jesse Nadeau @villain_thropic on X

Extract the RankingBatchPrep search-query fill into
fill_candidate_search_query_embeddings and pin the rows the per-item
write never touches: items whose input_buffer is None, rows past
length_of_input in a reused array, and an empty or wrong-width query.

Production behavior is unchanged; each test fails if the
sq_slice.fill(0.0) guard is removed.
@jnadeau207-collab
jnadeau207-collab force-pushed the fix/stale-search-query-embeddings branch from de9f975 to 1d55fc9 Compare September 2, 2026 23:06
@jnadeau207-collab jnadeau207-collab changed the title Prevent stale search-query embeddings in reused ranking batches Cover reused search-query batch rows with tests Sep 2, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant