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