diff --git a/README.md b/README.md index f13ad01..0b8fca5 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,7 @@ Based on the [liblance RFC](https://github.com/lance-format/lance/discussions/60 | [x] | Async scan | Callback-based `lance_scanner_scan_async()` for non-blocking scans | | [x] | Dataset metadata | `lance_dataset_version()`, `lance_dataset_count_rows()`, `lance_dataset_latest_version()` | | [x] | Filter pushdown | `lance_scanner_set_substrait_filter()` accepts a serialized Substrait `ExtendedExpression`; `lance_scanner_additional_sql_filter()` adds SQL predicates with AND before scanning starts | +| [x] | Row prefilter | `lance_scanner_set_row_addr_allowlist()` / `lance_scanner_set_row_addr_blocklist()` restrict scans, nearest-neighbor and full-text searches to (or exclude) a set of `_rowid`s before ranking | ## Segment-scoped array label filters diff --git a/include/lance/lance.h b/include/lance/lance.h index 6467a46..f1c69bc 100644 --- a/include/lance/lance.h +++ b/include/lance/lance.h @@ -1287,6 +1287,42 @@ int32_t lance_scanner_set_fragment_ids( size_t len ); +/** + * Restrict the scan or search to rows whose `_rowid` is in `row_ids`. + * + * The ids are in the dataset's `_rowid` space: the stable row id when stable + * row ids are enabled, otherwise the row address. The allowlist is applied + * before ranking: a nearest or full-text search (including one using a + * prepared FTS query context) computes its top-k over the allowed rows only, + * and a plain scan reads only those rows. It is combined with any filter + * expression (AND) and with a blocklist (allowed AND NOT blocked). It cannot be + * combined with lance_scanner_set_scalar_index_segment(). + * Calling it again replaces the previous allowlist. Must be called before + * iteration; the ids are copied. + * @param row_ids Array of `_rowid` values, in any order (may be NULL when len is 0) + * @param len Number of ids; 0 allows no rows + * @return 0 on success, -1 on error + */ +int32_t lance_scanner_set_row_addr_allowlist( + LanceScanner* scanner, + const uint64_t* row_ids, + size_t len +); + +/** + * Exclude rows whose `_rowid` is in `row_ids` from the scan or search. + * Same `_rowid` space, ordering, and restrictions as + * lance_scanner_set_row_addr_allowlist(). + * @param row_ids Array of `_rowid` values, in any order (may be NULL when len is 0) + * @param len Number of ids; 0 blocks no rows + * @return 0 on success, -1 on error + */ +int32_t lance_scanner_set_row_addr_blocklist( + LanceScanner* scanner, + const uint64_t* row_ids, + size_t len +); + /** * Set a Substrait filter on the scanner. * diff --git a/include/lance/lance.hpp b/include/lance/lance.hpp index 092929c..ba17040 100644 --- a/include/lance/lance.hpp +++ b/include/lance/lance.hpp @@ -1537,6 +1537,31 @@ class Scanner { return fragment_ids(ids.data(), ids.size()); } + /// Restrict the scan or search to rows whose `_rowid` is in `row_ids` + /// (applied before ranking). `len == 0` allows no rows. + Scanner& row_addr_allowlist(const uint64_t* row_ids, size_t len) { + if (lance_scanner_set_row_addr_allowlist(handle_.get(), row_ids, len) != 0) + check_error(); + return *this; + } + + /// Restrict the scan or search to rows whose `_rowid` is in `row_ids` (vector overload). + Scanner& row_addr_allowlist(const std::vector& row_ids) { + return row_addr_allowlist(row_ids.data(), row_ids.size()); + } + + /// Exclude rows whose `_rowid` is in `row_ids` (applied before ranking). + Scanner& row_addr_blocklist(const uint64_t* row_ids, size_t len) { + if (lance_scanner_set_row_addr_blocklist(handle_.get(), row_ids, len) != 0) + check_error(); + return *this; + } + + /// Exclude rows whose `_rowid` is in `row_ids` (vector overload). + Scanner& row_addr_blocklist(const std::vector& row_ids) { + return row_addr_blocklist(row_ids.data(), row_ids.size()); + } + /// Set a Substrait filter (serialized ExtendedExpression bytes). /// Wins over any SQL filter passed to the Scanner constructor. Scanner& substrait_filter(const uint8_t* bytes, size_t len) { diff --git a/src/scanner.rs b/src/scanner.rs index 2bbe7f4..e5ff5ca 100644 --- a/src/scanner.rs +++ b/src/scanner.rs @@ -16,7 +16,8 @@ use datafusion::physical_plan::{ExecutionPlan, empty::EmptyExec}; use futures::{FutureExt, Stream, StreamExt}; use lance::Dataset; use lance::dataset::scanner::{ - DatasetRecordBatchStream, ExecutionStatsCallback, ExecutionSummaryCounts, + DatasetRecordBatchStream, ExecutionStatsCallback, ExecutionSummaryCounts, RowAddrMask, + RowAddrTreeMap, }; use lance::io::exec::fts::{FlatMatchQueryExec, MatchQueryExec, PhraseQueryExec}; use lance_core::Result; @@ -110,6 +111,10 @@ pub struct LanceScanner { with_row_address: bool, include_deleted_rows: bool, fragment_ids: Option>, + // External row prefilter keyed in the dataset's `_rowid` space; combined into + // a single `RowAddrMask` (allowed AND NOT blocked) when the scanner is built. + row_addr_allowlist: Option, + row_addr_blocklist: Option, index_segments: Option>, scalar_index_segment: Option, nearest: Option, @@ -308,6 +313,8 @@ impl LanceScanner { with_row_address: false, include_deleted_rows: false, fragment_ids: None, + row_addr_allowlist: None, + row_addr_blocklist: None, index_segments: None, scalar_index_segment: None, nearest: None, @@ -370,6 +377,20 @@ impl LanceScanner { Ok(()) } + /// The external row prefilter, combined the same way as + /// `RowAddrMask::from_serialized_parts`: rows in the allowlist and not in the + /// blocklist. `None` when neither list is set (no mask, not "select nothing"). + fn row_addr_mask(&self) -> Option { + match (&self.row_addr_allowlist, &self.row_addr_blocklist) { + (Some(allow), Some(block)) => { + Some(RowAddrMask::from_allowed(allow.clone()).also_block(block.clone())) + } + (Some(allow), None) => Some(RowAddrMask::from_allowed(allow.clone())), + (None, Some(block)) => Some(RowAddrMask::from_block(block.clone())), + (None, None) => None, + } + } + fn apply_filter(&self, scanner: &mut lance::dataset::scanner::Scanner) -> Result<()> { if let Some(substrait) = &self.substrait_filter { scanner.filter_substrait(substrait)?; @@ -502,6 +523,18 @@ impl LanceScanner { "fragment_ids cannot be combined with an FTS query context; split the query by FTS index segment UUID instead".into(), )); } + // A segment-scoped scalar scan installs its own row prefilter from the segment's + // candidates, which would replace this one. + let row_addr_mask = self.row_addr_mask().map(Arc::new); + if row_addr_mask.is_some() && self.scalar_index_segment.is_some() { + return Err(lance_core::Error::invalid_input_source( + "row address allowlist/blocklist cannot be combined with scalar_index_segment" + .into(), + )); + } + if let Some(mask) = &row_addr_mask { + scanner.with_row_addr_prefilter(RowAddrMask::clone(mask)); + } if self.prefilter { scanner.prefilter(true); } @@ -577,6 +610,7 @@ impl LanceScanner { && (self.filter.is_some() || self.substrait_filter.is_some() || !self.additional_sql_filters.is_empty()), + row_addr_mask, batch_size: self.batch_size, scan_statistics_callback: self.scan_statistics_callback.clone(), }) @@ -631,6 +665,9 @@ struct PreparedFtsExecution { context: Arc, segments: Vec, scope_prefilter_to_fts_segments: bool, + // Re-applied to the rebuilt Match/Phrase execs: they are reconstructed from the + // planned ones, which do not expose the external mask Lance attached to them. + row_addr_mask: Option>, batch_size: Option, scan_statistics_callback: Option, } @@ -689,6 +726,7 @@ impl PreparedScanner { &distributed_fts.segments, &distributed_fts.context.prepared, selected_segments_have_current_fragments, + distributed_fts.row_addr_mask.as_ref(), )?; if rewritten.indexed_query_execs > 1 || rewritten.flat_match_query_execs > 1 @@ -785,6 +823,7 @@ fn rewrite_prepared_fts_plan( segments: &[IndexMetadata], prepared: &PreparedFtsQuery, selected_segments_have_current_fragments: bool, + row_addr_mask: Option<&Arc>, ) -> Result<(Arc, PreparedFtsPlanRewriteCounts)> { // Lance's ordinary FTS planner adds a flat-search branch for fragments not // covered by the logical index. A prepared INDEX_ONLY scan must omit that @@ -812,6 +851,7 @@ fn rewrite_prepared_fts_plan( segments, prepared, selected_segments_have_current_fragments, + row_addr_mask, )?; new_children.push(new_child); rewritten.indexed_query_execs += child_rewritten.indexed_query_execs; @@ -841,7 +881,8 @@ fn rewrite_prepared_fts_plan( exec.prefilter_source().clone(), segments.to_vec(), )? - .with_base_scorer(Arc::clone(scorer)); + .with_base_scorer(Arc::clone(scorer)) + .with_external_mask(row_addr_mask.cloned()); return Ok((Arc::new(replacement), rewritten)); } if let Some(exec) = rebuilt.downcast_ref::() { @@ -861,7 +902,8 @@ fn rewrite_prepared_fts_plan( exec.prefilter_source().clone(), segments.to_vec(), )? - .with_base_scorer(Arc::clone(scorer)); + .with_base_scorer(Arc::clone(scorer)) + .with_external_mask(row_addr_mask.cloned()); return Ok((Arc::new(replacement), rewritten)); } Ok((rebuilt, rewritten)) @@ -1654,6 +1696,87 @@ unsafe fn scanner_set_fragment_ids_inner( Ok(0) } +/// Restrict the scan or search to rows whose `_rowid` is in `row_ids`. +/// +/// The ids are in the dataset's `_rowid` space: the stable row id when stable +/// row ids are enabled, otherwise the row address. The allowlist is applied +/// before ranking: a nearest or full-text search (including a prepared FTS +/// query context) computes its top-k over the allowed rows only, and a plain +/// scan reads only those rows. Combined with any filter expression (AND) and +/// with a blocklist (allowed AND NOT blocked). `len == 0` allows no rows. +/// Calling it again replaces the previous allowlist. Must be called before the +/// scan starts; the ids are copied. +/// +/// Returns 0 on success, -1 on error. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_scanner_set_row_addr_allowlist( + scanner: *mut LanceScanner, + row_ids: *const u64, + len: usize, +) -> i32 { + scanner_poison_check!(scanner, -1); + scanner_ffi_try!(scanner, unsafe { + scanner_set_row_addr_list_inner(scanner, row_ids, len, RowAddrList::Allow) + }) +} + +/// Exclude rows whose `_rowid` is in `row_ids` from the scan or search. +/// +/// Same `_rowid` space, timing, and replacement semantics as +/// `lance_scanner_set_row_addr_allowlist`. `len == 0` blocks no rows. +/// +/// Returns 0 on success, -1 on error. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_scanner_set_row_addr_blocklist( + scanner: *mut LanceScanner, + row_ids: *const u64, + len: usize, +) -> i32 { + scanner_poison_check!(scanner, -1); + scanner_ffi_try!(scanner, unsafe { + scanner_set_row_addr_list_inner(scanner, row_ids, len, RowAddrList::Block) + }) +} + +enum RowAddrList { + Allow, + Block, +} + +unsafe fn scanner_set_row_addr_list_inner( + scanner: *mut LanceScanner, + row_ids: *const u64, + len: usize, + list: RowAddrList, +) -> Result { + if scanner.is_null() { + return Err(lance_core::Error::invalid_input_source( + "scanner is NULL".into(), + )); + } + if row_ids.is_null() && len > 0 { + return Err(lance_core::Error::invalid_input_source( + "row_ids is NULL but len > 0".into(), + )); + } + let s = unsafe { &mut *scanner }; + s.ensure_scan_not_started(match list { + RowAddrList::Allow => "row_addr_allowlist", + RowAddrList::Block => "row_addr_blocklist", + })?; + let ids = if len > 0 { + unsafe { std::slice::from_raw_parts(row_ids, len) } + } else { + &[] + }; + let tree_map: RowAddrTreeMap = ids.iter().collect(); + match list { + RowAddrList::Allow => s.row_addr_allowlist = Some(tree_map), + RowAddrList::Block => s.row_addr_blocklist = Some(tree_map), + } + Ok(0) +} + /// Set a Substrait filter on the scanner. /// /// `bytes` must point to a serialized Substrait @@ -3839,6 +3962,7 @@ mod tests { &segments, &distributed.context.prepared, has_current_fragments, + None, ) .unwrap(); assert_eq!(counts.indexed_query_execs, 1); diff --git a/tests/cpp/test_c_api.c b/tests/cpp/test_c_api.c index 77af1a6..a6842ac 100644 --- a/tests/cpp/test_c_api.c +++ b/tests/cpp/test_c_api.c @@ -263,6 +263,58 @@ static void test_scan_with_limit(const char *uri) { printf("OK\n"); } +static uint64_t count_stream_rows(struct ArrowArrayStream *stream) { + uint64_t total_rows = 0; + while (1) { + struct ArrowArray array; + memset(&array, 0, sizeof(array)); + ASSERT(stream->get_next(stream, &array) == 0, "get_next failed"); + if (array.release == NULL) break; + total_rows += (uint64_t)array.length; + array.release(&array); + } + if (stream->release) stream->release(stream); + return total_rows; +} + +static void test_scan_with_row_addr_lists(const char *uri) { + printf(" test_scan_with_row_addr_lists... "); + + LanceDataset *ds = lance_dataset_open(uri, NULL, 0); + ASSERT(ds != NULL, "dataset open failed"); + uint64_t expected_rows = lance_dataset_count_rows(ds); + CHECK_OK(); + + /* Rows 0, 2 and 4 of the first fragment, minus row 2. */ + const uint64_t allowed[] = {4, 0, 2}; + const uint64_t blocked[] = {2}; + LanceScanner *scanner = lance_scanner_new(ds, NULL, NULL); + ASSERT(scanner != NULL, "scanner creation failed"); + ASSERT(lance_scanner_set_row_addr_allowlist(scanner, allowed, 3) == 0, "set allowlist"); + ASSERT(lance_scanner_set_row_addr_blocklist(scanner, blocked, 1) == 0, "set blocklist"); + struct ArrowArrayStream stream; + memset(&stream, 0, sizeof(stream)); + ASSERT(lance_scanner_to_arrow_stream(scanner, &stream) == 0, "to_arrow_stream failed"); + ASSERT(count_stream_rows(&stream) == 2, "allowlist minus blocklist should keep 2 rows"); + ASSERT(lance_scanner_set_row_addr_allowlist(scanner, allowed, 3) == -1, + "allowlist after the scan started must fail"); + lance_scanner_close(scanner); + + /* An empty blocklist keeps every row; NULL is valid when len is 0. */ + scanner = lance_scanner_new(ds, NULL, NULL); + ASSERT(scanner != NULL, "scanner creation failed"); + ASSERT(lance_scanner_set_row_addr_blocklist(scanner, NULL, 0) == 0, "set empty blocklist"); + ASSERT(lance_scanner_set_row_addr_allowlist(scanner, NULL, 1) == -1, + "NULL row_ids with len > 0 must fail"); + memset(&stream, 0, sizeof(stream)); + ASSERT(lance_scanner_to_arrow_stream(scanner, &stream) == 0, "to_arrow_stream failed"); + ASSERT(count_stream_rows(&stream) == expected_rows, "empty blocklist should keep every row"); + lance_scanner_close(scanner); + + lance_dataset_close(ds); + printf("OK\n"); +} + /* Copy the Arrow C Data Interface format of the `blob` column of a stream's * schema into `out`; `out` is empty when the column is missing. */ static void blob_column_format(struct ArrowArrayStream *stream, char *out, size_t out_len) { @@ -1479,6 +1531,7 @@ int main(int argc, char **argv) { test_scan(uri); test_distance_range(uri); test_scan_with_limit(uri); + test_scan_with_row_addr_lists(uri); test_scanner_blob_handling(blob_uri); test_take_blobs(blob_uri); test_versions(uri); diff --git a/tests/cpp/test_cpp_api.cpp b/tests/cpp/test_cpp_api.cpp index 3bc313e..6e08d92 100644 --- a/tests/cpp/test_cpp_api.cpp +++ b/tests/cpp/test_cpp_api.cpp @@ -205,6 +205,43 @@ static void test_scanner_fluent(const std::string& uri) { PASS(); } +static void test_scanner_row_addr_lists(const std::string& uri) { + TEST(test_scanner_row_addr_lists); + + auto ds = lance::Dataset::open(uri); + + // Rows 0, 2 and 4 of the first fragment, minus row 2. + auto scanner = ds.scan(); + scanner.row_addr_allowlist(std::vector{4, 0, 2}) + .row_addr_blocklist(std::vector{2}); + ArrowArrayStream stream; + memset(&stream, 0, sizeof(stream)); + scanner.to_arrow_stream(&stream); + uint64_t total = 0; + while (true) { + ArrowArray arr; + memset(&arr, 0, sizeof(arr)); + int rc = stream.get_next(&stream, &arr); + assert(rc == 0); + if (!arr.release) break; + total += (uint64_t)arr.length; + arr.release(&arr); + } + if (stream.release) stream.release(&stream); + assert(total == 2); + + bool caught = false; + try { + ds.scan().row_addr_allowlist(nullptr, 1); + } catch (const lance::Error&) { + caught = true; + } + assert(caught); + printf("rows=%llu... ", (unsigned long long)total); + + PASS(); +} + static void test_scanner_async_stream_ownership(const std::string& uri) { TEST(test_scanner_async_stream_ownership); @@ -1367,6 +1404,7 @@ int main(int argc, char** argv) { test_shared_session(uri); test_dataset_schema(uri); test_scanner_fluent(uri); + test_scanner_row_addr_lists(uri); test_distance_range(uri); test_scanner_async_stream_ownership(uri); test_scanner_blob_handling(blob_uri); diff --git a/tests/row_addr_prefilter_test.rs b/tests/row_addr_prefilter_test.rs new file mode 100644 index 0000000..73e4e77 --- /dev/null +++ b/tests/row_addr_prefilter_test.rs @@ -0,0 +1,381 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Row address allowlist / blocklist prefilters (`lance_scanner_set_row_addr_*`). + +use std::collections::BTreeMap; +use std::ffi::{CStr, CString}; +use std::ptr; +use std::sync::Arc; + +use arrow::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream}; +use arrow::record_batch::RecordBatchIterator; +use arrow_array::{Array, FixedSizeListArray, Float32Array, Int32Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema}; +use lance::Dataset; +use lance::dataset::{WriteMode, WriteParams}; +use lance_c::*; + +const TEXTS: [&str; 10] = [ + "apple banana", + "apple", + "banana cherry", + "cherry", + "apple cherry", + "apple banana cherry", + "banana", + "apple apple", + "durian", + "banana apple", +]; + +fn c(s: &str) -> CString { + CString::new(s).unwrap() +} + +fn last_error() -> String { + unsafe { + let msg = lance_last_error_message(); + let message = CStr::from_ptr(msg).to_string_lossy().into_owned(); + lance_free_string(msg); + message + } +} + +fn ok(code: i32) { + assert_eq!(code, 0, "{}", last_error()); +} + +/// Two fragments of five rows with stable row ids, so `_rowid` (== `id`, 0..9) +/// differs from the row address of every row in the second fragment. Each row +/// has a text column with an inverted index and a 2-d vector `[id, 1]`. +fn fixture() -> (tempfile::TempDir, *mut LanceDataset) { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("row_addr_prefilter.lance"); + lance_c::runtime::block_on(async { + for fragment in 0..2 { + let ids: Vec = (fragment * 5..(fragment + 1) * 5).collect(); + let texts: Vec<&str> = ids.iter().map(|id| TEXTS[*id as usize]).collect(); + let vectors = FixedSizeListArray::try_new( + Arc::new(Field::new("item", DataType::Float32, false)), + 2, + Arc::new(Float32Array::from( + ids.iter() + .flat_map(|id| [*id as f32, 1.]) + .collect::>(), + )), + None, + ) + .unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("text", DataType::Utf8, false), + Field::new("embedding", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(texts)), + Arc::new(vectors), + ], + ) + .unwrap(); + Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + path.to_str().unwrap(), + Some(WriteParams { + mode: if fragment == 0 { + WriteMode::Create + } else { + WriteMode::Append + }, + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + } + }); + + let uri = c(path.to_str().unwrap()); + let dataset = unsafe { lance_dataset_open(uri.as_ptr(), ptr::null(), 0) }; + assert!(!dataset.is_null(), "{}", last_error()); + let column = c("text"); + let params = c(r#"{"base_tokenizer":"simple","language":"English"}"#); + ok(unsafe { + lance_dataset_create_scalar_index( + dataset, + column.as_ptr(), + ptr::null(), + LanceScalarIndexType::Inverted as i32, + params.as_ptr(), + false, + ) + }); + (dir, dataset) +} + +struct Scan(*mut LanceScanner); + +impl Scan { + fn new(dataset: *const LanceDataset, sql_filter: Option<&str>) -> Self { + let filter = sql_filter.map(c); + let scanner = unsafe { + lance_scanner_new( + dataset, + ptr::null(), + filter.as_ref().map_or(ptr::null(), |f| f.as_ptr()), + ) + }; + assert!(!scanner.is_null(), "{}", last_error()); + Self(scanner) + } + + fn allow(self, row_ids: &[u64]) -> Self { + ok(unsafe { + lance_scanner_set_row_addr_allowlist(self.0, row_ids.as_ptr(), row_ids.len()) + }); + self + } + + fn block(self, row_ids: &[u64]) -> Self { + ok(unsafe { + lance_scanner_set_row_addr_blocklist(self.0, row_ids.as_ptr(), row_ids.len()) + }); + self + } + + /// `id` → the named score column (`_score` / `_distance`), or 0 when absent. + fn collect(self, score_column: Option<&str>) -> BTreeMap { + let mut stream = FFI_ArrowArrayStream::empty(); + ok(unsafe { lance_scanner_to_arrow_stream(self.0, &mut stream) }); + let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream).unwrap() }; + let mut out = BTreeMap::new(); + for batch in reader { + let batch = batch.unwrap(); + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let scores = score_column.map(|name| { + batch + .column_by_name(name) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .clone() + }); + for row in 0..batch.num_rows() { + let score = scores.as_ref().map_or(0.0, |s| s.value(row)); + assert!(out.insert(ids.value(row), score).is_none()); + } + } + out + } +} + +impl Drop for Scan { + fn drop(&mut self) { + unsafe { lance_scanner_close(self.0) } + } +} + +fn ids(result: &BTreeMap) -> Vec { + result.keys().copied().collect() +} + +/// `full` restricted to `allowed`: what a masked search must return, scores unchanged. +fn restrict(full: &BTreeMap, allowed: &[u64]) -> BTreeMap { + full.iter() + .filter(|(id, _)| allowed.contains(&(**id as u64))) + .map(|(id, score)| (*id, *score)) + .collect() +} + +#[test] +fn test_row_addr_lists_restrict_plain_scan() { + let (_dir, dataset) = fixture(); + + // Ids 5 and 8 live in the second fragment, so this also checks the lists are + // keyed by stable row id rather than row address. + let all = |scan: Scan| ids(&scan.collect(None)); + assert_eq!( + all(Scan::new(dataset, None).allow(&[8, 1, 5])), + vec![1, 5, 8] + ); + assert_eq!( + all(Scan::new(dataset, None).block(&[0, 1, 2, 3, 4])), + vec![5, 6, 7, 8, 9] + ); + assert_eq!( + all(Scan::new(dataset, None).allow(&[1, 5, 8]).block(&[5])), + vec![1, 8] + ); + assert!(all(Scan::new(dataset, None).allow(&[])).is_empty()); + assert_eq!(all(Scan::new(dataset, None).block(&[])).len(), 10); + // A later call replaces the earlier list. + assert_eq!( + all(Scan::new(dataset, None).allow(&[1]).allow(&[2])), + vec![2] + ); + // ANDed with the filter expression. + assert_eq!( + all(Scan::new(dataset, Some("id > 2")).allow(&[1, 5, 8])), + vec![5, 8] + ); + + unsafe { lance_dataset_close(dataset) }; +} + +#[test] +fn test_row_addr_lists_reject_invalid_inputs() { + let (_dir, dataset) = fixture(); + + let scan = Scan::new(dataset, None); + assert_eq!( + unsafe { lance_scanner_set_row_addr_allowlist(scan.0, ptr::null(), 1) }, + -1 + ); + assert!(last_error().contains("row_ids is NULL but len > 0")); + assert_eq!( + unsafe { lance_scanner_set_row_addr_blocklist(scan.0, ptr::null(), 1) }, + -1 + ); + assert!(last_error().contains("row_ids is NULL but len > 0")); + assert_eq!( + unsafe { lance_scanner_set_row_addr_allowlist(ptr::null_mut(), ptr::null(), 0) }, + -1 + ); + // NULL is valid for an empty list. + ok(unsafe { lance_scanner_set_row_addr_allowlist(scan.0, ptr::null(), 0) }); + + let started = Scan::new(dataset, None); + let mut stream = FFI_ArrowArrayStream::empty(); + ok(unsafe { lance_scanner_to_arrow_stream(started.0, &mut stream) }); + drop(unsafe { ArrowArrayStreamReader::from_raw(&mut stream).unwrap() }); + assert_eq!( + unsafe { lance_scanner_set_row_addr_allowlist(started.0, [1u64].as_ptr(), 1) }, + -1 + ); + assert!(last_error().contains("row_addr_allowlist must be set before the scan starts")); + assert_eq!( + unsafe { lance_scanner_set_row_addr_blocklist(started.0, [1u64].as_ptr(), 1) }, + -1 + ); + assert!(last_error().contains("row_addr_blocklist must be set before the scan starts")); + + let segment = Scan::new(dataset, Some("id > 2")).allow(&[1]); + ok(unsafe { lance_scanner_set_scalar_index_segment(segment.0, [7u8; 16].as_ptr()) }); + let mut stream = FFI_ArrowArrayStream::empty(); + assert_eq!( + unsafe { lance_scanner_to_arrow_stream(segment.0, &mut stream) }, + -1 + ); + assert!(last_error().contains("cannot be combined with scalar_index_segment")); + + unsafe { lance_dataset_close(dataset) }; +} + +#[test] +fn test_row_addr_allowlist_full_text_search_scores_allowed_rows_only() { + let (_dir, dataset) = fixture(); + let query = c("apple banana"); + let search = |scan: Scan| { + ok(unsafe { lance_scanner_full_text_search(scan.0, query.as_ptr(), ptr::null(), 0) }); + scan.collect(Some("_score")) + }; + + let full = search(Scan::new(dataset, None)); + assert_eq!(ids(&full), vec![0, 1, 2, 4, 5, 6, 7, 9]); + let allowed = [0, 3, 6, 9]; + assert_eq!( + search(Scan::new(dataset, None).allow(&allowed)), + restrict(&full, &allowed) + ); + assert_eq!( + search(Scan::new(dataset, None).block(&[0, 1, 2, 3, 4])), + restrict(&full, &[5, 6, 7, 8, 9]) + ); + + unsafe { lance_dataset_close(dataset) }; +} + +#[test] +fn test_row_addr_allowlist_prepared_fts_scores_allowed_rows_only() { + let (_dir, dataset) = fixture(); + let column = c("text"); + let query = c("apple banana"); + + for (operator, expected_ids) in [ + (LanceFtsMatchOperator::Or, vec![0, 1, 2, 4, 5, 6, 7, 9]), + (LanceFtsMatchOperator::And, vec![0, 5, 9]), + ] { + let context = unsafe { + lance_dataset_prepare_fts_match_query( + dataset, + column.as_ptr(), + query.as_ptr(), + operator as i32, + 0, + LanceFtsCoverageMode::Strict as i32, + ) + }; + assert!(!context.is_null(), "{}", last_error()); + let search = |scan: Scan| { + ok(unsafe { lance_scanner_set_fts_query_context(scan.0, context) }); + scan.collect(Some("_score")) + }; + + // The prepared plan is rebuilt around the shared global scorer; the mask must + // survive that rewrite and leave the scores of the allowed rows unchanged. + let full = search(Scan::new(dataset, None)); + assert_eq!(ids(&full), expected_ids); + for allowed in [&[0u64, 5, 6][..], &[9], &[3, 8], &[]] { + assert_eq!( + search(Scan::new(dataset, None).allow(allowed)), + restrict(&full, allowed) + ); + } + assert_eq!( + search(Scan::new(dataset, Some("id >= 5")).allow(&[0, 5, 6])), + restrict(&full, &[5, 6]) + ); + unsafe { lance_fts_query_context_close(context) }; + } + + unsafe { lance_dataset_close(dataset) }; +} + +#[test] +fn test_row_addr_allowlist_nearest_ranks_allowed_rows_only() { + let (_dir, dataset) = fixture(); + let column = c("embedding"); + let query = [9.0f32, 1.0]; + let nearest = |scan: Scan| { + ok(unsafe { + lance_scanner_nearest( + scan.0, + column.as_ptr(), + query.as_ptr().cast(), + query.len(), + LanceDataType::Float32 as i32, + 2, + ) + }); + ids(&scan.collect(Some("_distance"))) + }; + + assert_eq!(nearest(Scan::new(dataset, None)), vec![8, 9]); + // Top-k is computed over the allowed rows, not filtered after ranking. + assert_eq!( + nearest(Scan::new(dataset, None).allow(&[1, 3, 6])), + vec![3, 6] + ); + + unsafe { lance_dataset_close(dataset) }; +}