fix: avoid overflow for large table state

crates/core/Cargo.toml

@@ -80,6 +80,7 @@ indexmap = "2.2.1"
 itertools = "0.14"
 parking_lot = "0.12"
 percent-encoding = "2"
+pin-project-lite = "^0.2.7"
 tracing = { workspace = true }
 rand = "0.8"
 sqlparser = { version = "0.59.0" }

crates/core/src/delta_datafusion/table_provider.rs

@@ -46,10 +46,8 @@ use datafusion::{
     prelude::Expr,
     scalar::ScalarValue,
 };
-use futures::TryStreamExt;
 use itertools::Itertools;
 use object_store::ObjectMeta;
-
 use serde::{Deserialize, Serialize};
 
 use crate::delta_datafusion::schema_adapter::DeltaSchemaAdapterFactory;
@@ -214,7 +212,7 @@ impl DataSink for DeltaDataSink {
             .map_err(|e| DataFusionError::External(Box::new(e)))?;
 
         let mut actions = if self.save_mode == SaveMode::Overwrite {
-            let current_files = self.snapshot.log_data().iter().map(|f| f.add_action());
+            let current_files = self.snapshot.log_data().into_iter().map(|f| f.add_action());
             current_files
                 .into_iter()
                 .map(|add| {
@@ -651,9 +649,21 @@ impl<'a> DeltaScanBuilder<'a> {
         //  Should we update datafusion_table_statistics to optionally take the mask?
         let stats = if let Some(mask) = pruning_mask {
             let es = self.snapshot.snapshot();
-            let pruned_stats =
-                filter_record_batch(&self.snapshot.files, &BooleanArray::from(mask))?;
-            LogDataHandler::new(&pruned_stats, es.table_configuration()).statistics()
+            let mut pruned_batches = Vec::new();
+            let mut mask_offset = 0;
+
+            for batch in &self.snapshot.files {
+                let batch_size = batch.num_rows();
+                let batch_mask = &mask[mask_offset..mask_offset + batch_size];
+                let batch_mask_array = BooleanArray::from(batch_mask.to_vec());
+                let pruned_batch = filter_record_batch(batch, &batch_mask_array)?;
+                if pruned_batch.num_rows() > 0 {
+                    pruned_batches.push(pruned_batch);
+                }
+                mask_offset += batch_size;
+            }
+
+            LogDataHandler::new(&pruned_batches, es.table_configuration()).statistics()
         } else {
             self.snapshot.log_data().statistics()
         };

crates/core/src/kernel/arrow/engine_ext.rs

@@ -3,7 +3,6 @@
 use std::borrow::Cow;
 use std::sync::Arc;
 
-use arrow_array::{Array, StructArray};
 use arrow_schema::{
     DataType as ArrowDataType, Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef,
 };
@@ -12,7 +11,6 @@ use delta_kernel::arrow::compute::filter_record_batch;
 use delta_kernel::arrow::record_batch::RecordBatch;
 use delta_kernel::engine::arrow_conversion::TryIntoArrow;
 use delta_kernel::engine::arrow_data::ArrowEngineData;
-use delta_kernel::engine::parse_json;
 use delta_kernel::expressions::{ColumnName, Scalar, StructData};
 use delta_kernel::scan::{Scan, ScanMetadata};
 use delta_kernel::schema::{
@@ -27,7 +25,6 @@ use delta_kernel::{
 use itertools::Itertools;
 
 use crate::errors::{DeltaResult as DeltaResultLocal, DeltaTableError};
-use crate::kernel::replay::parse_partitions;
 use crate::kernel::SCAN_ROW_ARROW_SCHEMA;
 
 /// [`ScanMetadata`] contains (1) a [`RecordBatch`] specifying data files to be scanned
@@ -118,9 +115,6 @@ pub(crate) trait SnapshotExt {
     /// computations by delta-rs. Specifically the `stats_parsed` and
     /// `partitionValues_parsed` fields are added.
     fn scan_row_parsed_schema_arrow(&self) -> DeltaResultLocal<ArrowSchemaRef>;
-
-    /// Parse stats column into a struct array.
-    fn parse_stats_column(&self, batch: &RecordBatch) -> DeltaResultLocal<RecordBatch>;
 }
 
 impl SnapshotExt for Snapshot {
@@ -150,14 +144,15 @@ impl SnapshotExt for Snapshot {
     /// scan row (file data).
     fn scan_row_parsed_schema_arrow(&self) -> DeltaResultLocal<ArrowSchemaRef> {
         let mut fields = SCAN_ROW_ARROW_SCHEMA.fields().to_vec();
+        let stats_idx = SCAN_ROW_ARROW_SCHEMA.index_of("stats").unwrap();
 
         let stats_schema = self.stats_schema()?;
         let stats_schema: ArrowSchema = stats_schema.as_ref().try_into_arrow()?;
-        fields.push(Arc::new(Field::new(
+        fields[stats_idx] = Arc::new(Field::new(
             "stats_parsed",
             ArrowDataType::Struct(stats_schema.fields().to_owned()),
             true,
-        )));
+        ));
 
         if let Some(partition_schema) = self.partitions_schema()? {
             let partition_schema: ArrowSchema = partition_schema.as_ref().try_into_arrow()?;
@@ -171,51 +166,6 @@ impl SnapshotExt for Snapshot {
         let schema = Arc::new(ArrowSchema::new(fields));
         Ok(schema)
     }
-
-    fn parse_stats_column(&self, batch: &RecordBatch) -> DeltaResultLocal<RecordBatch> {
-        let Some((stats_idx, _)) = batch.schema_ref().column_with_name("stats") else {
-            return Err(DeltaTableError::SchemaMismatch {
-                msg: "stats column not found".to_string(),
-            });
-        };
-
-        let mut columns = batch.columns().to_vec();
-        let mut fields = batch.schema().fields().to_vec();
-
-        let stats_schema = self.stats_schema()?;
-        let stats_batch = batch.project(&[stats_idx])?;
-        let stats_data = Box::new(ArrowEngineData::new(stats_batch));
-
-        let parsed = parse_json(stats_data, stats_schema)?;
-        let parsed: RecordBatch = ArrowEngineData::try_from_engine_data(parsed)?.into();
-
-        let stats_array: Arc<StructArray> = Arc::new(parsed.into());
-        fields.push(Arc::new(Field::new(
-            "stats_parsed",
-            stats_array.data_type().to_owned(),
-            true,
-        )));
-        columns.push(stats_array.clone());
-
-        if let Some(partition_schema) = self.partitions_schema()? {
-            let partition_array = parse_partitions(
-                batch,
-                partition_schema.as_ref(),
-                "fileConstantValues.partitionValues",
-            )?;
-            fields.push(Arc::new(Field::new(
-                "partitionValues_parsed",
-                partition_array.data_type().to_owned(),
-                false,
-            )));
-            columns.push(Arc::new(partition_array));
-        }
-
-        Ok(RecordBatch::try_new(
-            Arc::new(ArrowSchema::new(fields)),
-            columns,
-        )?)
-    }
 }
 
 fn partitions_schema(
@@ -916,4 +866,24 @@ mod tests {
         assert_eq!(None, result);
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_partition_schema2() {
+        let schema = StructType::try_new(vec![
+            StructField::new("id", DataType::LONG, true),
+            StructField::new("name", DataType::STRING, true),
+            StructField::new("date", DataType::DATE, true),
+        ])
+        .unwrap();
+
+        let partition_columns = vec!["date".to_string()];
+        let expected =
+            StructType::try_new(vec![StructField::new("date", DataType::DATE, true)]).unwrap();
+        assert_eq!(
+            partitions_schema(&schema, &partition_columns).unwrap(),
+            Some(expected)
+        );
+
+        assert_eq!(partitions_schema(&schema, &[]).unwrap(), None);
+    }
 }

crates/core/src/kernel/snapshot/iterators.rs

@@ -8,6 +8,7 @@ use arrow_array::types::Int64Type;
 use arrow_array::{Array, RecordBatch, StructArray};
 use arrow_schema::DataType as ArrowDataType;
 use chrono::{DateTime, Utc};
+use delta_kernel::engine::arrow_expression::evaluate_expression::to_json;
 use delta_kernel::expressions::{Scalar, StructData};
 use delta_kernel::scan::scan_row_schema;
 use delta_kernel::schema::DataType;
@@ -19,14 +20,15 @@ use crate::kernel::scalars::ScalarExt;
 use crate::kernel::{Add, DeletionVectorDescriptor, Remove};
 use crate::{DeltaResult, DeltaTableError};
 
+pub(crate) use self::scan_row::{scan_row_in_eval, ScanRowOutStream};
 pub use self::tombstones::TombstoneView;
 
+mod scan_row;
 mod tombstones;
 
 const FIELD_NAME_PATH: &str = "path";
 const FIELD_NAME_SIZE: &str = "size";
 const FIELD_NAME_MODIFICATION_TIME: &str = "modificationTime";
-const FIELD_NAME_STATS: &str = "stats";
 const FIELD_NAME_STATS_PARSED: &str = "stats_parsed";
 const FIELD_NAME_PARTITION_VALUES_PARSED: &str = "partitionValues_parsed";
 const FIELD_NAME_DELETION_VECTOR: &str = "deletionVector";
@@ -55,9 +57,6 @@ static FIELD_INDICES: LazyLock<HashMap<&'static str, usize>> = LazyLock::new(||
     let modification_time_idx = schema.index_of(FIELD_NAME_MODIFICATION_TIME).unwrap();
     indices.insert(FIELD_NAME_MODIFICATION_TIME, modification_time_idx);
 
-    let stats_idx = schema.index_of(FIELD_NAME_STATS).unwrap();
-    indices.insert(FIELD_NAME_STATS, stats_idx);
-
     indices
 });
 
@@ -158,12 +157,12 @@ impl LogicalFileView {
     }
 
     /// Returns the raw JSON statistics string for this file, if available.
-    pub fn stats(&self) -> Option<&str> {
-        get_string_value(
-            self.files
-                .column(*FIELD_INDICES.get(FIELD_NAME_STATS).unwrap()),
-            self.index,
-        )
+    pub fn stats(&self) -> Option<String> {
+        let stats = self.stats_parsed()?.slice(self.index, 1);
+        let value = to_json(&stats)
+            .ok()
+            .map(|arr| arr.as_string::<i32>().value(0).to_string());
+        value.and_then(|v| (!v.is_empty()).then_some(v))
     }
 
     /// Returns the parsed partition values as structured data.
@@ -275,7 +274,7 @@ impl LogicalFileView {
             size: self.size(),
             modification_time: self.modification_time(),
             data_change: true,
-            stats: self.stats().map(|v| v.to_string()),
+            stats: self.stats(),
             tags: None,
             deletion_vector: self.deletion_vector().map(|dv| dv.descriptor()),
             base_row_id: None,