diff --git a/Cargo.lock b/Cargo.lock index f7feae6..b08af8f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3918,6 +3918,7 @@ dependencies = [ "async-trait", "futures", "lance", + "lance-datafusion", "lance-file", "lance-index", "lance-io", diff --git a/crates/lance-context-ingestion/Cargo.toml b/crates/lance-context-ingestion/Cargo.toml index f037351..6df96bb 100644 --- a/crates/lance-context-ingestion/Cargo.toml +++ b/crates/lance-context-ingestion/Cargo.toml @@ -7,7 +7,7 @@ description = "Bounded session-ordered ingestion with recoverable WAL and indepe [features] default = [] -lance = ["dep:arrow-array", "dep:arrow-ipc", "dep:arrow-schema", "dep:lance", "dep:lance-file", "dep:lance-index", "dep:lance-io", "dep:lance-table", "dep:redb", "dep:tempfile"] +lance = ["dep:arrow-array", "dep:arrow-ipc", "dep:arrow-schema", "dep:lance", "dep:lance-datafusion", "dep:lance-file", "dep:lance-index", "dep:lance-io", "dep:lance-table", "dep:redb", "dep:tempfile"] [dependencies] async-trait = "0.1" @@ -15,6 +15,7 @@ arrow-array = { version = "58", optional = true } arrow-ipc = { version = "58", optional = true } arrow-schema = { version = "58", optional = true } lance = { version = "9.0.0", optional = true } +lance-datafusion = { version = "9.0.0", optional = true } lance-file = { version = "9.0.0", optional = true } lance-index = { version = "9.0.0", optional = true } lance-io = { version = "9.0.0", optional = true } diff --git a/crates/lance-context-ingestion/src/local_lance_reader.rs b/crates/lance-context-ingestion/src/local_lance_reader.rs index 6c293e3..3289813 100644 --- a/crates/lance-context-ingestion/src/local_lance_reader.rs +++ b/crates/lance-context-ingestion/src/local_lance_reader.rs @@ -2,7 +2,8 @@ use crate::{lance_sink::covered_sequence_at, Binding, Error, Result}; use arrow_array::{Array, RecordBatch, StructArray, UInt64Array}; use futures::TryStreamExt; -use lance::Dataset; +use lance::{dataset::scanner::DatasetRecordBatchStream, Dataset}; +use lance_datafusion::exec::{execute_plan, LanceExecutionOptions}; fn fail(error: impl std::fmt::Display) -> Error { Error::Stage(error.to_string()) @@ -37,6 +38,30 @@ pub enum OutputReadStage { Complete, } +/// Finer scan boundaries without extending the original `OutputReadStage` enum. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OutputReadPhase { + Read(OutputReadStage), + PlanCreation, + ExecutionInitialization, +} + +#[derive(Clone, Copy)] +enum ReadObserver<'a> { + Stages(Option<&'a (dyn Fn(OutputReadStage, u64) + Send + Sync)>), + Phases(&'a (dyn Fn(OutputReadPhase, u64) + Send + Sync)), +} + +impl ReadObserver<'_> { + fn notify(self, stage: OutputReadStage, batch: u64) { + match self { + Self::Stages(Some(observer)) => observer(stage, batch), + Self::Stages(None) => {} + Self::Phases(observer) => observer(OutputReadPhase::Read(stage), batch), + } + } +} + impl OutputRange { pub async fn open(dataset: Dataset, binding: &Binding, range: LogRange) -> Result { if range.base_version == 0 @@ -129,7 +154,7 @@ impl OutputRange { max_decoded_bytes, batch_rows, None, - None, + ReadObserver::Stages(None), ) .await } @@ -166,6 +191,46 @@ impl OutputRange { batch_rows: usize, fields: &[&str], observer: Option<&(dyn Fn(OutputReadStage, u64) + Send + Sync)>, + ) -> Result> { + self.read_projected_with_observer( + max_physical_bytes, + max_decoded_bytes, + batch_rows, + fields, + ReadObserver::Stages(observer), + ) + .await + } + + /// Observe planning separately from synchronous execution initialization. + /// Callbacks run inline and must not block or panic. Entry is not completion; + /// no task, deadline, retry, or cancellation is introduced. Batch size, stream + /// schema adaptation, errors and drop behavior match `read_projected`. + pub async fn read_projected_phases( + self, + max_physical_bytes: u64, + max_decoded_bytes: usize, + batch_rows: usize, + fields: &[&str], + observer: &(dyn Fn(OutputReadPhase, u64) + Send + Sync), + ) -> Result> { + self.read_projected_with_observer( + max_physical_bytes, + max_decoded_bytes, + batch_rows, + fields, + ReadObserver::Phases(observer), + ) + .await + } + + async fn read_projected_with_observer( + self, + max_physical_bytes: u64, + max_decoded_bytes: usize, + batch_rows: usize, + fields: &[&str], + observer: ReadObserver<'_>, ) -> Result> { if fields.is_empty() || fields.iter().any(|field| field.is_empty()) { return Err(Error::Invalid("empty Lance output projection".into())); @@ -198,7 +263,7 @@ impl OutputRange { max_decoded_bytes: usize, batch_rows: usize, columns: Option<&[String]>, - observer: Option<&(dyn Fn(OutputReadStage, u64) + Send + Sync)>, + observer: ReadObserver<'_>, ) -> Result> { if self.physical_bytes > max_physical_bytes || max_decoded_bytes == 0 || batch_rows == 0 { return Err(Error::Invalid("Lance log range exceeds read budget".into())); @@ -221,13 +286,27 @@ impl OutputRange { self.range.first_sequence, self.range.last_sequence )) .map_err(fail)?; - let notify = |stage, batch| { - if let Some(observer) = observer { - observer(stage, batch); - } - }; + let notify = |stage, batch| observer.notify(stage, batch); notify(OutputReadStage::StreamCreation, 0); - let mut stream = scan.try_into_stream().await.map_err(fail)?; + let mut stream = if let ReadObserver::Phases(observer) = observer { + observer(OutputReadPhase::PlanCreation, 0); + let plan = scan.create_plan().await.map_err(fail)?; + observer(OutputReadPhase::ExecutionInitialization, 0); + // Match Scanner::try_into_stream: this locally constructed scanner + // has batch_size set above and no execution statistics callback. + DatasetRecordBatchStream::new( + execute_plan( + plan, + LanceExecutionOptions { + batch_size: Some(batch_rows), + ..Default::default() + }, + ) + .map_err(fail)?, + ) + } else { + scan.try_into_stream().await.map_err(fail)? + }; let mut bytes = 0usize; let mut output = Vec::new(); loop { diff --git a/crates/lance-context-ingestion/tests/local_lance.rs b/crates/lance-context-ingestion/tests/local_lance.rs index 1356a31..ab47aed 100644 --- a/crates/lance-context-ingestion/tests/local_lance.rs +++ b/crates/lance-context-ingestion/tests/local_lance.rs @@ -760,7 +760,9 @@ async fn wide_shared_ipc_buffers_keep_small_calls_in_one_bounded_commit() { #[tokio::test] async fn projected_output_avoids_unneeded_payload_and_preserves_read_limits() { - use lance_context_ingestion::local_lance_reader::{LogRange, OutputRange, OutputReadStage}; + use lance_context_ingestion::local_lance_reader::{ + LogRange, OutputRange, OutputReadPhase, OutputReadStage, + }; let dir = tempfile::tempdir().unwrap(); let uri = dir.path().join("projected.lance"); @@ -770,10 +772,11 @@ async fn projected_output_avoids_unneeded_payload_and_preserves_read_limits() { Field::new("nullable", DataType::Utf8, true), ])); let large = "unused recovery metadata".repeat(4096); + let contents = (0..8).map(|i| format!("kept 世界 {i}")).collect::>(); let output = RecordBatch::try_new( output_schema.clone(), vec![ - Arc::new(StringArray::from(vec!["kept 世界"; 8])), + Arc::new(StringArray::from(contents.clone())), Arc::new(StringArray::from(vec![large.as_str(); 8])), Arc::new(StringArray::from(vec![None::<&str>; 8])), ], @@ -841,6 +844,73 @@ async fn projected_output_avoids_unneeded_payload_and_preserves_read_limits() { expected_stages.push((OutputReadStage::RowIds, observed.len() as u64)); expected_stages.push((OutputReadStage::Complete, observed.len() as u64)); assert_eq!(*stages.lock().unwrap(), expected_stages); + let phases = std::sync::Mutex::new(Vec::new()); + let phase_observer = |phase, batch| phases.lock().unwrap().push((phase, batch)); + for batch_rows in [1, 2, 1024] { + phases.lock().unwrap().clear(); + let expected = open_range() + .await + .unwrap() + .read_projected(1 << 20, 64 << 10, batch_rows, &["content", "nullable"]) + .await + .unwrap(); + let detailed = open_range() + .await + .unwrap() + .read_projected_phases( + 1 << 20, + 64 << 10, + batch_rows, + &["content", "nullable"], + &phase_observer, + ) + .await + .unwrap(); + assert_eq!(detailed, expected); + let actual: Vec<_> = detailed + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .map(|value| value.unwrap().to_owned()) + }) + .collect(); + assert_eq!(actual, contents); + let mut expected_phases = vec![ + (OutputReadPhase::Read(OutputReadStage::StreamCreation), 0), + (OutputReadPhase::PlanCreation, 0), + (OutputReadPhase::ExecutionInitialization, 0), + ]; + for batch in 0..detailed.len() as u64 { + expected_phases.push((OutputReadPhase::Read(OutputReadStage::RowIds), batch)); + expected_phases.push((OutputReadPhase::Read(OutputReadStage::TakeRows), batch)); + } + expected_phases.push(( + OutputReadPhase::Read(OutputReadStage::RowIds), + detailed.len() as u64, + )); + expected_phases.push(( + OutputReadPhase::Read(OutputReadStage::Complete), + detailed.len() as u64, + )); + assert_eq!(*phases.lock().unwrap(), expected_phases); + } + phases.lock().unwrap().clear(); + let phase_error = open_range() + .await + .unwrap() + .read_projected_phases(1 << 20, 1, 2, &["content"], &phase_observer) + .await + .unwrap_err(); + assert!(!phases + .lock() + .unwrap() + .iter() + .any(|(phase, _)| { *phase == OutputReadPhase::Read(OutputReadStage::Complete) })); stages.lock().unwrap().clear(); let observed_error = open_range() .await @@ -855,6 +925,7 @@ async fn projected_output_avoids_unneeded_payload_and_preserves_read_limits() { .await .unwrap_err(); assert_eq!(observed_error.to_string(), original_error.to_string()); + assert_eq!(phase_error.to_string(), original_error.to_string()); assert!(!stages .lock() .unwrap() @@ -889,12 +960,28 @@ async fn projected_output_avoids_unneeded_payload_and_preserves_read_limits() { (1 << 20, 64 << 10, vec![]), (1 << 20, 64 << 10, vec!["missing"]), ] { - assert!(open_range() + let expected_error = open_range() .await .unwrap() .read_projected(physical, decoded, 2, &fields) .await - .is_err()); + .unwrap_err(); + phases.lock().unwrap().clear(); + let detailed_error = open_range() + .await + .unwrap() + .read_projected_phases(physical, decoded, 2, &fields, &phase_observer) + .await + .unwrap_err(); + assert_eq!(detailed_error.to_string(), expected_error.to_string()); + assert!(!phases + .lock() + .unwrap() + .iter() + .any(|(phase, _)| { *phase == OutputReadPhase::Read(OutputReadStage::Complete) })); + if physical == 1 || fields.is_empty() || fields == ["missing"] { + assert!(phases.lock().unwrap().is_empty()); + } } // A null record is corrupt, even when all selected child values may be null. let mut scan = snapshot.scan();