16935267080

Committed 13 Aug 2025 11:00AM UTC coverage: 24.312% (-63.3%) from 87.658%

Build # 16935267080

Build Type

Pull #4226

github

Committed by

web-flow

Commit Message

Merge 81b48c7fb into baa6ea202

Pull Request Pull Request #4226: Support converting TimestampTZ to and from duckdb

Run Details

0 of 2 new or added lines in 1 file covered. (0.0%)

20666 existing lines in 469 files now uncovered.

8726 of 35892 relevant lines covered (24.31%)

147.74 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

34.67

/vortex-dtype/src/struct_.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use std::hash::Hash;
use std::sync::Arc;

use itertools::Itertools;
use vortex_error::{
    VortexExpect, VortexResult, VortexUnwrap, vortex_bail, vortex_err, vortex_panic,
};

use crate::flatbuffers::ViewedDType;
use crate::{DType, FieldName, FieldNames, PType};

/// DType of a struct's field, either owned or a pointer to an underlying flatbuffer.
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct FieldDType {
    inner: FieldDTypeInner,
}

impl From<ViewedDType> for FieldDType {
    fn from(value: ViewedDType) -> Self {
        Self {
            inner: FieldDTypeInner::View(value),
        }
    }
}

impl From<DType> for FieldDType {
    fn from(value: DType) -> Self {
        Self {
            inner: FieldDTypeInner::Owned(value),
        }
    }
}

impl From<PType> for FieldDType {
    fn from(value: PType) -> Self {
        Self {
            inner: FieldDTypeInner::Owned(DType::from(value)),
        }
    }
}

#[derive(Debug, Clone, Eq)]
enum FieldDTypeInner {
    /// Owned DType instance
    // TODO(ngates): we should consider making this an Arc<DType>.
    Owned(DType),
    /// A view over a flatbuffer, parsed only when accessed.
    View(ViewedDType),
}

impl PartialEq for FieldDTypeInner {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Self::Owned(lhs), Self::Owned(rhs)) => lhs == rhs,
            (Self::View(lhs), Self::View(rhs)) => {
                let lhs = DType::try_from(lhs.clone())
                    .vortex_expect("Failed to parse FieldDType into DType");
                let rhs = DType::try_from(rhs.clone())
                    .vortex_expect("Failed to parse FieldDType into DType");

                lhs == rhs
            }
            (Self::View(view), Self::Owned(owned)) | (Self::Owned(owned), Self::View(view)) => {
                let view = DType::try_from(view.clone())
                    .vortex_expect("Failed to parse FieldDType into DType");
                owned == &view
            }
        }
    }
}

impl Hash for FieldDTypeInner {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        match self {
            FieldDTypeInner::Owned(owned) => {
                owned.hash(state);
            }
            FieldDTypeInner::View(view) => {
                let owned = DType::try_from(view.clone())
                    .vortex_expect("Failed to parse FieldDType into DType");
                owned.hash(state);
            }
        }
    }
}

impl FieldDType {
    /// Returns the concrete DType, parsing it from the underlying buffer if necessary.
    pub fn value(&self) -> VortexResult<DType> {
        self.inner.value()
    }
}

impl FieldDTypeInner {
    fn value(&self) -> VortexResult<DType> {
        match &self {
            FieldDTypeInner::Owned(owned) => Ok(owned.clone()),
            FieldDTypeInner::View(view) => DType::try_from(view.clone()),
        }
    }
}

#[cfg(feature = "serde")]
impl serde::Serialize for FieldDTypeInner {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        use serde::ser::Error;

        let value = self.value().map_err(S::Error::custom)?;
        serializer.serialize_newtype_variant("FieldDType", 0, "Owned", &value)
    }
}

#[cfg(feature = "serde")]
impl<'de> serde::Deserialize<'de> for FieldDTypeInner {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        deserializer.deserialize_enum("FieldDType", &["Owned", "View"], FieldDTypeDeVisitor)
    }
}

#[cfg(feature = "serde")]
struct FieldDTypeDeVisitor;

#[cfg(feature = "serde")]
impl<'de> serde::de::Visitor<'de> for FieldDTypeDeVisitor {
    type Value = FieldDTypeInner;

    fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "variant identifier")
    }

    fn visit_enum<A>(self, data: A) -> Result<Self::Value, A::Error>
    where
        A: serde::de::EnumAccess<'de>,
    {
        use serde::de::{Error, VariantAccess};

        #[derive(serde::Deserialize, Debug)]
        enum FieldDTypeVariant {
            Owned,
            View,
        }
        let (variant, variant_data): (FieldDTypeVariant, _) = data.variant()?;

        match variant {
            FieldDTypeVariant::Owned => {
                let inner = variant_data.newtype_variant::<DType>()?;
                Ok(FieldDTypeInner::Owned(inner))
            }
            other => Err(A::Error::custom(format!("unsupported variant {other:?}"))),
        }
    }
}

/// Type information for a struct column.
///
/// The `StructFields` holds all field names and field types, and provides
/// access to them by index or by name.
///
/// ## Duplicate field names
///
/// In memory, it is not an error for a `StructFields` to contain duplicate
/// field names. In that case, any name-based access to fields will resolve
/// to the first such field with a given name.
///
/// ```rust
/// # use vortex_dtype::{DType, Nullability, PType, StructFields};
///
/// let fields = StructFields::from_iter([
///     ("string_col", DType::Utf8(Nullability::NonNullable)),
///     ("binary_col", DType::Binary(Nullability::NonNullable)),
///     ("int_col", DType::Primitive(PType::I32, Nullability::Nullable)),
///     ("int_col", DType::Primitive(PType::I64, Nullability::Nullable)),
/// ]);
///
/// // Accessing a field by name will yield the first
/// assert_eq!(fields.field("int_col").unwrap(), DType::Primitive(PType::I32, Nullability::Nullable));
/// ```
#[derive(Clone, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct StructFields(Arc<StructFieldsInner>);

impl std::fmt::Debug for StructFields {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StructFields")
            .field("names", &self.0.names)
            .field("dtypes", &self.0.dtypes)
            .finish()
    }
}

#[derive(PartialEq, Eq, Hash, Default)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
struct StructFieldsInner {
    names: FieldNames,
    dtypes: Arc<[FieldDType]>,
}

impl Default for StructFields {
    fn default() -> Self {
        Self::empty()
    }
}

impl StructFields {
    /// The fields of the empty struct.
    pub fn empty() -> Self {
        Self(Arc::new(StructFieldsInner {
            names: FieldNames::default(),
            dtypes: Arc::from([]),
        }))
    }

    /// Create a new [`StructFields`] from a list of names and dtypes
    pub fn new(names: FieldNames, dtypes: Vec<DType>) -> Self {
        if names.len() != dtypes.len() {
            vortex_panic!(
                "length mismatch between names ({}) and dtypes ({})",
                names.len(),
                dtypes.len()
            );
        }

        let dtypes = dtypes
            .into_iter()
            .map(|dt| FieldDType {
                inner: FieldDTypeInner::Owned(dt),
            })
            .collect::<Vec<_>>();

        Self::from_fields(names, dtypes)
    }

    /// Create a new [`StructFields`] from a  list of names and [`FieldDType`] which can be either lazily or eagerly serialized.
    pub fn from_fields(names: FieldNames, dtypes: Vec<FieldDType>) -> Self {
        if names.len() != dtypes.len() {
            vortex_panic!(
                "length mismatch between names ({}) and dtypes ({})",
                names.len(),
                dtypes.len()
            );
        }

        let inner = Arc::new(StructFieldsInner {
            names,
            dtypes: dtypes.into(),
        });

        Self(inner)
    }

    /// Get the names of the fields in the struct
    pub fn names(&self) -> &FieldNames {
        &self.0.names
    }

    /// Returns the number of fields in the struct
    pub fn nfields(&self) -> usize {
        self.0.names.len()
    }

    /// Returns the name of the field at the given index
    pub fn field_name(&self, index: usize) -> Option<&FieldName> {
        self.0.names.get(index)
    }

    /// Find the index of a field by name
    /// Returns `None` if the field is not found
    pub fn find(&self, name: impl AsRef<str>) -> Option<usize> {
        let name = name.as_ref();
        self.0.names.iter().position(|n| n.as_ref() == name)
    }

    /// Get the [`DType`] of a field.
    ///
    /// It is possible for there to be more than one field with
    /// the same name, in which case, this will return the DType
    /// of the first field encountered with a given name.
    pub fn field(&self, name: impl AsRef<str>) -> Option<DType> {
        let index = self.find(name)?;
        Some(self.0.dtypes[index].value().vortex_unwrap())
    }

    /// Get the [`DType`] of a field by index.
    pub fn field_by_index(&self, index: usize) -> Option<DType> {
        Some(self.0.dtypes.get(index)?.value().vortex_unwrap())
    }

    /// Returns an ordered iterator over the fields.
    pub fn fields(&self) -> impl ExactSizeIterator<Item = DType> + '_ {
        self.0.dtypes.iter().map(|dt| dt.value().vortex_unwrap())
    }

    /// Project a subset of fields from the struct
    ///
    /// If any of the fields are not found, this method will return
    /// an error.
    pub fn project(&self, projection: &[FieldName]) -> VortexResult<Self> {
        let mut names = Vec::with_capacity(projection.len());
        let mut dtypes = Vec::with_capacity(projection.len());

        for field in projection {
            let idx = self
                .find(field)
                .ok_or_else(|| vortex_err!("{field} not found"))?;
            names.push(self.0.names[idx].clone());
            dtypes.push(self.0.dtypes[idx].clone());
        }

        Ok(StructFields::from_fields(names.into(), dtypes))
    }

    /// Returns a new [`StructFields`] without the field at the given index.
    ///
    /// ## Errors
    /// Returns an error if the index is out of bounds for the struct fields.
    pub fn without_field(&self, index: usize) -> VortexResult<Self> {
        if index >= self.nfields() {
            vortex_bail!(
                "index {} out of bounds for struct with {} fields",
                index,
                self.nfields()
            );
        }

        let names = self
            .0
            .names
            .iter()
            .enumerate()
            .filter(|&(i, _)| i != index)
            .map(|(_, name)| name.clone())
            .collect::<FieldNames>();

        let dtypes = self
            .0
            .dtypes
            .iter()
            .enumerate()
            .filter(|&(i, _)| i != index)
            .map(|(_, dtype)| dtype.clone())
            .collect::<Vec<_>>();

        Ok(StructFields::from_fields(names, dtypes))
    }

    /// Merge two [`StructFields`] instances into a new one.
    /// Order of fields in arguments is preserved
    ///
    /// # Errors
    /// Returns an error if the merged struct would have duplicate field names.
    pub fn disjoint_merge(&self, other: &Self) -> VortexResult<Self> {
        let names = self
            .0
            .names
            .iter()
            .chain(other.0.names.iter())
            .cloned()
            .collect::<FieldNames>();

        if !names.iter().all_unique() {
            vortex_bail!("Can't merge struct fields with duplicate names");
        }

        let dtypes = self
            .0
            .dtypes
            .iter()
            .chain(other.0.dtypes.iter())
            .cloned()
            .collect::<Vec<_>>();

        Ok(Self::from_fields(names, dtypes))
    }
}

impl<T, V> FromIterator<(T, V)> for StructFields
where
    T: Into<FieldName>,
    V: Into<FieldDType>,
{
    fn from_iter<I: IntoIterator<Item = (T, V)>>(iter: I) -> Self {
        let (names, dtypes): (Vec<_>, Vec<_>) = iter
            .into_iter()
            .map(|(name, dtype)| (name.into(), dtype.into()))
            .unzip();
        StructFields::from_fields(names.into(), dtypes)
    }
}

#[cfg(test)]
mod test {
    use itertools::Itertools;

    use crate::dtype::DType;
    use crate::{FieldNames, Nullability, PType, StructFields};

    #[test]
    fn nullability() {
        assert!(
            !DType::Struct(
                StructFields::new(FieldNames::default(), Vec::new()),
                Nullability::NonNullable
            )
            .is_nullable()
        );

        let primitive = DType::Primitive(PType::U8, Nullability::Nullable);
        assert!(primitive.is_nullable());
        assert!(!primitive.as_nonnullable().is_nullable());
        assert!(primitive.as_nonnullable().as_nullable().is_nullable());
    }

    #[test]
    fn test_struct() {
        let a_type = DType::Primitive(PType::I32, Nullability::Nullable);
        let b_type = DType::Bool(Nullability::NonNullable);

        let dtype = DType::Struct(
            StructFields::from_iter([("A", a_type.clone()), ("B", b_type.clone())]),
            Nullability::Nullable,
        );
        assert!(dtype.is_nullable());
        assert!(dtype.as_struct().is_some());
        assert!(a_type.as_struct().is_none());

        let sdt = dtype.as_struct().unwrap();
        assert_eq!(sdt.names().len(), 2);
        assert_eq!(sdt.fields().len(), 2);
        assert_eq!(sdt.names()[0], "A".into());
        assert_eq!(sdt.names()[1], "B".into());
        assert_eq!(sdt.field_by_index(0).unwrap(), a_type);
        assert_eq!(sdt.field_by_index(1).unwrap(), b_type);

        let proj = sdt.project(&["B".into(), "A".into()]).unwrap();
        assert_eq!(proj.names()[0], "B".into());
        assert_eq!(proj.field_by_index(0).unwrap(), b_type);
        assert_eq!(proj.names()[1], "A".into());
        assert_eq!(proj.field_by_index(1).unwrap(), a_type);

        assert_eq!(sdt.find("A").unwrap(), 0);
        assert_eq!(sdt.find("B").unwrap(), 1);
        assert!(sdt.find("C").is_none());

        let without_a = sdt.without_field(0).unwrap();
        assert_eq!(without_a.names()[0], "B".into());
        assert_eq!(without_a.field_by_index(0).unwrap(), b_type);
        assert_eq!(without_a.nfields(), 1);
    }

    #[test]
    fn test_without_field_out_of_bounds() {
        let a_type = DType::Primitive(PType::I32, Nullability::Nullable);
        let b_type = DType::Bool(Nullability::NonNullable);
        let sdt = StructFields::from_iter([("A", a_type), ("B", b_type)]);

        let result = sdt.without_field(2);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("out of bounds"));

        let result = sdt.without_field(100);
        assert!(result.is_err());
    }

    #[test]
    fn test_without_field_deprecated() {
        let a_type = DType::Primitive(PType::I32, Nullability::Nullable);
        let b_type = DType::Bool(Nullability::NonNullable);
        let sdt = StructFields::from_iter([("A", a_type), ("B", b_type.clone())]);

        let without_a = sdt.without_field(0).unwrap();
        assert_eq!(without_a.names()[0], "B".into());
        assert_eq!(without_a.field_by_index(0).unwrap(), b_type);
        assert_eq!(without_a.nfields(), 1);
    }

    #[test]
    fn test_merge() {
        let child_a = DType::Primitive(PType::I32, Nullability::NonNullable);
        let child_b = DType::Bool(Nullability::Nullable);
        let child_c = DType::Utf8(Nullability::NonNullable);

        let sf1 = StructFields::from_iter([("A", child_a.clone()), ("B", child_b.clone())]);

        let sf2 = StructFields::from_iter([("C", child_c.clone())]);

        let merged = StructFields::disjoint_merge(&sf1, &sf2).unwrap();
        assert_eq!(merged.names(), &FieldNames::from_iter(["A", "B", "C"]));
        assert_eq!(
            merged.fields().collect_vec(),
            vec![child_a, child_b, child_c]
        );

        let err = StructFields::disjoint_merge(&sf1, &sf1).err().unwrap();
        assert!(err.to_string().contains("duplicate names"),);
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use std::hash::Hash;
5	use std::sync::Arc;
6
7	use itertools::Itertools;
8	use vortex_error::{
9	VortexExpect, VortexResult, VortexUnwrap, vortex_bail, vortex_err, vortex_panic,
10	};
11
12	use crate::flatbuffers::ViewedDType;
13	use crate::{DType, FieldName, FieldNames, PType};
14
15	/// DType of a struct's field, either owned or a pointer to an underlying flatbuffer.
16	#[derive(Debug, Clone, Eq, PartialEq, Hash)]
17	#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
18	pub struct FieldDType {
19	inner: FieldDTypeInner,
20	}
21
22	impl From<ViewedDType> for FieldDType {
23	fn from(value: ViewedDType) -> Self {	4✔
24	Self {	4✔
25	inner: FieldDTypeInner::View(value),	4✔
26	}	4✔
27	}	4✔
28	}
29
30	impl From<DType> for FieldDType {
31	fn from(value: DType) -> Self {	78✔
32	Self {	78✔
33	inner: FieldDTypeInner::Owned(value),	78✔
34	}	78✔
35	}	78✔
36	}
37
38	impl From<PType> for FieldDType {
UNCOV 39	fn from(value: PType) -> Self {	×
UNCOV 40	Self {	×
UNCOV 41	inner: FieldDTypeInner::Owned(DType::from(value)),	×
UNCOV 42	}	×
UNCOV 43	}	×
44	}
45
46	#[derive(Debug, Clone, Eq)]
47	enum FieldDTypeInner {
48	/// Owned DType instance
49	// TODO(ngates): we should consider making this an Arc<DType>.
50	Owned(DType),
51	/// A view over a flatbuffer, parsed only when accessed.
52	View(ViewedDType),
53	}
54
55	impl PartialEq for FieldDTypeInner {
56	fn eq(&self, other: &Self) -> bool {	282✔
57	match (self, other) {	282✔
58	(Self::Owned(lhs), Self::Owned(rhs)) => lhs == rhs,	282✔
59	(Self::View(lhs), Self::View(rhs)) => {	×
60	let lhs = DType::try_from(lhs.clone())	×
61	.vortex_expect("Failed to parse FieldDType into DType");	×
62	let rhs = DType::try_from(rhs.clone())	×
63	.vortex_expect("Failed to parse FieldDType into DType");	×
64
65	lhs == rhs	×
66	}
UNCOV 67	(Self::View(view), Self::Owned(owned)) \| (Self::Owned(owned), Self::View(view)) => {	×
UNCOV 68	let view = DType::try_from(view.clone())	×
UNCOV 69	.vortex_expect("Failed to parse FieldDType into DType");	×
UNCOV 70	owned == &view	×
71	}
72	}
73	}	282✔
74	}
75
76	impl Hash for FieldDTypeInner {
UNCOV 77	fn hash<H: std::hash::Hasher>(&self, state: &mut H) {	×
UNCOV 78	match self {	×
UNCOV 79	FieldDTypeInner::Owned(owned) => {	×
UNCOV 80	owned.hash(state);	×
UNCOV 81	}	×
82	FieldDTypeInner::View(view) => {	×
83	let owned = DType::try_from(view.clone())	×
84	.vortex_expect("Failed to parse FieldDType into DType");	×
85	owned.hash(state);	×
86	}	×
87	}
UNCOV 88	}	×
89	}
90
91	impl FieldDType {
92	/// Returns the concrete DType, parsing it from the underlying buffer if necessary.
93	pub fn value(&self) -> VortexResult<DType> {	596✔
94	self.inner.value()	596✔
95	}	596✔
96	}
97
98	impl FieldDTypeInner {
99	fn value(&self) -> VortexResult<DType> {	596✔
100	match &self {	596✔
101	FieldDTypeInner::Owned(owned) => Ok(owned.clone()),	584✔
102	FieldDTypeInner::View(view) => DType::try_from(view.clone()),	12✔
103	}
104	}	596✔
105	}
106
107	#[cfg(feature = "serde")]
108	impl serde::Serialize for FieldDTypeInner {
109	fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>	×
110	where	×
111	S: serde::Serializer,	×
112	{
113	use serde::ser::Error;
114
115	let value = self.value().map_err(S::Error::custom)?;	×
116	serializer.serialize_newtype_variant("FieldDType", 0, "Owned", &value)	×
117	}	×
118	}
119
120	#[cfg(feature = "serde")]
121	impl<'de> serde::Deserialize<'de> for FieldDTypeInner {
122	fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>	×
123	where	×
124	D: serde::Deserializer<'de>,	×
125	{
126	deserializer.deserialize_enum("FieldDType", &["Owned", "View"], FieldDTypeDeVisitor)	×
127	}	×
128	}
129
130	#[cfg(feature = "serde")]
131	struct FieldDTypeDeVisitor;
132
133	#[cfg(feature = "serde")]
134	impl<'de> serde::de::Visitor<'de> for FieldDTypeDeVisitor {
135	type Value = FieldDTypeInner;
136
137	fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {	×
138	write!(f, "variant identifier")	×
139	}	×
140
141	fn visit_enum<A>(self, data: A) -> Result<Self::Value, A::Error>	×
142	where	×
143	A: serde::de::EnumAccess<'de>,	×
144	{
145	use serde::de::{Error, VariantAccess};
146
147	#[derive(serde::Deserialize, Debug)]
148	enum FieldDTypeVariant {
149	Owned,
150	View,
151	}
152	let (variant, variant_data): (FieldDTypeVariant, _) = data.variant()?;	×
153
154	match variant {	×
155	FieldDTypeVariant::Owned => {
156	let inner = variant_data.newtype_variant::<DType>()?;	×
157	Ok(FieldDTypeInner::Owned(inner))	×
158	}
159	other => Err(A::Error::custom(format!("unsupported variant {other:?}"))),	×
160	}
161	}	×
162	}
163
164	/// Type information for a struct column.
165	///
166	/// The `StructFields` holds all field names and field types, and provides
167	/// access to them by index or by name.
168	///
169	/// ## Duplicate field names
170	///
171	/// In memory, it is not an error for a `StructFields` to contain duplicate
172	/// field names. In that case, any name-based access to fields will resolve
173	/// to the first such field with a given name.
174	///
175	/// ```rust
176	/// # use vortex_dtype::{DType, Nullability, PType, StructFields};
177	///
178	/// let fields = StructFields::from_iter([
179	/// ("string_col", DType::Utf8(Nullability::NonNullable)),
180	/// ("binary_col", DType::Binary(Nullability::NonNullable)),
181	/// ("int_col", DType::Primitive(PType::I32, Nullability::Nullable)),
182	/// ("int_col", DType::Primitive(PType::I64, Nullability::Nullable)),
183	/// ]);
184	///
185	/// // Accessing a field by name will yield the first
186	/// assert_eq!(fields.field("int_col").unwrap(), DType::Primitive(PType::I32, Nullability::Nullable));
187	/// ```
188	#[derive(Clone, PartialEq, Eq, Hash)]
189	#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
190	pub struct StructFields(Arc<StructFieldsInner>);
191
192	impl std::fmt::Debug for StructFields {
193	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {	×
194	f.debug_struct("StructFields")	×
195	.field("names", &self.0.names)	×
196	.field("dtypes", &self.0.dtypes)	×
197	.finish()	×
198	}	×
199	}
200
201	#[derive(PartialEq, Eq, Hash, Default)]
202	#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
203	struct StructFieldsInner {
204	names: FieldNames,
205	dtypes: Arc<[FieldDType]>,
206	}
207
208	impl Default for StructFields {
209	fn default() -> Self {	×
210	Self::empty()	×
211	}	×
212	}
213
214	impl StructFields {
215	/// The fields of the empty struct.
216	pub fn empty() -> Self {	×
217	Self(Arc::new(StructFieldsInner {	×
218	names: FieldNames::default(),	×
219	dtypes: Arc::from([]),	×
220	}))	×
221	}	×
222
223	/// Create a new [`StructFields`] from a list of names and dtypes
224	pub fn new(names: FieldNames, dtypes: Vec<DType>) -> Self {	276✔
225	if names.len() != dtypes.len() {	276✔
226	vortex_panic!(	×
227	"length mismatch between names ({}) and dtypes ({})",	×
228	names.len(),	×
229	dtypes.len()	×
230	);
231	}	276✔
232
233	let dtypes = dtypes	276✔
234	.into_iter()	276✔
235	.map(\|dt\| FieldDType {	276✔
236	inner: FieldDTypeInner::Owned(dt),	590✔
237	})	590✔
238	.collect::<Vec<_>>();	276✔
239
240	Self::from_fields(names, dtypes)	276✔
241	}	276✔
242
243	/// Create a new [`StructFields`] from a list of names and [`FieldDType`] which can be either lazily or eagerly serialized.
244	pub fn from_fields(names: FieldNames, dtypes: Vec<FieldDType>) -> Self {	290✔
245	if names.len() != dtypes.len() {	290✔
246	vortex_panic!(	×
247	"length mismatch between names ({}) and dtypes ({})",	×
248	names.len(),	×
249	dtypes.len()	×
250	);
251	}	290✔
252
253	let inner = Arc::new(StructFieldsInner {	290✔
254	names,	290✔
255	dtypes: dtypes.into(),	290✔
256	});	290✔
257
258	Self(inner)	290✔
259	}	290✔
260
261	/// Get the names of the fields in the struct
262	pub fn names(&self) -> &FieldNames {	66✔
263	&self.0.names	66✔
264	}	66✔
265
266	/// Returns the number of fields in the struct
267	pub fn nfields(&self) -> usize {	20✔
268	self.0.names.len()	20✔
269	}	20✔
270
271	/// Returns the name of the field at the given index
UNCOV 272	pub fn field_name(&self, index: usize) -> Option<&FieldName> {	×
UNCOV 273	self.0.names.get(index)	×
UNCOV 274	}	×
275
276	/// Find the index of a field by name
277	/// Returns `None` if the field is not found
278	pub fn find(&self, name: impl AsRef<str>) -> Option<usize> {	20✔
279	let name = name.as_ref();	20✔
280	self.0.names.iter().position(\|n\| n.as_ref() == name)	28✔
281	}	20✔
282
283	/// Get the [`DType`] of a field.
284	///
285	/// It is possible for there to be more than one field with
286	/// the same name, in which case, this will return the DType
287	/// of the first field encountered with a given name.
288	pub fn field(&self, name: impl AsRef<str>) -> Option<DType> {	12✔
289	let index = self.find(name)?;	12✔
290	Some(self.0.dtypes[index].value().vortex_unwrap())	12✔
291	}	12✔
292
293	/// Get the [`DType`] of a field by index.
294	pub fn field_by_index(&self, index: usize) -> Option<DType> {	296✔
295	Some(self.0.dtypes.get(index)?.value().vortex_unwrap())	296✔
296	}	296✔
297
298	/// Returns an ordered iterator over the fields.
299	pub fn fields(&self) -> impl ExactSizeIterator<Item = DType> + '_ {	144✔
300	self.0.dtypes.iter().map(\|dt\| dt.value().vortex_unwrap())	288✔
301	}	144✔
302
303	/// Project a subset of fields from the struct
304	///
305	/// If any of the fields are not found, this method will return
306	/// an error.
UNCOV 307	pub fn project(&self, projection: &[FieldName]) -> VortexResult<Self> {	×
UNCOV 308	let mut names = Vec::with_capacity(projection.len());	×
UNCOV 309	let mut dtypes = Vec::with_capacity(projection.len());	×
310
UNCOV 311	for field in projection {	×
UNCOV 312	let idx = self	×
UNCOV 313	.find(field)	×
UNCOV 314	.ok_or_else(\|\| vortex_err!("{field} not found"))?;	×
UNCOV 315	names.push(self.0.names[idx].clone());	×
UNCOV 316	dtypes.push(self.0.dtypes[idx].clone());	×
317	}
318
UNCOV 319	Ok(StructFields::from_fields(names.into(), dtypes))	×
UNCOV 320	}	×
321
322	/// Returns a new [`StructFields`] without the field at the given index.
323	///
324	/// ## Errors
325	/// Returns an error if the index is out of bounds for the struct fields.
UNCOV 326	pub fn without_field(&self, index: usize) -> VortexResult<Self> {	×
UNCOV 327	if index >= self.nfields() {	×
UNCOV 328	vortex_bail!(	×
UNCOV 329	"index {} out of bounds for struct with {} fields",	×
330	index,
UNCOV 331	self.nfields()	×
332	);
UNCOV 333	}	×
334
UNCOV 335	let names = self	×
UNCOV 336	.0	×
UNCOV 337	.names	×
UNCOV 338	.iter()	×
UNCOV 339	.enumerate()	×
UNCOV 340	.filter(\|&(i, _)\| i != index)	×
UNCOV 341	.map(\|(_, name)\| name.clone())	×
UNCOV 342	.collect::<FieldNames>();	×
343
UNCOV 344	let dtypes = self	×
UNCOV 345	.0	×
UNCOV 346	.dtypes	×
UNCOV 347	.iter()	×
UNCOV 348	.enumerate()	×
UNCOV 349	.filter(\|&(i, _)\| i != index)	×
UNCOV 350	.map(\|(_, dtype)\| dtype.clone())	×
UNCOV 351	.collect::<Vec<_>>();	×
352
UNCOV 353	Ok(StructFields::from_fields(names, dtypes))	×
UNCOV 354	}	×
355
356	/// Merge two [`StructFields`] instances into a new one.
357	/// Order of fields in arguments is preserved
358	///
359	/// # Errors
360	/// Returns an error if the merged struct would have duplicate field names.
UNCOV 361	pub fn disjoint_merge(&self, other: &Self) -> VortexResult<Self> {	×
UNCOV 362	let names = self	×
UNCOV 363	.0	×
UNCOV 364	.names	×
UNCOV 365	.iter()	×
UNCOV 366	.chain(other.0.names.iter())	×
UNCOV 367	.cloned()	×
UNCOV 368	.collect::<FieldNames>();	×
369
UNCOV 370	if !names.iter().all_unique() {	×
UNCOV 371	vortex_bail!("Can't merge struct fields with duplicate names");	×
UNCOV 372	}	×
373
UNCOV 374	let dtypes = self	×
UNCOV 375	.0	×
UNCOV 376	.dtypes	×
UNCOV 377	.iter()	×
UNCOV 378	.chain(other.0.dtypes.iter())	×
UNCOV 379	.cloned()	×
UNCOV 380	.collect::<Vec<_>>();	×
381
UNCOV 382	Ok(Self::from_fields(names, dtypes))	×
UNCOV 383	}	×
384	}
385
386	impl<T, V> FromIterator<(T, V)> for StructFields
387	where
388	T: Into<FieldName>,
389	V: Into<FieldDType>,
390	{
391	fn from_iter<I: IntoIterator<Item = (T, V)>>(iter: I) -> Self {	12✔
392	let (names, dtypes): (Vec<_>, Vec<_>) = iter	12✔
393	.into_iter()	12✔
394	.map(\|(name, dtype)\| (name.into(), dtype.into()))	78✔
395	.unzip();	12✔
396	StructFields::from_fields(names.into(), dtypes)	12✔
397	}	12✔
398	}
399
400	#[cfg(test)]
401	mod test {
402	use itertools::Itertools;
403
404	use crate::dtype::DType;
405	use crate::{FieldNames, Nullability, PType, StructFields};
406
407	#[test]
408	fn nullability() {
409	assert!(
410	!DType::Struct(
411	StructFields::new(FieldNames::default(), Vec::new()),
412	Nullability::NonNullable
413	)
414	.is_nullable()
415	);
416
417	let primitive = DType::Primitive(PType::U8, Nullability::Nullable);
418	assert!(primitive.is_nullable());
419	assert!(!primitive.as_nonnullable().is_nullable());
420	assert!(primitive.as_nonnullable().as_nullable().is_nullable());
421	}
422
423	#[test]
424	fn test_struct() {
425	let a_type = DType::Primitive(PType::I32, Nullability::Nullable);
426	let b_type = DType::Bool(Nullability::NonNullable);
427
428	let dtype = DType::Struct(
429	StructFields::from_iter([("A", a_type.clone()), ("B", b_type.clone())]),
430	Nullability::Nullable,
431	);
432	assert!(dtype.is_nullable());
433	assert!(dtype.as_struct().is_some());
434	assert!(a_type.as_struct().is_none());
435
436	let sdt = dtype.as_struct().unwrap();
437	assert_eq!(sdt.names().len(), 2);
438	assert_eq!(sdt.fields().len(), 2);
439	assert_eq!(sdt.names()[0], "A".into());
440	assert_eq!(sdt.names()[1], "B".into());
441	assert_eq!(sdt.field_by_index(0).unwrap(), a_type);
442	assert_eq!(sdt.field_by_index(1).unwrap(), b_type);
443
444	let proj = sdt.project(&["B".into(), "A".into()]).unwrap();
445	assert_eq!(proj.names()[0], "B".into());
446	assert_eq!(proj.field_by_index(0).unwrap(), b_type);
447	assert_eq!(proj.names()[1], "A".into());
448	assert_eq!(proj.field_by_index(1).unwrap(), a_type);
449
450	assert_eq!(sdt.find("A").unwrap(), 0);
451	assert_eq!(sdt.find("B").unwrap(), 1);
452	assert!(sdt.find("C").is_none());
453
454	let without_a = sdt.without_field(0).unwrap();
455	assert_eq!(without_a.names()[0], "B".into());
456	assert_eq!(without_a.field_by_index(0).unwrap(), b_type);
457	assert_eq!(without_a.nfields(), 1);
458	}
459
460	#[test]
461	fn test_without_field_out_of_bounds() {
462	let a_type = DType::Primitive(PType::I32, Nullability::Nullable);
463	let b_type = DType::Bool(Nullability::NonNullable);
464	let sdt = StructFields::from_iter([("A", a_type), ("B", b_type)]);
465
466	let result = sdt.without_field(2);
467	assert!(result.is_err());
468	assert!(result.unwrap_err().to_string().contains("out of bounds"));
469
470	let result = sdt.without_field(100);
471	assert!(result.is_err());
472	}
473
474	#[test]
475	fn test_without_field_deprecated() {
476	let a_type = DType::Primitive(PType::I32, Nullability::Nullable);
477	let b_type = DType::Bool(Nullability::NonNullable);
478	let sdt = StructFields::from_iter([("A", a_type), ("B", b_type.clone())]);
479
480	let without_a = sdt.without_field(0).unwrap();
481	assert_eq!(without_a.names()[0], "B".into());
482	assert_eq!(without_a.field_by_index(0).unwrap(), b_type);
483	assert_eq!(without_a.nfields(), 1);
484	}
485
486	#[test]
487	fn test_merge() {
488	let child_a = DType::Primitive(PType::I32, Nullability::NonNullable);
489	let child_b = DType::Bool(Nullability::Nullable);
490	let child_c = DType::Utf8(Nullability::NonNullable);
491
492	let sf1 = StructFields::from_iter([("A", child_a.clone()), ("B", child_b.clone())]);
493
494	let sf2 = StructFields::from_iter([("C", child_c.clone())]);
495
496	let merged = StructFields::disjoint_merge(&sf1, &sf2).unwrap();
497	assert_eq!(merged.names(), &FieldNames::from_iter(["A", "B", "C"]));
498	assert_eq!(
499	merged.fields().collect_vec(),
500	vec![child_a, child_b, child_c]
501	);
502
503	let err = StructFields::disjoint_merge(&sf1, &sf1).err().unwrap();
504	assert!(err.to_string().contains("duplicate names"),);
505	}
506	}

vortex-data / vortex / 16935267080

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous