diff --git a/include/lance/lance.h b/include/lance/lance.h index 9aa66d0..aac75a9 100644 --- a/include/lance/lance.h +++ b/include/lance/lance.h @@ -1065,6 +1065,38 @@ int32_t lance_scanner_set_include_deleted_rows( bool include_deleted_rows ); +/** How blob columns are materialized by a scan. Validated as an integer. */ +typedef enum { + /** + * Default: blob columns are returned as descriptor structs and every + * other binary column is returned as bytes. The descriptor layout + * depends on the storage format of the column: Blob v2 columns yield + * (kind, position, size, blob_id, blob_uri), while legacy blob columns + * (large_binary tagged `lance-encoding: blob`) yield (position, size). + */ + LANCE_BLOB_HANDLING_BLOBS_DESCRIPTIONS = 0, + /** Every blob column is materialized as bytes (LargeBinary). */ + LANCE_BLOB_HANDLING_ALL_BINARY = 1, + /** + * Requests descriptors for every binary column. On lance v11.0.0 only + * columns carrying blob metadata are affected; other binary columns keep + * their bytes, so this behaves like + * LANCE_BLOB_HANDLING_BLOBS_DESCRIPTIONS. + */ + LANCE_BLOB_HANDLING_ALL_DESCRIPTIONS = 2, +} LanceBlobHandling; + +/** + * Choose how blob columns are materialized by this scan. Default: + * LANCE_BLOB_HANDLING_BLOBS_DESCRIPTIONS. ALL_BINARY pulls the full payload + * into the batches, so keep descriptors for large values. Columns without + * blob metadata keep their bytes under every mode. + * + * Must be set before scanning starts; values outside the enum are rejected. + * @return 0 on success, -1 on error + */ +int32_t lance_scanner_set_blob_handling(LanceScanner* scanner, LanceBlobHandling handling); + /** * Restrict scan to the given fragment IDs. Must be called before iteration. * @param ids Array of fragment IDs diff --git a/include/lance/lance.hpp b/include/lance/lance.hpp index 027e6d5..4050cf6 100644 --- a/include/lance/lance.hpp +++ b/include/lance/lance.hpp @@ -1291,6 +1291,14 @@ class Scanner { return *this; } + /// Choose how blob columns are materialized (default: descriptors for blob + /// columns, bytes for every other binary column). + Scanner& blob_handling(LanceBlobHandling handling) { + if (lance_scanner_set_blob_handling(handle_.get(), handling) != 0) + check_error(); + return *this; + } + /// Enable/disable row ID in output. Scanner& with_row_id(bool enable = true) { if (lance_scanner_with_row_id(handle_.get(), enable) != 0) diff --git a/src/scanner.rs b/src/scanner.rs index 5f9844c..71d2dd2 100644 --- a/src/scanner.rs +++ b/src/scanner.rs @@ -20,6 +20,7 @@ use lance::dataset::scanner::{ }; use lance::io::exec::fts::{FlatMatchQueryExec, MatchQueryExec, PhraseQueryExec}; use lance_core::Result; +use lance_core::datatypes::BlobHandling; use lance_index::scalar::FullTextSearchQuery; use lance_index::vector::ApproxMode; use lance_io::stream::RecordBatchStream; @@ -92,6 +93,7 @@ pub struct LanceScanner { filter: Option, substrait_filter: Option>, additional_sql_filters: Vec, + blob_handling: Option, limit: Option, offset: Option, batch_size: Option, @@ -241,6 +243,7 @@ impl LanceScanner { filter: None, substrait_filter: None, additional_sql_filters: Vec::new(), + blob_handling: None, limit: None, offset: None, batch_size: None, @@ -363,6 +366,9 @@ impl LanceScanner { if let Some(cols) = &self.columns { scanner.project(cols)?; } + if let Some(handling) = &self.blob_handling { + scanner.blob_handling(handling.clone()); + } let multi_vector = self.nearest.as_ref().is_some_and(|query| { matches!( query.query.data_type(), @@ -1390,6 +1396,46 @@ unsafe fn scanner_set_use_stats_inner(scanner: *mut LanceScanner, use_stats: boo Ok(0) } +/// Set how blob columns are materialized. `handling` is the C enum +/// `LanceBlobHandling` as an integer: 0 descriptors for blob columns (the +/// default), 1 bytes for every blob column, 2 descriptors for every binary +/// column. Other values are rejected. Must be set before the scan starts. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_scanner_set_blob_handling( + scanner: *mut LanceScanner, + handling: i32, +) -> i32 { + scanner_poison_check!(scanner, -1); + scanner_ffi_try!(scanner, unsafe { + scanner_set_blob_handling_inner(scanner, handling) + }) +} + +unsafe fn scanner_set_blob_handling_inner( + scanner: *mut LanceScanner, + handling: i32, +) -> Result { + if scanner.is_null() { + return Err(lance_core::Error::invalid_input_source( + "scanner is NULL".into(), + )); + } + let scanner = unsafe { &mut *scanner }; + scanner.ensure_scan_not_started("blob_handling")?; + let parsed = match handling { + 0 => BlobHandling::BlobsDescriptions, + 1 => BlobHandling::AllBinary, + 2 => BlobHandling::AllDescriptions, + other => { + return Err(lance_core::Error::invalid_input(format!( + "blob_handling must be 0 (blobs as descriptions), 1 (all binary) or 2 (all descriptions); got {other}" + ))); + } + }; + scanner.blob_handling = Some(parsed); + Ok(0) +} + /// Enable or disable row ID in scan output. Returns 0. #[unsafe(no_mangle)] pub unsafe extern "C" fn lance_scanner_with_row_id( @@ -3369,6 +3415,48 @@ mod tests { } } + #[test] + fn set_blob_handling_stores_the_matching_upstream_variant() { + // A scan cannot tell AllDescriptions from BlobsDescriptions on lance + // v11, so check the stored variant directly. + let (_tmp, uri) = create_test_dataset(); + let (dataset, scanner) = open_dataset_and_scanner(&uri); + + assert_eq!( + unsafe { &*scanner }.blob_handling, + None, + "blob handling should be unset until the setter is called" + ); + + for (handling, expected) in [ + (0, BlobHandling::BlobsDescriptions), + (1, BlobHandling::AllBinary), + (2, BlobHandling::AllDescriptions), + ] { + assert_eq!( + unsafe { lance_scanner_set_blob_handling(scanner, handling) }, + 0 + ); + assert_eq!( + unsafe { &*scanner }.blob_handling, + Some(expected), + "blob handling {handling} stored the wrong variant" + ); + } + + // A rejected value leaves the last accepted mode in place. + assert_eq!(unsafe { lance_scanner_set_blob_handling(scanner, 3) }, -1); + assert_eq!( + unsafe { &*scanner }.blob_handling, + Some(BlobHandling::AllDescriptions) + ); + + unsafe { + lance_scanner_close(scanner); + lance_dataset_close(dataset); + } + } + #[test] fn null_poll_waker_is_rejected_and_clears_out() { let (_tmp, uri) = create_test_dataset(); diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs index 0424c97..55a04ac 100644 --- a/tests/c_api_test.rs +++ b/tests/c_api_test.rs @@ -17,7 +17,10 @@ use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; use arrow::ffi_stream::ArrowArrayStreamReader; use arrow::ffi_stream::FFI_ArrowArrayStream; use arrow::record_batch::RecordBatchReader; -use arrow_array::{Array, Float32Array, Int32Array, RecordBatch, StringArray, UInt64Array}; +use arrow_array::{ + Array, BinaryArray, Float32Array, Int32Array, LargeBinaryArray, RecordBatch, StringArray, + UInt32Array, UInt64Array, +}; use arrow_schema::{DataType, Field, Schema}; use lance::Dataset; use lance_c::*; @@ -13259,3 +13262,506 @@ fn test_scalar_segment_requires_explicit_domain_and_checks_uuid() { lance_dataset_close(ds); } } + +// --------------------------------------------------------------------------- +// Scanner blob handling +// --------------------------------------------------------------------------- + +// Mirror of the C enum `LanceBlobHandling`; the FFI parameter is an int32. +const BLOB_HANDLING_BLOBS_DESCRIPTIONS: i32 = 0; +const BLOB_HANDLING_ALL_BINARY: i32 = 1; +const BLOB_HANDLING_ALL_DESCRIPTIONS: i32 = 2; + +/// Sub-fields of a Blob v2 description struct, in schema order. +const BLOB_DESCRIPTION_FIELDS: [&str; 5] = ["kind", "position", "size", "blob_id", "blob_uri"]; + +/// Blob storage thresholds used by [`create_blob_v2_dataset`]. +const BLOB_INLINE_THRESHOLD: usize = 16; +const BLOB_DEDICATED_THRESHOLD: usize = 256; + +/// Blob sizes of the five rows in each fragment: inline, packed and dedicated +/// against the thresholds above, then an empty blob and a null. +const BLOB_ROW_SIZES: [Option; 5] = [Some(8), Some(128), Some(1024), Some(0), None]; + +/// First `id` of each fragment; also seeds its payloads. +const BLOB_FRAGMENT_BASE_IDS: [u32; 2] = [0, 100]; + +/// Blob payload: byte `i` is `(i * 7 + 3 + seed) as u8`. +fn blob_payload(len: usize, seed: usize) -> Vec { + (0..len).map(|i| (i * 7 + 3 + seed) as u8).collect() +} + +/// One fragment's batch: ids `base_id..base_id + 5`, blobs per +/// [`BLOB_ROW_SIZES`], `raw-` in the plain binary column (null where the +/// blob is null). +fn blob_batch(schema: &Arc, base_id: u32) -> RecordBatch { + let seed = base_id as usize; + let mut blobs = lance::BlobArrayBuilder::new(BLOB_ROW_SIZES.len()); + for size in BLOB_ROW_SIZES { + match size { + Some(0) => blobs.push_empty().unwrap(), + Some(len) => blobs.push_bytes(blob_payload(len, seed)).unwrap(), + None => blobs.push_null().unwrap(), + } + } + + let ids: Vec = (0..BLOB_ROW_SIZES.len() as u32) + .map(|row| base_id + row) + .collect(); + let raw: Vec> = ids + .iter() + .map(|id| format!("raw-{id}").into_bytes()) + .collect(); + let raw_array = BinaryArray::from_iter( + raw.iter() + .zip(BLOB_ROW_SIZES) + .map(|(value, size)| size.map(|_| value.as_slice())), + ); + + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(UInt32Array::from(ids)), + blobs.finish().unwrap(), + Arc::new(raw_array), + ], + ) + .unwrap() +} + +/// Two-fragment v2.2 dataset with a blob column, a plain binary column and an +/// id column; one [`blob_batch`] per entry of [`BLOB_FRAGMENT_BASE_IDS`]. +fn create_blob_v2_dataset() -> (tempfile::TempDir, String) { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().join("blob_ds").to_str().unwrap().to_string(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + lance::blob_field_with_options( + "blob", + true, + lance::BlobFieldOptions { + inline_size_threshold: Some(BLOB_INLINE_THRESHOLD), + dedicated_size_threshold: std::num::NonZeroUsize::new(BLOB_DEDICATED_THRESHOLD), + }, + ), + Field::new("raw", DataType::Binary, true), + ])); + + lance_c::runtime::block_on(async { + for (fragment, base_id) in BLOB_FRAGMENT_BASE_IDS.into_iter().enumerate() { + let params = lance::dataset::WriteParams { + mode: if fragment == 0 { + lance::dataset::WriteMode::Create + } else { + lance::dataset::WriteMode::Append + }, + // Blob v2 is a 2.2 storage feature. + data_storage_version: Some(lance_file::version::LanceFileVersion::V2_2), + ..Default::default() + }; + Dataset::write( + arrow::record_batch::RecordBatchIterator::new( + vec![Ok(blob_batch(&schema, base_id))], + schema.clone(), + ), + &uri, + Some(params), + ) + .await + .unwrap(); + } + }); + + (tmp, uri) +} + +/// Run the scanner through the C Arrow stream; return its schema and batches. +fn scan_stream(scanner: *mut LanceScanner) -> (Schema, Vec) { + let mut ffi_stream = FFI_ArrowArrayStream::empty(); + assert_eq!( + unsafe { lance_scanner_to_arrow_stream(scanner, &mut ffi_stream) }, + 0, + "to_arrow_stream should succeed" + ); + let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut ffi_stream) }.unwrap(); + let schema = reader.schema().as_ref().clone(); + let batches: Vec = reader.map(|batch| batch.unwrap()).collect(); + (schema, batches) +} + +/// Collect `(id, blob bytes)` pairs from batches whose blob column was +/// materialized as bytes, sorted by id. +fn collect_blob_bytes(batches: &[RecordBatch]) -> Vec<(u32, Option>)> { + let mut rows = Vec::new(); + for batch in batches { + let ids = batch + .column_by_name("id") + .expect("id column") + .as_any() + .downcast_ref::() + .expect("id is UInt32"); + let blobs = batch + .column_by_name("blob") + .expect("blob column") + .as_any() + .downcast_ref::() + .expect("blob is LargeBinary"); + for row in 0..batch.num_rows() { + let value = (!blobs.is_null(row)).then(|| blobs.value(row).to_vec()); + rows.push((ids.value(row), value)); + } + } + rows.sort_by_key(|(id, _)| *id); + rows +} + +/// Collect `(id, raw bytes)` pairs from the plain binary column, sorted by id. +fn collect_raw_bytes(batches: &[RecordBatch]) -> Vec<(u32, Option>)> { + let mut rows = Vec::new(); + for batch in batches { + let ids = batch + .column_by_name("id") + .expect("id column") + .as_any() + .downcast_ref::() + .expect("id is UInt32"); + let raw = batch + .column_by_name("raw") + .expect("raw column") + .as_any() + .downcast_ref::() + .expect("raw is Binary"); + for row in 0..batch.num_rows() { + let value = (!raw.is_null(row)).then(|| raw.value(row).to_vec()); + rows.push((ids.value(row), value)); + } + } + rows.sort_by_key(|(id, _)| *id); + rows +} + +/// Assert that the plain binary column of the fragment based at `base_id` +/// round-tripped: `raw-` bytes, and null in the last row. +fn assert_raw_bytes_of_fragment(rows: &[(u32, Option>)], base_id: u32) { + let row = |id: u32| -> &Option> { + &rows + .iter() + .find(|(row_id, _)| *row_id == id) + .unwrap_or_else(|| panic!("row {id} missing from scan output")) + .1 + }; + + for offset in 0..4 { + let id = base_id + offset; + assert_eq!( + row(id).as_deref(), + Some(format!("raw-{id}").as_bytes()), + "plain binary payload of row {id} must round-trip byte for byte" + ); + } + assert_eq!( + row(base_id + 4), + &None, + "null plain binary value must stay null" + ); +} + +/// Assert that the five rows written for `base_id` round-tripped byte for byte. +fn assert_blob_bytes_of_fragment(rows: &[(u32, Option>)], base_id: u32) { + let row = |id: u32| -> &Option> { + &rows + .iter() + .find(|(row_id, _)| *row_id == id) + .unwrap_or_else(|| panic!("row {id} missing from scan output")) + .1 + }; + let seed = base_id as usize; + + assert_eq!( + row(base_id).as_deref(), + Some(blob_payload(8, seed).as_slice()), + "inline blob (8 bytes) must round-trip byte for byte" + ); + assert_eq!( + row(base_id + 1).as_deref(), + Some(blob_payload(128, seed).as_slice()), + "packed blob (128 bytes) must round-trip byte for byte" + ); + assert_eq!( + row(base_id + 2).as_deref(), + Some(blob_payload(1024, seed).as_slice()), + "dedicated blob (1024 bytes) must round-trip byte for byte" + ); + assert_eq!( + row(base_id + 3).as_deref(), + Some([].as_slice()), + "empty blob must be a zero-length, non-null value" + ); + assert_eq!(row(base_id + 4), &None, "null blob must stay null"); +} + +/// Assert that the named field is a blob description struct. +fn assert_blob_description_field(schema: &Schema, name: &str) { + let field = schema.field_with_name(name).expect("field exists"); + match field.data_type() { + DataType::Struct(children) => { + let names: Vec<&str> = children.iter().map(|c| c.name().as_str()).collect(); + assert_eq!( + names, BLOB_DESCRIPTION_FIELDS, + "{name} should be a blob description struct" + ); + } + other => panic!("{name} should be a blob description struct, got {other:?}"), + } +} + +#[test] +fn test_scanner_blob_handling_all_binary_materializes_bytes() { + let (_tmp, uri) = create_blob_v2_dataset(); + let c_uri = c_str(&uri); + let ds = unsafe { lance_dataset_open(c_uri.as_ptr(), ptr::null(), 0) }; + assert!(!ds.is_null()); + + let scanner = unsafe { lance_scanner_new(ds, ptr::null(), ptr::null()) }; + assert!(!scanner.is_null()); + assert_eq!( + unsafe { lance_scanner_set_blob_handling(scanner, BLOB_HANDLING_ALL_BINARY) }, + 0 + ); + + let (schema, batches) = scan_stream(scanner); + let blob_field = schema.field_with_name("blob").expect("blob column"); + assert_eq!( + *blob_field.data_type(), + DataType::LargeBinary, + "ALL_BINARY should materialize the blob column as bytes" + ); + + // Neither blob marker survives materialization (lance v11), so a C caller + // cannot tell a materialized blob from a plain binary column by metadata. + let metadata = blob_field.metadata(); + assert!( + !metadata.contains_key("lance-encoding:blob"), + "the blob marker should not survive materialization: {metadata:?}" + ); + assert!( + !metadata.contains_key("ARROW:extension:name"), + "the blob v2 extension name should not survive materialization: {metadata:?}" + ); + + let rows = collect_blob_bytes(&batches); + assert_eq!(rows.len(), 10, "both fragments should be scanned"); + assert_blob_bytes_of_fragment(&rows, 0); + assert_blob_bytes_of_fragment(&rows, 100); + + unsafe { lance_scanner_close(scanner) }; + unsafe { lance_dataset_close(ds) }; +} + +#[test] +fn test_scanner_blob_handling_defaults_to_descriptions() { + let (_tmp, uri) = create_blob_v2_dataset(); + let c_uri = c_str(&uri); + let ds = unsafe { lance_dataset_open(c_uri.as_ptr(), ptr::null(), 0) }; + assert!(!ds.is_null()); + + // Without the setter, and with an explicit BLOBS_DESCRIPTIONS, the blob + // column is a description struct while plain binary columns stay bytes. + for handling in [None, Some(BLOB_HANDLING_BLOBS_DESCRIPTIONS)] { + let scanner = unsafe { lance_scanner_new(ds, ptr::null(), ptr::null()) }; + assert!(!scanner.is_null()); + if let Some(handling) = handling { + assert_eq!( + unsafe { lance_scanner_set_blob_handling(scanner, handling) }, + 0 + ); + } + + let (schema, batches) = scan_stream(scanner); + assert_blob_description_field(&schema, "blob"); + assert_eq!( + *schema + .field_with_name("raw") + .expect("raw column") + .data_type(), + DataType::Binary, + "a plain binary column stays bytes under {handling:?}" + ); + assert_eq!( + batches.iter().map(|b| b.num_rows()).sum::(), + 10, + "both fragments should be scanned under {handling:?}" + ); + + let raw_rows = collect_raw_bytes(&batches); + assert_raw_bytes_of_fragment(&raw_rows, 0); + assert_raw_bytes_of_fragment(&raw_rows, 100); + + unsafe { lance_scanner_close(scanner) }; + } + + unsafe { lance_dataset_close(ds) }; +} + +#[test] +fn test_scanner_blob_handling_all_descriptions() { + let (_tmp, uri) = create_blob_v2_dataset(); + let c_uri = c_str(&uri); + let ds = unsafe { lance_dataset_open(c_uri.as_ptr(), ptr::null(), 0) }; + assert!(!ds.is_null()); + + let scanner = unsafe { lance_scanner_new(ds, ptr::null(), ptr::null()) }; + assert!(!scanner.is_null()); + assert_eq!( + unsafe { lance_scanner_set_blob_handling(scanner, BLOB_HANDLING_ALL_DESCRIPTIONS) }, + 0 + ); + + let (schema, batches) = scan_stream(scanner); + assert_blob_description_field(&schema, "blob"); + // On lance v11 ALL_DESCRIPTIONS only rewrites fields with blob metadata + // (`Field::unloaded_mut` is gated on `is_blob`), so `raw` keeps its bytes. + assert_eq!( + *schema + .field_with_name("raw") + .expect("raw column") + .data_type(), + DataType::Binary, + "a column without blob metadata is not turned into a description" + ); + assert_eq!( + batches.iter().map(|b| b.num_rows()).sum::(), + 10, + "both fragments should be scanned" + ); + + let raw_rows = collect_raw_bytes(&batches); + assert_raw_bytes_of_fragment(&raw_rows, 0); + assert_raw_bytes_of_fragment(&raw_rows, 100); + + unsafe { lance_scanner_close(scanner) }; + unsafe { lance_dataset_close(ds) }; +} + +#[test] +fn test_scanner_blob_handling_rejected_after_scan_started() { + let (_tmp, uri) = create_blob_v2_dataset(); + let c_uri = c_str(&uri); + let ds = unsafe { lance_dataset_open(c_uri.as_ptr(), ptr::null(), 0) }; + assert!(!ds.is_null()); + + let scanner = unsafe { lance_scanner_new(ds, ptr::null(), ptr::null()) }; + assert!(!scanner.is_null()); + + let mut ffi_stream = FFI_ArrowArrayStream::empty(); + assert_eq!( + unsafe { lance_scanner_to_arrow_stream(scanner, &mut ffi_stream) }, + 0 + ); + // Release the stream; the scan has started either way. + drop(unsafe { ArrowArrayStreamReader::from_raw(&mut ffi_stream) }.unwrap()); + + assert_eq!( + unsafe { lance_scanner_set_blob_handling(scanner, BLOB_HANDLING_ALL_BINARY) }, + -1, + "blob handling must not change once the scan has started" + ); + let message = take_last_error_message(); + assert!( + message.contains("blob_handling must be set before the scan starts"), + "unexpected error: {message}" + ); + + unsafe { lance_scanner_close(scanner) }; + unsafe { lance_dataset_close(ds) }; +} + +#[test] +fn test_scanner_blob_handling_rejects_invalid_values() { + let (_tmp, uri) = create_blob_v2_dataset(); + let c_uri = c_str(&uri); + let ds = unsafe { lance_dataset_open(c_uri.as_ptr(), ptr::null(), 0) }; + assert!(!ds.is_null()); + + let scanner = unsafe { lance_scanner_new(ds, ptr::null(), ptr::null()) }; + assert!(!scanner.is_null()); + + for invalid in [3, -1] { + assert_eq!( + unsafe { lance_scanner_set_blob_handling(scanner, invalid) }, + -1, + "blob_handling {invalid} should be rejected" + ); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + let message = take_last_error_message(); + assert!( + message.contains(&format!("got {invalid}")), + "error for {invalid} should name the rejected value: {message}" + ); + } + + assert_eq!( + unsafe { lance_scanner_set_blob_handling(ptr::null_mut(), BLOB_HANDLING_ALL_BINARY) }, + -1, + "NULL scanner should be rejected" + ); + + // A rejected value leaves the default handling in place. + let (schema, _batches) = scan_stream(scanner); + assert_blob_description_field(&schema, "blob"); + + unsafe { lance_scanner_close(scanner) }; + unsafe { lance_dataset_close(ds) }; +} + +#[test] +fn test_scanner_blob_handling_all_binary_with_fragment_ids() { + let (_tmp, uri) = create_blob_v2_dataset(); + let c_uri = c_str(&uri); + let ds = unsafe { lance_dataset_open(c_uri.as_ptr(), ptr::null(), 0) }; + assert!(!ds.is_null()); + assert_eq!(unsafe { lance_dataset_fragment_count(ds) }, 2); + + let mut fragment_ids = vec![0u64; 2]; + assert_eq!( + unsafe { lance_dataset_fragment_ids(ds, fragment_ids.as_mut_ptr()) }, + 0 + ); + + let scanner = unsafe { lance_scanner_new(ds, ptr::null(), ptr::null()) }; + assert!(!scanner.is_null()); + assert_eq!( + unsafe { lance_scanner_set_fragment_ids(scanner, fragment_ids[1..].as_ptr(), 1) }, + 0 + ); + assert_eq!( + unsafe { lance_scanner_set_blob_handling(scanner, BLOB_HANDLING_ALL_BINARY) }, + 0 + ); + + let (schema, batches) = scan_stream(scanner); + assert_eq!( + *schema + .field_with_name("blob") + .expect("blob column") + .data_type(), + DataType::LargeBinary + ); + + let rows = collect_blob_bytes(&batches); + assert_eq!( + rows.len(), + 5, + "only the selected fragment should be scanned" + ); + assert!( + rows.iter().all(|(id, _)| (100..105).contains(id)), + "unexpected rows from the unselected fragment: {:?}", + rows.iter().map(|(id, _)| *id).collect::>() + ); + assert_blob_bytes_of_fragment(&rows, 100); + + unsafe { lance_scanner_close(scanner) }; + unsafe { lance_dataset_close(ds) }; +} diff --git a/tests/compile_and_run_test.rs b/tests/compile_and_run_test.rs index 8566d10..626e7a2 100644 --- a/tests/compile_and_run_test.rs +++ b/tests/compile_and_run_test.rs @@ -16,9 +16,14 @@ use std::path::{Path, PathBuf}; use std::process::Command; use std::sync::Arc; -use arrow_array::{FixedSizeListArray, Float32Array, Int32Array, RecordBatch, StringArray}; +use arrow_array::{ + BinaryArray, FixedSizeListArray, Float32Array, Int32Array, RecordBatch, StringArray, + UInt32Array, +}; use arrow_schema::{DataType, Field, Schema}; use lance::Dataset; +use lance::dataset::{WriteMode, WriteParams}; +use lance_file::version::LanceFileVersion; /// Build the lance-c cdylib and return the path to the shared library and include dir. fn build_lance_c() -> (PathBuf, PathBuf) { @@ -123,6 +128,76 @@ fn create_test_dataset_on_disk() -> (tempfile::TempDir, String) { (tmp, uri) } +/// Create a two-fragment Blob v2 dataset on disk and return (TempDir, path_string). +/// +/// Each fragment has five rows: blobs of 8, 128 and 1024 bytes (inline, packed +/// and dedicated under the 16 / 256 thresholds), an empty blob and a null, +/// next to `id` and a plain `raw` binary column. Payload byte `i` is +/// `(i * 7 + 3) as u8`. +fn create_blob_dataset_on_disk() -> (tempfile::TempDir, String) { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().join("blob_ds").to_str().unwrap().to_string(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + lance::blob_field_with_options( + "blob", + true, + lance::BlobFieldOptions { + inline_size_threshold: Some(16), + dedicated_size_threshold: std::num::NonZeroUsize::new(256), + }, + ), + Field::new("raw", DataType::Binary, true), + ])); + + let make_batch = |first_id: u32| { + let mut blobs = lance::BlobArrayBuilder::new(5); + for len in [8usize, 128, 1024] { + let payload: Vec = (0..len).map(|i| (i * 7 + 3) as u8).collect(); + blobs.push_bytes(payload).unwrap(); + } + blobs.push_empty().unwrap(); + blobs.push_null().unwrap(); + + let ids: Vec = (first_id..first_id + 5).collect(); + // The plain binary column is null in the same row as the blob column. + let raw = BinaryArray::from_iter((0..5).map(|row| (row < 4).then_some(&b"raw"[..]))); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(UInt32Array::from(ids)), + blobs.finish().unwrap(), + Arc::new(raw), + ], + ) + .unwrap() + }; + + lance_c::runtime::block_on(async { + for (first_id, mode) in [(0u32, WriteMode::Create), (100u32, WriteMode::Append)] { + let params = WriteParams { + mode, + // Blob v2 is a 2.2 storage feature. + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }; + Dataset::write( + arrow::record_batch::RecordBatchIterator::new( + vec![Ok(make_batch(first_id))], + schema.clone(), + ), + &uri, + Some(params), + ) + .await + .unwrap(); + } + }); + + (tmp, uri) +} + /// Compile a C source file, linking against lance-c. fn compile_c_test(source: &Path, output: &Path, include_dir: &Path, lib_path: &Path) -> bool { let lib_dir = lib_path.parent().unwrap(); @@ -174,12 +249,14 @@ fn compile_cpp_test(source: &Path, output: &Path, include_dir: &Path, lib_path: .success() } -/// Run a compiled test binary with the source dataset URI and a destination URI -/// for the write test. The destination path must not pre-exist. -fn run_test_binary(binary: &Path, dataset_uri: &str, write_uri: &str) { +/// Run a compiled test binary with the source dataset URI, a destination URI +/// for the write test and the URI of a Blob v2 dataset. The destination path +/// must not pre-exist. +fn run_test_binary(binary: &Path, dataset_uri: &str, write_uri: &str, blob_uri: &str) { let output = Command::new(binary) .arg(dataset_uri) .arg(write_uri) + .arg(blob_uri) .output() .unwrap_or_else(|e| panic!("Failed to run {}: {e}", binary.display())); @@ -207,6 +284,7 @@ fn test_c_compilation_and_execution() { let (lib_path, include_dir) = build_lance_c(); let (tmp, dataset_uri) = create_test_dataset_on_disk(); let write_uri = tmp.path().join("c_write_ds").to_str().unwrap().to_string(); + let (_blob_tmp, blob_uri) = create_blob_dataset_on_disk(); let build_dir = tempfile::tempdir().unwrap(); let source = PathBuf::from(env!("CARGO_MANIFEST_DIR")) @@ -220,7 +298,7 @@ fn test_c_compilation_and_execution() { "C test compilation failed" ); - run_test_binary(&binary, &dataset_uri, &write_uri); + run_test_binary(&binary, &dataset_uri, &write_uri, &blob_uri); } #[test] @@ -234,6 +312,7 @@ fn test_cpp_compilation_and_execution() { .to_str() .unwrap() .to_string(); + let (_blob_tmp, blob_uri) = create_blob_dataset_on_disk(); let build_dir = tempfile::tempdir().unwrap(); let source = PathBuf::from(env!("CARGO_MANIFEST_DIR")) @@ -247,7 +326,7 @@ fn test_cpp_compilation_and_execution() { "C++ test compilation failed" ); - run_test_binary(&binary, &dataset_uri, &write_uri); + run_test_binary(&binary, &dataset_uri, &write_uri, &blob_uri); } /// A fresh C executable must initialize OpenDAL even when archive constructors are omitted. diff --git a/tests/cpp/test_c_api.c b/tests/cpp/test_c_api.c index c49ecfa..0e524d7 100644 --- a/tests/cpp/test_c_api.c +++ b/tests/cpp/test_c_api.c @@ -8,7 +8,7 @@ * This file is compiled by the Rust integration test to verify that * lance.h is valid C and the API works end-to-end. * - * Usage: test_c_api + * Usage: test_c_api */ #include "lance/lance.h" @@ -218,6 +218,83 @@ static void test_scan_with_limit(const char *uri) { 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) { + struct ArrowSchema schema; + memset(&schema, 0, sizeof(schema)); + int rc = stream->get_schema(stream, &schema); + ASSERT(rc == 0, "get_schema from stream failed"); + out[0] = '\0'; + for (int64_t i = 0; i < schema.n_children; i++) { + if (strcmp(schema.children[i]->name, "blob") == 0) { + snprintf(out, out_len, "%s", schema.children[i]->format); + } + } + if (schema.release) schema.release(&schema); +} + +static void test_scanner_blob_handling(const char *blob_uri) { + printf(" test_scanner_blob_handling... "); + + LanceDataset *ds = lance_dataset_open(blob_uri, NULL, 0); + ASSERT(ds != NULL, "blob dataset open failed"); + uint64_t expected_rows = lance_dataset_count_rows(ds); + CHECK_OK(); + + char format[16]; + struct ArrowArrayStream stream; + + /* By default a blob column arrives as its description struct. */ + LanceScanner *scanner = lance_scanner_new(ds, NULL, NULL); + ASSERT(scanner != NULL, "scanner creation failed"); + memset(&stream, 0, sizeof(stream)); + int32_t rc = lance_scanner_to_arrow_stream(scanner, &stream); + ASSERT(rc == 0, "to_arrow_stream failed"); + blob_column_format(&stream, format, sizeof(format)); + ASSERT(strcmp(format, "+s") == 0, "default blob column should be a struct"); + if (stream.release) stream.release(&stream); + lance_scanner_close(scanner); + + /* ALL_BINARY materializes the bytes as LargeBinary and keeps every row. */ + scanner = lance_scanner_new(ds, NULL, NULL); + ASSERT(scanner != NULL, "scanner creation failed"); + rc = lance_scanner_set_blob_handling(scanner, LANCE_BLOB_HANDLING_ALL_BINARY); + ASSERT(rc == 0, "set_blob_handling failed"); + memset(&stream, 0, sizeof(stream)); + rc = lance_scanner_to_arrow_stream(scanner, &stream); + ASSERT(rc == 0, "to_arrow_stream failed"); + blob_column_format(&stream, format, sizeof(format)); + ASSERT(strcmp(format, "Z") == 0, "ALL_BINARY blob column should be LargeBinary"); + + uint64_t total_rows = 0; + while (1) { + struct ArrowArray array; + memset(&array, 0, sizeof(array)); + rc = stream.get_next(&stream, &array); + ASSERT(rc == 0, "get_next failed"); + if (array.release == NULL) { + break; + } + total_rows += (uint64_t)array.length; + array.release(&array); + } + ASSERT(total_rows == expected_rows, "row count mismatch"); + if (stream.release) stream.release(&stream); + + /* Once the scan has started the setting is rejected. */ + rc = lance_scanner_set_blob_handling(scanner, LANCE_BLOB_HANDLING_BLOBS_DESCRIPTIONS); + ASSERT(rc == -1, "set_blob_handling after the scan started should fail"); + ASSERT(lance_last_error_code() == LANCE_ERR_INVALID_ARGUMENT, "wrong error code"); + const char *msg = lance_last_error_message(); + if (msg) lance_free_string(msg); + + printf("rows=%llu... ", (unsigned long long)total_rows); + lance_scanner_close(scanner); + lance_dataset_close(ds); + printf("OK\n"); +} + static void test_versions(const char *uri) { printf(" test_versions... "); @@ -962,19 +1039,21 @@ static void test_delete(const char *write_uri) { } int main(int argc, char **argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s \n", argv[0]); + if (argc < 4) { + fprintf(stderr, "Usage: %s \n", argv[0]); return 1; } const char *uri = argv[1]; const char *write_uri = argv[2]; + const char *blob_uri = argv[3]; printf("Running C API tests with dataset: %s\n", uri); test_open_and_metadata(uri); test_shared_session(uri); test_scan(uri); test_scan_with_limit(uri); + test_scanner_blob_handling(blob_uri); test_versions(uri); test_restore_to_current(uri); test_error_handling(); diff --git a/tests/cpp/test_cpp_api.cpp b/tests/cpp/test_cpp_api.cpp index ff5ed1e..58e38bd 100644 --- a/tests/cpp/test_cpp_api.cpp +++ b/tests/cpp/test_cpp_api.cpp @@ -7,7 +7,7 @@ * * Tests the RAII wrappers, exception handling, and builder pattern. * - * Usage: test_cpp_api + * Usage: test_cpp_api */ #include "lance/lance.hpp" @@ -214,6 +214,73 @@ static void test_scanner_async_stream_ownership(const std::string& uri) { PASS(); } +/// Arrow C Data Interface format of the `blob` column in a stream's schema, +/// or an empty string when the column is missing. +static std::string blob_column_format(ArrowArrayStream& stream) { + ArrowSchema schema; + memset(&schema, 0, sizeof(schema)); + int rc = stream.get_schema(&stream, &schema); + assert(rc == 0); + std::string format; + for (int64_t i = 0; i < schema.n_children; i++) { + if (strcmp(schema.children[i]->name, "blob") == 0) { + format = schema.children[i]->format; + } + } + if (schema.release) schema.release(&schema); + return format; +} + +static void test_scanner_blob_handling(const std::string& blob_uri) { + TEST(test_scanner_blob_handling); + + auto ds = lance::Dataset::open(blob_uri); + + // By default a blob column arrives as its description struct ("+s"). + { + auto scanner = ds.scan(); + ArrowArrayStream stream; + memset(&stream, 0, sizeof(stream)); + scanner.to_arrow_stream(&stream); + assert(blob_column_format(stream) == "+s"); + if (stream.release) stream.release(&stream); + } + + // ALL_BINARY: LargeBinary ("Z"), and every row is still returned. + auto scanner = ds.scan(); + scanner.blob_handling(LANCE_BLOB_HANDLING_ALL_BINARY); + ArrowArrayStream stream; + memset(&stream, 0, sizeof(stream)); + scanner.to_arrow_stream(&stream); + assert(blob_column_format(stream) == "Z"); + + 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); + } + assert(total == ds.count_rows()); + if (stream.release) stream.release(&stream); + + // Once the scan has started the setting is rejected. + bool caught = false; + try { + scanner.blob_handling(LANCE_BLOB_HANDLING_BLOBS_DESCRIPTIONS); + } catch (const lance::Error& e) { + caught = true; + assert(e.code == LANCE_ERR_INVALID_ARGUMENT); + } + assert(caught); + + printf("rows=%llu... ", (unsigned long long)total); + PASS(); +} + static void test_dataset_take(const std::string& uri) { TEST(test_dataset_take); @@ -940,13 +1007,14 @@ static void test_delete_rows(const std::string& dst_uri) { } int main(int argc, char** argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s \n", argv[0]); + if (argc < 4) { + fprintf(stderr, "Usage: %s \n", argv[0]); return 1; } std::string uri(argv[1]); std::string write_uri(argv[2]); + std::string blob_uri(argv[3]); printf("Running C++ API tests with dataset: %s\n", uri.c_str()); test_dataset_open(uri); @@ -954,6 +1022,7 @@ int main(int argc, char** argv) { test_dataset_schema(uri); test_scanner_fluent(uri); test_scanner_async_stream_ownership(uri); + test_scanner_blob_handling(blob_uri); test_dataset_take(uri); test_dataset_take_rows(uri); test_raii_cleanup(uri);