Expose projected read planning and execution phases

crates/lance-context-ingestion/Cargo.toml

@@ -7,14 +7,15 @@ 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"
 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 }

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<Self> {
         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<Vec<RecordBatch>> {
+        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<Vec<RecordBatch>> {
+        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<Vec<RecordBatch>> {
         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<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()));
@@ -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 {

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::<Vec<_>>();
     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::<StringArray>()
+                    .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();