Skip to content

Index thinning - #226

Merged
fedorgrab merged 4 commits into
mainfrom
fg/issue-225
Aug 4, 2026
Merged

Index thinning#226
fedorgrab merged 4 commits into
mainfrom
fg/issue-225

Conversation

@fedorgrab

@fedorgrab fedorgrab commented Jul 16, 2026

Copy link
Copy Markdown
Collaborator

Add allow-list (thinning) support to create_vsindex

Lets you build a TileDB IVF_FLAT index over a subsample of an existing embedding set without regenerating embeddings. Previously create_vsindex indexed every cell in every
batch file; now an optional --allow-list-csv / allow_list_csv= (local path or gs://, single soma_joinid column with header) restricts indexing to the listed cells and
drops the rest.

Behavior

  • The filter is applied per batch in both phases — training (Phase 1) and streaming updates (Phase 2) — so the train/update split is preserved and every allow-listed cell
    is indexed exactly once. Training fills to --training-sample-size allowed cells; the boundary batch is appended whole and may slightly overshoot, matching the existing
    unfiltered path.
  • Phase 3 consolidation is now gated on vectors actually written rather than "remaining batches existed" — with thinning, remaining batches can filter down to zero rows.
    Empty chunks short-circuit instead of issuing zero-row update_batch calls.
  • An allow-list matching no input cell now fails fast with a ValueError instead of erroring deep inside index creation.
  • The allow-list is held as a pd.Index, not a Python set: isin() then reuses one pre-built C-level hash table instead of rebuilding a temporary object array on every
    per-batch call — material at millions of IDs across thousands of batches.
  • The script does not verify the allow-list is a genuine subsample of the embeddings; that stays the caller's responsibility.

Also in here

  • Extracted _pack_and_update() to de-duplicate the object-array packing + update_batch write.
  • Tests for CLI kwarg forwarding, _load_allowed_ids parsing, boundary overshoot, the empty-allow-list error, and a stubbed end-to-end check that the train and update ID
    sets are disjoint and together cover exactly the allow-list.
  • tests/unit/scripts/conftest.py: test-only tiledb / tiledb.vector_search stub so the module imports where the native tiledb-vector-search wheel is unavailable (e.g.
    linux/aarch64).
  • .gitignore entries for stray python / python3.1x artifacts; CHANGELOG entry under 1.8.5.

Closes #225

@fedorgrab fedorgrab self-assigned this Jul 16, 2026

@sjfleming sjfleming left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Some comments about the "carryover" concept and whether it can be eliminated. If it's necessary, that's okay. Just looking for a potential simplification

need = training_sample_size - total_training_rows
train_dfs.append(df.iloc[:need])
carry_df = df.iloc[need:]
total_training_rows += need

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

is need right, or is it len(df) - need ?

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

oh wait, it looks like it is need

Comment thread cellarium/cas_backend/scripts/create_vsindex.py Outdated
@sonarqubecloud

Copy link
Copy Markdown

@fedorgrab
fedorgrab requested a review from sjfleming July 29, 2026 18:45
@fedorgrab fedorgrab added the enhancement New feature or request label Jul 29, 2026

@sjfleming sjfleming left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Looks good to me!

@fedorgrab
fedorgrab merged commit cf910b5 into main Aug 4, 2026
3 checks passed
@fedorgrab
fedorgrab deleted the fg/issue-225 branch August 4, 2026 17:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Index thinning capability

2 participants