diff --git a/include/lance/lance.h b/include/lance/lance.h index 3bf291f..541134f 100644 --- a/include/lance/lance.h +++ b/include/lance/lance.h @@ -934,6 +934,74 @@ LanceScanner* lance_scanner_new( int32_t lance_scanner_set_limit(LanceScanner* scanner, int64_t limit); int32_t lance_scanner_set_offset(LanceScanner* scanner, int64_t offset); int32_t lance_scanner_set_batch_size(LanceScanner* scanner, int64_t batch_size); + +/** + * Set the target output batch size in bytes. + * + * When set, this takes precedence over the row-based batch size. The value + * must be greater than zero and must be set before scanning starts. + */ +int32_t lance_scanner_set_batch_size_bytes( + LanceScanner* scanner, + uint64_t batch_size_bytes +); + +/** + * Set the scanner I/O buffer size in bytes. + * + * The value must be between 1 and INT64_MAX, inclusive, and must be set before + * scanning starts. This bounds buffered I/O received from storage, but is not + * a hard limit on all memory used by the scanner. + * + * @param scanner Scanner handle. Must not be NULL. + * @param io_buffer_size_bytes I/O buffer size in bytes, in the range [1, INT64_MAX]. + * @return 0 on success, -1 on error. + */ +int32_t lance_scanner_set_io_buffer_size( + LanceScanner* scanner, + uint64_t io_buffer_size_bytes +); + +/** + * Set the maximum number of batches decoded concurrently. + * + * @param batch_readahead Number of in-flight batch decode tasks. Must be greater than zero. + */ +int32_t lance_scanner_set_batch_readahead( + LanceScanner* scanner, + size_t batch_readahead +); + +/** + * Set fragment readahead for unordered scans. + * + * This setting is only used when scan-in-order is disabled. The value must be + * greater than zero. + */ +int32_t lance_scanner_set_fragment_readahead( + LanceScanner* scanner, + size_t fragment_readahead +); + +/** + * Set the target number of physical execution partitions. + * + * This controls the partition count used by the physical optimizer and can be + * used to bound scan CPU parallelism. The value must be greater than zero and + * must be set before scanning starts. + */ +int32_t lance_scanner_set_target_parallelism( + LanceScanner* scanner, + size_t target_parallelism +); + +/** + * Configure whether batches are returned in storage order (default: true). + * + * Disabling ordering can improve throughput by returning batches as soon as + * they are ready. + */ +int32_t lance_scanner_set_scan_in_order(LanceScanner* scanner, bool scan_in_order); int32_t lance_scanner_with_row_id(LanceScanner* scanner, bool enable); /** @@ -1656,6 +1724,19 @@ int32_t lance_scanner_nearest( ); int32_t lance_scanner_set_nprobes(LanceScanner* scanner, uint32_t n); + +/** + * Set vector index partition-search concurrency for each query. + * + * A value of -1 uses the CPU pool size, 0 selects Lance's automatic policy, + * 1 uses the sequential path, and values greater than 1 request parallel + * partition search. The effective value is capped by available parallelism. + * Values below -1 are rejected. Must be set before scanning starts. + */ +int32_t lance_scanner_set_query_parallelism( + LanceScanner* scanner, + int32_t query_parallelism +); int32_t lance_scanner_set_refine_factor(LanceScanner* scanner, uint32_t f); int32_t lance_scanner_set_ef(LanceScanner* scanner, uint32_t e); int32_t lance_scanner_set_metric(LanceScanner* scanner, LanceMetricType metric); diff --git a/include/lance/lance.hpp b/include/lance/lance.hpp index 6cf245f..a60dcd4 100644 --- a/include/lance/lance.hpp +++ b/include/lance/lance.hpp @@ -1196,6 +1196,48 @@ class Scanner { return *this; } + /// Set the target output batch size in bytes. + Scanner& batch_size_bytes(uint64_t bytes) { + if (lance_scanner_set_batch_size_bytes(handle_.get(), bytes) != 0) + check_error(); + return *this; + } + + /// Set the scanner I/O buffer size in bytes, in the range [1, INT64_MAX]. + Scanner& io_buffer_size(uint64_t bytes) { + if (lance_scanner_set_io_buffer_size(handle_.get(), bytes) != 0) + check_error(); + return *this; + } + + /// Set the number of batches decoded concurrently. + Scanner& batch_readahead(size_t batches) { + if (lance_scanner_set_batch_readahead(handle_.get(), batches) != 0) + check_error(); + return *this; + } + + /// Set fragment readahead for unordered scans. + Scanner& fragment_readahead(size_t fragments) { + if (lance_scanner_set_fragment_readahead(handle_.get(), fragments) != 0) + check_error(); + return *this; + } + + /// Set the target number of physical execution partitions. + Scanner& target_parallelism(size_t partitions) { + if (lance_scanner_set_target_parallelism(handle_.get(), partitions) != 0) + check_error(); + return *this; + } + + /// Configure whether batches are returned in storage order. + Scanner& scan_in_order(bool ordered = true) { + if (lance_scanner_set_scan_in_order(handle_.get(), ordered) != 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) @@ -1313,6 +1355,11 @@ class Scanner { if (lance_scanner_set_nprobes(handle_.get(), n) != 0) check_error(); return *this; } + Scanner& query_parallelism(int32_t parallelism) { + if (lance_scanner_set_query_parallelism(handle_.get(), parallelism) != 0) + check_error(); + return *this; + } Scanner& refine_factor(uint32_t f) { if (lance_scanner_set_refine_factor(handle_.get(), f) != 0) check_error(); return *this; diff --git a/src/scanner.rs b/src/scanner.rs index 0c29b17..6414cc0 100644 --- a/src/scanner.rs +++ b/src/scanner.rs @@ -60,11 +60,18 @@ pub struct LanceScanner { limit: Option, offset: Option, batch_size: Option, + batch_size_bytes: Option, + io_buffer_size: Option, + batch_readahead: Option, + fragment_readahead: Option, + target_parallelism: Option, + scan_in_order: Option, with_row_id: bool, fragment_ids: Option>, index_segments: Option>, nearest: Option, nprobes: Option, + query_parallelism: Option, refine_factor: Option, ef: Option, metric_override: Option, @@ -129,11 +136,18 @@ impl LanceScanner { limit: None, offset: None, batch_size: None, + batch_size_bytes: None, + io_buffer_size: None, + batch_readahead: None, + fragment_readahead: None, + target_parallelism: None, + scan_in_order: None, with_row_id: false, fragment_ids: None, index_segments: None, nearest: None, nprobes: None, + query_parallelism: None, refine_factor: None, ef: None, metric_override: None, @@ -164,6 +178,15 @@ impl LanceScanner { Arc::clone(&self.poisoned) } + fn ensure_scan_not_started(&self, setting_name: &str) -> Result<()> { + if self.scan_started.load(Ordering::Acquire) { + return Err(lance_core::Error::invalid_input_source( + format!("{setting_name} must be set before the scan starts").into(), + )); + } + Ok(()) + } + /// Apply fragment selection to a scanner builder if fragment_ids is set. fn apply_fragment_filter(&self, scanner: &mut lance::dataset::scanner::Scanner) -> Result<()> { if let Some(ids) = &self.fragment_ids { @@ -231,6 +254,24 @@ impl LanceScanner { if let Some(bs) = self.batch_size { scanner.batch_size(bs); } + if let Some(batch_size_bytes) = self.batch_size_bytes { + scanner.batch_size_bytes(batch_size_bytes); + } + if let Some(io_buffer_size) = self.io_buffer_size { + scanner.io_buffer_size(io_buffer_size); + } + if let Some(batch_readahead) = self.batch_readahead { + scanner.batch_readahead(batch_readahead); + } + if let Some(fragment_readahead) = self.fragment_readahead { + scanner.fragment_readahead(fragment_readahead); + } + if let Some(target_parallelism) = self.target_parallelism { + scanner.target_parallelism(target_parallelism); + } + if let Some(scan_in_order) = self.scan_in_order { + scanner.scan_in_order(scan_in_order); + } if self.with_row_id { scanner.with_row_id(); } @@ -260,6 +301,9 @@ impl LanceScanner { if let Some(np) = self.nprobes { scanner.nprobes(np as usize); } + if let Some(query_parallelism) = self.query_parallelism { + scanner.query_parallelism(query_parallelism); + } if let Some(rf) = self.refine_factor { scanner.refine(rf); } @@ -757,6 +801,214 @@ unsafe fn scanner_set_batch_size_inner(scanner: *mut LanceScanner, batch_size: i Ok(0) } +/// Set the target output batch size in bytes. Returns 0 on success. +/// +/// The size must be greater than zero and must be set before the scan starts. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_scanner_set_batch_size_bytes( + scanner: *mut LanceScanner, + batch_size_bytes: u64, +) -> i32 { + scanner_poison_check!(scanner, -1); + scanner_ffi_try!(scanner, unsafe { + scanner_set_batch_size_bytes_inner(scanner, batch_size_bytes) + }) +} + +unsafe fn scanner_set_batch_size_bytes_inner( + scanner: *mut LanceScanner, + batch_size_bytes: u64, +) -> Result { + if scanner.is_null() { + return Err(lance_core::Error::invalid_input_source( + "scanner is NULL".into(), + )); + } + if batch_size_bytes == 0 { + return Err(lance_core::Error::invalid_input_source( + "batch_size_bytes must be greater than 0, got 0".into(), + )); + } + let scanner = unsafe { &mut *scanner }; + scanner.ensure_scan_not_started("batch_size_bytes")?; + scanner.batch_size_bytes = Some(batch_size_bytes); + Ok(0) +} + +/// Set the scanner I/O buffer size in bytes. Returns 0 on success. +/// +/// The size must be between 1 and [`i64::MAX`] and must be set before the scan starts. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_scanner_set_io_buffer_size( + scanner: *mut LanceScanner, + io_buffer_size_bytes: u64, +) -> i32 { + scanner_poison_check!(scanner, -1); + scanner_ffi_try!(scanner, unsafe { + scanner_set_io_buffer_size_inner(scanner, io_buffer_size_bytes) + }) +} + +unsafe fn scanner_set_io_buffer_size_inner( + scanner: *mut LanceScanner, + io_buffer_size_bytes: u64, +) -> Result { + if scanner.is_null() { + return Err(lance_core::Error::invalid_input_source( + "scanner is NULL".into(), + )); + } + if io_buffer_size_bytes == 0 { + return Err(lance_core::Error::invalid_input_source( + "io_buffer_size_bytes must be greater than 0, got 0".into(), + )); + } + if io_buffer_size_bytes > i64::MAX as u64 { + return Err(lance_core::Error::invalid_input_source( + format!( + "io_buffer_size_bytes must be at most {}, got {io_buffer_size_bytes}", + i64::MAX + ) + .into(), + )); + } + let scanner = unsafe { &mut *scanner }; + scanner.ensure_scan_not_started("io_buffer_size_bytes")?; + scanner.io_buffer_size = Some(io_buffer_size_bytes); + Ok(0) +} + +/// Set the number of batches to decode concurrently. Returns 0 on success. +/// +/// The value must be greater than zero and must be set before the scan starts. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_scanner_set_batch_readahead( + scanner: *mut LanceScanner, + batch_readahead: usize, +) -> i32 { + scanner_poison_check!(scanner, -1); + scanner_ffi_try!(scanner, unsafe { + scanner_set_batch_readahead_inner(scanner, batch_readahead) + }) +} + +unsafe fn scanner_set_batch_readahead_inner( + scanner: *mut LanceScanner, + batch_readahead: usize, +) -> Result { + if scanner.is_null() { + return Err(lance_core::Error::invalid_input_source( + "scanner is NULL".into(), + )); + } + if batch_readahead == 0 { + return Err(lance_core::Error::invalid_input_source( + "batch_readahead must be greater than 0, got 0".into(), + )); + } + let scanner = unsafe { &mut *scanner }; + scanner.ensure_scan_not_started("batch_readahead")?; + scanner.batch_readahead = Some(batch_readahead); + Ok(0) +} + +/// Set the number of fragments to read ahead for unordered scans. +/// +/// The value must be greater than zero and must be set before the scan starts. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_scanner_set_fragment_readahead( + scanner: *mut LanceScanner, + fragment_readahead: usize, +) -> i32 { + scanner_poison_check!(scanner, -1); + scanner_ffi_try!(scanner, unsafe { + scanner_set_fragment_readahead_inner(scanner, fragment_readahead) + }) +} + +unsafe fn scanner_set_fragment_readahead_inner( + scanner: *mut LanceScanner, + fragment_readahead: usize, +) -> Result { + if scanner.is_null() { + return Err(lance_core::Error::invalid_input_source( + "scanner is NULL".into(), + )); + } + if fragment_readahead == 0 { + return Err(lance_core::Error::invalid_input_source( + "fragment_readahead must be greater than 0, got 0".into(), + )); + } + let scanner = unsafe { &mut *scanner }; + scanner.ensure_scan_not_started("fragment_readahead")?; + scanner.fragment_readahead = Some(fragment_readahead); + Ok(0) +} + +/// Set the target number of physical execution partitions. +/// +/// The value must be greater than zero and must be set before the scan starts. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_scanner_set_target_parallelism( + scanner: *mut LanceScanner, + target_parallelism: usize, +) -> i32 { + scanner_poison_check!(scanner, -1); + scanner_ffi_try!(scanner, unsafe { + scanner_set_target_parallelism_inner(scanner, target_parallelism) + }) +} + +unsafe fn scanner_set_target_parallelism_inner( + scanner: *mut LanceScanner, + target_parallelism: usize, +) -> Result { + if scanner.is_null() { + return Err(lance_core::Error::invalid_input_source( + "scanner is NULL".into(), + )); + } + if target_parallelism == 0 { + return Err(lance_core::Error::invalid_input_source( + "target_parallelism must be greater than 0, got 0".into(), + )); + } + let scanner = unsafe { &mut *scanner }; + scanner.ensure_scan_not_started("target_parallelism")?; + scanner.target_parallelism = Some(target_parallelism); + Ok(0) +} + +/// Configure whether scan results are returned in storage order. +/// +/// Must be set before the scan starts. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_scanner_set_scan_in_order( + scanner: *mut LanceScanner, + scan_in_order: bool, +) -> i32 { + scanner_poison_check!(scanner, -1); + scanner_ffi_try!(scanner, unsafe { + scanner_set_scan_in_order_inner(scanner, scan_in_order) + }) +} + +unsafe fn scanner_set_scan_in_order_inner( + scanner: *mut LanceScanner, + scan_in_order: bool, +) -> 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("scan_in_order")?; + scanner.scan_in_order = Some(scan_in_order); + 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( @@ -1766,6 +2018,42 @@ scanner_set_u32!(lance_scanner_set_nprobes, nprobes); scanner_set_u32!(lance_scanner_set_refine_factor, refine_factor); scanner_set_u32!(lance_scanner_set_ef, ef); +/// Set vector index partition-search concurrency for each query. +/// +/// `-1` uses the CPU pool size, `0` selects Lance's automatic policy, and +/// positive values request that many workers. Values below `-1` are invalid. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_scanner_set_query_parallelism( + scanner: *mut LanceScanner, + query_parallelism: i32, +) -> i32 { + scanner_poison_check!(scanner, -1); + scanner_ffi_try!(scanner, unsafe { + scanner_set_query_parallelism_inner(scanner, query_parallelism) + }) +} + +unsafe fn scanner_set_query_parallelism_inner( + scanner: *mut LanceScanner, + query_parallelism: i32, +) -> Result { + if scanner.is_null() { + return Err(lance_core::Error::invalid_input_source( + "scanner is NULL".into(), + )); + } + if query_parallelism < -1 { + return Err(lance_core::Error::invalid_input_source( + format!("query_parallelism must be -1, 0, or greater than 0, got {query_parallelism}") + .into(), + )); + } + let scanner = unsafe { &mut *scanner }; + scanner.ensure_scan_not_started("query_parallelism")?; + scanner.query_parallelism = Some(query_parallelism); + Ok(0) +} + #[unsafe(no_mangle)] pub unsafe extern "C" fn lance_scanner_set_metric(scanner: *mut LanceScanner, metric: i32) -> i32 { scanner_poison_check!(scanner, -1); diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs index 8805764..ffdd915 100644 --- a/tests/c_api_test.rs +++ b/tests/c_api_test.rs @@ -1207,6 +1207,169 @@ fn test_scanner_batch_size() { unsafe { lance_dataset_close(ds) }; } +#[test] +fn test_scanner_execution_tuning_options() { + let (_tmp, uri) = create_multi_fragment_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_batch_size_bytes(scanner, 1024) }, + 0 + ); + assert_eq!( + unsafe { lance_scanner_set_io_buffer_size(scanner, 64 * 1024) }, + 0 + ); + assert_eq!(unsafe { lance_scanner_set_batch_readahead(scanner, 1) }, 0); + assert_eq!( + unsafe { lance_scanner_set_fragment_readahead(scanner, 1) }, + 0 + ); + assert_eq!( + unsafe { lance_scanner_set_target_parallelism(scanner, 1) }, + 0 + ); + assert_eq!( + unsafe { lance_scanner_set_scan_in_order(scanner, false) }, + 0 + ); + + let mut ffi_stream = FFI_ArrowArrayStream::empty(); + assert_eq!( + unsafe { lance_scanner_to_arrow_stream(scanner, &mut ffi_stream) }, + 0 + ); + let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut ffi_stream) }.unwrap(); + let total_rows: usize = reader.map(|batch| batch.unwrap().num_rows()).sum(); + assert_eq!(total_rows, 10); + + unsafe { lance_scanner_close(scanner) }; + unsafe { lance_dataset_close(ds) }; +} + +#[test] +fn test_scanner_execution_tuning_options_reject_invalid_values() { + let (_tmp, uri) = create_test_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_batch_size_bytes(scanner, 0) }, + -1 + ); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + assert!(take_last_error_message().contains("batch_size_bytes must be greater than 0, got 0")); + + assert_eq!(unsafe { lance_scanner_set_io_buffer_size(scanner, 0) }, -1); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + assert!( + take_last_error_message().contains("io_buffer_size_bytes must be greater than 0, got 0") + ); + + assert_eq!( + unsafe { lance_scanner_set_io_buffer_size(scanner, i64::MAX as u64) }, + 0 + ); + assert_eq!( + unsafe { lance_scanner_set_io_buffer_size(scanner, u64::MAX) }, + -1 + ); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + assert!(take_last_error_message().contains(&format!( + "io_buffer_size_bytes must be at most {}, got {}", + i64::MAX, + u64::MAX + ))); + + assert_eq!(unsafe { lance_scanner_set_batch_readahead(scanner, 0) }, -1); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + assert!(take_last_error_message().contains("batch_readahead must be greater than 0, got 0")); + + assert_eq!( + unsafe { lance_scanner_set_fragment_readahead(scanner, 0) }, + -1 + ); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + assert!(take_last_error_message().contains("fragment_readahead must be greater than 0, got 0")); + + assert_eq!( + unsafe { lance_scanner_set_target_parallelism(scanner, 0) }, + -1 + ); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + assert!(take_last_error_message().contains("target_parallelism must be greater than 0, got 0")); + + unsafe { lance_scanner_close(scanner) }; + unsafe { lance_dataset_close(ds) }; +} + +#[test] +fn test_scanner_execution_tuning_options_reject_after_scan_start() { + let (_tmp, uri) = create_test_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 + ); + + assert_eq!( + unsafe { lance_scanner_set_batch_size_bytes(scanner, 1024) }, + -1 + ); + assert!(take_last_error_message().contains("batch_size_bytes must be set before")); + + assert_eq!( + unsafe { lance_scanner_set_io_buffer_size(scanner, 64 * 1024) }, + -1 + ); + assert!(take_last_error_message().contains("io_buffer_size_bytes must be set before")); + + assert_eq!(unsafe { lance_scanner_set_batch_readahead(scanner, 1) }, -1); + assert!(take_last_error_message().contains("batch_readahead must be set before")); + + assert_eq!( + unsafe { lance_scanner_set_fragment_readahead(scanner, 1) }, + -1 + ); + assert!(take_last_error_message().contains("fragment_readahead must be set before")); + + assert_eq!( + unsafe { lance_scanner_set_target_parallelism(scanner, 1) }, + -1 + ); + assert!(take_last_error_message().contains("target_parallelism must be set before")); + + assert_eq!( + unsafe { lance_scanner_set_scan_in_order(scanner, false) }, + -1 + ); + assert!(take_last_error_message().contains("scan_in_order must be set before")); + + let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut ffi_stream) }.unwrap(); + assert_eq!( + reader.map(|batch| batch.unwrap().num_rows()).sum::(), + 5 + ); + + unsafe { lance_scanner_close(scanner) }; + unsafe { lance_dataset_close(ds) }; +} + // --------------------------------------------------------------------------- // Combined filter + projection + limit // --------------------------------------------------------------------------- @@ -1441,6 +1604,34 @@ fn test_null_safety_comprehensive() { unsafe { lance_scanner_set_batch_size(ptr::null_mut(), 10) }, -1 ); + assert_eq!( + unsafe { lance_scanner_set_batch_size_bytes(ptr::null_mut(), 1024) }, + -1 + ); + assert_eq!( + unsafe { lance_scanner_set_io_buffer_size(ptr::null_mut(), 64 * 1024) }, + -1 + ); + assert_eq!( + unsafe { lance_scanner_set_batch_readahead(ptr::null_mut(), 1) }, + -1 + ); + assert_eq!( + unsafe { lance_scanner_set_fragment_readahead(ptr::null_mut(), 1) }, + -1 + ); + assert_eq!( + unsafe { lance_scanner_set_target_parallelism(ptr::null_mut(), 1) }, + -1 + ); + assert_eq!( + unsafe { lance_scanner_set_query_parallelism(ptr::null_mut(), 1) }, + -1 + ); + assert_eq!( + unsafe { lance_scanner_set_scan_in_order(ptr::null_mut(), true) }, + -1 + ); assert_eq!( unsafe { lance_scanner_with_row_id(ptr::null_mut(), true) }, -1 @@ -5216,6 +5407,7 @@ fn test_scanner_nearest_with_ivf_pq_index() { 10, ); lance_scanner_set_nprobes(scanner, 4); + assert_eq!(lance_scanner_set_query_parallelism(scanner, 4), 0); } let mut stream = FFI_ArrowArrayStream::empty(); @@ -5234,6 +5426,59 @@ fn test_scanner_nearest_with_ivf_pq_index() { unsafe { lance_dataset_close(ds) }; } +#[test] +fn test_scanner_query_parallelism_validation_and_lifecycle() { + let (_tmp, uri) = create_test_dataset(); + let uri_c = c_str(&uri); + let ds = unsafe { lance_dataset_open(uri_c.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_query_parallelism(scanner, -1) }, + 0 + ); + assert_eq!( + unsafe { lance_scanner_set_query_parallelism(scanner, 0) }, + 0 + ); + assert_eq!( + unsafe { lance_scanner_set_query_parallelism(scanner, 2) }, + 0 + ); + + assert_eq!( + unsafe { lance_scanner_set_query_parallelism(scanner, -2) }, + -1 + ); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + assert!( + take_last_error_message() + .contains("query_parallelism must be -1, 0, or greater than 0, got -2") + ); + + let mut stream = FFI_ArrowArrayStream::empty(); + assert_eq!( + unsafe { lance_scanner_to_arrow_stream(scanner, &mut stream) }, + 0 + ); + assert_eq!( + unsafe { lance_scanner_set_query_parallelism(scanner, 1) }, + -1 + ); + assert!(take_last_error_message().contains("query_parallelism must be set before")); + + let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream) }.unwrap(); + assert_eq!( + reader.map(|batch| batch.unwrap().num_rows()).sum::(), + 5 + ); + + unsafe { lance_scanner_close(scanner) }; + unsafe { lance_dataset_close(ds) }; +} + #[test] fn test_scanner_nearest_dim_mismatch() { let (_tmp, uri) = create_vector_dataset(64, 8); diff --git a/tests/cpp/test_cpp_api.cpp b/tests/cpp/test_cpp_api.cpp index 17b1ab6..28762b8 100644 --- a/tests/cpp/test_cpp_api.cpp +++ b/tests/cpp/test_cpp_api.cpp @@ -134,6 +134,12 @@ static void test_scanner_fluent(const std::string& uri) { scanner.limit(5) .offset(0) .batch_size(2) + .batch_size_bytes(1024) + .io_buffer_size(64 * 1024) + .batch_readahead(1) + .fragment_readahead(1) + .target_parallelism(1) + .scan_in_order(false) .statistics_callback(capture_scan_statistics, &captured); ArrowArrayStream stream; @@ -370,6 +376,7 @@ static void test_nearest_smoke(const std::string& uri) { try { scanner.nearest("embedding", q, 8, 5) .nprobes(2) + .query_parallelism(2) .refine_factor(1) .ef(50) .metric(LANCE_METRIC_L2)