Expose progress boundaries for projected ingestion-log reads

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<Self> {
         if range.base_version == 0
@@ -115,8 +124,14 @@ impl OutputRange {
         max_decoded_bytes: usize,
         batch_rows: usize,
     ) -> Result<Vec<RecordBatch>> {
-        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<Vec<RecordBatch>> {
+        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<Vec<RecordBatch>> {
         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<Vec<RecordBatch>> {
         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::<UInt64Array>())
                     .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)
     }
 }

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()