zoned cleanup (#7634)

## Summary

Cleans up `ZonedLayout`. This is just a cosmetic change as I was trying
to understand it. Also adds some tests that I thought would be helpful.

## Testing

N/A

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

docs/concepts/scanning.md

@@ -68,10 +68,12 @@ independently. The scan tracks the selectivity of each conjunct using a probabil
 and dynamically reorders them so that the most selective predicates are evaluated first. This
 means that as a scan progresses, it learns the most efficient evaluation order for the filter.
 
-Filters are evaluated in two stages. First, pruning evaluation uses statistics stored in
-`ZonedLayout` (such as per-zone min/max values) to eliminate entire regions without reading any
-data. Second, filter evaluation materializes only the filter-referenced columns and computes a
-row mask of matching rows.
+Filters are evaluated in two stages. First, pruning evaluation uses statistics stored in a
+`ZonedLayout` auxiliary `zones` child to eliminate entire row zones without reading the underlying
+data child. These pruning predicates are falsification checks derived from the original filter, for
+example by comparing a zone's min/max values against the requested predicate. Second, filter
+evaluation materializes only the filter-referenced columns and computes a row mask of matching
+rows.
 
 ## Projection Pushdown
 
@@ -89,4 +91,3 @@ Query engines integrate with the Scan API by translating their internal plan rep
 scan requests and consuming the resulting array stream in their preferred format. Integrations
 exist for DuckDB, DataFusion, Spark, and Trino, with each engine converting its native filter
 and projection representations into Vortex [expressions](expressions.md).
-

vortex-layout/src/layouts/file_stats.rs

@@ -21,7 +21,7 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_panic;
 use vortex_session::VortexSession;
 
-use crate::layouts::zoned::zone_map::StatsAccumulator;
+use crate::layouts::zoned::StatsAccumulator;
 use crate::sequence::SendableSequentialStream;
 use crate::sequence::SequenceId;
 use crate::sequence::SequentialStreamAdapter;

vortex-layout/src/layouts/zoned/builder.rs

@@ -1,32 +1,137 @@
+//! Write-time accumulation and builders for zoned layout stats tables.
+
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::marker::PhantomData;
 
+use itertools::Itertools;
 use vortex_array::ArrayRef;
+use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
 use vortex_array::LEGACY_SESSION;
 use vortex_array::VortexSessionExecute;
 use vortex_array::arrays::ConstantArray;
+use vortex_array::arrays::StructArray;
 use vortex_array::builders::ArrayBuilder;
 use vortex_array::builders::BoolBuilder;
 use vortex_array::builders::builder_with_capacity;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::FieldName;
 use vortex_array::dtype::Nullability;
+use vortex_array::expr::stats::Precision;
 use vortex_array::expr::stats::Stat;
+use vortex_array::expr::stats::StatsProvider;
 use vortex_array::scalar::Scalar;
 use vortex_array::scalar::ScalarTruncation;
 use vortex_array::scalar::lower_bound;
 use vortex_array::scalar::upper_bound;
+use vortex_array::validity::Validity;
 use vortex_buffer::BufferString;
 use vortex_buffer::ByteBuffer;
+use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 
-pub const MAX_IS_TRUNCATED: &str = "max_is_truncated";
-pub const MIN_IS_TRUNCATED: &str = "min_is_truncated";
+use crate::layouts::zoned::schema::MAX_IS_TRUNCATED;
+use crate::layouts::zoned::schema::MIN_IS_TRUNCATED;
+use crate::layouts::zoned::zone_map::ZoneMap;
+
+/// Accumulates write-time statistics for each logical zone.
+pub struct StatsAccumulator {
+    builders: Vec<Box<dyn StatsArrayBuilder>>,
+    length: usize,
+}
+
+impl StatsAccumulator {
+    pub fn new(dtype: &DType, stats: &[Stat], max_variable_length_statistics_size: usize) -> Self {
+        let builders = stats
+            .iter()
+            .filter_map(|&stat| {
+                stat.dtype(dtype).map(|stat_dtype| {
+                    stats_builder_with_capacity(
+                        stat,
+                        &stat_dtype.as_nullable(),
+                        1024,
+                        max_variable_length_statistics_size,
+                    )
+                })
+            })
+            .collect::<Vec<_>>();
+
+        Self {
+            builders,
+            length: 0,
+        }
+    }
+
+    pub fn push_chunk_without_compute(&mut self, array: &ArrayRef) -> VortexResult<()> {
+        for builder in &mut self.builders {
+            if let Some(Precision::Exact(value)) = array.statistics().get(builder.stat()) {
+                builder.append_scalar(value.cast(&value.dtype().as_nullable())?)?;
+            } else {
+                builder.append_null();
+            }
+        }
+        self.length += 1;
+        Ok(())
+    }
 
-pub fn stats_builder_with_capacity(
+    pub fn push_chunk(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> {
+        for builder in &mut self.builders {
+            if let Some(value) = array.statistics().compute_stat(builder.stat(), ctx)? {
+                builder.append_scalar(value.cast(&value.dtype().as_nullable())?)?;
+            } else {
+                builder.append_null();
+            }
+        }
+        self.length += 1;
+        Ok(())
+    }
+
+    /// Finishes the accumulator into a [`ZoneMap`].
+    ///
+    /// Returns `None` if none of the requested statistics can be computed, for example they are
+    /// not applicable to the column's data type.
+    pub fn as_stats_table(&mut self) -> VortexResult<Option<ZoneMap>> {
+        let mut names = Vec::new();
+        let mut fields = Vec::new();
+        let mut stats = Vec::new();
+
+        for builder in self
+            .builders
+            .iter_mut()
+            // We sort the stats so the DType is deterministic based on which stats are present.
+            .sorted_unstable_by_key(|builder| builder.stat())
+        {
+            let values = builder.finish();
+
+            // We drop any all-null stats columns.
+            if values.all_invalid()? {
+                continue;
+            }
+
+            stats.push(builder.stat());
+            names.extend(values.names);
+            fields.extend(values.arrays);
+        }
+
+        if names.is_empty() {
+            return Ok(None);
+        }
+
+        let array = StructArray::try_new(names.into(), fields, self.length, Validity::NonNullable)
+            .vortex_expect("Failed to create zone map");
+        let stats = stats.into();
+
+        // SAFETY: `StatsAccumulator` builds the struct fields from `stats_builder_with_capacity`
+        // using the same field-ordering and truncation-column rules as `stats_table_dtype`.
+        // The `stats` list is collected in that same sorted order, so the resulting struct array
+        // matches the expected zoned stats-table dtype by construction.
+        Ok(Some(unsafe { ZoneMap::new_unchecked(array, stats) }))
+    }
+}
+
+fn stats_builder_with_capacity(
     stat: Stat,
     dtype: &DType,
     capacity: usize,
@@ -64,21 +169,20 @@ pub fn stats_builder_with_capacity(
     }
 }
 
-/// Arrays with their associated names, reduced version of a StructArray
-pub struct NamedArrays {
-    pub names: Vec<FieldName>,
-    pub arrays: Vec<ArrayRef>,
+/// Arrays with their associated names, reduced version of a `StructArray`.
+struct NamedArrays {
+    names: Vec<FieldName>,
+    arrays: Vec<ArrayRef>,
 }
 
 impl NamedArrays {
-    pub fn all_invalid(&self) -> VortexResult<bool> {
-        // by convention we assume that the first array is the one we care about for logical validity
+    fn all_invalid(&self) -> VortexResult<bool> {
+        // By convention the first array is the logical validity signal for the stat column.
         self.arrays[0].all_invalid(&mut LEGACY_SESSION.create_execution_ctx())
     }
 }
 
-/// Minimal array builder interface for use by StatsTable for building stats arrays
-pub trait StatsArrayBuilder: Send {
+trait StatsArrayBuilder: Send {
     fn stat(&self) -> Stat;
 
     fn append_scalar(&mut self, value: Scalar) -> VortexResult<()>;
@@ -88,13 +192,13 @@ pub trait StatsArrayBuilder: Send {
     fn finish(&mut self) -> NamedArrays;
 }
 
-pub struct StatNameArrayBuilder {
+struct StatNameArrayBuilder {
     stat: Stat,
     builder: Box<dyn ArrayBuilder>,
 }
 
 impl StatNameArrayBuilder {
-    pub fn new(stat: Stat, builder: Box<dyn ArrayBuilder>) -> Self {
+    fn new(stat: Stat, builder: Box<dyn ArrayBuilder>) -> Self {
         Self { stat, builder }
     }
 }
@@ -140,7 +244,7 @@ struct TruncatedMaxBinaryStatsBuilder<T: ScalarTruncation> {
 }
 
 impl<T: ScalarTruncation> TruncatedMaxBinaryStatsBuilder<T> {
-    pub fn new(
+    fn new(
         values: Box<dyn ArrayBuilder>,
         is_truncated: BoolBuilder,
         max_value_length: usize,
@@ -162,7 +266,7 @@ struct TruncatedMinBinaryStatsBuilder<T: ScalarTruncation> {
 }
 
 impl<T: ScalarTruncation> TruncatedMinBinaryStatsBuilder<T> {
-    pub fn new(
+    fn new(
         values: Box<dyn ArrayBuilder>,
         is_truncated: BoolBuilder,
         max_value_length: usize,
@@ -194,7 +298,6 @@ impl<T: ScalarTruncation> StatsArrayBuilder for TruncatedMaxBinaryStatsBuilder<T
         Ok(())
     }
 
-    #[inline]
     fn append_null(&mut self) {
         ArrayBuilder::append_null(self.values.as_mut());
         self.is_truncated.append_value(false);
@@ -229,7 +332,6 @@ impl<T: ScalarTruncation> StatsArrayBuilder for TruncatedMinBinaryStatsBuilder<T
         Ok(())
     }
 
-    #[inline]
     fn append_null(&mut self) {
         ArrayBuilder::append_null(self.values.as_mut());
         self.is_truncated.append_value(false);
@@ -245,3 +347,113 @@ impl<T: ScalarTruncation> StatsArrayBuilder for TruncatedMinBinaryStatsBuilder<T
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+    use vortex_array::IntoArray;
+    use vortex_array::LEGACY_SESSION;
+    use vortex_array::VortexSessionExecute;
+    use vortex_array::arrays::BoolArray;
+    use vortex_array::arrays::bool::BoolArrayExt;
+    use vortex_array::arrays::struct_::StructArrayExt;
+    use vortex_array::builders::ArrayBuilder;
+    use vortex_array::builders::VarBinViewBuilder;
+    use vortex_array::dtype::DType;
+    use vortex_array::dtype::Nullability;
+    use vortex_array::expr::stats::Stat;
+    use vortex_buffer::BitBuffer;
+    use vortex_buffer::buffer;
+    use vortex_error::VortexExpect;
+
+    use super::*;
+
+    #[rstest]
+    #[case(DType::Utf8(Nullability::NonNullable))]
+    #[case(DType::Binary(Nullability::NonNullable))]
+    fn truncates_accumulated_stats(#[case] dtype: DType) {
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        let mut builder = VarBinViewBuilder::with_capacity(dtype.clone(), 2);
+        builder.append_value("Value to be truncated");
+        builder.append_value("untruncated");
+        let mut builder2 = VarBinViewBuilder::with_capacity(dtype, 2);
+        builder2.append_value("Another");
+        builder2.append_value("wait a minute");
+        let mut acc =
+            StatsAccumulator::new(builder.dtype(), &[Stat::Max, Stat::Min, Stat::Sum], 12);
+        acc.push_chunk(&builder.finish(), &mut ctx)
+            .vortex_expect("push_chunk should succeed for test data");
+        acc.push_chunk(&builder2.finish(), &mut ctx)
+            .vortex_expect("push_chunk should succeed for test data");
+        let stats_table = acc
+            .as_stats_table()
+            .unwrap()
+            .expect("Must have stats table");
+        assert_eq!(
+            stats_table.array().names().as_ref(),
+            &[
+                Stat::Max.name(),
+                MAX_IS_TRUNCATED,
+                Stat::Min.name(),
+                MIN_IS_TRUNCATED,
+            ]
+        );
+        let field1_bool = stats_table
+            .array()
+            .unmasked_field(1)
+            .clone()
+            .execute::<BoolArray>(&mut ctx)
+            .unwrap();
+        assert_eq!(
+            field1_bool.to_bit_buffer(),
+            BitBuffer::from(vec![false, true])
+        );
+        let field3_bool = stats_table
+            .array()
+            .unmasked_field(3)
+            .clone()
+            .execute::<BoolArray>(&mut ctx)
+            .unwrap();
+        assert_eq!(
+            field3_bool.to_bit_buffer(),
+            BitBuffer::from(vec![true, false])
+        );
+    }
+
+    #[test]
+    fn always_adds_is_truncated_column() {
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        let array = buffer![0, 1, 2].into_array();
+        let mut acc = StatsAccumulator::new(array.dtype(), &[Stat::Max, Stat::Min, Stat::Sum], 12);
+        acc.push_chunk(&array, &mut ctx)
+            .vortex_expect("push_chunk should succeed for test array");
+        let stats_table = acc
+            .as_stats_table()
+            .unwrap()
+            .expect("Must have stats table");
+        assert_eq!(
+            stats_table.array().names().as_ref(),
+            &[
+                Stat::Max.name(),
+                MAX_IS_TRUNCATED,
+                Stat::Min.name(),
+                MIN_IS_TRUNCATED,
+                Stat::Sum.name(),
+            ]
+        );
+        let field1_bool = stats_table
+            .array()
+            .unmasked_field(1)
+            .clone()
+            .execute::<BoolArray>(&mut ctx)
+            .unwrap();
+        assert_eq!(field1_bool.to_bit_buffer(), BitBuffer::from(vec![false]));
+        let field3_bool = stats_table
+            .array()
+            .unmasked_field(3)
+            .clone()
+            .execute::<BoolArray>(&mut ctx)
+            .unwrap();
+        assert_eq!(field3_bool.to_bit_buffer(), BitBuffer::from(vec![false]));
+    }
+}