diff --git a/include/lance/lance.h b/include/lance/lance.h index 7630913..33194e0 100644 --- a/include/lance/lance.h +++ b/include/lance/lance.h @@ -1273,7 +1273,10 @@ typedef enum { int32_t lance_scanner_set_blob_handling(LanceScanner* scanner, LanceBlobHandling handling); /** - * Restrict scan to the given fragment IDs. Must be called before iteration. + * Restrict the scan/search domain to the given fragment IDs. For nearest-neighbor + * queries, this limits the candidates searched before Top-K regardless of the + * prefilter setting; it is not a predicate applied to ranked results. + * Must be called before iteration. * @param ids Array of fragment IDs * @param len Number of fragment IDs * @return 0 on success, -1 on error @@ -2198,6 +2201,13 @@ int32_t lance_scanner_set_refine_factor(LanceScanner* scanner, uint32_t f); int32_t lance_scanner_set_ef(LanceScanner* scanner, uint32_t e); int32_t lance_scanner_set_metric(LanceScanner* scanner, LanceMetricType metric); int32_t lance_scanner_set_use_index(LanceScanner* scanner, bool enable); + +/** + * Choose when scanner filter expressions are applied during nearest-neighbor search. + * When enabled, filters constrain ANN candidates; when disabled, they refine the + * candidates after ANN and may produce fewer than k rows. This does not change the + * fragment search domain set by lance_scanner_set_fragment_ids(). + */ int32_t lance_scanner_set_prefilter(LanceScanner* scanner, bool enable); /** diff --git a/include/lance/lance.hpp b/include/lance/lance.hpp index 286724e..2ba7a1b 100644 --- a/include/lance/lance.hpp +++ b/include/lance/lance.hpp @@ -1516,14 +1516,15 @@ class Scanner { return *this; } - /// Restrict scan to specific fragment IDs. + /// Restrict the scan/search domain to specific fragment IDs. For nearest + /// queries, they limit candidates before Top-K regardless of prefilter(). Scanner& fragment_ids(const uint64_t* ids, size_t len) { if (lance_scanner_set_fragment_ids(handle_.get(), ids, len) != 0) check_error(); return *this; } - /// Restrict scan to specific fragment IDs (vector overload). + /// Restrict the scan/search domain to specific fragment IDs (vector overload). Scanner& fragment_ids(const std::vector& ids) { return fragment_ids(ids.data(), ids.size()); } @@ -1672,6 +1673,9 @@ class Scanner { if (lance_scanner_set_use_index(handle_.get(), enable) != 0) check_error(); return *this; } + /// Apply scanner filters before ANN when enabled, or refine ANN candidates + /// afterward when disabled; postfiltering may return fewer than k rows. This + /// does not change the fragment search domain. Scanner& prefilter(bool enable) { if (lance_scanner_set_prefilter(handle_.get(), enable) != 0) check_error(); return *this; diff --git a/src/scanner.rs b/src/scanner.rs index 71d2dd2..32aedec 100644 --- a/src/scanner.rs +++ b/src/scanner.rs @@ -24,7 +24,7 @@ use lance_core::datatypes::BlobHandling; use lance_index::scalar::FullTextSearchQuery; use lance_index::vector::ApproxMode; use lance_io::stream::RecordBatchStream; -use lance_table::format::IndexMetadata; +use lance_table::format::{Fragment, IndexMetadata}; use uuid::Uuid; use crate::async_dispatcher::{self, LanceCallback}; @@ -305,7 +305,9 @@ impl LanceScanner { Ok(()) } - /// Apply fragment selection to a scanner builder if fragment_ids is set. + /// Restrict the scan/search domain to fragment_ids when they are set. + /// For nearest queries this limits candidates before Top-K regardless of prefilter; + /// it is not a predicate applied to the ranked results. fn apply_fragment_filter(&self, scanner: &mut lance::dataset::scanner::Scanner) -> Result<()> { if let Some(ids) = &self.fragment_ids { let all_fragments = self.dataset.get_fragments(); @@ -421,7 +423,14 @@ impl LanceScanner { if self.include_deleted_rows { scanner.include_deleted_rows(); } - self.apply_fragment_filter(&mut scanner)?; + // Lance rejects nearest() when fragments are already configured and prefilter is off. + // Deferring with_fragments only works around that configuration-time guard: fragments + // still restrict indexed and flat-search inputs before Top-K. The prefilter setting + // controls scanner filter expressions, not this search-domain scope. + let apply_fragment_filter_after_nearest = self.nearest.is_some() && !self.prefilter; + if !apply_fragment_filter_after_nearest { + self.apply_fragment_filter(&mut scanner)?; + } if self.index_segments.is_some() && self.nearest.is_none() { return Err(lance_core::Error::invalid_input_source( "index_segments requires nearest() to be configured".into(), @@ -437,8 +446,6 @@ impl LanceScanner { "fragment_ids cannot be combined with an FTS query context; split the query by FTS index segment UUID instead".into(), )); } - // nearest() checks the current prefilter setting before accepting a - // fragment-scoped search. Enable it before installing the query. if self.prefilter { scanner.prefilter(true); } @@ -495,6 +502,10 @@ impl LanceScanner { scanner.with_index_segments(segments.clone())?; } } + if apply_fragment_filter_after_nearest { + // This delayed builder call preserves the same fragment search domain. + self.apply_fragment_filter(&mut scanner)?; + } if let Some(fts) = &self.fts_query { scanner.full_text_search(fts.clone())?; } @@ -505,6 +516,10 @@ impl LanceScanner { Some(PreparedFtsExecution { context: Arc::clone(context), segments, + scope_prefilter_to_fts_segments: self.prefilter + && (self.filter.is_some() + || self.substrait_filter.is_some() + || !self.additional_sql_filters.is_empty()), batch_size: self.batch_size, scan_statistics_callback: self.scan_statistics_callback.clone(), }) @@ -558,6 +573,7 @@ impl LanceScanner { struct PreparedFtsExecution { context: Arc, segments: Vec, + scope_prefilter_to_fts_segments: bool, batch_size: Option, scan_statistics_callback: Option, } @@ -596,11 +612,21 @@ impl PreparedScanner { let Some(distributed_fts) = self.distributed_fts else { return self.scanner.try_into_stream().await; }; - let plan = self.scanner.create_plan().await?; - let selected_segments_have_current_fragments = segments_have_current_fragments( + let selected_fragments = selected_current_fts_fragments( &distributed_fts.context.dataset, &distributed_fts.segments, )?; + let selected_segments_have_current_fragments = !selected_fragments.is_empty(); + let mut scanner = self.scanner; + if distributed_fts.scope_prefilter_to_fts_segments + && selected_segments_have_current_fragments + { + // The scanner is already split by the selected FTS segment(s). Applying the same + // fragment scope before plan creation lets Lance restrict scalar-index segment loads + // for the TVF prefilter without changing unfiltered FTS scans. + scanner.with_fragments(selected_fragments); + } + let plan = scanner.create_plan().await?; let (plan, rewritten) = rewrite_prepared_fts_plan( plan, &distributed_fts.segments, @@ -661,22 +687,11 @@ fn select_fts_segments( Ok(selected) } -fn segments_have_current_fragments( +fn selected_current_fts_fragments( dataset: &lance::Dataset, segments: &[IndexMetadata], -) -> Result { - let current_fragment_ids = dataset - .get_fragments() - .into_iter() - .map(|fragment| { - u32::try_from(fragment.id()).map_err(|_| { - lance_core::Error::internal(format!( - "current fragment id {} exceeds the validated u32 FTS coverage range", - fragment.id() - )) - }) - }) - .collect::>>()?; +) -> Result> { + let mut selected_fragment_ids = std::collections::HashSet::new(); for segment in segments { let fragment_bitmap = segment.fragment_bitmap.as_ref().ok_or_else(|| { lance_core::Error::internal(format!( @@ -684,14 +699,22 @@ fn segments_have_current_fragments( segment.uuid )) })?; - if fragment_bitmap - .iter() - .any(|fragment_id| current_fragment_ids.contains(&fragment_id)) - { - return Ok(true); + selected_fragment_ids.extend(fragment_bitmap.iter()); + } + + let mut selected_fragments = Vec::new(); + for fragment in dataset.get_fragments() { + let fragment_id = u32::try_from(fragment.id()).map_err(|_| { + lance_core::Error::internal(format!( + "current fragment id {} exceeds the validated u32 FTS coverage range", + fragment.id() + )) + })?; + if selected_fragment_ids.contains(&fragment_id) { + selected_fragments.push(fragment.metadata().clone()); } } - Ok(false) + Ok(selected_fragments) } #[derive(Default)] @@ -3357,7 +3380,9 @@ mod tests { ); let has_current_fragments = - segments_have_current_fragments(&distributed.context.dataset, &segments).unwrap(); + !selected_current_fts_fragments(&distributed.context.dataset, &segments) + .unwrap() + .is_empty(); let (rewritten, counts) = rewrite_prepared_fts_plan( plan, &segments, diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs index 1a25e34..94c7a46 100644 --- a/tests/c_api_test.rs +++ b/tests/c_api_test.rs @@ -7997,6 +7997,101 @@ fn test_scanner_nearest_full_snapshot_prefilter_statistics() { } } +#[test] +fn test_scanner_nearest_indexed_fragment_scope_with_prefilter_disabled() { + let (_tmp, uri) = create_multi_fragment_vector_dataset(2, 32, 8, false); + let uri = c_str(&uri); + let dataset = unsafe { lance_dataset_open(uri.as_ptr(), ptr::null(), 0) }; + assert!(!dataset.is_null()); + + let column = c_str("embedding"); + let params = LanceVectorIndexParams { + index_type: LanceVectorIndexType::IvfFlat, + metric: LanceMetricType::L2, + num_partitions: 1, + num_sub_vectors: 0, + num_bits: 0, + max_iterations: 2, + hnsw_m: 0, + hnsw_ef_construction: 0, + sample_rate: 0, + }; + assert_eq!( + unsafe { + lance_dataset_create_vector_index(dataset, column.as_ptr(), ptr::null(), ¶ms, false) + }, + 0 + ); + + let mut fragments = [0_u64; 2]; + assert_eq!( + unsafe { lance_dataset_fragment_ids(dataset, fragments.as_mut_ptr()) }, + 0 + ); + let filter = c_str("id >= 32"); + let scanner = unsafe { lance_scanner_new(dataset, ptr::null(), filter.as_ptr()) }; + assert!(!scanner.is_null()); + // Model Doris: the split's fragments remain the search domain while the + // search filter is configured for postfiltering. + assert_eq!( + unsafe { lance_scanner_set_fragment_ids(scanner, fragments[1..].as_ptr(), 1) }, + 0 + ); + assert_eq!(unsafe { lance_scanner_set_prefilter(scanner, false) }, 0); + assert_eq!(unsafe { lance_scanner_set_use_index(scanner, true) }, 0); + + // The dataset-wide nearest rows are in fragment 0. Restricting the search + // domain to fragment 1 must return that fragment's local Top-K, not an empty + // result from filtering the dataset-wide Top-K afterward. + let query = [0.0_f32; 8]; + assert_eq!( + unsafe { + lance_scanner_nearest( + scanner, + column.as_ptr(), + query.as_ptr().cast(), + query.len(), + LanceDataType::Float32 as i32, + 5, + ) + }, + 0 + ); + let mut captured = CapturedScanStatistics::default(); + assert_eq!( + unsafe { + lance_scanner_set_statistics_callback( + scanner, + Some(capture_scan_statistics), + (&mut captured as *mut CapturedScanStatistics).cast(), + ) + }, + 0 + ); + + let batches = scan_all_rows_from_scanner(scanner); + let mut ids = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + ids.sort_unstable(); + assert_eq!(ids, vec![32, 33, 34, 35, 36]); + assert_eq!(captured.calls, 1); + assert!(captured.indices_loaded > 0, "expected the indexed ANN path"); + + unsafe { lance_scanner_close(scanner) }; + unsafe { lance_dataset_close(dataset) }; +} + #[test] fn test_scanner_nearest_segment_prefilter_statistics() { for stable_row_ids in [false, true] { @@ -9272,6 +9367,212 @@ fn test_prepared_fts_global_scorer_is_shared_across_segment_splits() { unsafe { lance_dataset_close(dataset) }; } +#[test] +fn test_prepared_fts_prefilter_is_scoped_to_selected_segments() { + use lance::index::DatasetIndexExt; + use lance_index::IndexCriteria; + use lance_index::optimize::OptimizeOptions; + + let (_tmp, uri) = create_test_dataset(); + let uri_c = c_str(&uri); + let column = c_str("name"); + let query = c_str("alice"); + let inverted_params = c_str(r#"{"base_tokenizer":"simple","language":"English"}"#); + + let dataset = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; + assert_eq!( + unsafe { + lance_dataset_create_scalar_index( + dataset, + column.as_ptr(), + ptr::null(), + LanceScalarIndexType::Inverted as i32, + inverted_params.as_ptr(), + false, + ) + }, + 0 + ); + unsafe { lance_dataset_close(dataset) }; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + ])); + append_batch( + &uri, + schema.clone(), + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![6, 7])), + Arc::new(StringArray::from(vec!["alice", "alice alice"])), + ], + ) + .unwrap(), + ); + lance_c::runtime::block_on(async { + let mut dataset = Dataset::open(&uri).await.unwrap(); + dataset + .optimize_indices(&OptimizeOptions::append()) + .await + .unwrap(); + }); + + let dataset = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; + let context = unsafe { + lance_dataset_prepare_fts_match_query( + dataset, + column.as_ptr(), + query.as_ptr(), + LanceFtsMatchOperator::Or as i32, + 0, + LanceFtsCoverageMode::Strict as i32, + ) + }; + assert!(!context.is_null(), "{}", take_last_error_message()); + + let mut fragment_ids = [0_u64; 2]; + assert_eq!(unsafe { lance_dataset_fragment_count(dataset) }, 2); + assert_eq!( + unsafe { lance_dataset_fragment_ids(dataset, fragment_ids.as_mut_ptr()) }, + 0 + ); + let selected_segment_uuid = lance_c::runtime::block_on(async { + let dataset = Dataset::open(&uri).await.unwrap(); + let logical_index = dataset + .load_scalar_index(IndexCriteria::default().for_column("name").supports_fts()) + .await + .unwrap() + .unwrap(); + let segments = dataset + .load_indices_by_name(&logical_index.name) + .await + .unwrap(); + assert_eq!(segments.len(), 2); + + let selected_fragment_id = u32::try_from(fragment_ids[1]).unwrap(); + let unselected_fragment_id = u32::try_from(fragment_ids[0]).unwrap(); + let selected_segments = segments + .iter() + .filter(|segment| { + segment + .fragment_bitmap + .as_ref() + .is_some_and(|bitmap| bitmap.iter().any(|id| id == selected_fragment_id)) + }) + .collect::>(); + assert_eq!(selected_segments.len(), 1); + assert!( + !selected_segments[0] + .fragment_bitmap + .as_ref() + .unwrap() + .iter() + .any(|id| id == unselected_fragment_id), + "the selected FTS segment should cover only the appended fragment" + ); + *selected_segments[0].uuid.as_bytes() + }); + + let full_scores = collect_context_fts_scores(dataset, context, None); + let mut full_ids = full_scores.keys().copied().collect::>(); + full_ids.sort_unstable(); + assert_eq!(full_ids, vec![1, 6, 7]); + + // Match rows in both fragments: four rows in the unselected fragment and + // both rows in the selected fragment. The prefilter scan must stay scoped + // to the selected segment's fragment. + let filter = c_str("id != 2"); + let scanner = unsafe { lance_scanner_new(dataset, ptr::null(), filter.as_ptr()) }; + assert!(!scanner.is_null()); + assert_eq!( + unsafe { lance_scanner_set_fts_query_context(scanner, context) }, + 0 + ); + assert_eq!( + unsafe { lance_scanner_set_fts_index_segments(scanner, selected_segment_uuid.as_ptr(), 1) }, + 0 + ); + assert_eq!(unsafe { lance_scanner_set_prefilter(scanner, true) }, 0); + let mut captured = CapturedScanStatistics::default(); + assert_eq!( + unsafe { + lance_scanner_set_statistics_callback( + scanner, + Some(capture_scan_statistics), + (&mut captured as *mut CapturedScanStatistics).cast(), + ) + }, + 0 + ); + + let mut stream = FFI_ArrowArrayStream::empty(); + assert_eq!( + unsafe { lance_scanner_to_arrow_stream(scanner, &mut stream) }, + 0, + "{}", + take_last_error_message() + ); + let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream).unwrap() }; + let mut filtered_scores = std::collections::HashMap::new(); + for batch in reader { + let batch = batch.unwrap(); + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let scores = batch + .column_by_name("_score") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + assert!( + filtered_scores + .insert(ids.value(row), scores.value(row)) + .is_none() + ); + } + } + let mut filtered_ids = filtered_scores.keys().copied().collect::>(); + filtered_ids.sort_unstable(); + assert_eq!(filtered_ids, vec![6, 7]); + for id in &filtered_ids { + let actual_score = filtered_scores[id]; + let global_score = full_scores[id]; + assert!( + (actual_score - global_score).abs() < 1e-6, + "id={id}, selected-segment score={actual_score}, global score={global_score}" + ); + } + + let metric_count = |name: &str| { + captured + .metrics + .iter() + .filter(|(metric, kind, _)| metric == name && *kind == LanceScanMetricKind::Count) + .map(|(_, _, value)| *value) + .sum::() + }; + assert_eq!(captured.calls, 1); + assert_eq!( + metric_count("prefilter_input_rows"), + 2, + "prefilter should read only the two rows in the selected segment's fragment" + ); + assert_eq!(metric_count("prefilter_row_ids"), 2); + + unsafe { + lance_scanner_close(scanner); + lance_fts_query_context_close(context); + lance_dataset_close(dataset); + } +} + #[test] fn test_prepare_fts_queries_reject_invalid_inputs() { let (_tmp, uri) = create_test_dataset();