From eea5fbb3504fbd8813d7059442773aa381397e4f Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Thu, 8 Oct 2026 11:45:16 +0000 Subject: [PATCH] Read projected immutable ingestion output without recovery payloads --- .../src/local_lance_reader.rs | 72 +++++++++- .../tests/local_lance.rs | 133 ++++++++++++++++++ 2 files changed, 203 insertions(+), 2 deletions(-) diff --git a/crates/lance-context-ingestion/src/local_lance_reader.rs b/crates/lance-context-ingestion/src/local_lance_reader.rs index bd276e2..6c06eab 100644 --- a/crates/lance-context-ingestion/src/local_lance_reader.rs +++ b/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> { + 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> { + 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::>(); + 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> { 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::()) + .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()))?; diff --git a/crates/lance-context-ingestion/tests/local_lance.rs b/crates/lance-context-ingestion/tests/local_lance.rs index 19cf20d..4a1f7b7 100644 --- a/crates/lance-context-ingestion/tests/local_lance.rs +++ b/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::(), + 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")); +}