Move nested stats to a new top level field. Legacy stats still populated

by default, but added a new writer option to omit them.

Author: thorfour

vortex-datafusion/src/persistent/format.rs

@@ -47,6 +47,7 @@ use object_store::ObjectStore;
 use vortex::VortexSessionDefault;
 use vortex::array::memory::MemorySessionExt;
 use vortex::dtype::DType;
+use vortex::dtype::FieldPath;
 use vortex::dtype::Nullability;
 use vortex::dtype::PType;
 use vortex::error::VortexExpect;
@@ -598,7 +599,7 @@ impl FileFormat for VortexFormat {
                         })
                 });
 
-            let (dtype, file_stats, row_count) = match cached_metadata {
+            let (_dtype, file_stats, row_count) = match cached_metadata {
                 Some(metadata) => metadata,
                 None => {
                     // Not entry - open the file
@@ -635,10 +636,6 @@ impl FileFormat for VortexFormat {
                 }
             };
 
-            let struct_dtype = dtype
-                .as_struct_fields_opt()
-                .vortex_expect("dtype is not a struct");
-
             // Evaluate the statistics for each column that we are able to return to DataFusion.
             let Some(file_stats) = file_stats else {
                 // If the file has no column stats, the best we can do is return a row count.
@@ -661,12 +658,13 @@ impl FileFormat for VortexFormat {
             for field in table_schema.fields().iter() {
                 // If the column does not exist, continue. This can happen if the schema has evolved
                 // but we have not yet updated the Vortex file.
-                let Some(col_idx) = struct_dtype.find(field.name()) else {
+                let Some((stats_set, stats_dtype)) =
+                    file_stats.get_by_path(&FieldPath::from_name(field.name().as_str()))
+                else {
                     // The default sets all statistics to `Precision<Absent>`.
                     column_statistics.push(ColumnStatistics::default());
                     continue;
                 };
-                let (stats_set, stats_dtype) = file_stats.get(col_idx);
 
                 // Update the total size in bytes.
                 let column_size =

vortex-duckdb/src/copy.rs

@@ -16,6 +16,7 @@ use vortex::array::ArrayRef;
 use vortex::array::stream::ArrayStreamAdapter;
 use vortex::dtype::DType;
 use vortex::dtype::FieldName;
+use vortex::dtype::FieldPath;
 use vortex::dtype::Nullability::NonNullable;
 use vortex::dtype::Nullability::Nullable;
 use vortex::dtype::StructFields;
@@ -211,8 +212,9 @@ pub(crate) fn written_column_stats(
 fn file_stats_from_summary(summary: &WriteSummary) -> WrittenFileStats {
     let num_columns = summary
         .footer()
-        .statistics()
-        .map_or(0, |s| s.stats_sets().len());
+        .dtype()
+        .as_struct_fields_opt()
+        .map_or(1, StructFields::nfields);
     WrittenFileStats {
         row_count: summary.row_count(),
         file_size_bytes: summary.size(),
@@ -223,11 +225,11 @@ fn file_stats_from_summary(summary: &WriteSummary) -> WrittenFileStats {
 }
 
 /// Per-column statistics from a finished write's summary and its precomputed compressed sizes
-/// (`column_sizes`, indexed the same as the footer's stats sets).
+/// (`column_sizes`, indexed by top-level field position).
 ///
-/// Only top-level columns are covered: the footer exposes one statistics set per top-level field,
-/// so nested struct/list leaf columns are not reported (parquet, by contrast, recurses to leaf
-/// paths). Flat tables - the common DuckLake case - are fully covered.
+/// Only top-level columns are covered: nested struct/list leaf columns are not reported (parquet,
+/// by contrast, recurses to leaf paths). Flat tables - the common DuckLake case - are fully
+/// covered.
 fn column_stats_from_summary(
     summary: &WriteSummary,
     column_index: usize,
@@ -237,15 +239,20 @@ fn column_stats_from_summary(
         .footer()
         .statistics()
         .ok_or_else(|| vortex_err!("written file has no statistics"))?;
-    let stats_sets = file_stats.stats_sets();
-    if column_index >= stats_sets.len() {
-        vortex_bail!(
-            "column index {column_index} out of range for {} statistics sets",
-            stats_sets.len()
-        );
-    }
-    let stats = &stats_sets[column_index];
-    let dtype = &file_stats.dtypes()[column_index];
+    let struct_fields = summary
+        .footer()
+        .dtype()
+        .as_struct_fields_opt()
+        .ok_or_else(|| vortex_err!("written file dtype is not a struct"))?;
+    let name = struct_fields.names().get(column_index).ok_or_else(|| {
+        vortex_err!(
+            "column index {column_index} out of range for {} top-level fields",
+            struct_fields.nfields()
+        )
+    })?;
+    let (stats, dtype) = file_stats
+        .get_by_path(&FieldPath::from_name(name.clone()))
+        .ok_or_else(|| vortex_err!("no statistics for column {name}"))?;
 
     Ok(WrittenColumnStats {
         min: exact_scalar_to_duckdb(stats.get(Stat::Min), dtype)?,
@@ -333,6 +340,7 @@ mod tests {
     use vortex::array::IntoArray;
     use vortex::array::arrays::StructArray;
     use vortex::array::stats::PRUNING_STATS;
+    use vortex::array::validity::Validity;
     use vortex::buffer::ByteBufferMut;
     use vortex::buffer::buffer;
 
@@ -367,4 +375,32 @@ mod tests {
         let summary = write_summary(vec![]);
         assert!(column_stats_from_summary(&summary, 0, &[]).is_err());
     }
+
+    #[test]
+    fn column_stats_resolves_by_name_past_a_nullable_struct_column() {
+        // Regression test: a nullable top-level struct field ("s") contributes both a `s.b` entry
+        // and a trailing entry for its own null count to the nested (post-order) stats layout, so
+        // that layout has one more entry than there are top-level fields. Resolving a later
+        // top-level column ("c", index 1) by flat position into that layout would silently return
+        // `s`'s own null-count-only entry (which has no Min) instead of `c`'s stats.
+        let summary = RUNTIME.block_on(async {
+            let b = buffer![10i32, 20, 30].into_array();
+            let inner = StructArray::new(["b"].into(), [b], 3, Validity::AllValid).into_array();
+            let c = buffer![7u32, 8, 9].into_array();
+            let outer = StructArray::new(["s", "c"].into(), [inner, c], 3, Validity::NonNullable)
+                .into_array();
+
+            let mut buf = ByteBufferMut::empty();
+            let mut writer = SESSION
+                .write_options()
+                .with_file_statistics(PRUNING_STATS.to_vec())
+                .writer(&mut buf, outer.dtype().clone());
+            writer.push(outer).await.unwrap();
+            writer.finish().await.unwrap()
+        });
+
+        let stats = column_stats_from_summary(&summary, 1, &[]).unwrap();
+        assert!(stats.min.is_some());
+        assert!(stats.max.is_some());
+    }
 }

vortex-duckdb/src/file_reader.rs

@@ -11,6 +11,7 @@ use url::Url;
 use vortex::array::VortexSessionExecute as _;
 use vortex::array::arrays::struct_::StructArrayExt as _;
 use vortex::dtype::DType;
+use vortex::dtype::FieldPath;
 use vortex::error::VortexExpect;
 use vortex::error::VortexResult;
 use vortex::error::vortex_panic;
@@ -274,16 +275,12 @@ pub fn reader_get_statistics(
         .as_any()
         .downcast_ref::<FileStatsLayoutReader>()?;
 
-    let DType::Struct(fields, _) = &file.reader.dtype() else {
-        return None;
-    };
-    let index = fields.find(column)?;
-    let stats_sets = reader.file_stats().stats_sets();
+    let (stats, dtype) = reader
+        .file_stats()
+        .get_by_path(&FieldPath::from_name(column))?;
 
-    let dtype = fields.field_by_index(index)?;
-
-    let stats = ColumnStatisticsAggregate::new(stats_sets.get(index)?);
-    match ColumnStatistics::try_from(stats, dtype) {
+    let stats = ColumnStatisticsAggregate::new(stats);
+    match ColumnStatistics::try_from(stats, dtype.clone()) {
         Ok(stats) => Some(stats),
         Err(e) => vortex_panic!(e),
     }
@@ -328,12 +325,56 @@ pub fn footer_get_statistics(footer: &Footer, index: usize) -> Option<ColumnStat
     let DType::Struct(fields, _) = footer.dtype() else {
         return None;
     };
+    let name = fields.names().get(index)?;
     let stats = footer.statistics()?;
-    let dtype = fields.field_by_index(index)?;
-    let stats = stats.stats_sets().get(index)?;
+    let (stats, dtype) = stats.get_by_path(&FieldPath::from_name(name.clone()))?;
     let stats = ColumnStatisticsAggregate::new(stats);
-    match ColumnStatistics::try_from(stats, dtype) {
+    match ColumnStatistics::try_from(stats, dtype.clone()) {
         Ok(stats) => Some(stats),
         Err(e) => vortex_panic!(e),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use vortex::array::IntoArray;
+    use vortex::array::arrays::StructArray;
+    use vortex::array::stats::PRUNING_STATS;
+    use vortex::array::validity::Validity;
+    use vortex::buffer::ByteBufferMut;
+    use vortex::buffer::buffer;
+    use vortex::file::WriteOptionsSessionExt;
+
+    use super::*;
+    use crate::RUNTIME;
+    use crate::SESSION;
+
+    #[test]
+    fn footer_get_statistics_resolves_by_name_past_a_nullable_struct_field() {
+        // Regression test: a nullable top-level struct field ("s") contributes both a `s.b` entry
+        // and a trailing entry for its own null count to the nested (post-order) stats layout, so
+        // that layout has one more entry than there are top-level fields. Resolving a later
+        // top-level column ("c", index 1) by flat position into that layout would silently return
+        // `s`'s own null-count-only entry (which has no Min) instead of `c`'s stats.
+        let footer = RUNTIME.block_on(async {
+            let b = buffer![10i32, 20, 30].into_array();
+            let inner = StructArray::new(["b"].into(), [b], 3, Validity::AllValid).into_array();
+            let c = buffer![7u32, 8, 9].into_array();
+            let outer = StructArray::new(["s", "c"].into(), [inner, c], 3, Validity::NonNullable)
+                .into_array();
+
+            let mut buf = ByteBufferMut::empty();
+            let mut writer = SESSION
+                .write_options()
+                .with_file_statistics(PRUNING_STATS.to_vec())
+                .writer(&mut buf, outer.dtype().clone());
+            writer.push(outer).await.unwrap();
+            writer.finish().await.unwrap().footer().clone()
+        });
+
+        let stats =
+            footer_get_statistics(&footer, 1).expect("stats for top-level column 1 (\"c\")");
+        assert!(stats.min.is_some());
+        assert!(stats.max.is_some());
+    }
+}

vortex-file/flatbuffers/vortex-file/footer.fbs

@@ -56,17 +56,22 @@ table PostscriptSegment {
 // [footer]
 /// The `FileStatistics` object contains file-level statistics for the Vortex file.
 table FileStatistics {
-    /// Statistics for each field in the root schema. If the root schema is not a struct, there will
-    /// be a single entry in this array.
+    /// Statistics for each top-level field in the root schema, in schema order. If the root schema
+    /// is not a struct, there will be a single entry in this array. Does not recurse into nested
+    /// struct fields; see `nested_field_stats` for that.
     ///
-    /// When `is_nested` is true, entries follow a post-order walk of the root `DType` tree: each
-    /// leaf field gets an entry, and each nullable struct additionally gets a trailing null-count
-    /// entry inserted after its children's entries.
+    /// Kept in this legacy shape for backward compatibility with readers that predate
+    /// `nested_field_stats`, which strictly validate this array's length against the number of
+    /// top-level struct fields.
     field_stats: [ArrayStats];
-    /// Whether `field_stats` follows the post-order nested-struct layout (true) or the legacy
-    /// top-level-fields-only layout (false). Defaults to false for backward compatibility with
-    /// files written before nested field stats were supported.
-    is_nested: bool = false;
+    /// Statistics following a post-order walk of the root `DType` tree: each leaf field gets an
+    /// entry, and each nullable struct additionally gets a trailing null-count entry inserted
+    /// after its children's entries.
+    ///
+    /// Readers that know about this field should prefer it over `field_stats` when present. It is
+    /// absent for files written before nested file stats were supported, and readers that don't
+    /// know about it fall back to `field_stats`'s legacy top-level-fields-only layout.
+    nested_field_stats: [ArrayStats];
 }
 
 /// The `Registry` object stores dictionary-encoded configuration for segments,