From a4d6e489c627fe4b0e49d9a2991c9436313debc0 Mon Sep 17 00:00:00 2001 From: zhangstar333 Date: Tue, 1 Sep 2026 11:15:00 +0800 Subject: [PATCH 1/3] update --- src/fts_query.rs | 2 ++ src/scanner.rs | 44 ++++++++++++++++++++++--- tests/c_api_test.rs | 78 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 120 insertions(+), 4 deletions(-) diff --git a/src/fts_query.rs b/src/fts_query.rs index cd194c7..e85ed99 100644 --- a/src/fts_query.rs +++ b/src/fts_query.rs @@ -54,6 +54,7 @@ pub(crate) struct FtsQueryContextInner { pub(crate) query: FullTextSearchQuery, pub(crate) segments: Vec, pub(crate) scorer: Arc, + pub(crate) has_unindexed_fragments: bool, } impl FtsQueryContextInner { @@ -216,6 +217,7 @@ async fn prepare_fts_query_context( query, segments, scorer, + has_unindexed_fragments: !unindexed_fragment_ids.is_empty(), }) } diff --git a/src/scanner.rs b/src/scanner.rs index 5b8c34e..f60f0c5 100644 --- a/src/scanner.rs +++ b/src/scanner.rs @@ -179,6 +179,45 @@ impl LanceScanner { Ok(()) } + /// Restrict an INDEX_ONLY prepared FTS scan to fragments covered by the + /// selected committed segments. This is deliberately separate from + /// `fast_search`: that option is scanner-wide, changes unrelated scalar + /// index fallback behavior, and also forces `_rowid` into the output. + fn apply_prepared_fts_fragment_filter( + &self, + scanner: &mut lance::dataset::scanner::Scanner, + context: &FtsQueryContextInner, + segments: &[IndexMetadata], + ) -> Result<()> { + if !context.has_unindexed_fragments { + return Ok(()); + } + + 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!( + "prepared FTS segment {} lost its validated fragment coverage", + segment.uuid + )) + })?; + selected_fragment_ids.extend(fragment_bitmap.iter()); + } + + let selected_fragments = self + .dataset + .get_fragments() + .into_iter() + .filter(|fragment| { + u32::try_from(fragment.id()) + .is_ok_and(|fragment_id| selected_fragment_ids.contains(&fragment_id)) + }) + .map(|fragment| fragment.metadata().clone()) + .collect(); + scanner.with_fragments(selected_fragments); + Ok(()) + } + fn apply_filter(&self, scanner: &mut lance::dataset::scanner::Scanner) -> Result<()> { if let Some(substrait) = &self.substrait_filter { scanner.filter_substrait(substrait)?; @@ -282,11 +321,8 @@ impl LanceScanner { let distributed_fts = if let Some(context) = &self.fts_context { context.validate_dataset_identity(&self.dataset)?; let segments = select_fts_segments(context, self.fts_index_segments.as_deref())?; + self.apply_prepared_fts_fragment_filter(&mut scanner, context, &segments)?; scanner.full_text_search(context.query.clone())?; - // Both STRICT and INDEX_ONLY context scans must use only the - // committed segments pinned in the context. In STRICT mode all - // current fragments were already proven covered during prepare. - scanner.fast_search(); Some(PreparedFtsExecution { context: Arc::clone(context), segments, diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs index 74b9f85..3627c4e 100644 --- a/tests/c_api_test.rs +++ b/tests/c_api_test.rs @@ -5762,6 +5762,84 @@ fn load_fts_segment_uuids(uri: &str, column: &str) -> Vec<[u8; 16]> { }) } +#[test] +fn test_prepared_fts_row_id_output_is_explicit() { + 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 + ); + let context = unsafe { + lance_dataset_prepare_fts_query( + dataset, + column.as_ptr(), + query.as_ptr(), + 0, + LanceFtsCoverageMode::Strict as i32, + ) + }; + assert!(!context.is_null(), "{}", unsafe { + std::ffi::CStr::from_ptr(lance_last_error_message()).to_string_lossy() + }); + + let id = c_str("id"); + let columns = [id.as_ptr(), ptr::null()]; + let scan_schema = |with_row_id: bool| { + let scanner = unsafe { lance_scanner_new(dataset, columns.as_ptr(), ptr::null()) }; + assert!(!scanner.is_null()); + if with_row_id { + assert_eq!(unsafe { lance_scanner_with_row_id(scanner, true) }, 0); + } + assert_eq!( + unsafe { lance_scanner_set_fts_query_context(scanner, context) }, + 0 + ); + let mut stream = FFI_ArrowArrayStream::empty(); + assert_eq!( + unsafe { lance_scanner_to_arrow_stream(scanner, &mut stream) }, + 0, + "{}", + unsafe { std::ffi::CStr::from_ptr(lance_last_error_message()).to_string_lossy() } + ); + let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream).unwrap() }; + let schema = reader.schema(); + let rows: usize = reader.map(|batch| batch.unwrap().num_rows()).sum(); + assert!(rows > 0); + unsafe { lance_scanner_close(scanner) }; + schema + }; + + let without_row_id = scan_schema(false); + assert_eq!(without_row_id.fields().len(), 2); + assert!(without_row_id.field_with_name("id").is_ok()); + assert!(without_row_id.field_with_name("_score").is_ok()); + assert!(without_row_id.field_with_name("_rowid").is_err()); + + let with_row_id = scan_schema(true); + assert_eq!(with_row_id.fields().len(), 3); + assert!(with_row_id.field_with_name("id").is_ok()); + assert!(with_row_id.field_with_name("_score").is_ok()); + assert!(with_row_id.field_with_name("_rowid").is_ok()); + + unsafe { lance_fts_query_context_close(context) }; + unsafe { lance_dataset_close(dataset) }; +} + #[test] fn test_prepare_fts_query_index_only_allows_unindexed_fragment() { let (_tmp, uri) = create_test_dataset(); From 6f0fae4cc51bf144685564b1d2e9f6f7afc71f8a Mon Sep 17 00:00:00 2001 From: zhangstar333 Date: Tue, 1 Sep 2026 12:34:05 +0800 Subject: [PATCH 2/3] update --- src/fts_query.rs | 2 - src/scanner.rs | 253 ++++++++++++++++++++++++++++++++++---------- tests/c_api_test.rs | 79 ++++++++++++++ 3 files changed, 279 insertions(+), 55 deletions(-) diff --git a/src/fts_query.rs b/src/fts_query.rs index e85ed99..cd194c7 100644 --- a/src/fts_query.rs +++ b/src/fts_query.rs @@ -54,7 +54,6 @@ pub(crate) struct FtsQueryContextInner { pub(crate) query: FullTextSearchQuery, pub(crate) segments: Vec, pub(crate) scorer: Arc, - pub(crate) has_unindexed_fragments: bool, } impl FtsQueryContextInner { @@ -217,7 +216,6 @@ async fn prepare_fts_query_context( query, segments, scorer, - has_unindexed_fragments: !unindexed_fragment_ids.is_empty(), }) } diff --git a/src/scanner.rs b/src/scanner.rs index f60f0c5..0c29b17 100644 --- a/src/scanner.rs +++ b/src/scanner.rs @@ -12,13 +12,13 @@ use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; use arrow::ffi_stream::FFI_ArrowArrayStream; use arrow_schema::SchemaRef; -use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::{ExecutionPlan, empty::EmptyExec}; use futures::{FutureExt, Stream, StreamExt}; use lance::Dataset; use lance::dataset::scanner::{ DatasetRecordBatchStream, ExecutionStatsCallback, ExecutionSummaryCounts, }; -use lance::io::exec::fts::MatchQueryExec; +use lance::io::exec::fts::{FlatMatchQueryExec, MatchQueryExec}; use lance_core::Result; use lance_index::scalar::FullTextSearchQuery; use lance_io::stream::RecordBatchStream; @@ -179,45 +179,6 @@ impl LanceScanner { Ok(()) } - /// Restrict an INDEX_ONLY prepared FTS scan to fragments covered by the - /// selected committed segments. This is deliberately separate from - /// `fast_search`: that option is scanner-wide, changes unrelated scalar - /// index fallback behavior, and also forces `_rowid` into the output. - fn apply_prepared_fts_fragment_filter( - &self, - scanner: &mut lance::dataset::scanner::Scanner, - context: &FtsQueryContextInner, - segments: &[IndexMetadata], - ) -> Result<()> { - if !context.has_unindexed_fragments { - return Ok(()); - } - - 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!( - "prepared FTS segment {} lost its validated fragment coverage", - segment.uuid - )) - })?; - selected_fragment_ids.extend(fragment_bitmap.iter()); - } - - let selected_fragments = self - .dataset - .get_fragments() - .into_iter() - .filter(|fragment| { - u32::try_from(fragment.id()) - .is_ok_and(|fragment_id| selected_fragment_ids.contains(&fragment_id)) - }) - .map(|fragment| fragment.metadata().clone()) - .collect(); - scanner.with_fragments(selected_fragments); - Ok(()) - } - fn apply_filter(&self, scanner: &mut lance::dataset::scanner::Scanner) -> Result<()> { if let Some(substrait) = &self.substrait_filter { scanner.filter_substrait(substrait)?; @@ -321,7 +282,6 @@ impl LanceScanner { let distributed_fts = if let Some(context) = &self.fts_context { context.validate_dataset_identity(&self.dataset)?; let segments = select_fts_segments(context, self.fts_index_segments.as_deref())?; - self.apply_prepared_fts_fragment_filter(&mut scanner, context, &segments)?; scanner.full_text_search(context.query.clone())?; Some(PreparedFtsExecution { context: Arc::clone(context), @@ -361,14 +321,24 @@ impl PreparedScanner { return self.scanner.try_into_stream().await; }; let plan = self.scanner.create_plan().await?; - let (plan, replaced) = replace_match_query_exec( + let selected_segments_have_current_fragments = segments_have_current_fragments( + &distributed_fts.context.dataset, + &distributed_fts.segments, + )?; + let (plan, rewritten) = rewrite_prepared_fts_plan( plan, &distributed_fts.segments, &distributed_fts.context.scorer, + selected_segments_have_current_fragments, )?; - if replaced != 1 { + if rewritten.match_query_execs > 1 + || rewritten.flat_match_query_execs > 1 + || rewritten.match_query_execs + rewritten.flat_match_query_execs == 0 + || (selected_segments_have_current_fragments && rewritten.match_query_execs != 1) + { return Err(lance_core::Error::internal(format!( - "expected exactly one MatchQueryExec in prepared FTS plan, replaced {replaced}" + "unexpected prepared FTS plan for selected segments with current fragment coverage {selected_segments_have_current_fragments}: rewrote {} MatchQueryExec node(s) and removed {} FlatMatchQueryExec node(s)", + rewritten.match_query_execs, rewritten.flat_match_query_execs ))); } let stream = lance_datafusion::exec::execute_plan( @@ -415,22 +385,81 @@ fn select_fts_segments( Ok(selected) } -fn replace_match_query_exec( +fn segments_have_current_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::>>()?; + for segment in segments { + let fragment_bitmap = segment.fragment_bitmap.as_ref().ok_or_else(|| { + lance_core::Error::internal(format!( + "prepared FTS segment {} lost its validated fragment coverage", + segment.uuid + )) + })?; + if fragment_bitmap + .iter() + .any(|fragment_id| current_fragment_ids.contains(&fragment_id)) + { + return Ok(true); + } + } + Ok(false) +} + +#[derive(Default)] +struct PreparedFtsPlanRewriteCounts { + match_query_execs: usize, + flat_match_query_execs: usize, +} + +fn rewrite_prepared_fts_plan( plan: Arc, segments: &[IndexMetadata], scorer: &Arc, -) -> Result<(Arc, usize)> { + selected_segments_have_current_fragments: bool, +) -> 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 + // branch, but using Scanner::with_fragments to do so would turn an + // otherwise unfiltered index search into a full row-id prefilter scan. + if plan.downcast_ref::().is_some() { + return Ok(( + Arc::new(EmptyExec::new(plan.schema())), + PreparedFtsPlanRewriteCounts { + match_query_execs: 0, + flat_match_query_execs: 1, + }, + )); + } + let children = plan.children(); - let mut replaced = 0; + let mut rewritten = PreparedFtsPlanRewriteCounts::default(); let rebuilt = if children.is_empty() { plan } else { let mut new_children = Vec::with_capacity(children.len()); for child in children { - let (new_child, child_replaced) = - replace_match_query_exec(Arc::clone(child), segments, scorer)?; + let (new_child, child_rewritten) = rewrite_prepared_fts_plan( + Arc::clone(child), + segments, + scorer, + selected_segments_have_current_fragments, + )?; new_children.push(new_child); - replaced += child_replaced; + rewritten.match_query_execs += child_rewritten.match_query_execs; + rewritten.flat_match_query_execs += child_rewritten.flat_match_query_execs; } plan.with_new_children(new_children).map_err(|error| { lance_core::Error::internal(format!( @@ -440,6 +469,10 @@ fn replace_match_query_exec( }; if let Some(exec) = rebuilt.downcast_ref::() { + rewritten.match_query_execs += 1; + if !selected_segments_have_current_fragments { + return Ok((Arc::new(EmptyExec::new(rebuilt.schema())), rewritten)); + } let replacement = MatchQueryExec::new_with_segments( Arc::clone(exec.dataset()), exec.query().clone(), @@ -448,9 +481,9 @@ fn replace_match_query_exec( segments.to_vec(), ) .with_base_scorer(Arc::clone(scorer)); - return Ok((Arc::new(replacement), replaced + 1)); + return Ok((Arc::new(replacement), rewritten)); } - Ok((rebuilt, replaced)) + Ok((rebuilt, rewritten)) } /// Type of a dynamically named scan metric. @@ -2118,6 +2151,9 @@ mod tests { use super::*; use crate::dataset::{lance_dataset_close, lance_dataset_open}; use crate::error::{lance_last_error_code, lance_last_error_message}; + use crate::fts_query::{ + LanceFtsCoverageMode, lance_dataset_prepare_fts_query, lance_fts_query_context_close, + }; use std::ffi::{CStr, CString}; use std::sync::atomic::{AtomicI32, AtomicUsize}; use std::sync::{Barrier, mpsc}; @@ -2125,6 +2161,9 @@ mod tests { use arrow_array::{Int32Array, RecordBatch, StringArray}; use arrow_schema::{DataType, Field, Schema}; + use lance::index::DatasetIndexExt; + use lance::io::exec::PreFilterSource; + use lance_index::{IndexType, scalar::InvertedIndexParams}; /// Write a 3-row dataset to a tempdir, returning (tempdir, uri). fn create_test_dataset() -> (tempfile::TempDir, String) { @@ -2169,6 +2208,114 @@ mod tests { .store(true, Ordering::SeqCst); } + fn prepared_fts_plan_shape(plan: &Arc) -> (usize, usize, usize) { + let mut match_query_execs = 0; + let mut flat_match_query_execs = 0; + let mut filtered_row_id_prefilters = 0; + if let Some(exec) = plan.downcast_ref::() { + match_query_execs += 1; + if matches!(exec.prefilter_source(), PreFilterSource::FilteredRowIds(_)) { + filtered_row_id_prefilters += 1; + } + } + if plan.downcast_ref::().is_some() { + flat_match_query_execs += 1; + } + for child in plan.children() { + let (child_match, child_flat, child_filtered) = + prepared_fts_plan_shape(&Arc::clone(child)); + match_query_execs += child_match; + flat_match_query_execs += child_flat; + filtered_row_id_prefilters += child_filtered; + } + ( + match_query_execs, + flat_match_query_execs, + filtered_row_id_prefilters, + ) + } + + #[test] + fn prepared_fts_index_only_plan_does_not_scan_indexed_fragment_row_ids() { + let (_tmp, uri) = create_test_dataset(); + block_on(async { + let mut dataset = Dataset::open(&uri).await.unwrap(); + dataset + .create_index( + &["name"], + IndexType::Inverted, + None, + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![4])), + Arc::new(StringArray::from(vec!["a"])), + ], + ) + .unwrap(); + dataset + .append( + arrow::record_batch::RecordBatchIterator::new(vec![Ok(batch)], schema), + None, + ) + .await + .unwrap(); + }); + + let (dataset, scanner) = open_dataset_and_scanner(&uri); + let column = CString::new("name").unwrap(); + let query = CString::new("a").unwrap(); + let context = unsafe { + lance_dataset_prepare_fts_query( + dataset, + column.as_ptr(), + query.as_ptr(), + 0, + LanceFtsCoverageMode::IndexOnly as i32, + ) + }; + assert!(!context.is_null()); + assert_eq!( + unsafe { lance_scanner_set_fts_query_context(scanner, context) }, + 0 + ); + + let prepared = unsafe { &*scanner }.build_scanner().unwrap(); + let distributed = prepared.distributed_fts.as_ref().unwrap(); + let segments = distributed.segments.clone(); + let scorer = Arc::clone(&distributed.context.scorer); + let plan = block_on(prepared.scanner.create_plan()).unwrap(); + assert_eq!( + prepared_fts_plan_shape(&plan), + (1, 1, 0), + "an unfiltered prepared FTS plan must not materialize selected fragment row IDs" + ); + + let has_current_fragments = + segments_have_current_fragments(&distributed.context.dataset, &segments).unwrap(); + let (rewritten, counts) = + rewrite_prepared_fts_plan(plan, &segments, &scorer, has_current_fragments).unwrap(); + assert_eq!(counts.match_query_execs, 1); + assert_eq!(counts.flat_match_query_execs, 1); + assert_eq!(prepared_fts_plan_shape(&rewritten), (1, 0, 0)); + + unsafe { + lance_scanner_close(scanner); + lance_fts_query_context_close(context); + lance_dataset_close(dataset); + } + } + /// Assert the pending thread-local error is `Panic` carrying the poison /// message; consumes it so the next assertion starts from a clean slate. fn assert_poison_error_pending() { diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs index 3627c4e..8805764 100644 --- a/tests/c_api_test.rs +++ b/tests/c_api_test.rs @@ -5943,6 +5943,85 @@ fn test_prepare_fts_query_index_only_allows_unindexed_fragment() { unsafe { lance_dataset_close(dataset) }; } +#[test] +fn test_prepared_fts_index_only_empty_segment_returns_empty_shard() { + use lance::index::DatasetIndexExt; + use lance_index::{IndexType, scalar::InvertedIndexParams}; + + let (_tmp, uri) = create_test_dataset(); + lance_c::runtime::block_on(async { + let mut dataset = Dataset::open(&uri).await.unwrap(); + let params = InvertedIndexParams::default(); + dataset + .create_index_builder(&["name"], IndexType::Inverted, ¶ms) + .name("empty_name_fts".to_string()) + .train(false) + .await + .unwrap(); + let segments = dataset + .load_indices_by_name("empty_name_fts") + .await + .unwrap(); + assert_eq!(segments.len(), 1); + assert!( + segments[0] + .fragment_bitmap + .as_ref() + .is_some_and(|fragment_bitmap| fragment_bitmap.is_empty()) + ); + }); + + let uri_c = c_str(&uri); + let column = c_str("name"); + let query = c_str("alice"); + let dataset = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; + let context = unsafe { + lance_dataset_prepare_fts_query( + dataset, + column.as_ptr(), + query.as_ptr(), + 0, + LanceFtsCoverageMode::IndexOnly as i32, + ) + }; + assert!(!context.is_null(), "{}", unsafe { + std::ffi::CStr::from_ptr(lance_last_error_message()).to_string_lossy() + }); + let segment_uuids = load_fts_segment_uuids(&uri, "name"); + assert_eq!(segment_uuids.len(), 1); + + let scanner = unsafe { lance_scanner_new(dataset, ptr::null(), ptr::null()) }; + assert_eq!( + unsafe { lance_scanner_set_fts_query_context(scanner, context) }, + 0 + ); + assert_eq!( + unsafe { + lance_scanner_set_fts_index_segments( + scanner, + segment_uuids.as_ptr().cast::(), + segment_uuids.len(), + ) + }, + 0 + ); + + let mut stream = FFI_ArrowArrayStream::empty(); + assert_eq!( + unsafe { lance_scanner_to_arrow_stream(scanner, &mut stream) }, + 0, + "{}", + unsafe { std::ffi::CStr::from_ptr(lance_last_error_message()).to_string_lossy() } + ); + let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream).unwrap() }; + let total_rows: usize = reader.map(|batch| batch.unwrap().num_rows()).sum(); + assert_eq!(total_rows, 0); + + unsafe { lance_scanner_close(scanner) }; + unsafe { lance_fts_query_context_close(context) }; + unsafe { lance_dataset_close(dataset) }; +} + #[test] fn test_prepared_fts_global_scorer_is_shared_across_segment_splits() { use lance::index::DatasetIndexExt; From 9499291da71ccea3855e4bddeeade355730a6738 Mon Sep 17 00:00:00 2001 From: zhangstar333 Date: Tue, 1 Sep 2026 16:48:41 +0800 Subject: [PATCH 3/3] data cache --- Cargo.lock | 3 + Cargo.toml | 3 + README.md | 38 +++ include/lance/lance.h | 99 ++++++- include/lance/lance.hpp | 49 +++ src/dataset.rs | 72 ++++- src/lib.rs | 2 + src/read_provider.rs | 590 +++++++++++++++++++++++++++++++++++++ tests/c_api_test.rs | 197 ++++++++++++- tests/cpp/test_c_api.c | 8 +- tests/cpp/test_cpp_api.cpp | 8 +- 11 files changed, 1061 insertions(+), 8 deletions(-) create mode 100644 src/read_provider.rs diff --git a/Cargo.lock b/Cargo.lock index 55e962a..b79205a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3703,6 +3703,8 @@ dependencies = [ "arrow", "arrow-array", "arrow-schema", + "async-trait", + "bytes", "chrono", "datafusion", "futures", @@ -3718,6 +3720,7 @@ dependencies = [ "lance-table", "libc", "log", + "object_store", "pin-project", "prost", "snafu", diff --git a/Cargo.toml b/Cargo.toml index 3fd7931..b0aab82 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,6 +18,8 @@ rust-version = "1.91.0" crate-type = ["cdylib", "staticlib", "rlib"] [dependencies] +async-trait = "0.1" +bytes = "1" lance = { git = "https://github.com/lance-format/lance.git", rev = "e934cc2c", features = ["substrait"] } lance-core = { git = "https://github.com/lance-format/lance.git", rev = "e934cc2c" } lance-file = { git = "https://github.com/lance-format/lance.git", rev = "e934cc2c" } @@ -39,6 +41,7 @@ tokio = { version = "1", features = ["rt-multi-thread", "sync"] } futures = "0.3" log = "0.4" libc = "0.2" +object_store = "0.13.2" pin-project = "1.0" prost = "0.14" snafu = "0.9" diff --git a/README.md b/README.md index 2056671..90225c0 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] | Host read provider | Query engines can route range reads through their native file reader and data cache while independently sharing Lance Session caches | ## Building @@ -197,6 +198,43 @@ auto ds = lance::Dataset::open_with_session(session, "data.lance"); auto stats = session.cache_stats(); ``` +### Combine a shared Session with a host read provider + +`LanceSession` and `LanceReadProvider` are intentionally independent. A +Session is long-lived and caches portable metadata/index state. A read provider +is bound to a Dataset and owns the current query engine's file readers, +credentials, cancellation state, data-cache policy, and I/O statistics. + +```c +LanceReadProviderOps provider_ops = { + .open = host_open, + .read_at = host_read_at, + .close_reader = host_close_reader, + .destroy_context = host_destroy_context, + .last_error_message = host_last_error_message, +}; +LanceReadProvider* provider = + lance_read_provider_new(&provider_ops, host_context, 16); + +LanceDatasetOpenOptions options = { + .uri = "s3://bucket/data.lance", + .storage_options = storage_options, + .version = 0, + .session = session, + .read_provider = provider, +}; +LanceDataset* ds = lance_dataset_open_with_options(&options); + +/* The Dataset retains both shared objects. */ +lance_read_provider_close(provider); +lance_session_close(session); +``` + +The provider receives object metadata from Lance's native object store and is +used only for object contents. Listing, metadata lookup, and writes continue to +use the native store. Returning `LANCE_READ_NOT_SUPPORTED` from `open` falls +back to the native read path for that object. + ### Open at a specific version `lance_dataset_open` takes a `version` argument — `0` means the latest, any diff --git a/include/lance/lance.h b/include/lance/lance.h index 3bf291f..99ab905 100644 --- a/include/lance/lance.h +++ b/include/lance/lance.h @@ -182,6 +182,7 @@ typedef struct LanceDataset LanceDataset; typedef struct LanceScanner LanceScanner; typedef struct LanceBatch LanceBatch; typedef struct LanceSession LanceSession; +typedef struct LanceReadProvider LanceReadProvider; typedef struct LanceVersions LanceVersions; typedef struct LanceDataStatistics LanceDataStatistics; typedef struct LanceIndexSegmentBuilder LanceIndexSegmentBuilder; @@ -233,8 +234,104 @@ int32_t lance_session_get_cache_stats( LanceSessionCacheStats* out_stats ); +/* ─── Host read provider ─── */ + +/** Status returned by host read-provider callbacks. */ +typedef enum LanceReadStatus { + LANCE_READ_OK = 0, + LANCE_READ_NOT_SUPPORTED = 1, + LANCE_READ_NOT_FOUND = 2, + LANCE_READ_CANCELLED = 3, + LANCE_READ_IO_ERROR = 4 +} LanceReadStatus; + +/** + * Stable identity of an object opened through a host read provider. + * + * All strings are borrowed and remain valid only for the duration of the + * `open` callback. `path` is relative to `store_prefix`. `e_tag` and + * `version` may be NULL. + */ +typedef struct LanceFileIdentity { + const char* store_prefix; + const char* path; + uint64_t size; + int64_t last_modified_millis; + const char* e_tag; + const char* version; +} LanceFileIdentity; + +/** + * Host callbacks for random-access reads. + * + * `open` and `read_at` may run concurrently on blocking worker threads and + * must be thread-safe. A successful `open` must set `out_reader` to a non-NULL + * value. `close_reader` is called exactly once for every successfully opened + * reader. Callbacks must not throw or unwind across the C ABI. + */ +typedef struct LanceReadProviderOps { + int32_t (*open)( + void* context, + const LanceFileIdentity* identity, + void** out_reader + ); + int32_t (*read_at)( + void* reader, + uint64_t offset, + uint8_t* buffer, + uint64_t length, + uint64_t* bytes_read + ); + void (*close_reader)(void* reader); + void (*destroy_context)(void* context); + /** + * Return the current thread's last host error. The string is borrowed and + * copied by lance-c immediately. May be NULL. + */ + const char* (*last_error_message)(void* context); +} LanceReadProviderOps; + +/** + * Create a reference-counted host read provider. + * + * `max_concurrency` bounds simultaneous blocking `open` and `read_at` + * callbacks and must be greater than zero. On success the provider owns + * `context` and, when supplied, eventually calls `destroy_context` exactly + * once. + */ +LanceReadProvider* lance_read_provider_new( + const LanceReadProviderOps* ops, + void* context, + uint32_t max_concurrency +); + +/** + * Close a provider handle. Datasets already opened with it remain valid and + * retain the provider until their outstanding reads are complete. + */ +void lance_read_provider_close(LanceReadProvider* provider); + /* ─── Dataset lifecycle ─── */ +/** + * Complete options for opening a dataset. + * + * `session` and `read_provider` are optional and borrowed for the duration of + * this call. The returned dataset retains shared ownership of both. + */ +typedef struct LanceDatasetOpenOptions { + const char* uri; + const char* const* storage_options; + uint64_t version; + const LanceSession* session; + const LanceReadProvider* read_provider; +} LanceDatasetOpenOptions; + +/** Open a dataset with an optional shared session and host read provider. */ +LanceDataset* lance_dataset_open_with_options( + const LanceDatasetOpenOptions* options +); + /** * Open a Lance dataset. * @@ -1016,7 +1113,7 @@ typedef struct { * best-effort and may be omitted if they cannot be materialized. `metrics` is * NULL when `metrics_len` is zero. */ -typedef struct { +typedef struct LanceScanStatistics { uint64_t iops; uint64_t requests; uint64_t bytes_read; diff --git a/include/lance/lance.hpp b/include/lance/lance.hpp index 6cf245f..310e4f9 100644 --- a/include/lance/lance.hpp +++ b/include/lance/lance.hpp @@ -190,6 +190,40 @@ class Session { const LanceSession* c_handle() const { return handle_.get(); } }; +// ─── Host Read Provider ───────────────────────────────────────────────────── + +/// Reference-counted host random-access reader. The provider takes ownership +/// of `context` after successful construction and, when supplied, invokes +/// `destroy_context` when the last Dataset/provider handle releases it. +class ReadProvider { + Handle handle_; + +public: + ReadProvider( + const LanceReadProviderOps& ops, + void* context, + uint32_t max_concurrency) + : handle_(lance_read_provider_new(&ops, context, max_concurrency)) { + if (!handle_) check_error(); + } + + ReadProvider(ReadProvider&&) noexcept = default; + ReadProvider& operator=(ReadProvider&&) noexcept = default; + ReadProvider(const ReadProvider&) = delete; + ReadProvider& operator=(const ReadProvider&) = delete; + + const LanceReadProvider* c_handle() const { return handle_.get(); } +}; + +/// Composable options for Dataset::open. Session caches and host reads have +/// independent lifetimes and may be enabled separately or together. +struct DatasetOpenOptions { + std::vector> storage_options; + uint64_t version = 0; + const Session* session = nullptr; + const ReadProvider* read_provider = nullptr; +}; + // ─── Process-local FTS query context ──────────────────────────────────────── /// Immutable, query-specific global BM25 scorer plus pinned FTS segment list. @@ -225,6 +259,21 @@ class Dataset { } public: + /// Open a dataset with composable Session and host-read-provider options. + static Dataset open(const std::string& uri, const DatasetOpenOptions& options) { + auto kv = to_c_storage_options(options.storage_options); + LanceDatasetOpenOptions c_options{ + uri.c_str(), + options.storage_options.empty() ? nullptr : kv.data(), + options.version, + options.session ? options.session->c_handle() : nullptr, + options.read_provider ? options.read_provider->c_handle() : nullptr, + }; + auto* ds = lance_dataset_open_with_options(&c_options); + if (!ds) check_error(); + return Dataset(ds); + } + /// Open a dataset at the given URI. Pass `version` = 0 (the default) for /// the latest, or a specific version id from `versions()` to check out /// that version, e.g. `lance::Dataset::open("data.lance", {}, /*version=*/42)`. diff --git a/src/dataset.rs b/src/dataset.rs index cc1f87c..c4fe59a 100644 --- a/src/dataset.rs +++ b/src/dataset.rs @@ -13,9 +13,11 @@ use arrow_schema::Schema as ArrowSchema; use lance::Dataset; use lance::dataset::builder::DatasetBuilder; use lance_core::Result; +use lance_io::object_store::{ObjectStoreParams, StorageOptionsAccessor}; use crate::error::{ffi_try, swallow_unwind}; use crate::helpers; +use crate::read_provider::{LanceReadProvider, ProviderWrapper}; use crate::runtime::block_on; use crate::session::LanceSession; use crate::stream_guard::guarded_ffi_stream_from_reader; @@ -25,6 +27,20 @@ pub struct LanceDataset { pub(crate) inner: RwLock>, } +/// Complete set of options for opening a dataset. +/// +/// `session` and `read_provider` are borrowed for this call. The returned +/// dataset retains shared ownership of their underlying state. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct LanceDatasetOpenOptions { + pub uri: *const c_char, + pub storage_options: *const *const c_char, + pub version: u64, + pub session: *const LanceSession, + pub read_provider: *const LanceReadProvider, +} + impl LanceDataset { /// Take a consistent snapshot of the inner dataset. /// Returns a cloned Arc so the caller can hold it without keeping the lock. @@ -122,7 +138,7 @@ pub unsafe extern "C" fn lance_dataset_open( version: u64, ) -> *mut LanceDataset { ffi_try!( - unsafe { open_dataset_inner(uri, storage_options, version, None) }, + unsafe { open_dataset_inner(uri, storage_options, version, None, None) }, null ) } @@ -156,7 +172,45 @@ unsafe fn open_dataset_with_session_inner( )); } let session = unsafe { &*session }; - unsafe { open_dataset_inner(uri, storage_options, version, Some(session)) } + unsafe { open_dataset_inner(uri, storage_options, version, Some(session), None) } +} + +/// Open a dataset with an optional shared session and host read provider. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_dataset_open_with_options( + options: *const LanceDatasetOpenOptions, +) -> *mut LanceDataset { + ffi_try!(unsafe { open_dataset_with_options_inner(options) }, null) +} + +unsafe fn open_dataset_with_options_inner( + options: *const LanceDatasetOpenOptions, +) -> Result<*mut LanceDataset> { + if options.is_null() { + return Err(lance_core::Error::invalid_input_source( + "options must not be NULL".into(), + )); + } + let options = unsafe { &*options }; + let session = if options.session.is_null() { + None + } else { + Some(unsafe { &*options.session }) + }; + let read_provider = if options.read_provider.is_null() { + None + } else { + Some(unsafe { &*options.read_provider }) + }; + unsafe { + open_dataset_inner( + options.uri, + options.storage_options, + options.version, + session, + read_provider, + ) + } } unsafe fn open_dataset_inner( @@ -164,6 +218,7 @@ unsafe fn open_dataset_inner( storage_options: *const *const c_char, version: u64, session: Option<&LanceSession>, + read_provider: Option<&LanceReadProvider>, ) -> Result<*mut LanceDataset> { let uri_str = unsafe { helpers::parse_c_string(uri)? } .ok_or_else(|| lance_core::Error::invalid_input_source("uri must not be NULL".into()))?; @@ -171,8 +226,17 @@ unsafe fn open_dataset_inner( let opts = unsafe { helpers::parse_storage_options(storage_options)? }; let mut builder = DatasetBuilder::from_uri(uri_str); - if !opts.is_empty() { - builder = builder.with_storage_options(opts); + if !opts.is_empty() || read_provider.is_some() { + let mut store_params = ObjectStoreParams::default(); + if !opts.is_empty() { + store_params.storage_options_accessor = + Some(Arc::new(StorageOptionsAccessor::with_static_options(opts))); + } + if let Some(read_provider) = read_provider { + store_params.object_store_wrapper = + Some(Arc::new(ProviderWrapper::new(&read_provider.inner))); + } + builder = builder.with_store_params(store_params); } if version != 0 { builder = builder.with_version(version); diff --git a/src/lib.rs b/src/lib.rs index 8b212f5..3452a61 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -37,6 +37,7 @@ mod index; mod index_model; mod index_segment; mod merge_insert; +mod read_provider; mod restore; pub mod runtime; mod scanner; @@ -64,6 +65,7 @@ pub use index::*; pub use index_model::*; pub use index_segment::*; pub use merge_insert::*; +pub use read_provider::*; pub use restore::*; pub use scanner::*; pub use session::*; diff --git a/src/read_provider.rs b/src/read_provider.rs new file mode 100644 index 0000000..61b6608 --- /dev/null +++ b/src/read_provider.rs @@ -0,0 +1,590 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Host-provided random-access reads for query-engine integration. +//! +//! A read provider is deliberately independent from [`crate::LanceSession`]. +//! Sessions may be shared across queries to cache portable metadata and index +//! state, while each dataset binds the provider that owns its current storage +//! credentials, cancellation state, data-cache policy, and I/O statistics. + +use std::collections::HashMap; +use std::ffi::{CStr, CString, c_char, c_void}; +use std::fmt::{Debug, Display, Formatter}; +use std::ops::Range; +use std::sync::{Arc, Mutex, Weak}; + +use async_trait::async_trait; +use bytes::Bytes; +use futures::StreamExt; +use futures::stream::{self, BoxStream}; +use lance_core::Result; +use lance_io::object_store::WrappingObjectStore; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta, + ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, + Result as ObjectStoreResult, +}; +use tokio::sync::Semaphore; + +use crate::error::{ffi_try, swallow_unwind}; + +/// The host callback completed successfully. +pub const LANCE_READ_OK: i32 = 0; +/// The provider does not handle this object; use Lance's native object store. +pub const LANCE_READ_NOT_SUPPORTED: i32 = 1; +/// The requested object does not exist. +pub const LANCE_READ_NOT_FOUND: i32 = 2; +/// The request was cancelled by the host. +pub const LANCE_READ_CANCELLED: i32 = 3; +/// The provider encountered an I/O error. +pub const LANCE_READ_IO_ERROR: i32 = 4; + +/// Stable identity of an object passed to the host's `open` callback. +/// +/// All string pointers are borrowed and remain valid only for the duration of +/// the callback. `path` is relative to `store_prefix`. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct LanceFileIdentity { + pub store_prefix: *const c_char, + pub path: *const c_char, + pub size: u64, + pub last_modified_millis: i64, + pub e_tag: *const c_char, + pub version: *const c_char, +} + +/// C callbacks used by a host-provided random-access reader. +/// +/// `open` and `read_at` may be called concurrently from blocking worker +/// threads. `close_reader` is called exactly once for every reader returned by +/// a successful `open`. When supplied, `destroy_context` is called exactly once +/// when the last provider reference is released. Callbacks must not unwind +/// across the C ABI. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct LanceReadProviderOps { + pub open: Option< + unsafe extern "C" fn( + context: *mut c_void, + identity: *const LanceFileIdentity, + out_reader: *mut *mut c_void, + ) -> i32, + >, + pub read_at: Option< + unsafe extern "C" fn( + reader: *mut c_void, + offset: u64, + buffer: *mut u8, + length: u64, + bytes_read: *mut u64, + ) -> i32, + >, + pub close_reader: Option, + pub destroy_context: Option, + /// Return the current thread's last host error. The returned string is + /// borrowed; lance-c copies it immediately. + pub last_error_message: Option *const c_char>, +} + +/// Opaque, reference-counted host read provider. +pub struct LanceReadProvider { + pub(crate) inner: Arc, +} + +/// Create a host read provider. +/// +/// `max_concurrency` bounds simultaneous blocking host callbacks. The provider +/// owns `context` after this function succeeds. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_read_provider_new( + ops: *const LanceReadProviderOps, + context: *mut c_void, + max_concurrency: u32, +) -> *mut LanceReadProvider { + ffi_try!( + unsafe { read_provider_new_inner(ops, context, max_concurrency) }, + null + ) +} + +unsafe fn read_provider_new_inner( + ops: *const LanceReadProviderOps, + context: *mut c_void, + max_concurrency: u32, +) -> Result<*mut LanceReadProvider> { + if ops.is_null() { + return Err(lance_core::Error::invalid_input_source( + "ops must not be NULL".into(), + )); + } + if max_concurrency == 0 { + return Err(lance_core::Error::invalid_input_source( + "max_concurrency must be greater than zero".into(), + )); + } + + let ops = unsafe { *ops }; + if ops.open.is_none() || ops.read_at.is_none() || ops.close_reader.is_none() { + return Err(lance_core::Error::invalid_input_source( + "ops.open, ops.read_at, and ops.close_reader must not be NULL".into(), + )); + } + + Ok(Box::into_raw(Box::new(LanceReadProvider { + inner: Arc::new(ForeignReadProvider { + ops, + context: context as usize, + permits: Arc::new(Semaphore::new(max_concurrency as usize)), + }), + }))) +} + +/// Close a provider handle. Datasets already opened with the provider retain +/// shared ownership and remain valid. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_read_provider_close(provider: *mut LanceReadProvider) { + if !provider.is_null() { + swallow_unwind("lance_read_provider_close", || unsafe { + let _ = Box::from_raw(provider); + }); + } +} + +pub(crate) struct ForeignReadProvider { + ops: LanceReadProviderOps, + context: usize, + permits: Arc, +} + +impl Debug for ForeignReadProvider { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("ForeignReadProvider") + .field("max_concurrency", &self.permits.available_permits()) + .finish_non_exhaustive() + } +} + +impl Drop for ForeignReadProvider { + fn drop(&mut self) { + if let Some(destroy_context) = self.ops.destroy_context { + unsafe { destroy_context(self.context as *mut c_void) }; + } + } +} + +impl ForeignReadProvider { + fn error_message(&self, operation: &str, status: i32) -> String { + let host_message = self.ops.last_error_message.and_then(|last_error_message| { + let message = unsafe { last_error_message(self.context as *mut c_void) }; + if message.is_null() { + None + } else { + Some( + unsafe { CStr::from_ptr(message) } + .to_string_lossy() + .into_owned(), + ) + } + }); + host_message.unwrap_or_else(|| format!("host {operation} failed with status {status}")) + } + + fn error(&self, operation: &str, path: &Path, status: i32) -> object_store::Error { + let message = self.error_message(operation, status); + let source = || Box::new(std::io::Error::other(message.clone())); + match status { + LANCE_READ_NOT_FOUND => object_store::Error::NotFound { + path: path.to_string(), + source: source(), + }, + LANCE_READ_NOT_SUPPORTED => object_store::Error::NotSupported { source: source() }, + _ => object_store::Error::Generic { + store: "host read provider", + source: source(), + }, + } + } + + async fn open_file( + self: &Arc, + store_prefix: Arc, + path: Path, + meta: ObjectMeta, + attributes: object_store::Attributes, + ) -> ObjectStoreResult>> { + let provider = Arc::clone(self); + let callback = self.ops.open.expect("validated at provider construction"); + let permit = Arc::clone(&self.permits) + .acquire_owned() + .await + .map_err(|error| object_store::Error::Generic { + store: "host read provider", + source: Box::new(error), + })?; + let callback_path = path.clone(); + let callback_meta = meta.clone(); + + tokio::task::spawn_blocking(move || { + let _permit = permit; + let store_prefix = CString::new(store_prefix.as_ref()).map_err(|error| { + object_store::Error::Generic { + store: "host read provider", + source: Box::new(error), + } + })?; + let path_string = CString::new(callback_path.as_ref()).map_err(|error| { + object_store::Error::Generic { + store: "host read provider", + source: Box::new(error), + } + })?; + let e_tag = callback_meta + .e_tag + .as_deref() + .map(CString::new) + .transpose() + .map_err(|error| object_store::Error::Generic { + store: "host read provider", + source: Box::new(error), + })?; + let version = callback_meta + .version + .as_deref() + .map(CString::new) + .transpose() + .map_err(|error| object_store::Error::Generic { + store: "host read provider", + source: Box::new(error), + })?; + let identity = LanceFileIdentity { + store_prefix: store_prefix.as_ptr(), + path: path_string.as_ptr(), + size: callback_meta.size, + last_modified_millis: callback_meta.last_modified.timestamp_millis(), + e_tag: e_tag + .as_ref() + .map_or(std::ptr::null(), |value| value.as_ptr()), + version: version + .as_ref() + .map_or(std::ptr::null(), |value| value.as_ptr()), + }; + let mut reader = std::ptr::null_mut(); + let status = + unsafe { callback(provider.context as *mut c_void, &identity, &mut reader) }; + if status == LANCE_READ_NOT_SUPPORTED { + return Ok(None); + } + if status != LANCE_READ_OK { + return Err(provider.error("open", &callback_path, status)); + } + if reader.is_null() { + return Err(object_store::Error::Generic { + store: "host read provider", + source: Box::new(std::io::Error::other( + "host open returned success with a NULL reader", + )), + }); + } + + Ok(Some(Arc::new(HostFile { + provider, + reader: reader as usize, + meta, + attributes, + }))) + }) + .await + .map_err(|error| object_store::Error::Generic { + store: "host read provider", + source: Box::new(error), + })? + } +} + +#[derive(Debug)] +pub(crate) struct ProviderWrapper { + // ObjectStoreRegistry stores ObjectStoreParams as map keys. Keeping only a + // Weak reference here prevents a long-lived Session registry entry from + // retaining a Dataset-scoped provider after the wrapped store is gone. + provider: Weak, +} + +impl ProviderWrapper { + pub(crate) fn new(provider: &Arc) -> Self { + Self { + provider: Arc::downgrade(provider), + } + } +} + +impl WrappingObjectStore for ProviderWrapper { + fn wrap(&self, store_prefix: &str, original: Arc) -> Arc { + let Some(provider) = self.provider.upgrade() else { + return original; + }; + Arc::new(HostReadObjectStore { + provider, + original, + store_prefix: Arc::from(store_prefix), + files: Arc::new(Mutex::new(HashMap::new())), + }) + } +} + +#[derive(Debug, Hash, PartialEq, Eq)] +struct FileKey { + path: Path, + requested_version: Option, +} + +#[derive(Debug)] +struct HostReadObjectStore { + provider: Arc, + original: Arc, + store_prefix: Arc, + files: Arc>>>, +} + +impl Display for HostReadObjectStore { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + write!(formatter, "HostReadObjectStore({})", self.store_prefix) + } +} + +impl HostReadObjectStore { + fn cached_file(&self, key: &FileKey) -> Option> { + self.files + .lock() + .unwrap_or_else(|error| error.into_inner()) + .get(key) + .cloned() + } + + fn cache_file(&self, key: FileKey, file: &Arc) { + self.files + .lock() + .unwrap_or_else(|error| error.into_inner()) + .insert(key, Arc::clone(file)); + } + + fn invalidate(&self, location: &Path) { + self.files + .lock() + .unwrap_or_else(|error| error.into_inner()) + .retain(|key, _| key.path != *location); + } + + async fn file( + &self, + location: &Path, + options: &GetOptions, + ) -> ObjectStoreResult>> { + let key = FileKey { + path: location.clone(), + requested_version: options.version.clone(), + }; + if let Some(file) = self.cached_file(&key) { + options.check_preconditions(&file.meta)?; + return Ok(Some(file)); + } + + let mut head_options = options.clone(); + head_options.head = true; + head_options.range = None; + let head = self.original.get_opts(location, head_options).await?; + options.check_preconditions(&head.meta)?; + let file = self + .provider + .open_file( + Arc::clone(&self.store_prefix), + location.clone(), + head.meta, + head.attributes, + ) + .await?; + if let Some(file) = &file { + self.cache_file(key, file); + } + Ok(file) + } +} + +#[derive(Debug)] +struct HostFile { + provider: Arc, + reader: usize, + meta: ObjectMeta, + attributes: object_store::Attributes, +} + +impl HostFile { + async fn read(self: &Arc, range: Range) -> ObjectStoreResult { + let length = range.end - range.start; + if length == 0 { + return Ok(Bytes::new()); + } + let length_usize = + usize::try_from(length).map_err(|error| object_store::Error::Generic { + store: "host read provider", + source: Box::new(error), + })?; + let file = Arc::clone(self); + let callback = self + .provider + .ops + .read_at + .expect("validated at provider construction"); + let permit = Arc::clone(&self.provider.permits) + .acquire_owned() + .await + .map_err(|error| object_store::Error::Generic { + store: "host read provider", + source: Box::new(error), + })?; + + tokio::task::spawn_blocking(move || { + let _permit = permit; + let mut buffer = vec![0_u8; length_usize]; + let mut bytes_read = 0_u64; + let status = unsafe { + callback( + file.reader as *mut c_void, + range.start, + buffer.as_mut_ptr(), + length, + &mut bytes_read, + ) + }; + if status != LANCE_READ_OK { + return Err(file.provider.error("read_at", &file.meta.location, status)); + } + if bytes_read != length { + return Err(object_store::Error::Generic { + store: "host read provider", + source: Box::new(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + format!( + "host read_at returned {bytes_read} bytes for requested range {}..{}", + range.start, range.end + ), + )), + }); + } + Ok(Bytes::from(buffer)) + }) + .await + .map_err(|error| object_store::Error::Generic { + store: "host read provider", + source: Box::new(error), + })? + } +} + +impl Drop for HostFile { + fn drop(&mut self) { + let close_reader = self + .provider + .ops + .close_reader + .expect("validated at provider construction"); + unsafe { close_reader(self.reader as *mut c_void) }; + } +} + +#[async_trait] +impl ObjectStore for HostReadObjectStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + options: PutOptions, + ) -> ObjectStoreResult { + self.invalidate(location); + self.original.put_opts(location, payload, options).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + options: PutMultipartOptions, + ) -> ObjectStoreResult> { + self.invalidate(location); + self.original.put_multipart_opts(location, options).await + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> ObjectStoreResult { + if options.head { + return self.original.get_opts(location, options).await; + } + let Some(file) = self.file(location, &options).await? else { + return self.original.get_opts(location, options).await; + }; + let range = options + .range + .as_ref() + .map(|range| range.as_range(file.meta.size)) + .transpose() + .map_err(|error| object_store::Error::Generic { + store: "host read provider", + source: Box::new(error), + })? + .unwrap_or(0..file.meta.size); + let bytes = file.read(range.clone()).await?; + let payload = stream::once(async move { Ok(bytes) }).boxed(); + Ok(GetResult { + payload: GetResultPayload::Stream(payload), + meta: file.meta.clone(), + range, + attributes: file.attributes.clone(), + }) + } + + fn delete_stream( + &self, + locations: BoxStream<'static, ObjectStoreResult>, + ) -> BoxStream<'static, ObjectStoreResult> { + let files = Arc::clone(&self.files); + self.original + .delete_stream(locations) + .map(move |result| { + if let Ok(location) = &result { + files + .lock() + .unwrap_or_else(|error| error.into_inner()) + .retain(|key, _| key.path != *location); + } + result + }) + .boxed() + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, ObjectStoreResult> { + self.original.list(prefix) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, ObjectStoreResult> { + self.original.list_with_offset(prefix, offset) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> ObjectStoreResult { + self.original.list_with_delimiter(prefix).await + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + options: CopyOptions, + ) -> ObjectStoreResult<()> { + self.invalidate(to); + self.original.copy_opts(from, to, options).await + } +} diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs index 8805764..12a8310 100644 --- a/tests/c_api_test.rs +++ b/tests/c_api_test.rs @@ -7,10 +7,13 @@ //! validating the C API contract without needing a C compiler. use std::ffi::{CString, c_char, c_void}; +use std::fs::File; +use std::io::{Read, Seek, SeekFrom}; +use std::path::PathBuf; use std::process::Command; use std::ptr; -use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering as AtomicOrdering}; +use std::sync::{Arc, Mutex}; use arrow::ffi::from_ffi; use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; @@ -313,6 +316,198 @@ fn test_shared_session_rejects_null_inputs() { } } +#[derive(Default)] +struct ReadProviderCounters { + opens: AtomicUsize, + reads: AtomicUsize, + closes: AtomicUsize, + destroys: AtomicUsize, +} + +struct TestReadProviderContext { + dataset_root: PathBuf, + counters: Arc, +} + +struct TestRandomAccessReader { + file: Mutex, + counters: Arc, +} + +unsafe extern "C" fn test_read_provider_open( + context: *mut c_void, + identity: *const LanceFileIdentity, + out_reader: *mut *mut c_void, +) -> i32 { + if context.is_null() || identity.is_null() || out_reader.is_null() { + return LANCE_READ_IO_ERROR; + } + let context = unsafe { &*context.cast::() }; + let identity = unsafe { &*identity }; + if identity.path.is_null() { + return LANCE_READ_IO_ERROR; + } + let relative_path = unsafe { std::ffi::CStr::from_ptr(identity.path) }; + let Ok(relative_path) = relative_path.to_str() else { + return LANCE_READ_IO_ERROR; + }; + let relative_path = PathBuf::from(relative_path); + let absolute_path = PathBuf::from("/").join(&relative_path); + let path = if absolute_path.exists() { + absolute_path + } else { + context.dataset_root.join(relative_path) + }; + let Ok(file) = File::open(path) else { + return LANCE_READ_NOT_FOUND; + }; + context.counters.opens.fetch_add(1, AtomicOrdering::SeqCst); + let reader = Box::new(TestRandomAccessReader { + file: Mutex::new(file), + counters: Arc::clone(&context.counters), + }); + unsafe { *out_reader = Box::into_raw(reader).cast() }; + LANCE_READ_OK +} + +unsafe extern "C" fn test_read_provider_read_at( + reader: *mut c_void, + offset: u64, + buffer: *mut u8, + length: u64, + bytes_read: *mut u64, +) -> i32 { + if reader.is_null() || buffer.is_null() || bytes_read.is_null() { + return LANCE_READ_IO_ERROR; + } + let Ok(length) = usize::try_from(length) else { + return LANCE_READ_IO_ERROR; + }; + let reader = unsafe { &*reader.cast::() }; + let mut file = reader.file.lock().unwrap(); + if file.seek(SeekFrom::Start(offset)).is_err() { + return LANCE_READ_IO_ERROR; + } + let output = unsafe { std::slice::from_raw_parts_mut(buffer, length) }; + if file.read_exact(output).is_err() { + return LANCE_READ_IO_ERROR; + } + unsafe { *bytes_read = length as u64 }; + reader.counters.reads.fetch_add(1, AtomicOrdering::SeqCst); + LANCE_READ_OK +} + +unsafe extern "C" fn test_read_provider_close_reader(reader: *mut c_void) { + if !reader.is_null() { + let reader = unsafe { Box::from_raw(reader.cast::()) }; + reader.counters.closes.fetch_add(1, AtomicOrdering::SeqCst); + } +} + +unsafe extern "C" fn test_read_provider_destroy_context(context: *mut c_void) { + if !context.is_null() { + let context = unsafe { Box::from_raw(context.cast::()) }; + context + .counters + .destroys + .fetch_add(1, AtomicOrdering::SeqCst); + } +} + +#[test] +fn test_read_provider_is_dataset_scoped_with_shared_session() { + let (_tmp, uri) = create_test_dataset(); + let counters = Arc::new(ReadProviderCounters::default()); + let context = Box::new(TestReadProviderContext { + dataset_root: PathBuf::from(&uri), + counters: Arc::clone(&counters), + }); + let ops = LanceReadProviderOps { + open: Some(test_read_provider_open), + read_at: Some(test_read_provider_read_at), + close_reader: Some(test_read_provider_close_reader), + destroy_context: Some(test_read_provider_destroy_context), + last_error_message: None, + }; + let provider = unsafe { lance_read_provider_new(&ops, Box::into_raw(context).cast(), 4) }; + assert!(!provider.is_null()); + let session = lance_session_new(0, 16 * 1024 * 1024); + assert!(!session.is_null()); + let c_uri = c_str(&uri); + let options = LanceDatasetOpenOptions { + uri: c_uri.as_ptr(), + storage_options: ptr::null(), + version: 0, + session, + read_provider: provider, + }; + let dataset = unsafe { lance_dataset_open_with_options(&options) }; + if dataset.is_null() { + let message = lance_last_error_message(); + let message = if message.is_null() { + "unknown error".to_string() + } else { + let owned = unsafe { std::ffi::CStr::from_ptr(message) } + .to_string_lossy() + .into_owned(); + unsafe { lance_free_string(message) }; + owned + }; + panic!("dataset open with provider should succeed: {message}"); + } + + // The Dataset owns the provider after open. + unsafe { lance_read_provider_close(provider) }; + assert_eq!(counters.destroys.load(AtomicOrdering::SeqCst), 0); + assert_eq!( + scan_all_rows(dataset) + .iter() + .map(|batch| batch.num_rows()) + .sum::(), + 5 + ); + assert!(counters.opens.load(AtomicOrdering::SeqCst) > 0); + assert!(counters.reads.load(AtomicOrdering::SeqCst) > 0); + + unsafe { lance_dataset_close(dataset) }; + assert_eq!( + counters.opens.load(AtomicOrdering::SeqCst), + counters.closes.load(AtomicOrdering::SeqCst) + ); + assert_eq!(counters.destroys.load(AtomicOrdering::SeqCst), 1); + + // Keeping the Session alive must not retain the Dataset-scoped provider. + // The same Session remains usable for a subsequent Dataset binding. + let second = + unsafe { lance_dataset_open_with_session(c_uri.as_ptr(), ptr::null(), 0, session) }; + assert!(!second.is_null()); + assert_eq!(unsafe { lance_dataset_count_rows(second) }, 5); + unsafe { + lance_dataset_close(second); + lance_session_close(session); + } +} + +#[test] +fn test_read_provider_rejects_invalid_inputs() { + assert!(unsafe { lance_read_provider_new(ptr::null(), ptr::null_mut(), 1) }.is_null()); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + + let ops = LanceReadProviderOps { + open: Some(test_read_provider_open), + read_at: Some(test_read_provider_read_at), + close_reader: Some(test_read_provider_close_reader), + destroy_context: None, + last_error_message: None, + }; + assert!(unsafe { lance_read_provider_new(&ops, ptr::null_mut(), 0) }.is_null()); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + assert!(unsafe { lance_dataset_open_with_options(ptr::null()) }.is_null()); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + + unsafe { lance_read_provider_close(ptr::null_mut()) }; +} + #[test] fn test_open_nonexistent() { let c_uri = c_str("memory://nonexistent_dataset_xyz"); diff --git a/tests/cpp/test_c_api.c b/tests/cpp/test_c_api.c index c49ecfa..9a27239 100644 --- a/tests/cpp/test_c_api.c +++ b/tests/cpp/test_c_api.c @@ -45,6 +45,12 @@ typedef struct { int invalid; } ScanStatisticsCapture; +static uint64_t tagged_scan_statistics_bytes_read( + const struct LanceScanStatistics *statistics +) { + return statistics->bytes_read; +} + static void capture_scan_statistics( void *callback_ctx, const LanceScanStatistics *statistics @@ -66,7 +72,7 @@ static void capture_scan_statistics( } } captured->calls += 1; - captured->bytes_read = statistics->bytes_read; + captured->bytes_read = tagged_scan_statistics_bytes_read(statistics); } static void test_open_and_metadata(const char *uri) { diff --git a/tests/cpp/test_cpp_api.cpp b/tests/cpp/test_cpp_api.cpp index 17b1ab6..c13918d 100644 --- a/tests/cpp/test_cpp_api.cpp +++ b/tests/cpp/test_cpp_api.cpp @@ -34,6 +34,10 @@ struct ScanStatisticsCapture { bool invalid = false; }; +static_assert(std::is_same_v< + LanceScanStatistics, + struct LanceScanStatistics>); + static void capture_scan_statistics( void* callback_ctx, const LanceScanStatistics* statistics) noexcept { @@ -88,7 +92,9 @@ static void test_shared_session(const std::string& uri) { TEST(test_shared_session); auto session = std::make_unique(0, 16 * 1024 * 1024); - auto ds = lance::Dataset::open_with_session(*session, uri); + lance::DatasetOpenOptions options; + options.session = session.get(); + auto ds = lance::Dataset::open(uri, options); auto stats = session->cache_stats(); session.reset();