perf: Collect parquet statistics in one contiguous buffer

crates/polars-arrow/src/array/fixed_size_binary/mutable.rs

@@ -222,6 +222,14 @@ impl MutableFixedSizeBinaryArray {
             validity.shrink_to_fit()
         }
     }
+
+    pub fn freeze(self) -> FixedSizeBinaryArray {
+        FixedSizeBinaryArray::new(
+            ArrowDataType::FixedSizeBinary(self.size),
+            self.values.into(),
+            self.validity.map(|x| x.into()),
+        )
+    }
 }
 
 /// Accessors

crates/polars-io/src/parquet/read/predicates.rs

@@ -1,69 +1,45 @@
 use polars_core::config;
 use polars_core::prelude::*;
-use polars_parquet::read::statistics::{deserialize, Statistics};
+use polars_parquet::read::statistics::{
+    deserialize, deserialize_all, ArrowColumnStatisticsArrays, Statistics,
+};
 use polars_parquet::read::RowGroupMetadata;
 
 use crate::predicates::{BatchStats, ColumnStats, ScanIOPredicate};
 
 /// Collect the statistics in a row-group
 pub fn collect_statistics_with_live_columns(
-    md: &RowGroupMetadata,
+    row_groups: &[RowGroupMetadata],
     schema: &ArrowSchema,
-    pl_schema: &SchemaRef,
     live_columns: &PlIndexSet<PlSmallStr>,
-) -> PolarsResult<Option<BatchStats>> {
-    // TODO! fix this performance. This is a full sequential scan.
-    let stats = live_columns
+) -> PolarsResult<Vec<Option<ArrowColumnStatisticsArrays>>> {
+    if row_groups.is_empty() {
+        return Ok((0..live_columns.len()).map(|_| None).collect());
+    }
+
+    let md = &row_groups[0];
+    live_columns
         .iter()
         .map(|c| {
             let field = schema.get(c).unwrap();
 
-            let default_fn = || ColumnStats::new(field.into(), None, None, None);
-
             // This can be None in the allow_missing_columns case.
-            let Some(mut iter) = md.columns_under_root_iter(&field.name) else {
-                return Ok(default_fn());
+            let Some(idxs) = md.columns_idxs_under_root_iter(&field.name) else {
+                return Ok(None);
             };
 
-            let statistics = deserialize(field, &mut iter)?;
-            assert!(iter.next().is_none());
-
-            // We don't support reading nested statistics for now. It does not really make any
-            // sense at the moment with how we structure statistics.
-            let Some(Statistics::Column(stats)) = statistics else {
-                return Ok(default_fn());
-            };
-
-            let stats = stats.into_arrow()?;
-
-            let null_count = stats
-                .null_count
-                .map(|x| Scalar::from(x).into_series(PlSmallStr::EMPTY));
-            let min_value = stats
-                .min_value
-                .map(|x| Series::try_from((PlSmallStr::EMPTY, x)).unwrap());
-            let max_value = stats
-                .max_value
-                .map(|x| Series::try_from((PlSmallStr::EMPTY, x)).unwrap());
+            // 0 is possible for possible for empty structs.
+            //
+            // 2+ is for structs. We don't support reading nested statistics for now. It does not
+            // really make any sense at the moment with how we structure statistics.
+            if idxs.is_empty() || idxs.len() > 1 {
+                return Ok(None);
+            }
 
-            Ok(ColumnStats::new(
-                field.into(),
-                null_count,
-                min_value,
-                max_value,
-            ))
+            let idx = idxs[0];
+            Ok(deserialize_all(field, row_groups, idx)?)
         })
-        .collect::<PolarsResult<Vec<_>>>()?;
-
-    if stats.is_empty() {
-        return Ok(None);
-    }
-
-    Ok(Some(BatchStats::new(
-        pl_schema.clone(),
-        stats,
-        Some(md.num_rows()),
-    )))
+        .collect::<PolarsResult<Vec<_>>>()
 }
 
 /// Collect the statistics in a row-group