vortex_tui/
segment_tree.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Shared segment tree collection logic used by both the TUI browse view and the CLI segments command.
5
6use std::sync::Arc;
7
8use vortex::dtype::FieldName;
9use vortex::error::VortexResult;
10use vortex::file::SegmentSpec;
11use vortex::layout::Layout;
12use vortex::layout::LayoutChildType;
13use vortex::utils::aliases::hash_map::HashMap;
14
15/// Information about a single segment for display purposes.
16pub struct SegmentDisplay {
17    /// Name of the segment (e.g., "data", `[0]`, "zones")
18    pub name: FieldName,
19    /// The underlying segment specification
20    pub spec: SegmentSpec,
21    /// Row offset within the file
22    pub row_offset: u64,
23    /// Number of rows in this segment
24    pub row_count: u64,
25}
26
27/// A tree of segments organized by field name.
28pub struct SegmentTree {
29    /// Map from field name to list of segments for that field
30    pub segments: HashMap<FieldName, Vec<SegmentDisplay>>,
31    /// Ordered list of field names (columns) in display order
32    pub segment_ordering: Vec<FieldName>,
33}
34
35/// Collect segment tree from a layout and segment map.
36pub fn collect_segment_tree(
37    root_layout: &dyn Layout,
38    segments: &Arc<[SegmentSpec]>,
39) -> SegmentTree {
40    let mut tree = SegmentTree {
41        segments: HashMap::new(),
42        segment_ordering: Vec::new(),
43    };
44    // Ignore errors during traversal - we want to collect as much as possible
45    drop(segments_by_name_impl(
46        root_layout,
47        None,
48        None,
49        Some(0),
50        segments,
51        &mut tree,
52    ));
53    tree
54}
55
56fn segments_by_name_impl(
57    root: &dyn Layout,
58    group_name: Option<FieldName>,
59    name: Option<FieldName>,
60    row_offset: Option<u64>,
61    segments: &Arc<[SegmentSpec]>,
62    segment_tree: &mut SegmentTree,
63) -> VortexResult<()> {
64    // Recurse into children
65    for (child, child_type) in root.children()?.into_iter().zip(root.child_types()) {
66        match child_type {
67            LayoutChildType::Transparent(sub_name) => segments_by_name_impl(
68                child.as_ref(),
69                group_name.clone(),
70                Some(
71                    name.as_ref()
72                        .map(|n| format!("{n}.{sub_name}").into())
73                        .unwrap_or_else(|| sub_name.into()),
74                ),
75                row_offset,
76                segments,
77                segment_tree,
78            )?,
79            LayoutChildType::Auxiliary(aux_name) => segments_by_name_impl(
80                child.as_ref(),
81                group_name.clone(),
82                Some(
83                    name.as_ref()
84                        .map(|n| format!("{n}.{aux_name}").into())
85                        .unwrap_or_else(|| aux_name.into()),
86                ),
87                Some(0),
88                segments,
89                segment_tree,
90            )?,
91            LayoutChildType::Chunk((idx, chunk_row_offset)) => segments_by_name_impl(
92                child.as_ref(),
93                group_name.clone(),
94                Some(
95                    name.as_ref()
96                        .map(|n| format!("{n}.[{idx}]"))
97                        .unwrap_or_else(|| format!("[{idx}]"))
98                        .into(),
99                ),
100                // Compute absolute row offset.
101                Some(chunk_row_offset + row_offset.unwrap_or(0)),
102                segments,
103                segment_tree,
104            )?,
105            LayoutChildType::Field(field_name) => {
106                // Step into a new group name
107                let new_group_name = group_name
108                    .as_ref()
109                    .map(|n| format!("{n}.{field_name}").into())
110                    .unwrap_or_else(|| field_name);
111                segment_tree.segment_ordering.push(new_group_name.clone());
112
113                segments_by_name_impl(
114                    child.as_ref(),
115                    Some(new_group_name),
116                    None,
117                    row_offset,
118                    segments,
119                    segment_tree,
120                )?
121            }
122        }
123    }
124
125    let current_segments = segment_tree
126        .segments
127        .entry(group_name.unwrap_or_else(|| FieldName::from("root")))
128        .or_default();
129
130    for segment_id in root.segment_ids() {
131        let segment_spec = segments[*segment_id as usize];
132        current_segments.push(SegmentDisplay {
133            name: name.clone().unwrap_or_else(|| "<unnamed>".into()),
134            spec: segment_spec,
135            row_count: root.row_count(),
136            row_offset: row_offset.unwrap_or(0),
137        })
138    }
139
140    Ok(())
141}
vortex_tui/segment_tree.rs

vortex_tui/
segment_tree.rs
