Statistices nested columns

vortex-file/src/file.rs

@@ -227,19 +227,14 @@ impl VortexFile {
     /// Row-count-aware pruning predicates are evaluated with the file's total
     /// row count as their scope.
     pub fn can_prune(&self, filter: &Expression) -> VortexResult<bool> {
-        let Some((stats, fields)) = self
-            .footer
-            .statistics()
-            .zip(self.footer.dtype().as_struct_fields_opt())
-        else {
+        let Some(stats) = self.footer.statistics() else {
             return Ok(false);
         };
 
         can_prune_file_stats(
             &filter.bind(self.footer.dtype())?,
             self.footer.row_count(),
             stats,
-            fields,
             &self.session,
         )
     }

vortex-file/src/footer/file_statistics.rs

@@ -12,6 +12,7 @@ use flatbuffers::FlatBufferBuilder;
 use flatbuffers::WIPOffset;
 use itertools::Itertools;
 use vortex_array::dtype::DType;
+use vortex_array::dtype::FieldPath;
 use vortex_array::stats::StatsSet;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
@@ -20,6 +21,7 @@ use vortex_flatbuffers::FlatBufferRoot;
 use vortex_flatbuffers::WriteFlatBuffer;
 use vortex_flatbuffers::array::ArrayStats;
 use vortex_flatbuffers::footer as fb;
+use vortex_layout::layouts::file_stats::postorder_stats_layout;
 use vortex_session::VortexSession;
 
 /// Contains statistical information about the data in a Vortex file.
@@ -33,61 +35,63 @@ pub struct FileStatistics {
     stats: Arc<[StatsSet]>,
     /// An array of `DType`s, one for each field or column in the file.
     dtypes: Arc<[DType]>,
+    /// An array of field paths, one for each field or column in the file. Parallel to `stats` and
+    /// `dtypes`. For files written before nested field stats, every path has depth 1 (or is the
+    /// root path, for a non-struct file dtype).
+    paths: Arc<[FieldPath]>,
 }
 
 impl FileStatistics {
-    /// Creates a new [`FileStatistics`] from the given statistics and data types.
+    /// Creates a new [`FileStatistics`] from the given statistics, data types, and field paths.
     ///
     /// # Panics
     ///
-    /// Panics if `stats` and `dtypes` have different lengths.
-    pub fn new(stats: Arc<[StatsSet]>, dtypes: Arc<[DType]>) -> Self {
+    /// Panics if `stats`, `dtypes`, and `paths` have different lengths.
+    pub fn new(stats: Arc<[StatsSet]>, dtypes: Arc<[DType]>, paths: Arc<[FieldPath]>) -> Self {
         assert_eq!(
             stats.len(),
             dtypes.len(),
             "stats and dtypes must have the same length"
         );
+        assert_eq!(
+            stats.len(),
+            paths.len(),
+            "stats and paths must have the same length"
+        );
 
-        Self { stats, dtypes }
+        Self {
+            stats,
+            dtypes,
+            paths,
+        }
     }
 
     /// Creates a new [`FileStatistics`] from the given statistics and file dtype.
     ///
-    /// If the [`DType`] of the file is a [`DType::Struct`], then there must be the same number of
-    /// stats as struct fields. Otherwise, there must be only 1 statistic.
+    /// `stats` must follow the post-order nested-struct layout produced by
+    /// [`postorder_stats_layout`] for `file_dtype`.
     ///
     /// # Panics
     ///
     /// Panics if the number of stats doesn't match the expected number based on the dtype.
     pub fn new_with_dtype(stats: Arc<[StatsSet]>, file_dtype: &DType) -> Self {
-        if let DType::Struct(struct_fields, _) = file_dtype {
-            assert_eq!(
-                stats.len(),
-                struct_fields.nfields(),
-                "stats length must match number of struct fields"
-            );
+        let layout = postorder_stats_layout(file_dtype);
+        assert_eq!(
+            stats.len(),
+            layout.len(),
+            "stats length must match the post-order stats layout for the file dtype"
+        );
 
-            let dtypes = struct_fields.fields().collect();
+        let (paths, dtypes): (Vec<FieldPath>, Vec<DType>) = layout.into_iter().unzip();
 
-            Self { stats, dtypes }
-        } else {
-            assert_eq!(
-                stats.len(),
-                1,
-                "non-struct dtype must have exactly 1 statistic"
-            );
-
-            Self {
-                stats,
-                dtypes: Arc::new([file_dtype.clone()]),
-            }
+        Self {
+            stats,
+            dtypes: dtypes.into(),
+            paths: paths.into(),
         }
     }
 
     /// Creates [`FileStatistics`] from a flatbuffers [`fb::FileStatistics<'a>`].
-    ///
-    /// If the [`DType`] of the file is a [`DType::Struct`], then there must be the same number of
-    /// file stats in the flatbuffer. Otherwise, there must be only 1 statistic.
     pub fn from_flatbuffer<'a>(
         fb: &fb::FileStatistics<'a>,
         file_dtype: &DType,
@@ -96,6 +100,28 @@ impl FileStatistics {
         let field_stats = fb.field_stats().unwrap_or_default();
         let mut array_stats: Vec<ArrayStats> = field_stats.iter().collect();
 
+        if fb.is_nested() {
+            let layout = postorder_stats_layout(file_dtype);
+            vortex_ensure_eq!(array_stats.len(), layout.len());
+
+            let mut stats_sets = Vec::with_capacity(array_stats.len());
+            let mut dtypes = Vec::with_capacity(layout.len());
+            let mut paths = Vec::with_capacity(layout.len());
+            for (array_stat, (path, dtype)) in array_stats.into_iter().zip(layout) {
+                stats_sets.push(StatsSet::from_flatbuffer(&array_stat, &dtype, session)?);
+                dtypes.push(dtype);
+                paths.push(path);
+            }
+
+            return Ok(Self {
+                stats: stats_sets.into(),
+                dtypes: dtypes.into(),
+                paths: paths.into(),
+            });
+        }
+
+        // Legacy (pre-nested-stats) layout: top-level struct fields only, or a single entry for a
+        // non-struct root dtype.
         if let DType::Struct(struct_fields, _) = file_dtype {
             vortex_ensure_eq!(array_stats.len(), struct_fields.nfields());
 
@@ -108,10 +134,16 @@ impl FileStatistics {
                 .try_collect()?;
 
             let dtypes = struct_fields.fields().collect();
+            let paths = struct_fields
+                .names()
+                .iter()
+                .map(|name| FieldPath::from_name(name.clone()))
+                .collect();
 
             Ok(Self {
                 stats: stats_sets,
                 dtypes,
+                paths,
             })
         } else {
             vortex_ensure_eq!(array_stats.len(), 1);
@@ -124,6 +156,7 @@ impl FileStatistics {
             Ok(Self {
                 stats: Arc::new([stats_set]),
                 dtypes: Arc::new([file_dtype.clone()]),
+                paths: Arc::new([FieldPath::root()]),
             })
         }
     }
@@ -138,6 +171,11 @@ impl FileStatistics {
         &self.dtypes
     }
 
+    /// Returns a reference to the field paths.
+    pub fn paths(&self) -> &Arc<[FieldPath]> {
+        &self.paths
+    }
+
     /// Returns the statistics and data type for a specific field.
     ///
     /// # Panics
@@ -146,6 +184,14 @@ impl FileStatistics {
     pub fn get(&self, field_idx: usize) -> (&StatsSet, &DType) {
         (&self.stats[field_idx], &self.dtypes[field_idx])
     }
+
+    /// Returns the statistics and data type for the field at the given path, if present.
+    pub fn get_by_path(&self, path: &FieldPath) -> Option<(&StatsSet, &DType)> {
+        self.paths
+            .iter()
+            .position(|p| p == path)
+            .map(|idx| (&self.stats[idx], &self.dtypes[idx]))
+    }
 }
 
 impl<'a> IntoIterator for &'a FileStatistics {
@@ -177,7 +223,102 @@ impl WriteFlatBuffer for FileStatistics {
             fbb,
             &fb::FileStatisticsArgs {
                 field_stats: Some(field_stats),
+                is_nested: true,
             },
         ))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use flatbuffers::FlatBufferBuilder;
+    use vortex_array::array_session;
+    use vortex_array::dtype::FieldPath;
+    use vortex_array::dtype::Nullability;
+    use vortex_array::dtype::PType;
+    use vortex_array::expr::stats::Precision;
+    use vortex_array::expr::stats::Stat;
+    use vortex_array::scalar::ScalarValue;
+    use vortex_flatbuffers::WriteFlatBuffer;
+    use vortex_flatbuffers::WriteFlatBufferExt;
+
+    use super::*;
+
+    fn i32_dtype() -> DType {
+        DType::Primitive(PType::I32, Nullability::NonNullable)
+    }
+
+    #[test]
+    fn nested_round_trip_resolves_by_path() -> VortexResult<()> {
+        let session = array_session();
+        let inner = DType::struct_([("b", i32_dtype())], Nullability::Nullable);
+        let file_dtype = DType::struct_([("a", inner)], Nullability::NonNullable);
+
+        // Layout: [a.b, a] (a's own null-count entry trails its child).
+        let mut b_stats = StatsSet::default();
+        b_stats.set(Stat::Min, Precision::exact(ScalarValue::from(1i32)));
+        let mut a_stats = StatsSet::default();
+        a_stats.set(Stat::NullCount, Precision::exact(ScalarValue::from(1u64)));
+
+        let file_stats = FileStatistics::new_with_dtype(Arc::from([b_stats, a_stats]), &file_dtype);
+
+        let bytes = file_stats.write_flatbuffer_bytes()?;
+        let fb = flatbuffers::root::<fb::FileStatistics>(bytes.as_ref())
+            .vortex_expect("valid flatbuffer");
+        assert!(fb.is_nested());
+
+        let read_back = FileStatistics::from_flatbuffer(&fb, &file_dtype, &session)?;
+
+        let (b, _) = read_back
+            .get_by_path(&FieldPath::from_name("a").push("b"))
+            .expect("a.b stats");
+        assert_eq!(b.get(Stat::Min).as_exact(), Some(ScalarValue::from(1i32)));
+
+        let (a, _) = read_back
+            .get_by_path(&FieldPath::from_name("a"))
+            .expect("a's own null-count stats");
+        assert_eq!(
+            a.get(Stat::NullCount).as_exact(),
+            Some(ScalarValue::from(1u64))
+        );
+
+        assert!(read_back.get_by_path(&FieldPath::root()).is_none());
+
+        Ok(())
+    }
+
+    #[test]
+    fn legacy_non_nested_footer_still_parses() -> VortexResult<()> {
+        // Simulates a footer written before nested field stats existed: `is_nested` is absent
+        // (defaults to false), and `field_stats` holds one entry per top-level struct field.
+        let session = array_session();
+        let file_dtype = DType::struct_([("col", i32_dtype())], Nullability::NonNullable);
+
+        let mut stats = StatsSet::default();
+        stats.set(Stat::Min, Precision::exact(ScalarValue::from(7i32)));
+
+        let mut fbb = FlatBufferBuilder::new();
+        let array_stats = stats.write_flatbuffer(&mut fbb)?;
+        let field_stats = fbb.create_vector(&[array_stats]);
+        let root = fb::FileStatistics::create(
+            &mut fbb,
+            &fb::FileStatisticsArgs {
+                field_stats: Some(field_stats),
+                is_nested: false,
+            },
+        );
+        fbb.finish_minimal(root);
+        let bytes = fbb.finished_data().to_vec();
+
+        let fb = flatbuffers::root::<fb::FileStatistics>(&bytes).vortex_expect("valid flatbuffer");
+        assert!(!fb.is_nested());
+
+        let read_back = FileStatistics::from_flatbuffer(&fb, &file_dtype, &session)?;
+        let (col, _) = read_back
+            .get_by_path(&FieldPath::from_name("col"))
+            .expect("col stats");
+        assert_eq!(col.get(Stat::Min).as_exact(), Some(ScalarValue::from(7i32)));
+
+        Ok(())
+    }
+}

vortex-file/src/pruning.rs

@@ -9,7 +9,6 @@ use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::NullArray;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::FieldPath;
-use vortex_array::dtype::StructFields;
 use vortex_array::expr::BoundExpression;
 use vortex_array::expr::bound::lit;
 use vortex_array::expr::stats::Stat;
@@ -29,17 +28,13 @@ pub(crate) fn can_prune_file_stats(
     expr: &BoundExpression,
     row_count: u64,
     file_stats: &FileStatistics,
-    struct_fields: &StructFields,
     session: &VortexSession,
 ) -> VortexResult<bool> {
     let Some(pruning_expr) = expr.falsify(session)? else {
         return Ok(false);
     };
 
-    let binder = FileStatsBinder {
-        file_stats,
-        struct_fields,
-    };
+    let binder = FileStatsBinder { file_stats };
     let pruning_expr = bind_stats(pruning_expr, &binder)?;
 
     if let Some(result) = pruning_expr.as_opt::<Literal>() {
@@ -62,7 +57,6 @@ pub(crate) fn can_prune_file_stats(
 
 struct FileStatsBinder<'a> {
     file_stats: &'a FileStatistics,
-    struct_fields: &'a StructFields,
 }
 
 impl StatBinder for FileStatsBinder<'_> {
@@ -84,18 +78,10 @@ impl StatBinder for FileStatsBinder<'_> {
 
 impl FileStatsBinder<'_> {
     fn stat_ref(&self, field_path: &FieldPath, stat: Stat) -> Option<BoundExpression> {
-        // FileStats currently only holds top-level field statistics.
-        if field_path.parts().len() != 1 {
-            return None;
-        }
-
-        let field_name = field_path.parts()[0].as_name()?;
-        let field_idx = self.struct_fields.find(field_name)?;
-        let field_stats = self.file_stats.stats_sets().get(field_idx)?;
+        let (field_stats, field_dtype) = self.file_stats.get_by_path(field_path)?;
 
         let stat_value = field_stats.get(stat).as_exact()?;
-        let field_dtype = self.struct_fields.field_by_index(field_idx)?;
-        let stat_dtype = stat.dtype(&field_dtype)?;
+        let stat_dtype = stat.dtype(field_dtype)?;
         let stat_scalar = Scalar::try_new(stat_dtype, Some(stat_value)).ok()?;
 
         Some(lit(stat_scalar))