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 (child, child_type) in layout.children()?.into_iter().zip(layout.child_types()) {
60                let child_path = match child_type {
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((child, 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        crate::enable_all_registered_array_encodings(&session);
145        session
146    });
147
148    async fn write_nested() -> VortexResult<WriteSummary> {
149        let inner = StructArray::from_fields(&[
150            ("c", buffer![10i64, 20, 30, 40].into_array()),
151            (
152                "d",
153                StructArray::from_fields(&[("e", buffer![1.0f64, 2.0, 3.0, 4.0].into_array())])?
154                    .into_array(),
155            ),
156        ])?;
157        let root = StructArray::from_fields(&[
158            ("a", buffer![1u32, 2, 3, 4].into_array()),
159            ("b", inner.into_array()),
160        ])?;
161
162        let mut buf = ByteBufferMut::empty();
163        SESSION
164            .write_options()
165            .write(&mut buf, root.into_array().to_array_stream())
166            .await
167    }
168
169    #[tokio::test]
170    async fn nested_field_sizes() -> VortexResult<()> {
171        let summary = write_nested().await?;
172        let sizes = summary.footer().compressed_field_sizes()?;
173
174        for path in [
175            field_path!(a),
176            field_path!(b),
177            field_path!(b.c),
178            field_path!(b.d),
179            field_path!(b.d.e),
180        ] {
181            assert!(
182                sizes.get(&path).is_some_and(|size| size > 0),
183                "expected non-zero size for {path}, got {:?}",
184                sizes.get(&path)
185            );
186        }
187        assert_eq!(sizes.get(&field_path!(missing)), None);
188        assert_eq!(sizes.get(&field_path!(b.d.e.missing)), None);
189
190        // Attribution partitions segments: a field's size is the sum of its children plus its own
191        // directly-attributed segments (none here, as the structs are non-nullable).
192        let size = |path: FieldPath| sizes.get(&path).unwrap_or_default();
193        assert_eq!(
194            size(field_path!(b)),
195            size(field_path!(b.c)) + size(field_path!(b.d))
196        );
197        assert_eq!(size(field_path!(b.d)), size(field_path!(b.d.e)));
198        assert_eq!(sizes.total(), size(field_path!(a)) + size(field_path!(b)));
199
200        // Every segment in the file is attributed exactly once.
201        let segment_total: u64 = summary
202            .footer()
203            .segment_map()
204            .iter()
205            .map(|segment| u64::from(segment.length))
206            .sum();
207        assert_eq!(sizes.total(), segment_total);
208
209        Ok(())
210    }
211
212    #[tokio::test]
213    async fn column_sizes_in_schema_order() -> VortexResult<()> {
214        let summary = write_nested().await?;
215        let sizes = summary.footer().compressed_field_sizes()?;
216
217        assert_eq!(
218            summary.compressed_column_sizes()?,
219            vec![
220                sizes.get(&field_path!(a)).unwrap_or_default(),
221                sizes.get(&field_path!(b)).unwrap_or_default(),
222            ]
223        );
224        Ok(())
225    }
226
227    #[tokio::test]
228    async fn non_struct_file_reports_single_column() -> VortexResult<()> {
229        let mut buf = ByteBufferMut::empty();
230        let summary = SESSION
231            .write_options()
232            .write(
233                &mut buf,
234                buffer![1i32, 2, 3, 4].into_array().to_array_stream(),
235            )
236            .await?;
237
238        let sizes = summary.footer().compressed_field_sizes()?;
239        assert!(sizes.total() > 0);
240        assert_eq!(sizes.get(&FieldPath::root()), Some(sizes.total()));
241        assert_eq!(summary.compressed_column_sizes()?, vec![sizes.total()]);
242        Ok(())
243    }
244}
vortex_file/footer/field_sizes.rs

vortex_file/footer/
field_sizes.rs
