feat(index): support segmented ngram index#7244
Conversation
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
18b23df to
5420a54
Compare
|
Hi @BubbleCal ~ This is a ngram segmented index adoption. Would u have time to take a look? Really appreciate it! |
|
Hi @Xuanwo ~ This is a ngram segmented index adoption. Would u mind to take look? Thanks in advance! |
114669a to
f2cfee8
Compare
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe PR adds NGram segmented-index support across Python and Rust, preserves dataset-version and source-field metadata, implements compaction-aware merging and remapping, strengthens index compatibility and CreateIndex conflict handling, and adds end-to-end coverage for overlays, commits, queries, and deferred compaction. ChangesNGram segmented index lifecycle
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Python
participant DatasetIndex
participant NGramMerge
participant Commit
participant Query
Python->>DatasetIndex: create_index_uncommitted
DatasetIndex->>NGramMerge: merge_existing_index_segments
NGramMerge-->>DatasetIndex: merged NGram segment and dataset version
DatasetIndex->>Commit: commit_existing_index_segments
Commit-->>Python: committed index metadata
Python->>Query: contains(text, 'needle')
Query-->>Python: matching row count
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rust/lance/src/dataset/optimize/remapping.rs (1)
252-296: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAdd coverage for legacy remapped-index semantics.
test_ngram_remap_excludes_newer_overlay_fragmentsexercises the known-coverage branch, including overlay exclusion anddataset_versionbumping. Add a parallel legacy path that starts withfragment_bitmap: None, compacts/remaps with an olderdataset_versionthan newer overlays, and asserts that:
dataset_versionremains at the legacy index baseline, and- newer-overlaid fragments are still treated as part of the legacy index coverage/stale mask.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance/src/dataset/optimize/remapping.rs` around lines 252 - 296, Add a parallel legacy-coverage test alongside test_ngram_remap_excludes_newer_overlay_fragments, initializing the index with fragment_bitmap: None and an older dataset_version before compaction/remapping with newer overlays. Assert that the resulting dataset_version remains the legacy index baseline and that newer-overlaid fragments remain included in the legacy index coverage/stale mask.
🧹 Nitpick comments (3)
rust/lance/src/index/scalar_logical.rs (1)
1558-1567: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExplain why "compaction_guard" is required for this test to work.
Creating this unrelated, fully-committed NGram index right before
compact_files(..., defer_index_remap: true)is what causes the compaction to retain a fragment-reuse mapping at all — without any committed index needing deferred remap,compact_fileslikely has nothing to defer, so the mapping wouldn't exist and the latermerge_existing_index_segmentscall (which succeeds here) would instead fail like it does intest_ngram_segment_merge_rejects_retired_coverage_without_remap(line 1687-1713 below). This causal link isn't obvious from reading the test and is easy to accidentally break (e.g. if someone removes what looks like a redundant guard index). A short comment would help.As per coding guidelines, "Add only meaningful comments and tests; comments should explain non-obvious reasoning rather than restating code."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance/src/index/scalar_logical.rs` around lines 1558 - 1567, The test’s committed “compaction_guard” NGram index establishes the fragment-reuse mapping required by compact_files with defer_index_remap enabled, allowing the later merge_existing_index_segments call to succeed. Add a concise comment immediately before dataset.create_index explaining this dependency and warning that removing the seemingly unrelated index changes the compaction behavior.Source: Coding guidelines
rust/lance/src/index.rs (1)
1408-1460: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftBoolean-flag dispatch across 8 index types is fragile — consider an enum.
merge_existing_index_segmentsnow has 8all_Xflags (vector/inverted/bitmap/btree/fmindex/zonemap/label_list/ngram) feeding both the dispatchif/elsechain and the separatedataset_versionoverride list at line 1458. The override list deliberately excludesall_ngram(correctly —ngram::merge_segmentsalready sets the rebuild-awaredataset_versionitself, and overriding it here would silently revert a fresh rebuild's version back to the stale min source version), but that correctness is easy to lose the next time an index type is added, since nothing enforces that the override list and the dispatch list stay in sync.Consider deriving a single
SegmentKindenum fromindex_detailsonce, then using an exhaustivematchfor both the merge dispatch and the dataset_version handling, so adding a new index type is a compiler-enforced decision rather than an easily-missed list update.As per coding guidelines: "Use enums instead of magic numbers for format versions, variant types, and discriminators, and rely on exhaustive
match."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance/src/index.rs` around lines 1408 - 1460, Replace the independent all_X flags in merge_existing_index_segments with a single SegmentKind derived from each segment’s index_details, validating that all source segments share the same kind. Use an exhaustive match on SegmentKind for merge dispatch and dataset_version handling, preserving ngram’s special behavior by not overriding its merged version; ensure new index kinds require explicit compiler-enforced handling.Source: Coding guidelines
rust/lance-index/src/scalar/ngram.rs (1)
895-901: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
RowAddressinstead of raw bitwise fragment extraction.
(row_id >> 32) as u32manually decodes the fragment id from a row address. The repo convention is to route this throughRowAddressso address encoding stays centralized.♻️ Proposed change
fn old_filter_keeps(filter: &super::OldIndexDataFilter, row_id: u64) -> bool { match filter { super::OldIndexDataFilter::Fragments { to_keep, .. } => { - to_keep.contains((row_id >> 32) as u32) + to_keep.contains(RowAddress::from(row_id).fragment_id()) } super::OldIndexDataFilter::RowIds(valid) => valid.contains(row_id), } }Add the import at the top of the file:
use lance_core::utils::address::RowAddress;As per coding guidelines: "Use
RowAddressfromlance-core/src/utils/address.rsinstead of raw bitwise operations on row addresses."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance-index/src/scalar/ngram.rs` around lines 895 - 901, Update old_filter_keeps to use lance_core::utils::address::RowAddress for fragment extraction instead of manually shifting row_id, adding the RowAddress import and preserving the existing to_keep membership check.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@rust/lance/src/dataset/optimize/remapping.rs`:
- Around line 252-296: Add a parallel legacy-coverage test alongside
test_ngram_remap_excludes_newer_overlay_fragments, initializing the index with
fragment_bitmap: None and an older dataset_version before compaction/remapping
with newer overlays. Assert that the resulting dataset_version remains the
legacy index baseline and that newer-overlaid fragments remain included in the
legacy index coverage/stale mask.
---
Nitpick comments:
In `@rust/lance-index/src/scalar/ngram.rs`:
- Around line 895-901: Update old_filter_keeps to use
lance_core::utils::address::RowAddress for fragment extraction instead of
manually shifting row_id, adding the RowAddress import and preserving the
existing to_keep membership check.
In `@rust/lance/src/index.rs`:
- Around line 1408-1460: Replace the independent all_X flags in
merge_existing_index_segments with a single SegmentKind derived from each
segment’s index_details, validating that all source segments share the same
kind. Use an exhaustive match on SegmentKind for merge dispatch and
dataset_version handling, preserving ngram’s special behavior by not overriding
its merged version; ensure new index kinds require explicit compiler-enforced
handling.
In `@rust/lance/src/index/scalar_logical.rs`:
- Around line 1558-1567: The test’s committed “compaction_guard” NGram index
establishes the fragment-reuse mapping required by compact_files with
defer_index_remap enabled, allowing the later merge_existing_index_segments call
to succeed. Add a concise comment immediately before dataset.create_index
explaining this dependency and warning that removing the seemingly unrelated
index changes the compaction behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 8f4356ee-ac64-4714-85a5-610baa6ca13b
📒 Files selected for processing (17)
java/lance-jni/src/blocking_dataset.rspython/python/lance/dataset.pypython/python/tests/test_scalar_index.pypython/src/dataset.rspython/src/indices.rsrust/lance-index/src/scalar/ngram.rsrust/lance/src/dataset/index.rsrust/lance/src/dataset/optimize/remapping.rsrust/lance/src/dataset/tests/dataset_overlay_index_masking.rsrust/lance/src/index.rsrust/lance/src/index/api.rsrust/lance/src/index/append.rsrust/lance/src/index/create.rsrust/lance/src/index/scalar.rsrust/lance/src/index/scalar/ngram.rsrust/lance/src/index/scalar_logical.rsrust/lance/src/io/commit/conflict_resolver.rs
Xuanwo
left a comment
Xuanwo
left a comment
There was a problem hiding this comment.
commit_existing_index_segments can attach a staged segment to a different column snapshot and silently hide matching rows.
The segment preserves dataset_version, but IntoIndexSegment for IndexMetadata drops the source fields; commit resolves the current column and uses the coordinator's current manifest version as the transaction read_version. If the indexed column is dropped/recreated (the same issue applies to an indexed-field Merge or DataReplacement) after the segment was built and the coordinator has refreshed, those intervening transactions are never checked. The old NGram postings are then advertised as covering the current fragment, so the query has no flat fallback.
On f2cfee815, the public-API reproducer below builds the segment at V1, replaces text at V2/V3, and commits the V1 segment. contains(text, 'dog') returns 3 before the commit and 0 afterward.
Reproducer
use std::sync::Arc;
use arrow_array::{ArrayRef, Int32Array, RecordBatch, RecordBatchIterator, StringArray};
use arrow_schema::{DataType, Field, Schema};
use lance::dataset::NewColumnTransform;
use lance::index::{CreateIndexBuilder, DatasetIndexExt};
use lance::Dataset;
use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams};
use lance_index::IndexType;
#[tokio::main]
async fn main() -> lance::Result<()> {
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("text", DataType::Utf8, false),
]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![0, 1, 2])) as ArrayRef,
Arc::new(StringArray::from(vec!["cat-zero", "cat-one", "cat-two"])) as ArrayRef,
],
)?;
let reader = RecordBatchIterator::new(vec![Ok(batch)], schema);
let tmp = tempfile::tempdir()?;
let mut dataset = Dataset::write(reader, tmp.path().to_str().unwrap(), None).await?;
let params = ScalarIndexParams::for_builtin(BuiltinIndexType::NGram);
let segment = CreateIndexBuilder::new(
&mut dataset,
&["text"],
IndexType::NGram,
¶ms,
)
.name("text_ngram".to_string())
.fragments(vec![0])
.execute_uncommitted()
.await?;
let source_version = segment.dataset_version;
let source_field_id = segment.fields[0];
dataset.drop_columns(&["text"]).await?;
dataset
.add_columns(
NewColumnTransform::SqlExpressions(vec![(
"text".to_string(),
"concat('dog-', cast(id as varchar))".to_string(),
)]),
Some(vec!["id".to_string()]),
None,
)
.await?;
let replacement_version = dataset.version().version;
let replacement_field_id = dataset.schema().field("text").unwrap().id;
let before_commit = dataset
.count_rows(Some("contains(text, 'dog')".to_string()))
.await?;
dataset
.commit_existing_index_segments("text_ngram", "text", vec![segment])
.await?;
let committed = dataset.load_indices_by_name("text_ngram").await?;
let after_commit = dataset
.count_rows(Some("contains(text, 'dog')".to_string()))
.await?;
assert_eq!(source_version, 1);
assert_eq!(replacement_version, 3);
assert_ne!(source_field_id, replacement_field_id);
assert_eq!(committed[0].dataset_version, source_version);
assert_eq!(before_commit, 3);
assert_eq!(
after_commit, before_commit,
"committing an index built for the replaced column must not hide current matches"
);
Ok(())
}The final assertion fails with left: 0, right: 3.
The delayed commit needs to reject or rebuild a segment when its source field/snapshot no longer matches the current indexed data, rather than publishing its old coverage.
f2cfee8 to
e75ed76
Compare
There was a problem hiding this comment.
Note
Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.
🟡 Other comments (1)
rust/lance/src/index.rs-1469-1478 (1)
1469-1478: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick winStamp FMIndex merges with source-min
dataset_version.
fmindex::merge_segmentssetsdataset_version: dataset.manifest.version, so merged FMIndex segments lose the source staleness provenance that NGram preserves with its explicit source-min. Includeall_fmindexin the post-merge override or have FMIndex merge return a source-mindataset_versiondirectly.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance/src/index.rs` around lines 1469 - 1478, Update the post-merge dataset version override in the merge flow to include the all_fmindex condition, ensuring FMIndex merged segments receive source_dataset_version like NGram and other index types. Alternatively, change fmindex::merge_segments to return the source-min dataset_version directly, but preserve the existing behavior for all other index types.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Other comments:
In `@rust/lance/src/index.rs`:
- Around line 1469-1478: Update the post-merge dataset version override in the
merge flow to include the all_fmindex condition, ensuring FMIndex merged
segments receive source_dataset_version like NGram and other index types.
Alternatively, change fmindex::merge_segments to return the source-min
dataset_version directly, but preserve the existing behavior for all other index
types.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 22e5030e-9c60-429e-bad7-27d227235ef5
📒 Files selected for processing (17)
java/lance-jni/src/blocking_dataset.rspython/python/lance/dataset.pypython/python/tests/test_scalar_index.pypython/src/dataset.rspython/src/indices.rsrust/lance-index/src/scalar/ngram.rsrust/lance/src/dataset/index.rsrust/lance/src/dataset/optimize/remapping.rsrust/lance/src/dataset/tests/dataset_overlay_index_masking.rsrust/lance/src/index.rsrust/lance/src/index/api.rsrust/lance/src/index/append.rsrust/lance/src/index/create.rsrust/lance/src/index/scalar.rsrust/lance/src/index/scalar/ngram.rsrust/lance/src/index/scalar_logical.rsrust/lance/src/io/commit/conflict_resolver.rs
There was a problem hiding this comment.
🧹 Nitpick comments (1)
rust/lance/src/io/commit/conflict_resolver.rs (1)
40-46: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the non-obvious semantics of
index_build_version.The
dataset_version == 0sentinel and themax(dataset_version, transaction_read_version)fallback drive segment-source keying and the conflict-filter (< other_version) across the whole CreateIndex rebase path, but the reasoning isn't stated. A short doc comment explaining that0means "unversioned/legacy segment → treat as built at the transaction read version" and why the newer of the two versions is chosen would prevent future misuse.As per coding guidelines: "Add doc comments to magic constants, thresholds, and non-obvious transformation functions, explaining what the value represents and why it was chosen."
📝 Suggested doc comment
+/// Returns the dataset version a staged index segment was built against. +/// +/// A `dataset_version` of `0` denotes a legacy/unversioned segment, which is +/// treated as built at the transaction's read version. Otherwise we take the +/// newer of the segment's recorded version and the read version so segments +/// staged ahead of the read version are keyed to their true source version. fn index_build_version(index: &IndexMetadata, transaction_read_version: u64) -> u64 { if index.dataset_version == 0 { transaction_read_version } else { index.dataset_version.max(transaction_read_version) } }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance/src/io/commit/conflict_resolver.rs` around lines 40 - 46, Add a concise Rust doc comment directly above index_build_version explaining that dataset_version == 0 represents an unversioned or legacy segment treated as built at transaction_read_version, and that otherwise the newer of dataset_version and transaction_read_version is selected for segment-source keying and conflict filtering.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@rust/lance/src/io/commit/conflict_resolver.rs`:
- Around line 40-46: Add a concise Rust doc comment directly above
index_build_version explaining that dataset_version == 0 represents an
unversioned or legacy segment treated as built at transaction_read_version, and
that otherwise the newer of dataset_version and transaction_read_version is
selected for segment-source keying and conflict filtering.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 4148cdb6-f84d-47b8-85ae-6961c6d50260
📒 Files selected for processing (8)
python/src/indices.rsrust/lance/src/index.rsrust/lance/src/index/api.rsrust/lance/src/index/create.rsrust/lance/src/index/scalar/inverted.rsrust/lance/src/index/scalar_logical.rsrust/lance/src/io/commit.rsrust/lance/src/io/commit/conflict_resolver.rs
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/lance/src/io/commit/conflict_resolver.rs`:
- Around line 2161-2165: Add a regression test in the deferred-rewrite conflict
test suite that rebases an index through a rewrite, invokes finish_create_index
through the normal finish path, and verifies the committed index retains its
original fragment bitmap coverage after remapping. Reuse the existing test setup
and assertions around check_txn, extending coverage through the final commit
rather than stopping at transaction validation.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: e5d207cd-3abf-4e0e-96df-27fb08798888
📒 Files selected for processing (1)
rust/lance/src/io/commit/conflict_resolver.rs
| for index in new_indices.iter_mut() { | ||
| if let Some(original) = self.original_index_fragment_bitmaps.get(&index.uuid) { | ||
| index.fragment_bitmap = Some(original.clone()); | ||
| } | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Add a finish-path regression test for bitmap restoration.
The deferred-rewrite tests stop at check_txn; none verifies that finish_create_index restores the original bitmap after conflict-time remapping. Add a test that rebases through a rewrite, calls finish, and asserts the committed index retains its pre-remap coverage.
As per coding guidelines, “Every bugfix and feature must have corresponding tests.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance/src/io/commit/conflict_resolver.rs` around lines 2161 - 2165, Add
a regression test in the deferred-rewrite conflict test suite that rebases an
index through a rewrite, invokes finish_create_index through the normal finish
path, and verifies the committed index retains its original fragment bitmap
coverage after remapping. Reuse the existing test setup and assertions around
check_txn, extending coverage through the final commit rather than stopping at
transaction validation.
Source: Coding guidelines
3340acc to
81877be
Compare
81877be to
f69bbe8
Compare
3 participants
No description provided.