From 47b3c6095cb4af1edbf00e211b802a7859dd9267 Mon Sep 17 00:00:00 2001 From: "zhanghaobo@kanzhun.com" Date: Thu, 1 Oct 2026 20:33:59 +0800 Subject: [PATCH 1/7] feat(java): count indexed rows on explicit scalar segments Pin countIndexedRows to the named scalar index and add an overload that opens only the requested physical segments. Selected segment coverage defines the count scope and must match an explicit fragment list. --- java/lance-jni/src/blocking_dataset.rs | 149 +++++---- java/src/main/java/org/lance/Dataset.java | 48 +++ .../java/org/lance/index/ScalarIndexTest.java | 305 ++++++++++++++++++ rust/lance/src/dataset/scanner.rs | 234 +++++++++++++- rust/lance/src/index.rs | 14 + .../src/index/frag_reuse_reader_tests.rs | 14 +- rust/lance/src/index/scalar_logical.rs | 14 +- rust/lance/src/io/exec/scalar_index.rs | 31 +- 8 files changed, 733 insertions(+), 76 deletions(-) diff --git a/java/lance-jni/src/blocking_dataset.rs b/java/lance-jni/src/blocking_dataset.rs index 320475def0a..3e8d62eef5c 100644 --- a/java/lance-jni/src/blocking_dataset.rs +++ b/java/lance-jni/src/blocking_dataset.rs @@ -4018,73 +4018,112 @@ pub extern "system" fn Java_org_lance_Dataset_nativeCountIndexedRows( fn inner_count_indexed_rows( env: &mut JNIEnv, java_dataset: JObject, - _jindex_name: JString, + jindex_name: JString, jfilter: JString, jfragment_ids: JObject, // Optional> ) -> Result { + let index_name: String = jindex_name.extract(env)?; let filter: String = jfilter.extract(env)?; + let fragment_ids = extract_count_fragment_ids(env, &jfragment_ids)?; - // Extract optional fragment IDs - let fragment_ids: Option> = if env - .call_method(&jfragment_ids, "isPresent", "()Z", &[])? - .z()? - { - let list_obj = env - .call_method(&jfragment_ids, "get", "()Ljava/lang/Object;", &[])? - .l()?; - let list = env.get_list(&list_obj)?; - let mut ids = Vec::new(); - let mut iter = list.iter(env)?; - while let Some(elem) = iter.next(env)? { - let int_val = env.call_method(&elem, "intValue", "()I", &[])?.i()?; - ids.push(int_val as u32); - } - Some(ids) - } else { - None - }; - - let count = { - let dataset_guard = - unsafe { env.get_rust_field::<_, _, BlockingDataset>(java_dataset, NATIVE_DATASET) }?; - - // Use a scanner with fragment filtering to count rows - // This ensures we only count rows in the specified fragments - let inner = dataset_guard.inner.clone(); + count_indexed_rows( + env, + java_dataset, + index_name, + filter, + None, + fragment_ids, + ) +} - block_on(async { - let mut scanner = inner.scan(); +#[unsafe(no_mangle)] +pub extern "system" fn Java_org_lance_Dataset_nativeCountIndexedRowsWithSegments( + mut env: JNIEnv, + java_dataset: JObject, + jindex_name: JString, + jfilter: JString, + jsegment_uuids: JObject, // List + jfragment_ids: JObject, // Optional> +) -> jlong { + ok_or_throw_with_return!( + env, + inner_count_indexed_rows_with_segments( + &mut env, + java_dataset, + jindex_name, + jfilter, + jsegment_uuids, + jfragment_ids, + ), + -1 + ) +} - // Apply filter - if !filter.is_empty() { - scanner.filter(&filter)?; - } +fn inner_count_indexed_rows_with_segments( + env: &mut JNIEnv, + java_dataset: JObject, + jindex_name: JString, + jfilter: JString, + jsegment_uuids: JObject, + jfragment_ids: JObject, +) -> Result { + let index_name: String = jindex_name.extract(env)?; + let filter: String = jfilter.extract(env)?; + let segment_uuids = + import_vec_to_rust(env, &jsegment_uuids, |env, obj| obj.extract_object(env))?; + let fragment_ids = extract_count_fragment_ids(env, &jfragment_ids)?; - // Empty projection and enable row_id for count_rows to work - // count_rows() requires metadata-only projection - scanner.project::(&[])?; - scanner.with_row_id(); - - // Apply fragment filter if specified - if let Some(frag_ids) = fragment_ids { - // Convert FileFragment to Fragment by extracting metadata - let filtered_fragments: Vec<_> = inner - .get_fragments() - .into_iter() - .filter(|f| frag_ids.contains(&(f.id() as u32))) - .map(|f| f.metadata().clone()) - .collect(); - scanner.with_fragments(filtered_fragments); - } + count_indexed_rows( + env, + java_dataset, + index_name, + filter, + Some(segment_uuids), + fragment_ids, + ) +} - // Use the scanner's count_rows method - let count = scanner.count_rows().await?; +fn extract_count_fragment_ids( + env: &mut JNIEnv, + jfragment_ids: &JObject, +) -> Result>> { + env.get_ints_opt(jfragment_ids)? + .map(|fragment_ids| { + fragment_ids + .into_iter() + .map(|fragment_id| { + u32::try_from(fragment_id).map_err(|_| { + Error::input_error(format!( + "fragment ID must be non-negative, got {fragment_id}" + )) + }) + }) + .collect() + }) + .transpose() +} - Ok::(count as i64) - })? +fn count_indexed_rows( + env: &mut JNIEnv, + java_dataset: JObject, + index_name: String, + filter: String, + segment_uuids: Option>, + fragment_ids: Option>, +) -> Result { + let count = { + let dataset_guard = + unsafe { env.get_rust_field::<_, _, BlockingDataset>(java_dataset, NATIVE_DATASET) }?; + block_on(dataset_guard.inner.count_indexed_rows( + &index_name, + &filter, + segment_uuids.as_deref(), + fragment_ids.as_deref(), + ))? }; - Ok(count) + i64::try_from(count) + .map_err(|_| Error::input_error(format!("indexed row count {count} exceeds Java long"))) } ////////////////////////////// diff --git a/java/src/main/java/org/lance/Dataset.java b/java/src/main/java/org/lance/Dataset.java index 53fe78c9d65..64134b8c8a5 100644 --- a/java/src/main/java/org/lance/Dataset.java +++ b/java/src/main/java/org/lance/Dataset.java @@ -71,6 +71,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.UUID; import java.util.concurrent.locks.ReentrantLock; import java.util.stream.Collectors; @@ -1482,6 +1483,10 @@ public Session session() { * counts matching row addresses, which is more efficient than scanning when the index covers the * filter column. * + *

Planning is pinned to {@code indexName}. A filter that cannot be answered by that scalar + * index is rejected instead of scanning the table or selecting another index. Deleted rows are + * excluded. + * * @param indexName the name of the scalar index to use * @param filter the filter expression (e.g., "column = 5") * @param fragmentIds optional list of fragment IDs to restrict the count to @@ -1499,9 +1504,52 @@ public long countIndexedRows( } } + /** + * Count rows matching a filter using explicit physical segments of a scalar index. + * + *

Only {@code segmentUuids} are opened. Their current fragment coverage defines the count + * scope: matching deleted rows inside that scope are excluded, and rows outside it are not + * counted. When {@code fragmentIds} is omitted, the scope is derived from that coverage. When it + * is present, its set must equal the coverage; order does not matter. A mismatch is rejected with + * an error that reports both sets. + * + *

An empty segment list, duplicate segment UUIDs, an unknown UUID, a UUID from a different + * index, or a segment without fragment coverage is rejected. The existing three-argument method + * remains available and does not take a segment list. + * + * @param indexName the logical scalar index name that every selected segment must belong to + * @param filter the filter expression (e.g., "column = 5") + * @param segmentUuids physical segment UUIDs to open; must be non-empty and contain no duplicates + * @param fragmentIds optional fragment IDs that must match the selected segments' current + * coverage + * @return count of matching rows in the selected segment scope + */ + public long countIndexedRows( + String indexName, + String filter, + List segmentUuids, + Optional> fragmentIds) { + try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) { + Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed"); + Preconditions.checkArgument( + indexName != null && !indexName.isEmpty(), "indexName cannot be null or empty"); + Preconditions.checkArgument( + filter != null && !filter.isEmpty(), "filter cannot be null or empty"); + Preconditions.checkNotNull(segmentUuids, "segmentUuids cannot be null"); + Preconditions.checkArgument(!segmentUuids.isEmpty(), "segmentUuids cannot be empty"); + return nativeCountIndexedRowsWithSegments(indexName, filter, segmentUuids, fragmentIds); + } + } + private native long nativeCountIndexedRows( String indexName, String filter, Optional> fragmentIds); + private native long nativeCountIndexedRowsWithSegments( + String indexName, + String filter, + List segmentUuids, + Optional> fragmentIds); + /** * Calculate the size of the dataset. * diff --git a/java/src/test/java/org/lance/index/ScalarIndexTest.java b/java/src/test/java/org/lance/index/ScalarIndexTest.java index 5c911447f94..35feb1b031d 100644 --- a/java/src/test/java/org/lance/index/ScalarIndexTest.java +++ b/java/src/test/java/org/lance/index/ScalarIndexTest.java @@ -49,10 +49,13 @@ import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; +import java.io.UncheckedIOException; +import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.Comparator; import java.util.List; import java.util.Map; import java.util.Optional; @@ -63,6 +66,7 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.Stream; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; @@ -762,6 +766,307 @@ public void testCreateBTreeIndexDistributively(@TempDir Path tempDir) throws Exc } } + @Test + public void testCountIndexedRowsSelectsPhysicalSegments(@TempDir Path tempDir) throws Exception { + assertPhysicalSegmentSelection(tempDir.resolve("segments"), false); + } + + @Test + public void testCountIndexedRowsStableRowIdsExcludeDeletions(@TempDir Path tempDir) + throws Exception { + assertPhysicalSegmentSelection(tempDir.resolve("stable-segments"), true); + } + + @Test + public void testCountIndexedRowsDoesNotOpenUnselectedSegment(@TempDir Path tempDir) + throws Exception { + String filter = "name >= 'Person 0'"; + Path datasetPath = tempDir.resolve("unopened"); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + UUID segmentA; + UUID segmentB; + String indexName; + try (NamedSegments segments = createNamedSegments(allocator, datasetPath, false, 10, 4)) { + segmentA = segments.segmentA.uuid(); + segmentB = segments.segmentB.uuid(); + indexName = segments.indexName; + } + // Reopen after deleting B so a session cache cannot satisfy the query. + deleteIndexSegment(datasetPath, segmentB); + try (Dataset dataset = Dataset.open(datasetPath.toString(), allocator)) { + assertEquals( + 10, + dataset.countIndexedRows( + indexName, filter, Collections.singletonList(segmentA), Optional.empty())); + Exception missingSegment = + Assertions.assertThrows( + Exception.class, + () -> + dataset.countIndexedRows( + indexName, filter, Collections.singletonList(segmentB), Optional.empty())); + assertTrue( + missingSegment.getMessage() != null && !missingSegment.getMessage().isEmpty(), + "opening the deleted segment should fail, got: " + missingSegment); + } + } + } + + private void assertPhysicalSegmentSelection(Path datasetPath, boolean stableRowIds) + throws Exception { + String filter = "name >= 'Person 0'"; + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + NamedSegments segments = createNamedSegments(allocator, datasetPath, stableRowIds, 10, 4)) { + Dataset dataset = segments.dataset; + List segmentA = Collections.singletonList(segments.segmentA.uuid()); + List segmentB = Collections.singletonList(segments.segmentB.uuid()); + List both = Arrays.asList(segments.segmentB.uuid(), segments.segmentA.uuid()); + int fragmentA = fragmentId(segments.segmentA); + int fragmentB = fragmentId(segments.segmentB); + + assertEquals( + 10, dataset.countIndexedRows(segments.indexName, filter, segmentA, Optional.empty())); + assertEquals( + 4, dataset.countIndexedRows(segments.indexName, filter, segmentB, Optional.empty())); + assertEquals( + 14, dataset.countIndexedRows(segments.indexName, filter, both, Optional.empty())); + assertEquals(14, dataset.countIndexedRows(segments.indexName, filter, Optional.empty())); + assertEquals( + 14, + dataset.countIndexedRows( + segments.indexName, filter, both, Optional.of(Arrays.asList(fragmentB, fragmentA)))); + assertEquals( + 10, + dataset.countIndexedRows( + segments.indexName, filter, Optional.of(Collections.singletonList(fragmentA)))); + + Exception unindexedFilter = + Assertions.assertThrows( + IllegalArgumentException.class, + () -> dataset.countIndexedRows(segments.indexName, "id = 0", Optional.empty())); + assertTrue( + unindexedFilter.getMessage().contains(segments.indexName), unindexedFilter.getMessage()); + assertTrue( + unindexedFilter.getMessage().contains("cannot be planned"), unindexedFilter.getMessage()); + + Exception coverageMismatch = + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + dataset.countIndexedRows( + segments.indexName, + filter, + segmentA, + Optional.of(Collections.singletonList(fragmentB)))); + assertTrue( + coverageMismatch.getMessage().contains("do not match selected segment coverage"), + coverageMismatch.getMessage()); + assertTrue( + coverageMismatch.getMessage().contains(Integer.toString(fragmentA)), + coverageMismatch.getMessage()); + assertTrue( + coverageMismatch.getMessage().contains(Integer.toString(fragmentB)), + coverageMismatch.getMessage()); + + Exception unknownFragment = + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + dataset.countIndexedRows( + segments.indexName, + filter, + segmentA, + Optional.of(Collections.singletonList(99)))); + assertTrue( + unknownFragment.getMessage().contains("unknown fragment IDs"), + unknownFragment.getMessage()); + + Exception duplicateFragments = + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + dataset.countIndexedRows( + segments.indexName, + filter, + segmentA, + Optional.of(Arrays.asList(fragmentA, fragmentA)))); + assertTrue( + duplicateFragments.getMessage().contains("duplicate fragment IDs"), + duplicateFragments.getMessage()); + + Exception unknownSegment = + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + dataset.countIndexedRows( + segments.indexName, + filter, + Collections.singletonList(UUID.randomUUID()), + Optional.empty())); + assertTrue( + unknownSegment.getMessage().contains("does not exist"), unknownSegment.getMessage()); + + Index otherIndex = + dataset.createIndex( + IndexOptions.builder( + Collections.singletonList("id"), IndexType.BTREE, btreeIndexParams()) + .withIndexName("other_index") + .build()); + Exception wrongIndex = + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + dataset.countIndexedRows( + segments.indexName, + filter, + Collections.singletonList(otherIndex.uuid()), + Optional.empty())); + assertTrue(wrongIndex.getMessage().contains("other_index"), wrongIndex.getMessage()); + assertTrue(wrongIndex.getMessage().contains(segments.indexName), wrongIndex.getMessage()); + + Exception duplicateSegments = + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + dataset.countIndexedRows( + segments.indexName, + filter, + Arrays.asList(segments.segmentA.uuid(), segments.segmentA.uuid()), + Optional.empty())); + assertTrue( + duplicateSegments.getMessage().contains("duplicate UUIDs"), + duplicateSegments.getMessage()); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + dataset.countIndexedRows( + segments.indexName, filter, Collections.emptyList(), Optional.empty())); + + dataset.delete("name = 'Person 0'"); + assertEquals( + 9, dataset.countIndexedRows(segments.indexName, filter, segmentA, Optional.empty())); + assertEquals( + 3, dataset.countIndexedRows(segments.indexName, filter, segmentB, Optional.empty())); + assertEquals( + 12, dataset.countIndexedRows(segments.indexName, filter, both, Optional.empty())); + assertEquals(12, dataset.countRows(filter)); + } + } + + private NamedSegments createNamedSegments( + BufferAllocator allocator, Path datasetPath, boolean stableRowIds, int rowsA, int rowsB) + throws Exception { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath.toString()); + WriteParams.Builder createParams = new WriteParams.Builder(); + if (stableRowIds) { + createParams.withEnableStableRowIds(true); + } + testDataset.createDatasetWithWriteParams(createParams.build()).close(); + testDataset.write(1, rowsA).close(); + Dataset dataset = testDataset.write(2, rowsB); + List fragments = dataset.getFragments(); + assertEquals(2, fragments.size()); + Fragment fragmentA = null; + Fragment fragmentB = null; + for (Fragment fragment : fragments) { + long physicalRows = fragment.metadata().getPhysicalRows(); + if (physicalRows == rowsA) { + fragmentA = fragment; + } else if (physicalRows == rowsB) { + fragmentB = fragment; + } + } + assertNotNull(fragmentA); + assertNotNull(fragmentB); + + String indexName = "name_idx"; + List committed = + commitBtreeSegments(dataset, indexName, "name", Arrays.asList(fragmentA, fragmentB)); + return new NamedSegments( + dataset, + indexName, + indexCovering(committed, fragmentA.getId()), + indexCovering(committed, fragmentB.getId())); + } + + private static IndexParams btreeIndexParams() { + return IndexParams.builder() + .setScalarIndexParams(ScalarIndexParams.create("btree", "{\"zone_size\": 2048}")) + .build(); + } + + private static List commitBtreeSegments( + Dataset dataset, String indexName, String column, List fragments) { + List segments = new ArrayList<>(); + for (Fragment fragment : fragments) { + segments.add( + dataset.createIndex( + IndexOptions.builder( + Collections.singletonList(column), IndexType.BTREE, btreeIndexParams()) + .withIndexName(indexName) + .withFragmentIds(Collections.singletonList(fragment.getId())) + .build())); + } + List committed = dataset.commitExistingIndexSegments(indexName, column, segments); + assertEquals(fragments.size(), committed.size()); + return committed; + } + + private static Index indexCovering(List indexes, int fragmentId) { + for (Index index : indexes) { + if (index.fragments().orElse(Collections.emptyList()).contains(fragmentId)) { + return index; + } + } + throw new AssertionError("no committed segment covers fragment " + fragmentId); + } + + private static int fragmentId(Index segment) { + List fragments = + segment + .fragments() + .orElseThrow( + () -> + new AssertionError("segment " + segment.uuid() + " has no fragment coverage")); + assertEquals(1, fragments.size(), "expected one fragment in " + segment.uuid()); + return fragments.get(0); + } + + private static void deleteIndexSegment(Path datasetPath, UUID uuid) throws IOException { + Path indexDir = datasetPath.resolve("_indices").resolve(uuid.toString()); + assertTrue(Files.isDirectory(indexDir), "expected index directory " + indexDir); + try (Stream walk = Files.walk(indexDir)) { + walk.sorted(Comparator.reverseOrder()) + .forEach( + path -> { + try { + Files.delete(path); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + }); + } + } + + private static final class NamedSegments implements AutoCloseable { + private final Dataset dataset; + private final String indexName; + private final Index segmentA; + private final Index segmentB; + + private NamedSegments(Dataset dataset, String indexName, Index segmentA, Index segmentB) { + this.dataset = dataset; + this.indexName = indexName; + this.segmentA = segmentA; + this.segmentB = segmentB; + } + + @Override + public void close() { + dataset.close(); + } + } + @Test public void testRangedBTreeIndex(@TempDir Path tempDir) throws Exception { String datasetPath = tempDir.resolve("ranged_btree_map").toString(); diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 281f34b23da..610b3a6140b 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -101,7 +101,7 @@ use crate::dataset::row_offsets_to_row_addresses; use crate::dataset::rowids::{live_row_addrs_to_row_ids, translate_addr_treemap_to_row_ids}; use crate::dataset::utils::SchemaAdapter; use crate::index::DatasetIndexInternalExt; -use crate::index::scalar::fetch_index_details; +use crate::index::scalar::{IndexDetails, fetch_index_details}; use crate::index::scalar::inverted::{ fts_index_fragment_bitmap, load_segment_details, load_segments, normalize_inverted_details, resolve_fts_field, resolve_query_document_granularity, validate_combined_fields_target_column, @@ -1100,6 +1100,13 @@ impl AggregateExprBuilder { /// .buffered(16) /// .sum() /// ``` +#[derive(Clone)] +struct ScalarIndexSelection { + index_name: String, + segments: Arc>, + segment_uuids: Option>>, +} + #[derive(Clone)] pub struct Scanner { dataset: Arc, @@ -1196,6 +1203,9 @@ pub struct Scanner { /// If set, this scanner serves only these fragments. fragments: Option>, + /// If set, scalar-index planning is restricted to this logical index and optional segment set. + scalar_index_selection: Option, + /// If set, this scanner will only search the specified vector index segments. index_segments: Option>, @@ -1426,6 +1436,170 @@ impl TakeOperation { } } +impl Dataset { + /// Count rows matching `filter` while pinning scalar-index planning to `index_name`. + /// + /// When `segment_uuids` is present, only those physical index segments are opened and their + /// current fragment coverage defines the count scope. In that mode, a supplied + /// `fragment_ids` list must exactly equal the selected segments' current coverage. + #[instrument(skip_all)] + pub async fn count_indexed_rows( + &self, + index_name: &str, + filter: &str, + segment_uuids: Option<&[Uuid]>, + fragment_ids: Option<&[u32]>, + ) -> Result { + if index_name.is_empty() { + return Err(Error::invalid_input("index_name must not be empty")); + } + if filter.is_empty() { + return Err(Error::invalid_input("filter must not be empty")); + } + + let requested_fragments = fragment_ids + .map(|fragment_ids| { + let fragment_set = fragment_ids.iter().copied().collect::(); + if fragment_set.len() != fragment_ids.len() as u64 { + return Err(Error::invalid_input(format!( + "fragment_ids contains duplicate fragment IDs: {fragment_ids:?}" + ))); + } + let unknown = &fragment_set - self.fragment_bitmap.as_ref(); + if !unknown.is_empty() { + return Err(Error::invalid_input(format!( + "fragment_ids contains unknown fragment IDs: {:?}", + unknown.iter().collect::>() + ))); + } + Ok(fragment_set) + }) + .transpose()?; + + let all_indices = self.load_indices().await?; + let (selected_segments, selected_segment_uuids, fragment_scope) = + if let Some(segment_uuids) = segment_uuids { + if segment_uuids.is_empty() { + return Err(Error::invalid_input( + "segment_uuids must contain at least one UUID", + )); + } + let selected_uuid_set = segment_uuids.iter().copied().collect::>(); + if selected_uuid_set.len() != segment_uuids.len() { + return Err(Error::invalid_input(format!( + "segment_uuids contains duplicate UUIDs: {segment_uuids:?}" + ))); + } + + let mut selected_segments = Vec::with_capacity(segment_uuids.len()); + let mut selected_coverage = RoaringBitmap::new(); + for segment_uuid in segment_uuids { + let segment = all_indices + .iter() + .find(|segment| segment.uuid == *segment_uuid) + .ok_or_else(|| { + Error::invalid_input(format!( + "Scalar index segment {segment_uuid} does not exist" + )) + })?; + if segment.name != index_name { + return Err(Error::invalid_input(format!( + "Scalar index segment {segment_uuid} belongs to index '{}', not requested index '{index_name}'", + segment.name + ))); + } + let Some(field_id) = segment.keyed_field() else { + return Err(Error::invalid_input(format!( + "Scalar index '{index_name}' segment {segment_uuid} has no keyed field" + ))); + }; + let field = self.schema().field_by_id(field_id).ok_or_else(|| { + Error::internal(format!( + "Scalar index '{index_name}' segment {segment_uuid} references missing field ID {field_id}" + )) + })?; + let field_path = if let Some(ancestors) = self.schema().field_ancestry_by_id(field.id) { + let field_refs = ancestors + .iter() + .map(|field| field.name.as_str()) + .collect::>(); + lance_core::datatypes::format_field_path(&field_refs) + } else { + field.name.clone() + }; + let details = + IndexDetails(fetch_index_details(self, &field_path, segment).await?); + if details.is_vector() || details.get_plugin().is_err() { + return Err(Error::invalid_input(format!( + "Index '{index_name}' segment {segment_uuid} is not a usable scalar index segment" + ))); + } + let coverage = segment.fragment_bitmap.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "Scalar index '{index_name}' segment {segment_uuid} is missing fragment coverage" + )) + })?; + let current_coverage = coverage & self.fragment_bitmap.as_ref(); + if current_coverage.is_empty() { + return Err(Error::invalid_input(format!( + "Scalar index '{index_name}' segment {segment_uuid} does not cover any current dataset fragments" + ))); + } + selected_coverage |= current_coverage; + selected_segments.push(segment.clone()); + } + + if let Some(requested_fragments) = requested_fragments.as_ref() + && requested_fragments != &selected_coverage + { + return Err(Error::invalid_input(format!( + "fragment_ids {:?} do not match selected segment coverage {:?} for scalar index '{index_name}'", + requested_fragments.iter().collect::>(), + selected_coverage.iter().collect::>() + ))); + } + + ( + selected_segments, + Some(Arc::new(selected_uuid_set)), + Some(selected_coverage), + ) + } else { + let selected_segments = all_indices + .iter() + .filter(|segment| segment.name == index_name) + .cloned() + .collect::>(); + if selected_segments.is_empty() { + return Err(Error::invalid_input(format!( + "Scalar index '{index_name}' does not exist or has no usable segments" + ))); + } + (selected_segments, None, requested_fragments) + }; + + let mut scanner = self.scan(); + scanner.filter(filter)?; + scanner.project::(&[])?; + scanner.with_row_id(); + scanner.with_scalar_index_selection( + index_name.to_string(), + selected_segments, + selected_segment_uuids, + ); + if let Some(fragment_scope) = fragment_scope { + let fragments = self + .fragments() + .iter() + .filter(|fragment| fragment_scope.contains(fragment.id as u32)) + .cloned() + .collect(); + scanner.with_fragments(fragments); + } + scanner.count_rows().await + } +} + impl Scanner { pub fn new(dataset: Arc) -> Self { let projection_plan = ProjectionPlan::full(dataset.clone()).unwrap(); @@ -1455,6 +1629,7 @@ impl Scanner { use_stats: true, ordered: true, fragments: None, + scalar_index_selection: None, index_segments: None, fast_search: false, use_scalar_index: true, @@ -1504,6 +1679,20 @@ impl Scanner { self } + fn with_scalar_index_selection( + &mut self, + index_name: String, + segments: Vec, + segment_uuids: Option>>, + ) -> &mut Self { + self.scalar_index_selection = Some(ScalarIndexSelection { + index_name, + segments: Arc::new(segments), + segment_uuids, + }); + self + } + /// Restrict vector index search to the specified index segments. /// /// This setting is only supported for vector search. @@ -3037,7 +3226,14 @@ impl Scanner { // Check expr filter let filter_plan = if let Some(filter) = self.filter.expr_filter.as_ref() { let expr = filter.to_datafusion(self.dataset.schema(), filter_schema.as_ref())?; - let index_info = self.dataset.scalar_index_info().await?; + let index_info = match self.scalar_index_selection.as_ref() { + Some(selection) => { + self.dataset + .scalar_index_info_for_segments(selection.segments.as_ref()) + .await? + } + None => self.dataset.scalar_index_info().await?, + }; let filter_plan = planner.create_filter_plan(expr.clone(), &index_info, use_scalar_index)?; @@ -3072,6 +3268,15 @@ impl Scanner { FilterPlan::new(query_filter, ExprFilterPlan::default()) }; + if let Some(selection) = self.scalar_index_selection.as_ref() + && filter_plan.expr_filter_plan.index_query.is_none() + { + return Err(Error::invalid_input(format!( + "Filter cannot be planned using scalar index '{}'", + selection.index_name + ))); + } + // Check query filter if filter_plan.query_filter.is_some() && self.nearest.is_none() @@ -3738,10 +3943,20 @@ impl Scanner { ) .await? & target_fragments; - Some(Arc::new( - ScalarIndexExec::new(self.dataset.clone(), index_query, result_format) - .with_fragment_scope(fragment_scope), - ) as Arc) + let mut exec = ScalarIndexExec::new( + self.dataset.clone(), + index_query, + result_format, + ) + .with_fragment_scope(fragment_scope); + if let Some(segment_uuids) = self + .scalar_index_selection + .as_ref() + .and_then(|selection| selection.segment_uuids.clone()) + { + exec = exec.with_segment_uuids(segment_uuids); + } + Some(Arc::new(exec) as Arc) } else { None } @@ -7875,6 +8090,13 @@ impl Scanner { // unless fast_search allows skipping uncovered fragments. let mut exec = ScalarIndexExec::new(self.dataset.clone(), index_query.clone(), result_format); + if let Some(segment_uuids) = self + .scalar_index_selection + .as_ref() + .and_then(|selection| selection.segment_uuids.clone()) + { + exec = exec.with_segment_uuids(segment_uuids); + } if missing_frags.is_empty() && !relevant_frags.is_empty() { exec = exec.with_fragment_scope( relevant_frags diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 70209f12ccf..b2411e45f9f 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -3266,6 +3266,12 @@ pub trait DatasetIndexInternalExt: DatasetIndexExt { /// Loads information about all the available scalar indices on the dataset async fn scalar_index_info(&self) -> Result; + /// Loads scalar-index planning information from an explicitly selected metadata set. + async fn scalar_index_info_for_segments( + &self, + segments: &[IndexMetadata], + ) -> Result; + /// Return the fragments that are not covered by any of the deltas of the index. async fn unindexed_fragments(&self, idx_name: &str) -> Result>; @@ -3849,6 +3855,14 @@ impl DatasetIndexInternalExt for Dataset { #[instrument(level = "trace", skip_all)] async fn scalar_index_info(&self) -> Result { let indices = self.load_indices().await?; + self.scalar_index_info_for_segments(indices.as_ref()).await + } + + #[instrument(level = "trace", skip_all)] + async fn scalar_index_info_for_segments( + &self, + indices: &[IndexMetadata], + ) -> Result { let schema = self.schema(); let mut indexed_fields = Vec::new(); // (column, index_name) → union of every contributing IndexMetadata's diff --git a/rust/lance/src/index/frag_reuse_reader_tests.rs b/rust/lance/src/index/frag_reuse_reader_tests.rs index bb28c1f1946..94cc8cd0800 100644 --- a/rust/lance/src/index/frag_reuse_reader_tests.rs +++ b/rust/lance/src/index/frag_reuse_reader_tests.rs @@ -2303,9 +2303,16 @@ async fn fragment_scope_prunes_on_fri_effective_coverage() { // not intersect {10}). let scope_f10 = RoaringBitmap::from_iter([10]); let metrics = lance_index::metrics::LocalMetricsCollector::default(); - let scoped = open_scalar_index_segments(&dataset, "i", "i_idx", Some(&scope_f10), &metrics) - .await - .unwrap(); + let scoped = open_scalar_index_segments( + &dataset, + "i", + "i_idx", + Some(&scope_f10), + None, + &metrics, + ) + .await + .unwrap(); assert_eq!( metrics .index_loads @@ -2399,6 +2406,7 @@ async fn fragment_scope_prunes_on_fri_effective_coverage() { "i", "i_idx", Some(&RoaringBitmap::from_iter([999])), + None, &lance_index::metrics::NoOpMetricsCollector, ) .await diff --git a/rust/lance/src/index/scalar_logical.rs b/rust/lance/src/index/scalar_logical.rs index 18cc840b823..1d9b65455fd 100644 --- a/rust/lance/src/index/scalar_logical.rs +++ b/rust/lance/src/index/scalar_logical.rs @@ -5,6 +5,7 @@ use lance_core::utils::row_addr_remap::RowAddrRemap; use std::any::Any; +use std::collections::HashSet; use std::sync::Arc; use async_trait::async_trait; @@ -341,20 +342,25 @@ pub async fn open_named_scalar_index( index_name: &str, metrics: &dyn MetricsCollector, ) -> Result> { - open_scalar_index_segments(dataset, column, index_name, None, metrics).await + open_scalar_index_segments(dataset, column, index_name, None, None, metrics).await } -/// Open scalar index segments whose coverage intersects `fragments`. +/// Open selected scalar index segments whose coverage intersects `fragments`. /// -/// `None` preserves the unscoped behavior and opens every usable segment. +/// A `None` scope preserves the unscoped behavior for that dimension. UUID selection happens +/// before any segment is opened. pub async fn open_scalar_index_segments( dataset: &Dataset, column: &str, index_name: &str, fragments: Option<&RoaringBitmap>, + segment_uuids: Option<&HashSet>, metrics: &dyn MetricsCollector, ) -> Result> { let mut indices = load_named_scalar_segments(dataset, column, index_name).await?; + if let Some(segment_uuids) = segment_uuids { + indices.retain(|index| segment_uuids.contains(&index.uuid)); + } if let Some(fragments) = fragments { indices.retain(|index| { index @@ -510,6 +516,7 @@ mod tests { "value", "value_btree", Some(&scope), + None, &NoOpMetricsCollector, ) .await @@ -587,6 +594,7 @@ mod tests { "value", "value_btree_pairs", Some(&RoaringBitmap::from_iter([target_fragment])), + None, &NoOpMetricsCollector, ) .await diff --git a/rust/lance/src/io/exec/scalar_index.rs b/rust/lance/src/io/exec/scalar_index.rs index 79c383c00fd..a4ac183f976 100644 --- a/rust/lance/src/io/exec/scalar_index.rs +++ b/rust/lance/src/io/exec/scalar_index.rs @@ -81,13 +81,14 @@ impl ScalarIndexLoader for Dataset { } } -struct FragmentScopedScalarIndexLoader<'a> { +struct ScopedScalarIndexLoader<'a> { dataset: &'a Dataset, - fragments: &'a RoaringBitmap, + fragments: Option<&'a RoaringBitmap>, + segment_uuids: Option<&'a HashSet>, } #[async_trait] -impl ScalarIndexLoader for FragmentScopedScalarIndexLoader<'_> { +impl ScalarIndexLoader for ScopedScalarIndexLoader<'_> { async fn load_index( &self, column: &str, @@ -98,7 +99,8 @@ impl ScalarIndexLoader for FragmentScopedScalarIndexLoader<'_> { self.dataset, column, index_name, - Some(self.fragments), + self.fragments, + self.segment_uuids, metrics, ) .await @@ -155,6 +157,7 @@ pub struct ScalarIndexExec { metrics: ExecutionPlanMetricsSet, result_format: IndexExprResultWireFormat, fragment_scope: Option>, + segment_uuids: Option>>, } impl DisplayAs for ScalarIndexExec { @@ -189,6 +192,7 @@ impl ScalarIndexExec { metrics: ExecutionPlanMetricsSet::new(), result_format, fragment_scope: None, + segment_uuids: None, } } @@ -201,6 +205,12 @@ impl ScalarIndexExec { self } + /// Restrict index loading to these physical scalar index segments. + pub(crate) fn with_segment_uuids(mut self, segment_uuids: Arc>) -> Self { + self.segment_uuids = Some(segment_uuids); + self + } + pub fn dataset(&self) -> &Arc { &self.dataset } @@ -252,26 +262,28 @@ impl ScalarIndexExec { plan_metrics: ExecutionPlanMetricsSet, result_format: IndexExprResultWireFormat, fragment_scope: Option>, + segment_uuids: Option>>, ) -> Result { let metrics = IndexMetrics::new(&plan_metrics, 0); let query_result = { let search_time = plan_metrics.new_time(SCALAR_INDEX_SEARCH_TIME_METRIC, 0); let _timer = search_time.timer(); - match fragment_scope.as_deref() { - Some(fragments) if fragments.is_empty() => { + match (fragment_scope.as_deref(), segment_uuids.as_deref()) { + (Some(fragments), _) if fragments.is_empty() => { IndexExprResult::exact(RowAddrMask::allow_nothing()) } - Some(fragments) => { + (fragments, segment_uuids) if fragments.is_some() || segment_uuids.is_some() => { expr.evaluate( - &FragmentScopedScalarIndexLoader { + &ScopedScalarIndexLoader { dataset: dataset.as_ref(), fragments, + segment_uuids, }, &metrics, ) .await? } - None => expr.evaluate(dataset.as_ref(), &metrics).await?, + (None, None) => expr.evaluate(dataset.as_ref(), &metrics).await?, } }; let mut fragments_covered_by_result = @@ -324,6 +336,7 @@ impl ExecutionPlan for ScalarIndexExec { self.metrics.clone(), self.result_format, self.fragment_scope.clone(), + self.segment_uuids.clone(), ); let stream = futures::stream::iter(vec![batch_fut]) .then(|batch_fut| batch_fut.map_err(|err| err.into())) From 91f748fe71f2b150ebbf7aa586e7dc11d734d1b0 Mon Sep 17 00:00:00 2001 From: "zhanghaobo@kanzhun.com" Date: Thu, 1 Oct 2026 20:46:31 +0800 Subject: [PATCH 2/7] fix(index): make scalar segment loading exhaustive Match guards are not exhaustive, so Rust 1.91 rejected the segment loader. Format the new indexed-count path with the repository rustfmt rules. --- java/lance-jni/src/blocking_dataset.rs | 9 +- rust/lance/src/dataset/scanner.rs | 167 +++++++++--------- .../src/index/frag_reuse_reader_tests.rs | 14 +- rust/lance/src/io/exec/scalar_index.rs | 15 +- 4 files changed, 101 insertions(+), 104 deletions(-) diff --git a/java/lance-jni/src/blocking_dataset.rs b/java/lance-jni/src/blocking_dataset.rs index 3e8d62eef5c..d15a4073345 100644 --- a/java/lance-jni/src/blocking_dataset.rs +++ b/java/lance-jni/src/blocking_dataset.rs @@ -4026,14 +4026,7 @@ fn inner_count_indexed_rows( let filter: String = jfilter.extract(env)?; let fragment_ids = extract_count_fragment_ids(env, &jfragment_ids)?; - count_indexed_rows( - env, - java_dataset, - index_name, - filter, - None, - fragment_ids, - ) + count_indexed_rows(env, java_dataset, index_name, filter, None, fragment_ids) } #[unsafe(no_mangle)] diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 610b3a6140b..b05cc3cd78c 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -101,11 +101,11 @@ use crate::dataset::row_offsets_to_row_addresses; use crate::dataset::rowids::{live_row_addrs_to_row_ids, translate_addr_treemap_to_row_ids}; use crate::dataset::utils::SchemaAdapter; use crate::index::DatasetIndexInternalExt; -use crate::index::scalar::{IndexDetails, fetch_index_details}; use crate::index::scalar::inverted::{ fts_index_fragment_bitmap, load_segment_details, load_segments, normalize_inverted_details, resolve_fts_field, resolve_query_document_granularity, validate_combined_fields_target_column, }; +use crate::index::scalar::{IndexDetails, fetch_index_details}; use crate::index::scalar_logical::{load_named_scalar_segments, scalar_index_fragment_bitmap}; use crate::index::vector::utils::{ default_distance_type_for, get_vector_dim, get_vector_type, validate_distance_type_for, @@ -1477,48 +1477,51 @@ impl Dataset { .transpose()?; let all_indices = self.load_indices().await?; - let (selected_segments, selected_segment_uuids, fragment_scope) = - if let Some(segment_uuids) = segment_uuids { - if segment_uuids.is_empty() { - return Err(Error::invalid_input( - "segment_uuids must contain at least one UUID", - )); - } - let selected_uuid_set = segment_uuids.iter().copied().collect::>(); - if selected_uuid_set.len() != segment_uuids.len() { + let (selected_segments, selected_segment_uuids, fragment_scope) = if let Some( + segment_uuids, + ) = segment_uuids + { + if segment_uuids.is_empty() { + return Err(Error::invalid_input( + "segment_uuids must contain at least one UUID", + )); + } + let selected_uuid_set = segment_uuids.iter().copied().collect::>(); + if selected_uuid_set.len() != segment_uuids.len() { + return Err(Error::invalid_input(format!( + "segment_uuids contains duplicate UUIDs: {segment_uuids:?}" + ))); + } + + let mut selected_segments = Vec::with_capacity(segment_uuids.len()); + let mut selected_coverage = RoaringBitmap::new(); + for segment_uuid in segment_uuids { + let segment = all_indices + .iter() + .find(|segment| segment.uuid == *segment_uuid) + .ok_or_else(|| { + Error::invalid_input(format!( + "Scalar index segment {segment_uuid} does not exist" + )) + })?; + if segment.name != index_name { return Err(Error::invalid_input(format!( - "segment_uuids contains duplicate UUIDs: {segment_uuids:?}" + "Scalar index segment {segment_uuid} belongs to index '{}', not requested index '{index_name}'", + segment.name ))); } - - let mut selected_segments = Vec::with_capacity(segment_uuids.len()); - let mut selected_coverage = RoaringBitmap::new(); - for segment_uuid in segment_uuids { - let segment = all_indices - .iter() - .find(|segment| segment.uuid == *segment_uuid) - .ok_or_else(|| { - Error::invalid_input(format!( - "Scalar index segment {segment_uuid} does not exist" - )) - })?; - if segment.name != index_name { - return Err(Error::invalid_input(format!( - "Scalar index segment {segment_uuid} belongs to index '{}', not requested index '{index_name}'", - segment.name - ))); - } - let Some(field_id) = segment.keyed_field() else { - return Err(Error::invalid_input(format!( - "Scalar index '{index_name}' segment {segment_uuid} has no keyed field" - ))); - }; - let field = self.schema().field_by_id(field_id).ok_or_else(|| { + let Some(field_id) = segment.keyed_field() else { + return Err(Error::invalid_input(format!( + "Scalar index '{index_name}' segment {segment_uuid} has no keyed field" + ))); + }; + let field = self.schema().field_by_id(field_id).ok_or_else(|| { Error::internal(format!( "Scalar index '{index_name}' segment {segment_uuid} references missing field ID {field_id}" )) })?; - let field_path = if let Some(ancestors) = self.schema().field_ancestry_by_id(field.id) { + let field_path = + if let Some(ancestors) = self.schema().field_ancestry_by_id(field.id) { let field_refs = ancestors .iter() .map(|field| field.name.as_str()) @@ -1527,56 +1530,55 @@ impl Dataset { } else { field.name.clone() }; - let details = - IndexDetails(fetch_index_details(self, &field_path, segment).await?); - if details.is_vector() || details.get_plugin().is_err() { - return Err(Error::invalid_input(format!( - "Index '{index_name}' segment {segment_uuid} is not a usable scalar index segment" - ))); - } - let coverage = segment.fragment_bitmap.as_ref().ok_or_else(|| { + let details = IndexDetails(fetch_index_details(self, &field_path, segment).await?); + if details.is_vector() || details.get_plugin().is_err() { + return Err(Error::invalid_input(format!( + "Index '{index_name}' segment {segment_uuid} is not a usable scalar index segment" + ))); + } + let coverage = segment.fragment_bitmap.as_ref().ok_or_else(|| { Error::invalid_input(format!( "Scalar index '{index_name}' segment {segment_uuid} is missing fragment coverage" )) })?; - let current_coverage = coverage & self.fragment_bitmap.as_ref(); - if current_coverage.is_empty() { - return Err(Error::invalid_input(format!( - "Scalar index '{index_name}' segment {segment_uuid} does not cover any current dataset fragments" - ))); - } - selected_coverage |= current_coverage; - selected_segments.push(segment.clone()); - } - - if let Some(requested_fragments) = requested_fragments.as_ref() - && requested_fragments != &selected_coverage - { + let current_coverage = coverage & self.fragment_bitmap.as_ref(); + if current_coverage.is_empty() { return Err(Error::invalid_input(format!( - "fragment_ids {:?} do not match selected segment coverage {:?} for scalar index '{index_name}'", - requested_fragments.iter().collect::>(), - selected_coverage.iter().collect::>() + "Scalar index '{index_name}' segment {segment_uuid} does not cover any current dataset fragments" ))); } + selected_coverage |= current_coverage; + selected_segments.push(segment.clone()); + } - ( - selected_segments, - Some(Arc::new(selected_uuid_set)), - Some(selected_coverage), - ) - } else { - let selected_segments = all_indices - .iter() - .filter(|segment| segment.name == index_name) - .cloned() - .collect::>(); - if selected_segments.is_empty() { - return Err(Error::invalid_input(format!( - "Scalar index '{index_name}' does not exist or has no usable segments" - ))); - } - (selected_segments, None, requested_fragments) - }; + if let Some(requested_fragments) = requested_fragments.as_ref() + && requested_fragments != &selected_coverage + { + return Err(Error::invalid_input(format!( + "fragment_ids {:?} do not match selected segment coverage {:?} for scalar index '{index_name}'", + requested_fragments.iter().collect::>(), + selected_coverage.iter().collect::>() + ))); + } + + ( + selected_segments, + Some(Arc::new(selected_uuid_set)), + Some(selected_coverage), + ) + } else { + let selected_segments = all_indices + .iter() + .filter(|segment| segment.name == index_name) + .cloned() + .collect::>(); + if selected_segments.is_empty() { + return Err(Error::invalid_input(format!( + "Scalar index '{index_name}' does not exist or has no usable segments" + ))); + } + (selected_segments, None, requested_fragments) + }; let mut scanner = self.scan(); scanner.filter(filter)?; @@ -3943,12 +3945,9 @@ impl Scanner { ) .await? & target_fragments; - let mut exec = ScalarIndexExec::new( - self.dataset.clone(), - index_query, - result_format, - ) - .with_fragment_scope(fragment_scope); + let mut exec = + ScalarIndexExec::new(self.dataset.clone(), index_query, result_format) + .with_fragment_scope(fragment_scope); if let Some(segment_uuids) = self .scalar_index_selection .as_ref() diff --git a/rust/lance/src/index/frag_reuse_reader_tests.rs b/rust/lance/src/index/frag_reuse_reader_tests.rs index 94cc8cd0800..a30eb76663a 100644 --- a/rust/lance/src/index/frag_reuse_reader_tests.rs +++ b/rust/lance/src/index/frag_reuse_reader_tests.rs @@ -2303,16 +2303,10 @@ async fn fragment_scope_prunes_on_fri_effective_coverage() { // not intersect {10}). let scope_f10 = RoaringBitmap::from_iter([10]); let metrics = lance_index::metrics::LocalMetricsCollector::default(); - let scoped = open_scalar_index_segments( - &dataset, - "i", - "i_idx", - Some(&scope_f10), - None, - &metrics, - ) - .await - .unwrap(); + let scoped = + open_scalar_index_segments(&dataset, "i", "i_idx", Some(&scope_f10), None, &metrics) + .await + .unwrap(); assert_eq!( metrics .index_loads diff --git a/rust/lance/src/io/exec/scalar_index.rs b/rust/lance/src/io/exec/scalar_index.rs index a4ac183f976..ff87aefb961 100644 --- a/rust/lance/src/io/exec/scalar_index.rs +++ b/rust/lance/src/io/exec/scalar_index.rs @@ -272,17 +272,28 @@ impl ScalarIndexExec { (Some(fragments), _) if fragments.is_empty() => { IndexExprResult::exact(RowAddrMask::allow_nothing()) } - (fragments, segment_uuids) if fragments.is_some() || segment_uuids.is_some() => { + (Some(fragments), segment_uuids) => { expr.evaluate( &ScopedScalarIndexLoader { dataset: dataset.as_ref(), - fragments, + fragments: Some(fragments), segment_uuids, }, &metrics, ) .await? } + (None, Some(segment_uuids)) => { + expr.evaluate( + &ScopedScalarIndexLoader { + dataset: dataset.as_ref(), + fragments: None, + segment_uuids: Some(segment_uuids), + }, + &metrics, + ) + .await? + } (None, None) => expr.evaluate(dataset.as_ref(), &metrics).await?, } }; From 296780775fefc519010ab3212481772bb0dfa031 Mon Sep 17 00:00:00 2001 From: "zhanghaobo@kanzhun.com" Date: Thu, 1 Oct 2026 20:54:26 +0800 Subject: [PATCH 3/7] fix(index): reject incomplete scalar segment selections After fragment reuse, one source segment advertises every destination fragment while still depending on its siblings. Require every contributor for that scope so a partial UUID set cannot under-count. --- rust/lance/src/dataset/scanner.rs | 49 +++++++++++- .../src/index/frag_reuse_reader_tests.rs | 77 +++++++++++++++++++ 2 files changed, 124 insertions(+), 2 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index b05cc3cd78c..8fd55eca00c 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -1440,8 +1440,11 @@ impl Dataset { /// Count rows matching `filter` while pinning scalar-index planning to `index_name`. /// /// When `segment_uuids` is present, only those physical index segments are opened and their - /// current fragment coverage defines the count scope. In that mode, a supplied - /// `fragment_ids` list must exactly equal the selected segments' current coverage. + /// current fragment coverage defines the count scope. That coverage is accepted only when the + /// selection includes every segment that contributes to it. After a fragment-reuse rewrite, one + /// source segment can advertise every destination fragment while still depending on sibling + /// segments; an incomplete selection is rejected. A supplied `fragment_ids` list must exactly + /// equal the selected segments' current coverage. #[instrument(skip_all)] pub async fn count_indexed_rows( &self, @@ -1550,6 +1553,15 @@ impl Dataset { selected_coverage |= current_coverage; selected_segments.push(segment.clone()); } + // Rewritten coverage is the contributor closure: a sibling whose bitmap + // overlaps this scope still owns rows in those fragments. Opening only + // the requested UUIDs would publish that scope and then under-count. + reject_incomplete_segment_contributors( + index_name, + all_indices.as_ref(), + &selected_uuid_set, + &selected_coverage, + )?; if let Some(requested_fragments) = requested_fragments.as_ref() && requested_fragments != &selected_coverage @@ -1602,6 +1614,39 @@ impl Dataset { } } +/// Reject a UUID selection whose advertised fragments still depend on an unselected segment. +fn reject_incomplete_segment_contributors( + index_name: &str, + indices: &[IndexMetadata], + selected_uuids: &HashSet, + selected_coverage: &RoaringBitmap, +) -> Result<()> { + let mut missing = Vec::new(); + let mut shared_fragments = RoaringBitmap::new(); + for segment in indices { + if segment.name != index_name || selected_uuids.contains(&segment.uuid) { + continue; + } + let Some(coverage) = segment.fragment_bitmap.as_ref() else { + continue; + }; + let overlap = coverage & selected_coverage; + if overlap.is_empty() { + continue; + } + shared_fragments |= overlap; + missing.push(segment.uuid); + } + if missing.is_empty() { + return Ok(()); + } + missing.sort_unstable(); + Err(Error::invalid_input(format!( + "Scalar index '{index_name}' selection omits contributing segments {missing:?} for fragments {:?}. Select every segment that covers those fragments", + shared_fragments.iter().collect::>() + ))) +} + impl Scanner { pub fn new(dataset: Arc) -> Self { let projection_plan = ProjectionPlan::full(dataset.clone()).unwrap(); diff --git a/rust/lance/src/index/frag_reuse_reader_tests.rs b/rust/lance/src/index/frag_reuse_reader_tests.rs index a30eb76663a..6a4e4e800e8 100644 --- a/rust/lance/src/index/frag_reuse_reader_tests.rs +++ b/rust/lance/src/index/frag_reuse_reader_tests.rs @@ -2409,6 +2409,83 @@ async fn fragment_scope_prunes_on_fri_effective_coverage() { ); } +/// One reclustered source advertises both destinations, but it only stores half +/// of their rows. Selecting that UUID alone must fail so the caller can pass +/// every contributor; the complete UUID set still counts the destinations. +#[tokio::test] +async fn selected_segment_recluster_requires_every_contributor() { + let mut dataset = lance_datagen::gen_batch() + .col("i", lance_datagen::array::step::()) + .into_ram_dataset(FragmentCount::from(2), FragmentRowCount::from(4)) + .await + .unwrap(); + let params = ScalarIndexParams::default(); + let source_fragments: Vec = dataset.fragments().iter().map(|f| f.id as u32).collect(); + let mut segments = Vec::new(); + for fragment in &source_fragments { + segments.push( + CreateIndexBuilder::new(&mut dataset, &["i"], IndexType::BTree, ¶ms) + .name("i_idx".into()) + .fragments(vec![*fragment]) + .execute_uncommitted() + .await + .unwrap(), + ); + } + dataset + .commit_existing_index_segments("i_idx", "i", segments) + .await + .unwrap(); + let stored = dataset.load_indices_by_name("i_idx").await.unwrap(); + let selected_uuid = stored[0].uuid; + let sibling_uuid = stored[1].uuid; + + let (transition, destinations) = prepare(&dataset).await; + let content = InlineContent { + legacy_versions: vec![], + transitions: vec![transition], + } + .encode_to_vec(); + install(&mut dataset, content, destinations, false).await; + + let current = dataset.load_indices_by_name("i_idx").await.unwrap(); + let selected = current + .iter() + .find(|segment| segment.uuid == selected_uuid) + .unwrap(); + assert_eq!( + selected.fragment_bitmap.as_ref(), + Some(dataset.fragment_bitmap.as_ref()) + ); + + let message = dataset + .count_indexed_rows("i_idx", "i >= 0", Some(&[selected_uuid]), None) + .await + .unwrap_err() + .to_string(); + assert!(message.contains(&sibling_uuid.to_string()), "{message}"); + assert!(message.contains("[10, 11]"), "{message}"); + + let both = [selected_uuid, sibling_uuid]; + let positive = dataset + .count_indexed_rows("i_idx", "i >= 0", Some(&both), None) + .await + .unwrap(); + let greater = dataset + .count_indexed_rows("i_idx", "i > 0", Some(&both), None) + .await + .unwrap(); + let negative = dataset + .count_indexed_rows("i_idx", "NOT (i = 0)", Some(&both), None) + .await + .unwrap(); + let logical = dataset + .count_indexed_rows("i_idx", "i >= 0", None, None) + .await + .unwrap(); + assert_eq!((positive, greater, negative, logical), (8, 7, 7, 8)); +} + // Combination: the derived-listing cache holds the FULL effective listing per // snapshot; each scoped query filters an independent copy. Querying one // destination then the other on the same snapshot must both be correct: the From 2abe625acafde5f7fbbc23e75cf50036dd0dc001 Mon Sep 17 00:00:00 2001 From: "zhanghaobo@kanzhun.com" Date: Thu, 1 Oct 2026 21:06:08 +0800 Subject: [PATCH 4/7] docs(java): document incomplete segment selection rejection Callers need to see that a fragment-reuse scope is accepted only when every contributing segment UUID is supplied. --- java/src/main/java/org/lance/Dataset.java | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/java/src/main/java/org/lance/Dataset.java b/java/src/main/java/org/lance/Dataset.java index 64134b8c8a5..7817d4aa1a1 100644 --- a/java/src/main/java/org/lance/Dataset.java +++ b/java/src/main/java/org/lance/Dataset.java @@ -1513,6 +1513,12 @@ public long countIndexedRows( * is present, its set must equal the coverage; order does not matter. A mismatch is rejected with * an error that reports both sets. * + *

The selection is accepted only when it includes every segment that contributes to that + * coverage. After fragment reuse, one source segment can advertise every destination fragment + * while still depending on its siblings. An incomplete selection is rejected, and the error names + * the missing segment UUIDs. A segment whose coverage does not overlap the selection can be + * queried alone. + * *

An empty segment list, duplicate segment UUIDs, an unknown UUID, a UUID from a different * index, or a segment without fragment coverage is rejected. The existing three-argument method * remains available and does not take a segment list. From da4ef24eae72854c8dfe0642fdbb0c3b82a8f3a0 Mon Sep 17 00:00:00 2001 From: "zhanghaobo@kanzhun.com" Date: Fri, 2 Oct 2026 16:44:39 +0800 Subject: [PATCH 5/7] ci: retrigger dependency check From 2c573fdad4729c7b0404debccb43987c73b00675 Mon Sep 17 00:00:00 2001 From: "zhanghaobo@kanzhun.com" Date: Fri, 2 Oct 2026 22:17:16 +0800 Subject: [PATCH 6/7] fix(index): honor segment selection on legacy storage Legacy datasets count through MaterializeIndexExec, which opened every segment of the logical index. Pass the requested UUIDs into that loader. --- rust/lance/src/dataset/scanner.rs | 9 ++- rust/lance/src/index/scalar_logical.rs | 77 ++++++++++++++++++++++++++ rust/lance/src/io/exec/scalar_index.rs | 20 ++++++- 3 files changed, 104 insertions(+), 2 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 8fd55eca00c..9a289a19b44 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -7303,11 +7303,18 @@ impl Scanner { // Build the MaterializeIndexExec, blocking stale row addresses so the index never // emits them. Stale rows are re-scored separately via a targeted take below. - let mat_exec = MaterializeIndexExec::new( + let mut mat_exec = MaterializeIndexExec::new( self.dataset.clone(), index_expr.clone(), Arc::new(relevant_frags), ); + if let Some(segment_uuids) = self + .scalar_index_selection + .as_ref() + .and_then(|selection| selection.segment_uuids.clone()) + { + mat_exec = mat_exec.with_segment_uuids(segment_uuids); + } let mat_exec = match self.stale_rows_block_mask(&stale_rows).await? { Some(block) => mat_exec.with_overlay_block(block), None => mat_exec, diff --git a/rust/lance/src/index/scalar_logical.rs b/rust/lance/src/index/scalar_logical.rs index 1d9b65455fd..23907f54f10 100644 --- a/rust/lance/src/index/scalar_logical.rs +++ b/rust/lance/src/index/scalar_logical.rs @@ -2021,4 +2021,81 @@ mod tests { vec![1] ); } + + /// Legacy storage counts through MaterializeIndexExec. Deleting the unselected + /// segment directory must still leave the selected segment countable. + #[tokio::test] + async fn legacy_storage_count_opens_only_selected_segment() { + let test_dir = TempStrDir::default(); + let mut dataset = lance_datagen::gen_batch() + .col("i", array::step::()) + .into_dataset_with_params( + test_dir.as_str(), + FragmentCount::from(2), + FragmentRowCount::from(4), + Some(WriteParams { + max_rows_per_file: 4, + data_storage_version: Some(lance_file::version::LanceFileVersion::Legacy), + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!( + dataset.manifest.data_storage_format.lance_file_format(), + lance_file::version::ConcreteFileVersion::V1 + ); + + let params = ScalarIndexParams::for_builtin(BuiltinIndexType::BTree); + let fragment_ids: Vec = dataset + .get_fragments() + .iter() + .map(|fragment| fragment.id() as u32) + .collect(); + let mut segments = Vec::new(); + for fragment_id in &fragment_ids { + segments.push( + CreateIndexBuilder::new(&mut dataset, &["i"], IndexType::BTree, ¶ms) + .name("i_idx".to_string()) + .fragments(vec![*fragment_id]) + .execute_uncommitted() + .await + .unwrap(), + ); + } + dataset + .commit_existing_index_segments("i_idx", "i", segments) + .await + .unwrap(); + let committed = dataset.load_indices_by_name("i_idx").await.unwrap(); + let selected = committed[0].uuid; + let omitted = committed[1].uuid; + assert_eq!( + dataset + .count_indexed_rows("i_idx", "i >= 0", Some(&[selected]), None) + .await + .unwrap(), + 4 + ); + + let omitted_dir = std::path::Path::new(test_dir.as_str()) + .join("_indices") + .join(omitted.to_string()); + assert!(omitted_dir.is_dir(), "{}", omitted_dir.display()); + std::fs::remove_dir_all(&omitted_dir).unwrap(); + drop(dataset); + + let dataset = Dataset::open(test_dir.as_str()).await.unwrap(); + assert_eq!( + dataset + .count_indexed_rows("i_idx", "i >= 0", Some(&[selected]), None) + .await + .unwrap(), + 4 + ); + dataset + .count_indexed_rows("i_idx", "i >= 0", Some(&[omitted]), None) + .await + .unwrap_err(); + } } diff --git a/rust/lance/src/io/exec/scalar_index.rs b/rust/lance/src/io/exec/scalar_index.rs index ff87aefb961..d90e26f5816 100644 --- a/rust/lance/src/io/exec/scalar_index.rs +++ b/rust/lance/src/io/exec/scalar_index.rs @@ -884,6 +884,8 @@ pub struct MaterializeIndexExec { dataset: Arc, expr: ScalarIndexExpr, fragments: Arc>, + /// Physical segments to open. `None` loads the whole logical index. + segment_uuids: Option>>, /// Row addresses blocked from the index result due to data overlay files committed after the /// index was built. ANDead into the candidate mask before row ID materialisation so that stale /// index entries never reach downstream operators. @@ -960,12 +962,19 @@ impl MaterializeIndexExec { dataset, expr, fragments, + segment_uuids: None, overlay_block: None, properties, metrics: ExecutionPlanMetricsSet::new(), } } + /// Restrict index loading to these physical scalar index segments. + pub(crate) fn with_segment_uuids(mut self, segment_uuids: Arc>) -> Self { + self.segment_uuids = Some(segment_uuids); + self + } + /// Block specific row addresses (see the `overlay_block` field) from the index result. pub fn with_overlay_block(mut self, block: RowAddrMask) -> Self { self.overlay_block = Some(block); @@ -977,10 +986,18 @@ impl MaterializeIndexExec { expr: ScalarIndexExpr, dataset: Arc, fragments: Arc>, + segment_uuids: Option>>, overlay_block: Option, metrics: Arc, ) -> Result { - let expr_result = expr.evaluate(dataset.as_ref(), metrics.as_ref()); + // Legacy storage materializes index hits here. UUID selection has to be applied at load + // time; evaluating against the dataset opens every segment of the logical index. + let loader = ScopedScalarIndexLoader { + dataset: dataset.as_ref(), + fragments: None, + segment_uuids: segment_uuids.as_deref(), + }; + let expr_result = expr.evaluate(&loader, metrics.as_ref()); let span = debug_span!("create_prefilter"); let prefilter = span.in_scope(|| { let fragment_bitmap = @@ -1145,6 +1162,7 @@ impl ExecutionPlan for MaterializeIndexExec { self.expr.clone(), self.dataset.clone(), self.fragments.clone(), + self.segment_uuids.clone(), self.overlay_block.clone(), metrics, ); From 8a671c9dbe609b0310f7d26161e2eb73f4f20bc2 Mon Sep 17 00:00:00 2001 From: "zhanghaobo@kanzhun.com" Date: Sat, 3 Oct 2026 21:59:22 +0800 Subject: [PATCH 7/7] ci: retrigger pipeline