Project immutable ingestion output before decoding recovery payloads

crates/lance-context-ingestion/src/local_lance_reader.rs

@@ -1,6 +1,6 @@
 //! Version-pinned reads for an independent consumer of a local-state Lance log.
 use crate::{lance_sink::covered_sequence_at, Binding, Error, Result};
-use arrow_array::{Array, RecordBatch, StructArray};
+use arrow_array::{Array, RecordBatch, StructArray, UInt64Array};
 use futures::TryStreamExt;
 use lance::Dataset;
 
@@ -114,6 +114,53 @@ impl OutputRange {
         max_physical_bytes: u64,
         max_decoded_bytes: usize,
         batch_rows: usize,
+    ) -> Result<Vec<RecordBatch>> {
+        self.read_projection(max_physical_bytes, max_decoded_bytes, batch_rows, None)
+            .await
+    }
+
+    /// Read selected field paths within the output `record` struct. Projection
+    /// first selects matching row IDs, then reads only the projected schema,
+    /// preserving the record struct's validity before charging retained output.
+    /// Physical-file, range and sequence checks are unchanged. Empty projections
+    /// are rejected; omitted fields remain in the immutable log for recovery.
+    pub async fn read_projected(
+        self,
+        max_physical_bytes: u64,
+        max_decoded_bytes: usize,
+        batch_rows: usize,
+        fields: &[&str],
+    ) -> Result<Vec<RecordBatch>> {
+        if fields.is_empty() || fields.iter().any(|field| field.is_empty()) {
+            return Err(Error::Invalid("empty Lance output projection".into()));
+        }
+        let columns = fields
+            .iter()
+            .map(|field| format!("record.{field}"))
+            .collect::<Vec<_>>();
+        if columns
+            .iter()
+            .any(|column| self.dataset.schema().field(column).is_none())
+        {
+            return Err(Error::Invalid(
+                "unknown Lance output projection field".into(),
+            ));
+        }
+        self.read_projection(
+            max_physical_bytes,
+            max_decoded_bytes,
+            batch_rows,
+            Some(&columns),
+        )
+        .await
+    }
+
+    async fn read_projection(
+        self,
+        max_physical_bytes: u64,
+        max_decoded_bytes: usize,
+        batch_rows: usize,
+        columns: Option<&[String]>,
     ) -> 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()));
@@ -122,7 +169,15 @@ impl OutputRange {
         scan.with_fragments(self.fragments)
             .batch_size(batch_rows)
             .scan_in_order(true);
-        scan.project(&["record"]).map_err(fail)?;
+        let projection = columns
+            .map(|columns| self.dataset.schema().project(columns))
+            .transpose()
+            .map_err(fail)?;
+        if projection.is_some() {
+            scan.empty_project().map_err(fail)?.with_row_id();
+        } else {
+            scan.project(&["record"]).map_err(fail)?;
+        }
         scan.filter(&format!(
             "kind = 3 AND sequence >= {} AND sequence <= {}",
             self.range.first_sequence, self.range.last_sequence
@@ -132,6 +187,19 @@ impl OutputRange {
         let mut bytes = 0usize;
         let mut output = Vec::new();
         while let Some(batch) = stream.try_next().await.map_err(fail)? {
+            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()))?;
+                self.dataset
+                    .take_rows(row_ids.values(), projection.clone())
+                    .await
+                    .map_err(fail)?
+            } else {
+                batch
+            };
             bytes = bytes
                 .checked_add(batch.get_array_memory_size())
                 .ok_or_else(|| Error::Invalid("decoded Lance log size overflow".into()))?;

crates/lance-context-ingestion/tests/local_lance.rs

@@ -757,3 +757,136 @@ async fn wide_shared_ipc_buffers_keep_small_calls_in_one_bounded_commit() {
     assert_eq!(writer.dataset().version().version, base + 1);
     assert_eq!(writer.through_sequence(), 128);
 }
+
+#[tokio::test]
+async fn projected_output_avoids_unneeded_payload_and_preserves_read_limits() {
+    use lance_context_ingestion::local_lance_reader::{LogRange, OutputRange};
+
+    let dir = tempfile::tempdir().unwrap();
+    let uri = dir.path().join("projected.lance");
+    let output_schema = Arc::new(Schema::new(vec![
+        Field::new("content", DataType::Utf8, false),
+        Field::new("source_metadata", DataType::Utf8, false),
+        Field::new("nullable", DataType::Utf8, true),
+    ]));
+    let large = "unused recovery metadata".repeat(4096);
+    let output = RecordBatch::try_new(
+        output_schema.clone(),
+        vec![
+            Arc::new(StringArray::from(vec!["kept 世界"; 8])),
+            Arc::new(StringArray::from(vec![large.as_str(); 8])),
+            Arc::new(StringArray::from(vec![None::<&str>; 8])),
+        ],
+    )
+    .unwrap();
+    let schema = Arc::new(log_schema(&output_schema, &binding()));
+    let dataset = Dataset::write(
+        RecordBatchIterator::new(vec![Ok(RecordBatch::new_empty(schema.clone()))], schema),
+        uri.to_str().unwrap(),
+        Some(WriteParams {
+            data_storage_version: Some(LanceFileVersion::V2_2),
+            ..Default::default()
+        }),
+    )
+    .await
+    .unwrap();
+    let base_version = dataset.version().version;
+    let mut writer = open(&dir.path().join("projected.redb"), dataset).await;
+    writer
+        .commit(&[AlignedCall {
+            sequence: 1,
+            session: "s".into(),
+            receipt: "receipt".into(),
+            input_digest: "digest".into(),
+            mutations: vec![],
+            records: output,
+        }])
+        .await
+        .unwrap();
+    let snapshot = writer.dataset().clone();
+    let range = LogRange {
+        base_version,
+        version: snapshot.version().version,
+        first_sequence: 1,
+        last_sequence: 1,
+    };
+    let expected_binding = binding();
+    let open_range = || OutputRange::open(snapshot.clone(), &expected_binding, range.clone());
+    assert!(open_range()
+        .await
+        .unwrap()
+        .read(1 << 20, 64 << 10, 2)
+        .await
+        .is_err());
+    let projected = open_range()
+        .await
+        .unwrap()
+        .read_projected(1 << 20, 64 << 10, 2, &["content"])
+        .await
+        .unwrap();
+    let full = open_range()
+        .await
+        .unwrap()
+        .read(1 << 20, 8 << 20, 2)
+        .await
+        .unwrap();
+    assert_eq!(projected.len(), full.len());
+    assert_eq!(
+        projected.iter().map(RecordBatch::num_rows).sum::<usize>(),
+        8
+    );
+    for (selected, complete) in projected.iter().zip(&full) {
+        assert_eq!(selected, &complete.project(&[0]).unwrap());
+    }
+    let nullable = open_range()
+        .await
+        .unwrap()
+        .read_projected(1 << 20, 64 << 10, 2, &["nullable"])
+        .await
+        .unwrap();
+    for (selected, complete) in nullable.iter().zip(&full) {
+        assert_eq!(selected, &complete.project(&[2]).unwrap());
+    }
+    for (physical, decoded, fields) in [
+        (1, 64 << 10, vec!["content"]),
+        (1 << 20, 1, vec!["content"]),
+        (1 << 20, 64 << 10, vec![]),
+        (1 << 20, 64 << 10, vec!["missing"]),
+    ] {
+        assert!(open_range()
+            .await
+            .unwrap()
+            .read_projected(physical, decoded, 2, &fields)
+            .await
+            .is_err());
+    }
+    // A null record is corrupt, even when all selected child values may be null.
+    let mut scan = snapshot.scan();
+    scan.filter("kind = 3")
+        .unwrap()
+        .limit(Some(1), None)
+        .unwrap();
+    let batch = scan.try_into_batch().await.unwrap();
+    let mut columns = batch.columns().to_vec();
+    columns[8] = arrow_array::new_null_array(batch.schema().field(8).data_type(), 1);
+    let corrupt = RecordBatch::try_new(batch.schema(), columns).unwrap();
+    let corrupt = Dataset::write(
+        RecordBatchIterator::new(vec![Ok(corrupt)], batch.schema()),
+        uri.to_str().unwrap(),
+        Some(WriteParams {
+            mode: lance::dataset::WriteMode::Append,
+            ..Default::default()
+        }),
+    )
+    .await
+    .unwrap();
+    let mut corrupt_range = range;
+    corrupt_range.version = corrupt.version().version;
+    let error = OutputRange::open(corrupt, &expected_binding, corrupt_range)
+        .await
+        .unwrap()
+        .read_projected(1 << 20, 64 << 10, 2, &["nullable"])
+        .await
+        .unwrap_err();
+    assert!(error.to_string().contains("null Lance log output row"));
+}