vortex_dtype/
field.rs

1//! Selectors for fields in (possibly nested) `StructDType`s
2//!
3//! A `Field` can either be a direct child field of the top-level struct (selected by name or index),
4//! or a nested field (selected by a sequence of such selectors)
5
6use core::fmt;
7use std::fmt::{Display, Formatter};
8use std::sync::Arc;
9
10use itertools::Itertools;
11
12/// A selector for a field in a struct
13#[derive(Clone, Debug, PartialEq, Eq, Hash)]
14#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
15pub enum Field {
16    /// Address a field of a [`crate::DType::Struct`].
17    Name(Arc<str>),
18    /// Address the element type of a [`crate::DType::List`].
19    ElementType,
20}
21
22impl From<&str> for Field {
23    fn from(value: &str) -> Self {
24        Field::Name(value.into())
25    }
26}
27
28impl From<Arc<str>> for Field {
29    fn from(value: Arc<str>) -> Self {
30        Self::Name(value)
31    }
32}
33
34impl From<&Arc<str>> for Field {
35    fn from(value: &Arc<str>) -> Self {
36        Self::Name(value.clone())
37    }
38}
39
40impl From<String> for Field {
41    fn from(value: String) -> Self {
42        Field::Name(value.into())
43    }
44}
45
46impl Display for Field {
47    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
48        match self {
49            Field::Name(name) => write!(f, "${name}"),
50            Field::ElementType => write!(f, "[]"),
51        }
52    }
53}
54
55impl Field {
56    /// Returns true if the field is defined by Name
57    pub fn is_named(&self) -> bool {
58        matches!(self, Field::Name(_))
59    }
60}
61
62/// A path through a (possibly nested) struct, composed of a sequence of field selectors
63// TODO(ngates): we should probably reverse the path. Or better yet, store a Arc<[Field]> along
64//  with a positional index to allow cheap step_into.
65#[derive(Clone, Debug, PartialEq, Eq, Hash)]
66#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
67pub struct FieldPath(Vec<Field>);
68
69impl FieldPath {
70    /// The selector for the root (i.e., the top-level struct itself)
71    pub fn root() -> Self {
72        Self(vec![])
73    }
74
75    /// Constructs a new `FieldPath` from a single field selector (i.e., a direct child field of the top-level struct)
76    pub fn from_name<F: Into<Field>>(name: F) -> Self {
77        Self(vec![name.into()])
78    }
79
80    /// Returns the sequence of field selectors that make up this path
81    pub fn path(&self) -> &[Field] {
82        &self.0
83    }
84
85    /// Returns whether this path is a root path.
86    pub fn is_root(&self) -> bool {
87        self.0.is_empty()
88    }
89
90    /// Pushes a new field selector to the end of this path
91    pub fn push<F: Into<Field>>(mut self, field: F) -> Self {
92        self.0.push(field.into());
93        self
94    }
95
96    /// Whether the path starts with the given field name
97    /// TODO(joe): handle asserts better.
98    pub fn starts_with_field(&self, field: &Field) -> bool {
99        assert!(matches!(field, Field::Name(_)));
100        let first = self.0.first();
101        assert!(matches!(first, Some(Field::Name(_))));
102        first.is_some_and(|f| f == field)
103    }
104
105    /// Steps into the next field in the path
106    pub fn step_into(mut self) -> Option<Self> {
107        if self.0.is_empty() {
108            return None;
109        }
110        self.0 = self.0.iter().skip(1).cloned().collect();
111        Some(self)
112    }
113}
114
115impl FromIterator<Field> for FieldPath {
116    fn from_iter<T: IntoIterator<Item = Field>>(iter: T) -> Self {
117        FieldPath(iter.into_iter().collect())
118    }
119}
120
121impl From<Field> for FieldPath {
122    fn from(value: Field) -> Self {
123        FieldPath(vec![value])
124    }
125}
126
127impl From<Vec<Field>> for FieldPath {
128    fn from(value: Vec<Field>) -> Self {
129        FieldPath(value)
130    }
131}
132
133impl Display for FieldPath {
134    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
135        Display::fmt(&self.0.iter().format("."), f)
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    #[test]
144    fn test_field_path() {
145        let path = FieldPath::from_name("A").push("B").push("C");
146        assert_eq!(path.to_string(), "$A.$B.$C");
147
148        let fields = vec!["A", "B", "C"]
149            .into_iter()
150            .map(Field::from)
151            .collect_vec();
152        assert_eq!(path.path(), &fields);
153
154        let vec_path = FieldPath::from(fields);
155        assert_eq!(vec_path.to_string(), "$A.$B.$C");
156        assert_eq!(path, vec_path);
157    }
158}