diff --git a/java/lance-jni/src/blocking_dataset.rs b/java/lance-jni/src/blocking_dataset.rs index 320475def0a..d15a4073345 100644 --- a/java/lance-jni/src/blocking_dataset.rs +++ b/java/lance-jni/src/blocking_dataset.rs @@ -4018,73 +4018,105 @@ 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..7817d4aa1a1 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,58 @@ 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. + * + *

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. + * + * @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..9a289a19b44 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::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, @@ -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,217 @@ 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. 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, + 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()); + } + // 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 + { + 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 + } +} + +/// 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(); @@ -1455,6 +1676,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 +1726,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 +3273,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 +3315,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 +3990,17 @@ impl Scanner { ) .await? & target_fragments; - Some(Arc::new( + let mut exec = ScalarIndexExec::new(self.dataset.clone(), index_query, result_format) - .with_fragment_scope(fragment_scope), - ) as Arc) + .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 } @@ -7044,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, @@ -7875,6 +8141,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..6a4e4e800e8 100644 --- a/rust/lance/src/index/frag_reuse_reader_tests.rs +++ b/rust/lance/src/index/frag_reuse_reader_tests.rs @@ -2303,9 +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), &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 +2400,7 @@ async fn fragment_scope_prunes_on_fri_effective_coverage() { "i", "i_idx", Some(&RoaringBitmap::from_iter([999])), + None, &lance_index::metrics::NoOpMetricsCollector, ) .await @@ -2407,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 diff --git a/rust/lance/src/index/scalar_logical.rs b/rust/lance/src/index/scalar_logical.rs index 18cc840b823..23907f54f10 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 @@ -2013,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 79c383c00fd..d90e26f5816 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,39 @@ 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) => { + (Some(fragments), segment_uuids) => { + expr.evaluate( + &ScopedScalarIndexLoader { + dataset: dataset.as_ref(), + fragments: Some(fragments), + segment_uuids, + }, + &metrics, + ) + .await? + } + (None, Some(segment_uuids)) => { expr.evaluate( - &FragmentScopedScalarIndexLoader { + &ScopedScalarIndexLoader { dataset: dataset.as_ref(), - fragments, + fragments: None, + segment_uuids: Some(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 +347,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())) @@ -860,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. @@ -936,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); @@ -953,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 = @@ -1121,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, );