vortex_file/footer/
field_sizes.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::dtype::Field;
5use vortex_array::dtype::FieldPath;
6use vortex_error::VortexResult;
7use vortex_error::vortex_err;
8use vortex_layout::LayoutChildType;
9use vortex_layout::LayoutRef;
10use vortex_layout::segments::SegmentId;
11use vortex_utils::aliases::hash_map::Entry;
12use vortex_utils::aliases::hash_map::HashMap;
13
14use crate::footer::SegmentSpec;
15
16/// Compressed sizes of a file's fields, keyed by field path.
17///
18/// Sizes are derived from the file's layout tree and segment map: walking the tree from the root,
19/// each physical segment is attributed to the deepest [`FieldPath`] whose layout subtree
20/// references it. A segment referenced by more than one field is attributed to their longest
21/// common ancestor, so a segment is counted exactly once and the size of a field always equals
22/// the sum of its child field sizes plus the bytes attributed directly to it (e.g. struct
23/// validity, or auxiliary segments such as zone maps and dictionaries written above the field
24/// split).
25///
26/// Fields nested inside non-struct types (e.g. structs within lists) are not split into separate
27/// layout subtrees by the writer, so their bytes are attributed to the enclosing field.
28///
29/// File-level metadata, alignment padding, and the footer are not part of any segment and are
30/// never attributed.
31#[derive(Debug, Clone)]
32pub struct CompressedFieldSizes {
33    sizes: HashMap<FieldPath, u64>,
34}
35
36impl CompressedFieldSizes {
37    pub(crate) fn try_new(root: &LayoutRef, segments: &[SegmentSpec]) -> VortexResult<Self> {
38        // Every segment is attributed to the deepest field path referencing it, collapsing to the
39        // longest common ancestor if multiple fields reference the same segment.
40        let mut attribution = HashMap::<SegmentId, FieldPath>::default();
41        // Seed every field path present in the layout tree so empty fields report a zero size.
42        let mut sizes = HashMap::<FieldPath, u64>::default();
43        sizes.insert(FieldPath::root(), 0);
44
45        let mut stack = vec![(LayoutRef::clone(root), FieldPath::root())];
46        while let Some((layout, path)) = stack.pop() {
47            for segment_id in layout.segment_ids() {
48                match attribution.entry(segment_id) {
49                    Entry::Occupied(mut entry) => {
50                        let ancestor = common_prefix(entry.get(), &path);
51                        entry.insert(ancestor);
52                    }
53                    Entry::Vacant(entry) => {
54                        entry.insert(path.clone());
55                    }
56                }
57            }
58
59            for idx in 0..layout.nchildren() {
60                let child_path = match layout.child_type(idx) {
61                    LayoutChildType::Field(name) => {
62                        let child_path = path.clone().push(name);
63                        sizes.entry(child_path.clone()).or_insert(0);
64                        child_path
65                    }
66                    _ => path.clone(),
67                };
68                stack.push((layout.child(idx)?, child_path));
69            }
70        }
71
72        for (segment_id, path) in attribution {
73            let segment = segments.get(*segment_id as usize).ok_or_else(|| {
74                vortex_err!(
75                    "layout references missing segment {} (segment count: {})",
76                    *segment_id,
77                    segments.len()
78                )
79            })?;
80            // A segment counts towards its attributed field and every enclosing field.
81            for len in 0..=path.parts().len() {
82                *sizes
83                    .entry(FieldPath::from(path.parts()[..len].to_vec()))
84                    .or_insert(0) += u64::from(segment.length);
85            }
86        }
87
88        Ok(Self { sizes })
89    }
90
91    /// Returns the compressed size in bytes of the field at the given path, including all of its
92    /// descendants, or `None` if the layout tree contains no such field.
93    ///
94    /// [`FieldPath::root`] returns the same value as [`total`][Self::total].
95    pub fn get(&self, path: &FieldPath) -> Option<u64> {
96        self.sizes.get(path).copied()
97    }
98
99    /// Returns the total compressed size in bytes of all segments referenced by the layout tree.
100    pub fn total(&self) -> u64 {
101        self.get(&FieldPath::root()).unwrap_or_default()
102    }
103
104    /// Iterates over all field paths in the layout tree and their compressed sizes.
105    pub fn iter(&self) -> impl Iterator<Item = (&FieldPath, u64)> {
106        self.sizes.iter().map(|(path, size)| (path, *size))
107    }
108}
109
110fn common_prefix(left: &FieldPath, right: &FieldPath) -> FieldPath {
111    left.parts()
112        .iter()
113        .zip(right.parts())
114        .take_while(|(l, r)| l == r)
115        .map(|(l, _)| l.clone())
116        .collect::<Vec<Field>>()
117        .into()
118}
119
120#[cfg(test)]
121mod tests {
122    use std::sync::LazyLock;
123
124    use vortex_array::IntoArray;
125    use vortex_array::array_session;
126    use vortex_array::arrays::StructArray;
127    use vortex_array::dtype::FieldPath;
128    use vortex_array::field_path;
129    use vortex_buffer::ByteBufferMut;
130    use vortex_buffer::buffer;
131    use vortex_error::VortexResult;
132    use vortex_io::session::RuntimeSession;
133    use vortex_layout::session::LayoutSession;
134    use vortex_session::VortexSession;
135
136    use crate::WriteOptionsSessionExt;
137    use crate::writer::WriteSummary;
138
139    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
140        let session = array_session()
141            .with::<LayoutSession>()
142            .with::<RuntimeSession>();
143        crate::register_default_encodings(&session);
144        session
145    });
146
147    async fn write_nested() -> VortexResult<WriteSummary> {
148        let inner = StructArray::from_fields(&[
149            ("c", buffer![10i64, 20, 30, 40].into_array()),
150            (
151                "d",
152                StructArray::from_fields(&[("e", buffer![1.0f64, 2.0, 3.0, 4.0].into_array())])?
153                    .into_array(),
154            ),
155        ])?;
156        let root = StructArray::from_fields(&[
157            ("a", buffer![1u32, 2, 3, 4].into_array()),
158            ("b", inner.into_array()),
159        ])?;
160
161        let mut buf = ByteBufferMut::empty();
162        SESSION
163            .write_options()
164            .write(&mut buf, root.into_array().to_array_stream())
165            .await
166    }
167
168    #[tokio::test]
169    async fn nested_field_sizes() -> VortexResult<()> {
170        let summary = write_nested().await?;
171        let sizes = summary.footer().compressed_field_sizes()?;
172
173        for path in [
174            field_path!(a),
175            field_path!(b),
176            field_path!(b.c),
177            field_path!(b.d),
178            field_path!(b.d.e),
179        ] {
180            assert!(
181                sizes.get(&path).is_some_and(|size| size > 0),
182                "expected non-zero size for {path}, got {:?}",
183                sizes.get(&path)
184            );
185        }
186        assert_eq!(sizes.get(&field_path!(missing)), None);
187        assert_eq!(sizes.get(&field_path!(b.d.e.missing)), None);
188
189        // Attribution partitions segments: a field's size is the sum of its children plus its own
190        // directly-attributed segments (none here, as the structs are non-nullable).
191        let size = |path: FieldPath| sizes.get(&path).unwrap_or_default();
192        assert_eq!(
193            size(field_path!(b)),
194            size(field_path!(b.c)) + size(field_path!(b.d))
195        );
196        assert_eq!(size(field_path!(b.d)), size(field_path!(b.d.e)));
197        assert_eq!(sizes.total(), size(field_path!(a)) + size(field_path!(b)));
198
199        // Every segment in the file is attributed exactly once.
200        let segment_total: u64 = summary
201            .footer()
202            .segment_map()
203            .iter()
204            .map(|segment| u64::from(segment.length))
205            .sum();
206        assert_eq!(sizes.total(), segment_total);
207
208        Ok(())
209    }
210
211    #[tokio::test]
212    async fn column_sizes_in_schema_order() -> VortexResult<()> {
213        let summary = write_nested().await?;
214        let sizes = summary.footer().compressed_field_sizes()?;
215
216        assert_eq!(
217            summary.compressed_column_sizes()?,
218            vec![
219                sizes.get(&field_path!(a)).unwrap_or_default(),
220                sizes.get(&field_path!(b)).unwrap_or_default(),
221            ]
222        );
223        Ok(())
224    }
225
226    #[tokio::test]
227    async fn non_struct_file_reports_single_column() -> VortexResult<()> {
228        let mut buf = ByteBufferMut::empty();
229        let summary = SESSION
230            .write_options()
231            .write(
232                &mut buf,
233                buffer![1i32, 2, 3, 4].into_array().to_array_stream(),
234            )
235            .await?;
236
237        let sizes = summary.footer().compressed_field_sizes()?;
238        assert!(sizes.total() > 0);
239        assert_eq!(sizes.get(&FieldPath::root()), Some(sizes.total()));
240        assert_eq!(summary.compressed_column_sizes()?, vec![sizes.total()]);
241        Ok(())
242    }
243}
vortex_file/footer/field_sizes.rs

vortex_file/footer/
field_sizes.rs
