diff --git a/rust/lancedb/src/data/scannable.rs b/rust/lancedb/src/data/scannable.rs index 8248a3202..67afecb2e 100644 --- a/rust/lancedb/src/data/scannable.rs +++ b/rust/lancedb/src/data/scannable.rs @@ -9,13 +9,6 @@ use std::sync::Arc; -use arrow_array::{ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader}; -use arrow_schema::{ArrowError, SchemaRef}; -use async_trait::async_trait; -use futures::stream::once; -use futures::StreamExt; -use lance_datafusion::utils::StreamingWriteSource; - use crate::arrow::{ SendableRecordBatchStream, SendableRecordBatchStreamExt, SimpleRecordBatchStream, }; @@ -25,6 +18,12 @@ use crate::embeddings::{ }; use crate::table::{ColumnDefinition, ColumnKind, TableDefinition}; use crate::{Error, Result}; +use arrow_array::{ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader}; +use arrow_schema::{ArrowError, SchemaRef}; +use async_trait::async_trait; +use futures::stream::once; +use futures::StreamExt; +use lance_datafusion::utils::StreamingWriteSource; pub trait Scannable: Send { /// Returns the schema of the data. @@ -349,6 +348,133 @@ pub fn scannable_with_embeddings( Ok(inner) } +/// A wrapper that buffers the first RecordBatch from a Scannable so we can +/// inspect it (e.g. to estimate data size) without losing it. +pub(crate) struct PeekedScannable { + inner: Box, + peeked: Option, + /// The first item from the stream, if it was an error. Stored so we can + /// re-emit it from `scan_as_stream` instead of silently dropping it. + first_error: Option, + stream: Option, +} + +impl PeekedScannable { + pub fn new(inner: Box) -> Self { + Self { + inner, + peeked: None, + first_error: None, + stream: None, + } + } + + /// Reads and buffers the first batch from the inner scannable. + /// Returns a clone of it. Subsequent calls return the same batch. + /// + /// Returns `None` if the stream is empty or the first item is an error. + /// Errors are preserved and re-emitted by `scan_as_stream`. + pub async fn peek(&mut self) -> Option { + if self.peeked.is_some() { + return self.peeked.clone(); + } + // Already peeked and got an error or empty stream. + if self.stream.is_some() || self.first_error.is_some() { + return None; + } + let mut stream = self.inner.scan_as_stream(); + match stream.next().await { + Some(Ok(batch)) => { + self.peeked = Some(batch.clone()); + self.stream = Some(stream); + Some(batch) + } + Some(Err(e)) => { + self.first_error = Some(e); + self.stream = Some(stream); + None + } + None => { + self.stream = Some(stream); + None + } + } + } +} + +impl Scannable for PeekedScannable { + fn schema(&self) -> SchemaRef { + self.inner.schema() + } + + fn num_rows(&self) -> Option { + self.inner.num_rows() + } + + fn rescannable(&self) -> bool { + self.inner.rescannable() + } + + fn scan_as_stream(&mut self) -> SendableRecordBatchStream { + let schema = self.inner.schema(); + + // If peek() hit an error, prepend it so downstream sees the error. + let error_item = self.first_error.take().map(Err); + + match (self.peeked.take(), self.stream.take()) { + (Some(batch), Some(rest)) => { + let prepend = futures::stream::once(std::future::ready(Ok(batch))); + Box::pin(SimpleRecordBatchStream { + schema, + stream: prepend.chain(rest), + }) + } + (Some(batch), None) => Box::pin(SimpleRecordBatchStream { + schema, + stream: futures::stream::once(std::future::ready(Ok(batch))), + }), + (None, Some(rest)) => { + if let Some(err) = error_item { + let stream = futures::stream::once(std::future::ready(err)); + Box::pin(SimpleRecordBatchStream { schema, stream }) + } else { + rest + } + } + (None, None) => { + // peek() was never called — just delegate + self.inner.scan_as_stream() + } + } + } +} + +/// Compute the number of write partitions based on data size estimates. +/// +/// `sample_bytes` and `sample_rows` come from a representative batch and are +/// used to estimate per-row size. `total_rows_hint` is the total row count +/// when known; otherwise `sample_rows` row count is used as a lower bound +/// estimate. +/// +/// Targets roughly 1 million rows or 2 GB per partition, capped at +/// `max_partitions` (typically the number of available CPU cores). +pub(crate) fn estimate_write_partitions( + sample_bytes: usize, + sample_rows: usize, + total_rows_hint: Option, + max_partitions: usize, +) -> usize { + if sample_rows == 0 { + return 1; + } + let bytes_per_row = sample_bytes / sample_rows; + let total_rows = total_rows_hint.unwrap_or(sample_rows); + let total_bytes = total_rows * bytes_per_row; + let by_rows = total_rows.div_ceil(1_000_000); + let by_bytes = total_bytes.div_ceil(2 * 1024 * 1024 * 1024); + by_rows.max(by_bytes).max(1).min(max_partitions) +} + #[cfg(test)] mod tests { use super::*; @@ -445,6 +571,231 @@ mod tests { assert!(result2.unwrap().is_err()); } + mod peeked_scannable_tests { + use crate::test_utils::TestCustomError; + + use super::*; + + #[tokio::test] + async fn test_peek_returns_first_batch() { + let batch = record_batch!(("id", Int64, [1, 2, 3])).unwrap(); + let mut peeked = PeekedScannable::new(Box::new(batch.clone())); + + let first = peeked.peek().await.unwrap(); + assert_eq!(first, batch); + } + + #[tokio::test] + async fn test_peek_is_idempotent() { + let batch = record_batch!(("id", Int64, [1, 2, 3])).unwrap(); + let mut peeked = PeekedScannable::new(Box::new(batch.clone())); + + let first = peeked.peek().await.unwrap(); + let second = peeked.peek().await.unwrap(); + assert_eq!(first, second); + } + + #[tokio::test] + async fn test_scan_after_peek_returns_all_data() { + let batches = vec![ + record_batch!(("id", Int64, [1, 2])).unwrap(), + record_batch!(("id", Int64, [3, 4, 5])).unwrap(), + ]; + let mut peeked = PeekedScannable::new(Box::new(batches.clone())); + + let first = peeked.peek().await.unwrap(); + assert_eq!(first, batches[0]); + + let result: Vec = peeked.scan_as_stream().try_collect().await.unwrap(); + assert_eq!(result.len(), 2); + assert_eq!(result[0], batches[0]); + assert_eq!(result[1], batches[1]); + } + + #[tokio::test] + async fn test_scan_without_peek_passes_through() { + let batch = record_batch!(("id", Int64, [1, 2, 3])).unwrap(); + let mut peeked = PeekedScannable::new(Box::new(batch.clone())); + + let result: Vec = peeked.scan_as_stream().try_collect().await.unwrap(); + assert_eq!(result.len(), 1); + assert_eq!(result[0], batch); + } + + #[tokio::test] + async fn test_delegates_num_rows() { + let batches = vec![ + record_batch!(("id", Int64, [1, 2])).unwrap(), + record_batch!(("id", Int64, [3])).unwrap(), + ]; + let peeked = PeekedScannable::new(Box::new(batches)); + assert_eq!(peeked.num_rows(), Some(3)); + } + + #[tokio::test] + async fn test_non_rescannable_stream_data_preserved() { + let batches = vec![ + record_batch!(("id", Int64, [1, 2])).unwrap(), + record_batch!(("id", Int64, [3])).unwrap(), + ]; + let schema = batches[0].schema(); + let inner = futures::stream::iter(batches.clone().into_iter().map(Ok)); + let stream: SendableRecordBatchStream = Box::pin(SimpleRecordBatchStream { + schema, + stream: inner, + }); + + let mut peeked = PeekedScannable::new(Box::new(stream)); + assert!(!peeked.rescannable()); + assert_eq!(peeked.num_rows(), None); + + let first = peeked.peek().await.unwrap(); + assert_eq!(first, batches[0]); + + // All data is still available via scan_as_stream + let result: Vec = peeked.scan_as_stream().try_collect().await.unwrap(); + assert_eq!(result.len(), 2); + assert_eq!(result[0], batches[0]); + assert_eq!(result[1], batches[1]); + } + + #[tokio::test] + async fn test_error_in_first_batch_propagates() { + let schema = Arc::new(arrow_schema::Schema::new(vec![arrow_schema::Field::new( + "id", + arrow_schema::DataType::Int64, + false, + )])); + let inner = futures::stream::iter(vec![Err(Error::External { + source: Box::new(TestCustomError), + })]); + let stream: SendableRecordBatchStream = Box::pin(SimpleRecordBatchStream { + schema, + stream: inner, + }); + + let mut peeked = PeekedScannable::new(Box::new(stream)); + + // peek returns None for errors + assert!(peeked.peek().await.is_none()); + + // But the error should come through when scanning + let mut stream = peeked.scan_as_stream(); + let first = stream.next().await.unwrap(); + assert!(first.is_err()); + let err = first.unwrap_err(); + assert!( + matches!(&err, Error::External { source } if source.downcast_ref::().is_some()), + "Expected TestCustomError to be preserved, got: {err}" + ); + } + + #[tokio::test] + async fn test_error_in_later_batch_propagates() { + let good_batch = record_batch!(("id", Int64, [1, 2])).unwrap(); + let schema = good_batch.schema(); + let inner = futures::stream::iter(vec![ + Ok(good_batch.clone()), + Err(Error::External { + source: Box::new(TestCustomError), + }), + ]); + let stream: SendableRecordBatchStream = Box::pin(SimpleRecordBatchStream { + schema, + stream: inner, + }); + + let mut peeked = PeekedScannable::new(Box::new(stream)); + + // peek succeeds with the first batch + let first = peeked.peek().await.unwrap(); + assert_eq!(first, good_batch); + + // scan_as_stream should yield the first batch, then the error + let mut stream = peeked.scan_as_stream(); + let batch1 = stream.next().await.unwrap().unwrap(); + assert_eq!(batch1, good_batch); + + let batch2 = stream.next().await.unwrap(); + assert!(batch2.is_err()); + let err = batch2.unwrap_err(); + assert!( + matches!(&err, Error::External { source } if source.downcast_ref::().is_some()), + "Expected TestCustomError to be preserved, got: {err}" + ); + } + + #[tokio::test] + async fn test_empty_stream_returns_none() { + let schema = Arc::new(arrow_schema::Schema::new(vec![arrow_schema::Field::new( + "id", + arrow_schema::DataType::Int64, + false, + )])); + let inner = futures::stream::empty(); + let stream: SendableRecordBatchStream = Box::pin(SimpleRecordBatchStream { + schema, + stream: inner, + }); + + let mut peeked = PeekedScannable::new(Box::new(stream)); + assert!(peeked.peek().await.is_none()); + + // Scanning an empty (post-peek) stream should yield nothing + let result: Vec = peeked.scan_as_stream().try_collect().await.unwrap(); + assert!(result.is_empty()); + } + } + + mod estimate_write_partitions_tests { + use super::*; + + #[test] + fn test_small_data_single_partition() { + // 100 rows * 24 bytes/row = 2400 bytes — well under both thresholds + assert_eq!(estimate_write_partitions(2400, 100, Some(100), 8), 1); + } + + #[test] + fn test_scales_by_row_count() { + // 2.5M rows at 24 bytes/row — row threshold dominates + // ceil(2_500_000 / 1_000_000) = 3 + assert_eq!(estimate_write_partitions(72, 3, Some(2_500_000), 8), 3); + } + + #[test] + fn test_scales_by_byte_size() { + // 100k rows at 40KB/row = ~4GB total → ceil(4GB / 2GB) = 2 + let sample_bytes = 40_000 * 10; + assert_eq!( + estimate_write_partitions(sample_bytes, 10, Some(100_000), 8), + 2 + ); + } + + #[test] + fn test_capped_at_max_partitions() { + // 10M rows would want 10 partitions, but capped at 4 + assert_eq!(estimate_write_partitions(72, 3, Some(10_000_000), 4), 4); + } + + #[test] + fn test_zero_sample_rows_returns_one() { + assert_eq!(estimate_write_partitions(0, 0, Some(1_000_000), 8), 1); + } + + #[test] + fn test_no_row_hint_uses_sample_size() { + // Without a hint, uses sample_rows (3), which is small + assert_eq!(estimate_write_partitions(72, 3, None, 8), 1); + } + + #[test] + fn test_always_at_least_one() { + assert_eq!(estimate_write_partitions(24, 1, Some(1), 8), 1); + } + } + mod embedding_tests { use super::*; use crate::embeddings::MemoryRegistry; diff --git a/rust/lancedb/src/error.rs b/rust/lancedb/src/error.rs index d04566d7d..7bfbc6f2e 100644 --- a/rust/lancedb/src/error.rs +++ b/rust/lancedb/src/error.rs @@ -97,10 +97,7 @@ pub type Result = std::result::Result; impl From for Error { fn from(source: ArrowError) -> Self { match source { - ArrowError::ExternalError(source) => match source.downcast::() { - Ok(e) => *e, - Err(source) => Self::External { source }, - }, + ArrowError::ExternalError(source) => Self::from_box_error(source), _ => Self::Arrow { source }, } } @@ -110,15 +107,7 @@ impl From for Error { fn from(source: DataFusionError) -> Self { match source { DataFusionError::ArrowError(source, _) => (*source).into(), - DataFusionError::External(source) => match source.downcast::() { - Ok(e) => *e, - Err(source) => match source.downcast::() { - Ok(arrow_error) => Self::Arrow { - source: *arrow_error, - }, - Err(source) => Self::External { source }, - }, - }, + DataFusionError::External(source) => Self::from_box_error(source), other => Self::External { source: Box::new(other), }, @@ -130,15 +119,52 @@ impl From for Error { fn from(source: lance::Error) -> Self { // Try to unwrap external errors that were wrapped by lance match source { - lance::Error::Wrapped { error, .. } => match error.downcast::() { - Ok(e) => *e, - Err(source) => Self::External { source }, - }, + lance::Error::Wrapped { error, .. } => Self::from_box_error(error), + lance::Error::External { source } => Self::from_box_error(source), _ => Self::Lance { source }, } } } +impl Error { + fn from_box_error(mut source: Box) -> Self { + source = match source.downcast::() { + Ok(e) => match *e { + Self::External { source } => return Self::from_box_error(source), + other => return other, + }, + Err(source) => source, + }; + + source = match source.downcast::() { + Ok(e) => match *e { + lance::Error::Wrapped { error, .. } => return Self::from_box_error(error), + other => return other.into(), + }, + Err(source) => source, + }; + + source = match source.downcast::() { + Ok(e) => match *e { + ArrowError::ExternalError(source) => return Self::from_box_error(source), + other => return other.into(), + }, + Err(source) => source, + }; + + source = match source.downcast::() { + Ok(e) => match *e { + DataFusionError::ArrowError(source, _) => return (*source).into(), + DataFusionError::External(source) => return Self::from_box_error(source), + other => return other.into(), + }, + Err(source) => source, + }; + + Self::External { source } + } +} + impl From for Error { fn from(source: object_store::Error) -> Self { Self::ObjectStore { source } diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 31bb8a1b3..e92a817b6 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -6,11 +6,12 @@ use arrow_array::{RecordBatch, RecordBatchReader}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use async_trait::async_trait; +use datafusion_execution::TaskContext; use datafusion_expr::Expr; use datafusion_physical_plan::display::DisplayableExecutionPlan; use datafusion_physical_plan::ExecutionPlan; +use futures::stream::FuturesUnordered; use futures::StreamExt; -use futures::TryStreamExt; use lance::dataset::builder::DatasetBuilder; pub use lance::dataset::ColumnAlteration; pub use lance::dataset::NewColumnTransform; @@ -21,7 +22,6 @@ use lance::dataset::{InsertBuilder, WriteParams}; use lance::index::vector::utils::infer_vector_dim; use lance::index::vector::VectorIndexParams; use lance::io::{ObjectStoreParams, WrappingObjectStore}; -use lance_datafusion::exec::execute_plan; use lance_datafusion::utils::StreamingWriteSource; use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams}; use lance_index::vector::bq::RQBuildParams; @@ -43,7 +43,7 @@ use std::format; use std::path::Path; use std::sync::Arc; -use crate::data::scannable::Scannable; +use crate::data::scannable::{estimate_write_partitions, PeekedScannable, Scannable}; use crate::database::Database; use crate::embeddings::{EmbeddingDefinition, EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; @@ -2113,7 +2113,7 @@ impl BaseTable for NativeTable { } } - async fn add(&self, add: AddDataBuilder) -> Result { + async fn add(&self, mut add: AddDataBuilder) -> Result { let table_def = self.table_definition().await?; self.dataset.ensure_mutable()?; @@ -2122,6 +2122,22 @@ impl BaseTable for NativeTable { let table_schema = Schema::from(&ds.schema().clone()); + // Peek at the first batch to estimate a good partition count for + // write parallelism. + let mut peeked = PeekedScannable::new(add.data); + let num_partitions = if let Some(first_batch) = peeked.peek().await { + let max_partitions = lance_core::utils::tokio::get_num_compute_intensive_cpus(); + estimate_write_partitions( + first_batch.get_array_memory_size(), + first_batch.num_rows(), + peeked.num_rows(), + max_partitions, + ) + } else { + 1 + }; + add.data = Box::new(peeked); + let output = add.into_plan(&table_schema, &table_def)?; let lance_params = output @@ -2135,18 +2151,41 @@ impl BaseTable for NativeTable { ..Default::default() }); - let plan = Arc::new(InsertExec::new( - ds_wrapper.clone(), - ds, - output.plan, - lance_params, - )); + // Repartition for write parallelism if beneficial. + let plan = if num_partitions > 1 { + Arc::new( + datafusion_physical_plan::repartition::RepartitionExec::try_new( + output.plan, + datafusion_physical_plan::Partitioning::RoundRobinBatch(num_partitions), + )?, + ) as Arc + } else { + output.plan + }; - let stream = execute_plan(plan, Default::default())?; - stream - .try_collect::>() - .await - .map_err(crate::Error::from)?; + let insert_exec = Arc::new(InsertExec::new(ds_wrapper.clone(), ds, plan, lance_params)); + + // Execute all partitions in parallel. + let task_ctx = Arc::new(TaskContext::default()); + let handles = FuturesUnordered::new(); + for partition in 0..num_partitions { + let exec = insert_exec.clone(); + let ctx = task_ctx.clone(); + handles.push(tokio::spawn(async move { + let mut stream = exec + .execute(partition, ctx) + .map_err(|e| -> Error { e.into() })?; + while let Some(batch) = stream.next().await { + batch.map_err(|e| -> Error { e.into() })?; + } + Ok::<_, Error>(()) + })); + } + for handle in handles { + handle.await.map_err(|e| Error::Runtime { + message: format!("Insert task panicked: {}", e), + })??; + } let version = ds_wrapper.get().await?.manifest().version; Ok(AddResult { version }) diff --git a/rust/lancedb/src/table/add_data.rs b/rust/lancedb/src/table/add_data.rs index ad64e36d7..ea5f111c2 100644 --- a/rust/lancedb/src/table/add_data.rs +++ b/rust/lancedb/src/table/add_data.rs @@ -219,6 +219,7 @@ mod tests { use crate::table::add_data::NaNVectorBehavior; use crate::table::{ColumnDefinition, ColumnKind, Table, TableDefinition, WriteOptions}; use crate::test_utils::embeddings::MockEmbed; + use crate::test_utils::TestCustomError; use crate::Error; use super::AddDataMode; @@ -283,17 +284,20 @@ mod tests { test_add_with_data(stream).await; } - #[derive(Debug)] - struct MyError; - - impl std::fmt::Display for MyError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "MyError occurred") - } + fn assert_preserves_external_error(err: &Error) { + assert!( + matches!(err, Error::External { source } if source.downcast_ref::().is_some()), + "Expected Error::External, got: {err:?}" + ); + // The original TestCustomError message should be preserved through the + // error chain, even if the error gets wrapped multiple times by + // lance's insert pipeline. + assert!( + err.to_string().contains("TestCustomError occurred"), + "Expected original error message to be preserved, got: {err}" + ); } - impl std::error::Error for MyError {} - #[tokio::test] async fn test_add_preserves_reader_error() { let table = create_test_table().await; @@ -301,7 +305,7 @@ mod tests { let schema = first_batch.schema(); let iterator = vec![ Ok(first_batch), - Err(ArrowError::ExternalError(Box::new(MyError))), + Err(ArrowError::ExternalError(Box::new(TestCustomError))), ]; let reader: Box = Box::new( RecordBatchIterator::new(iterator.into_iter(), schema.clone()), @@ -309,7 +313,7 @@ mod tests { let result = table.add(reader).execute().await; - assert!(result.is_err()); + assert_preserves_external_error(&result.unwrap_err()); } #[tokio::test] @@ -320,7 +324,7 @@ mod tests { let iterator = vec![ Ok(first_batch), Err(Error::External { - source: Box::new(MyError), + source: Box::new(TestCustomError), }), ]; let stream = futures::stream::iter(iterator); @@ -331,7 +335,7 @@ mod tests { let result = table.add(stream).execute().await; - assert!(result.is_err()); + assert_preserves_external_error(&result.unwrap_err()); } #[tokio::test] diff --git a/rust/lancedb/src/test_utils.rs b/rust/lancedb/src/test_utils.rs index 1e7fc742b..3210d0568 100644 --- a/rust/lancedb/src/test_utils.rs +++ b/rust/lancedb/src/test_utils.rs @@ -4,3 +4,14 @@ pub mod connection; pub mod datagen; pub mod embeddings; + +#[derive(Debug)] +pub struct TestCustomError; + +impl std::fmt::Display for TestCustomError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "TestCustomError occurred") + } +} + +impl std::error::Error for TestCustomError {}