From 89025240444d39a5c2c7005ee59fd1c4263bb7de Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Fri, 9 Oct 2026 08:12:47 +0000 Subject: [PATCH] Expose projected WAL read progress without changing read semantics --- .../src/local_lance_reader.rs | 57 ++++++++++++++++++- .../tests/local_lance.rs | 38 ++++++++++++- 2 files changed, 91 insertions(+), 4 deletions(-) diff --git a/crates/lance-context-ingestion/src/local_lance_reader.rs b/crates/lance-context-ingestion/src/local_lance_reader.rs index 6c06eab..6c293e3 100644 --- a/crates/lance-context-ingestion/src/local_lance_reader.rs +++ b/crates/lance-context-ingestion/src/local_lance_reader.rs @@ -28,6 +28,15 @@ pub struct OutputRange { range: LogRange, } +/// Current projected-read boundary. Notifications contain no row values or paths. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OutputReadStage { + StreamCreation, + RowIds, + TakeRows, + Complete, +} + impl OutputRange { pub async fn open(dataset: Dataset, binding: &Binding, range: LogRange) -> Result { if range.base_version == 0 @@ -115,8 +124,14 @@ impl OutputRange { max_decoded_bytes: usize, batch_rows: usize, ) -> Result> { - self.read_projection(max_physical_bytes, max_decoded_bytes, batch_rows, None) - .await + self.read_projection( + max_physical_bytes, + max_decoded_bytes, + batch_rows, + None, + None, + ) + .await } /// Read selected field paths within the output `record` struct. Projection @@ -130,6 +145,27 @@ impl OutputRange { max_decoded_bytes: usize, batch_rows: usize, fields: &[&str], + ) -> Result> { + self.read_projected_observed( + max_physical_bytes, + max_decoded_bytes, + batch_rows, + fields, + None, + ) + .await + } + + /// Observe stage entry with a zero-based batch ordinal. The callback runs + /// inline and must not block or panic. Errors and cancellation retain the + /// same semantics as `read_projected`; `Complete` is emitted only on success. + pub async fn read_projected_observed( + self, + max_physical_bytes: u64, + max_decoded_bytes: usize, + batch_rows: usize, + fields: &[&str], + observer: Option<&(dyn Fn(OutputReadStage, u64) + Send + Sync)>, ) -> Result> { if fields.is_empty() || fields.iter().any(|field| field.is_empty()) { return Err(Error::Invalid("empty Lance output projection".into())); @@ -151,6 +187,7 @@ impl OutputRange { max_decoded_bytes, batch_rows, Some(&columns), + observer, ) .await } @@ -161,6 +198,7 @@ impl OutputRange { max_decoded_bytes: usize, batch_rows: usize, columns: Option<&[String]>, + observer: Option<&(dyn Fn(OutputReadStage, u64) + Send + Sync)>, ) -> 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())); @@ -183,16 +221,28 @@ 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); + } + }; + notify(OutputReadStage::StreamCreation, 0); let mut stream = scan.try_into_stream().await.map_err(fail)?; let mut bytes = 0usize; let mut output = Vec::new(); - while let Some(batch) = stream.try_next().await.map_err(fail)? { + loop { + let batch_ordinal = output.len() as u64; + notify(OutputReadStage::RowIds, batch_ordinal); + let Some(batch) = stream.try_next().await.map_err(fail)? else { + break; + }; let batch = if let Some(projection) = &projection { let row_ids = batch .column_by_name("_rowid") .and_then(|column| column.as_any().downcast_ref::()) .filter(|ids| ids.null_count() == 0) .ok_or_else(|| Error::Invalid("invalid Lance output row IDs".into()))?; + notify(OutputReadStage::TakeRows, batch_ordinal); self.dataset .take_rows(row_ids.values(), projection.clone()) .await @@ -218,6 +268,7 @@ impl OutputRange { } output.push(RecordBatch::from(records.clone())); } + notify(OutputReadStage::Complete, output.len() as u64); Ok(output) } } diff --git a/crates/lance-context-ingestion/tests/local_lance.rs b/crates/lance-context-ingestion/tests/local_lance.rs index 4a1f7b7..1356a31 100644 --- a/crates/lance-context-ingestion/tests/local_lance.rs +++ b/crates/lance-context-ingestion/tests/local_lance.rs @@ -760,7 +760,7 @@ 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}; + use lance_context_ingestion::local_lance_reader::{LogRange, OutputRange, OutputReadStage}; let dir = tempfile::tempdir().unwrap(); let uri = dir.path().join("projected.lance"); @@ -824,6 +824,42 @@ async fn projected_output_avoids_unneeded_payload_and_preserves_read_limits() { .read_projected(1 << 20, 64 << 10, 2, &["content"]) .await .unwrap(); + let stages = std::sync::Mutex::new(Vec::new()); + let observer = |stage, batch| stages.lock().unwrap().push((stage, batch)); + let observed = open_range() + .await + .unwrap() + .read_projected_observed(1 << 20, 64 << 10, 2, &["content"], Some(&observer)) + .await + .unwrap(); + assert_eq!(observed, projected); + let mut expected_stages = vec![(OutputReadStage::StreamCreation, 0)]; + for batch in 0..observed.len() as u64 { + expected_stages.push((OutputReadStage::RowIds, batch)); + expected_stages.push((OutputReadStage::TakeRows, batch)); + } + 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); + stages.lock().unwrap().clear(); + let observed_error = open_range() + .await + .unwrap() + .read_projected_observed(1 << 20, 1, 2, &["content"], Some(&observer)) + .await + .unwrap_err(); + let original_error = open_range() + .await + .unwrap() + .read_projected(1 << 20, 1, 2, &["content"]) + .await + .unwrap_err(); + assert_eq!(observed_error.to_string(), original_error.to_string()); + assert!(!stages + .lock() + .unwrap() + .iter() + .any(|(stage, _)| *stage == OutputReadStage::Complete)); let full = open_range() .await .unwrap()