arrow_schema/
fields.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use std::ops::Deref;
6use std::sync::Arc;
7
8use crate::{ArrowError, DataType, Field, FieldRef};
9
10/// A cheaply cloneable, owned slice of [`FieldRef`]
11///
12/// Similar to `Arc<Vec<FieldRef>>` or `Arc<[FieldRef]>`
13///
14/// Can be constructed in a number of ways
15///
16/// ```
17/// # use std::sync::Arc;
18/// # use arrow_schema::{DataType, Field, Fields, SchemaBuilder};
19/// // Can be constructed from Vec<Field>
20/// Fields::from(vec![Field::new("a", DataType::Boolean, false)]);
21/// // Can be constructed from Vec<FieldRef>
22/// Fields::from(vec![Arc::new(Field::new("a", DataType::Boolean, false))]);
23/// // Can be constructed from an iterator of Field
24/// std::iter::once(Field::new("a", DataType::Boolean, false)).collect::<Fields>();
25/// // Can be constructed from an iterator of FieldRef
26/// std::iter::once(Arc::new(Field::new("a", DataType::Boolean, false))).collect::<Fields>();
27/// ```
28///
29/// See [`SchemaBuilder`] for mutating or updating [`Fields`]
30///
31/// ```
32/// # use arrow_schema::{DataType, Field, SchemaBuilder};
33/// let mut builder = SchemaBuilder::new();
34/// builder.push(Field::new("a", DataType::Boolean, false));
35/// builder.push(Field::new("b", DataType::Boolean, false));
36/// let fields = builder.finish().fields;
37///
38/// let mut builder = SchemaBuilder::from(&fields);
39/// builder.remove(0);
40/// let new = builder.finish().fields;
41/// ```
42///
43/// [`SchemaBuilder`]: crate::SchemaBuilder
44#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
45#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
46#[cfg_attr(feature = "serde", serde(transparent))]
47pub struct Fields(Arc<[FieldRef]>);
48
49impl std::fmt::Debug for Fields {
50    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
51        self.0.as_ref().fmt(f)
52    }
53}
54
55impl Fields {
56    /// Returns a new empty [`Fields`]
57    pub fn empty() -> Self {
58        Self(Arc::new([]))
59    }
60
61    /// Return size of this instance in bytes.
62    pub fn size(&self) -> usize {
63        self.iter()
64            .map(|field| field.size() + std::mem::size_of::<FieldRef>())
65            .sum()
66    }
67
68    /// Searches for a field by name, returning it along with its index if found
69    pub fn find(&self, name: &str) -> Option<(usize, &FieldRef)> {
70        self.0.iter().enumerate().find(|(_, b)| b.name() == name)
71    }
72
73    /// Check to see if `self` is a superset of `other`
74    ///
75    /// In particular returns true if both have the same number of fields, and [`Field::contains`]
76    /// for each field across self and other
77    ///
78    /// In other words, any record that conforms to `other` should also conform to `self`
79    pub fn contains(&self, other: &Fields) -> bool {
80        if Arc::ptr_eq(&self.0, &other.0) {
81            return true;
82        }
83        self.len() == other.len()
84            && self
85                .iter()
86                .zip(other.iter())
87                .all(|(a, b)| Arc::ptr_eq(a, b) || a.contains(b))
88    }
89
90    /// Returns a copy of this [`Fields`] containing only those [`FieldRef`] passing a predicate
91    ///
92    /// Performs a depth-first scan of [`Fields`] invoking `filter` for each [`FieldRef`]
93    /// containing no child [`FieldRef`], a leaf field, along with a count of the number
94    /// of such leaves encountered so far. Only [`FieldRef`] for which `filter`
95    /// returned `true` will be included in the result.
96    ///
97    /// This can therefore be used to select a subset of fields from nested types
98    /// such as [`DataType::Struct`] or [`DataType::List`].
99    ///
100    /// ```
101    /// # use arrow_schema::{DataType, Field, Fields};
102    /// let fields = Fields::from(vec![
103    ///     Field::new("a", DataType::Int32, true), // Leaf 0
104    ///     Field::new("b", DataType::Struct(Fields::from(vec![
105    ///         Field::new("c", DataType::Float32, false), // Leaf 1
106    ///         Field::new("d", DataType::Float64, false), // Leaf 2
107    ///         Field::new("e", DataType::Struct(Fields::from(vec![
108    ///             Field::new("f", DataType::Int32, false),   // Leaf 3
109    ///             Field::new("g", DataType::Float16, false), // Leaf 4
110    ///         ])), true),
111    ///     ])), false)
112    /// ]);
113    /// let filtered = fields.filter_leaves(|idx, _| [0, 2, 3, 4].contains(&idx));
114    /// let expected = Fields::from(vec![
115    ///     Field::new("a", DataType::Int32, true),
116    ///     Field::new("b", DataType::Struct(Fields::from(vec![
117    ///         Field::new("d", DataType::Float64, false),
118    ///         Field::new("e", DataType::Struct(Fields::from(vec![
119    ///             Field::new("f", DataType::Int32, false),
120    ///             Field::new("g", DataType::Float16, false),
121    ///         ])), true),
122    ///     ])), false)
123    /// ]);
124    /// assert_eq!(filtered, expected);
125    /// ```
126    pub fn filter_leaves<F: FnMut(usize, &FieldRef) -> bool>(&self, mut filter: F) -> Self {
127        self.try_filter_leaves(|idx, field| Ok(filter(idx, field)))
128            .unwrap()
129    }
130
131    /// Returns a copy of this [`Fields`] containing only those [`FieldRef`] passing a predicate
132    /// or an error if the predicate fails.
133    ///
134    /// See [`Fields::filter_leaves`] for more information.
135    pub fn try_filter_leaves<F: FnMut(usize, &FieldRef) -> Result<bool, ArrowError>>(
136        &self,
137        mut filter: F,
138    ) -> Result<Self, ArrowError> {
139        fn filter_field<F: FnMut(&FieldRef) -> Result<bool, ArrowError>>(
140            f: &FieldRef,
141            filter: &mut F,
142        ) -> Result<Option<FieldRef>, ArrowError> {
143            use DataType::*;
144
145            let v = match f.data_type() {
146                Dictionary(_, v) => v.as_ref(),       // Key must be integer
147                RunEndEncoded(_, v) => v.data_type(), // Run-ends must be integer
148                d => d,
149            };
150            let d = match v {
151                List(child) => {
152                    let fields = filter_field(child, filter)?;
153                    if let Some(fields) = fields {
154                        List(fields)
155                    } else {
156                        return Ok(None);
157                    }
158                }
159                LargeList(child) => {
160                    let fields = filter_field(child, filter)?;
161                    if let Some(fields) = fields {
162                        LargeList(fields)
163                    } else {
164                        return Ok(None);
165                    }
166                }
167                ListView(child) => {
168                    let fields = filter_field(child, filter)?;
169                    if let Some(fields) = fields {
170                        ListView(fields)
171                    } else {
172                        return Ok(None);
173                    }
174                }
175                LargeListView(child) => {
176                    let fields = filter_field(child, filter)?;
177                    if let Some(fields) = fields {
178                        LargeListView(fields)
179                    } else {
180                        return Ok(None);
181                    }
182                }
183                Map(child, ordered) => {
184                    let fields = filter_field(child, filter)?;
185                    if let Some(fields) = fields {
186                        Map(fields, *ordered)
187                    } else {
188                        return Ok(None);
189                    }
190                }
191                FixedSizeList(child, size) => {
192                    let fields = filter_field(child, filter)?;
193                    if let Some(fields) = fields {
194                        FixedSizeList(fields, *size)
195                    } else {
196                        return Ok(None);
197                    }
198                }
199                Struct(fields) => {
200                    let filtered: Result<Vec<_>, _> =
201                        fields.iter().map(|f| filter_field(f, filter)).collect();
202                    let filtered: Fields = filtered?.iter().filter_map(|f| f.clone()).collect();
203
204                    if filtered.is_empty() {
205                        return Ok(None);
206                    }
207
208                    Struct(filtered)
209                }
210                Union(fields, mode) => {
211                    let filtered: Result<Vec<_>, _> = fields
212                        .iter()
213                        .map(|(id, f)| filter_field(f, filter).map(|f| f.map(|f| (id, f))))
214                        .collect();
215                    let filtered: UnionFields =
216                        filtered?.iter().filter_map(|f| f.clone()).collect();
217
218                    if filtered.is_empty() {
219                        return Ok(None);
220                    }
221
222                    Union(filtered, *mode)
223                }
224                _ => {
225                    let filtered = filter(f)?;
226                    return Ok(filtered.then(|| f.clone()));
227                }
228            };
229            let d = match f.data_type() {
230                Dictionary(k, _) => Dictionary(k.clone(), Box::new(d)),
231                RunEndEncoded(v, f) => {
232                    RunEndEncoded(v.clone(), Arc::new(f.as_ref().clone().with_data_type(d)))
233                }
234                _ => d,
235            };
236            Ok(Some(Arc::new(f.as_ref().clone().with_data_type(d))))
237        }
238
239        let mut leaf_idx = 0;
240        let mut filter = |f: &FieldRef| {
241            let t = filter(leaf_idx, f)?;
242            leaf_idx += 1;
243            Ok(t)
244        };
245
246        let filtered: Result<Vec<_>, _> = self
247            .0
248            .iter()
249            .map(|f| filter_field(f, &mut filter))
250            .collect();
251        let filtered = filtered?.iter().filter_map(|f| f.clone()).collect();
252        Ok(filtered)
253    }
254}
255
256impl Default for Fields {
257    fn default() -> Self {
258        Self::empty()
259    }
260}
261
262impl FromIterator<Field> for Fields {
263    fn from_iter<T: IntoIterator<Item = Field>>(iter: T) -> Self {
264        iter.into_iter().map(Arc::new).collect()
265    }
266}
267
268impl FromIterator<FieldRef> for Fields {
269    fn from_iter<T: IntoIterator<Item = FieldRef>>(iter: T) -> Self {
270        Self(iter.into_iter().collect())
271    }
272}
273
274impl From<Vec<Field>> for Fields {
275    fn from(value: Vec<Field>) -> Self {
276        value.into_iter().collect()
277    }
278}
279
280impl From<Vec<FieldRef>> for Fields {
281    fn from(value: Vec<FieldRef>) -> Self {
282        Self(value.into())
283    }
284}
285
286impl From<&[FieldRef]> for Fields {
287    fn from(value: &[FieldRef]) -> Self {
288        Self(value.into())
289    }
290}
291
292impl<const N: usize> From<[FieldRef; N]> for Fields {
293    fn from(value: [FieldRef; N]) -> Self {
294        Self(Arc::new(value))
295    }
296}
297
298impl Deref for Fields {
299    type Target = [FieldRef];
300
301    fn deref(&self) -> &Self::Target {
302        self.0.as_ref()
303    }
304}
305
306impl<'a> IntoIterator for &'a Fields {
307    type Item = &'a FieldRef;
308    type IntoIter = std::slice::Iter<'a, FieldRef>;
309
310    fn into_iter(self) -> Self::IntoIter {
311        self.0.iter()
312    }
313}
314
315/// A cheaply cloneable, owned collection of [`FieldRef`] and their corresponding
316/// type ids.
317///
318/// Invariants:
319/// - All type ids are unique
320/// - All type ids are non-negative
321#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
322#[cfg_attr(feature = "serde", derive(serde::Serialize))]
323#[cfg_attr(feature = "serde", serde(transparent))]
324pub struct UnionFields(Arc<[(i8, FieldRef)]>);
325
326// Manual implementation to enforce invariants per docstring above
327#[cfg(feature = "serde")]
328impl<'de> serde::Deserialize<'de> for UnionFields {
329    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
330    where
331        D: serde::Deserializer<'de>,
332    {
333        let entries = Vec::<(i8, FieldRef)>::deserialize(deserializer)?;
334        let (type_ids, fields): (Vec<i8>, Vec<FieldRef>) = entries.into_iter().unzip();
335        Self::try_new(type_ids, fields).map_err(serde::de::Error::custom)
336    }
337}
338
339impl std::fmt::Debug for UnionFields {
340    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
341        self.0.as_ref().fmt(f)
342    }
343}
344
345/// Allows direct indexing into [`UnionFields`] to access fields by position.
346///
347/// # Panics
348///
349/// Panics if the index is out of bounds. Note that [`UnionFields`] supports
350/// a maximum of 128 fields, as type IDs are represented as `i8` values.
351///
352/// For a non-panicking alternative, use [`UnionFields::get`].
353impl std::ops::Index<usize> for UnionFields {
354    type Output = (i8, FieldRef);
355
356    fn index(&self, index: usize) -> &Self::Output {
357        &self.0[index]
358    }
359}
360
361impl UnionFields {
362    /// Create a new [`UnionFields`] with no fields
363    pub fn empty() -> Self {
364        Self(Arc::from([]))
365    }
366
367    /// Create a new [`UnionFields`] from a [`Fields`] and array of type_ids
368    ///
369    /// See <https://arrow.apache.org/docs/format/Columnar.html#union-layout>
370    ///
371    /// # Errors
372    ///
373    /// This function returns an error if:
374    /// - Iterators are not the same length
375    /// - Any type_id appears more than once (duplicate type ids)
376    /// - There is a negative type_id
377    ///
378    /// # Examples
379    ///
380    /// ```
381    /// use arrow_schema::{DataType, Field, UnionFields};
382    /// // Create a new UnionFields with type id mapping
383    /// // 1 -> DataType::UInt8
384    /// // 3 -> DataType::Utf8
385    /// let result = UnionFields::try_new(
386    ///     vec![1, 3],
387    ///     vec![
388    ///         Field::new("field1", DataType::UInt8, false),
389    ///         Field::new("field3", DataType::Utf8, false),
390    ///     ],
391    /// );
392    /// assert!(result.is_ok());
393    ///
394    /// // This will fail due to duplicate type ids
395    /// let result = UnionFields::try_new(
396    ///     vec![1, 1],
397    ///     vec![
398    ///         Field::new("field1", DataType::UInt8, false),
399    ///         Field::new("field2", DataType::Utf8, false),
400    ///     ],
401    /// );
402    /// assert!(result.is_err());
403    /// ```
404    pub fn try_new<F, T>(type_ids: T, fields: F) -> Result<Self, ArrowError>
405    where
406        F: IntoIterator,
407        F::Item: Into<FieldRef>,
408        T: IntoIterator<Item = i8>,
409    {
410        let mut type_ids_iter = type_ids.into_iter();
411        let mut fields_iter = fields.into_iter().map(Into::into);
412
413        let mut seen_type_ids = 0u128;
414
415        let mut out = Vec::new();
416
417        loop {
418            match (type_ids_iter.next(), fields_iter.next()) {
419                (None, None) => return Ok(Self(out.into())),
420                (Some(type_id), Some(field)) => {
421                    // check type id is non-negative
422                    if type_id < 0 {
423                        return Err(ArrowError::InvalidArgumentError(format!(
424                            "type ids must be non-negative: {type_id}"
425                        )));
426                    }
427
428                    // check type id uniqueness
429                    let mask = 1_u128 << type_id;
430                    if (seen_type_ids & mask) != 0 {
431                        return Err(ArrowError::InvalidArgumentError(format!(
432                            "duplicate type id: {type_id}"
433                        )));
434                    }
435
436                    seen_type_ids |= mask;
437
438                    out.push((type_id, field));
439                }
440                (None, Some(_)) => {
441                    return Err(ArrowError::InvalidArgumentError(
442                        "fields iterator has more elements than type_ids iterator".to_string(),
443                    ));
444                }
445                (Some(_), None) => {
446                    return Err(ArrowError::InvalidArgumentError(
447                        "type_ids iterator has more elements than fields iterator".to_string(),
448                    ));
449                }
450            }
451        }
452    }
453
454    /// Create a new [`UnionFields`] from a collection of fields with automatically
455    /// assigned type IDs starting from 0.
456    ///
457    /// The type IDs are assigned in increasing order: 0, 1, 2, 3, etc.
458    ///
459    /// See <https://arrow.apache.org/docs/format/Columnar.html#union-layout>
460    ///
461    /// # Panics
462    ///
463    /// Panics if the number of fields exceeds 127 (the maximum value for i8 type IDs).
464    ///
465    /// If you want to avoid panics, use [`UnionFields::try_from_fields`] instead, which
466    /// returns a `Result`.
467    ///
468    /// # Examples
469    ///
470    /// ```
471    /// use arrow_schema::{DataType, Field, UnionFields};
472    /// // Create a new UnionFields with automatic type id assignment
473    /// // 0 -> DataType::UInt8
474    /// // 1 -> DataType::Utf8
475    /// let union_fields = UnionFields::from_fields(vec![
476    ///     Field::new("field1", DataType::UInt8, false),
477    ///     Field::new("field2", DataType::Utf8, false),
478    /// ]);
479    /// assert_eq!(union_fields.len(), 2);
480    /// ```
481    pub fn from_fields<F>(fields: F) -> Self
482    where
483        F: IntoIterator,
484        F::Item: Into<FieldRef>,
485    {
486        fields
487            .into_iter()
488            .enumerate()
489            .map(|(i, field)| {
490                let id = i8::try_from(i).expect("UnionFields cannot contain more than 128 fields");
491
492                (id, field.into())
493            })
494            .collect()
495    }
496
497    /// Create a new [`UnionFields`] from a collection of fields with automatically
498    /// assigned type IDs starting from 0.
499    ///
500    /// The type IDs are assigned in increasing order: 0, 1, 2, 3, etc.
501    ///
502    /// This is the non-panicking version of [`UnionFields::from_fields`].
503    ///
504    /// See <https://arrow.apache.org/docs/format/Columnar.html#union-layout>
505    ///
506    /// # Errors
507    ///
508    /// Returns an error if the number of fields exceeds 127 (the maximum value for i8 type IDs).
509    ///
510    /// # Examples
511    ///
512    /// ```
513    /// use arrow_schema::{DataType, Field, UnionFields};
514    /// // Create a new UnionFields with automatic type id assignment
515    /// // 0 -> DataType::UInt8
516    /// // 1 -> DataType::Utf8
517    /// let result = UnionFields::try_from_fields(vec![
518    ///     Field::new("field1", DataType::UInt8, false),
519    ///     Field::new("field2", DataType::Utf8, false),
520    /// ]);
521    /// assert!(result.is_ok());
522    /// assert_eq!(result.unwrap().len(), 2);
523    ///
524    /// // This will fail with too many fields
525    /// let many_fields: Vec<_> = (0..200)
526    ///     .map(|i| Field::new(format!("field{}", i), DataType::Int32, false))
527    ///     .collect();
528    /// let result = UnionFields::try_from_fields(many_fields);
529    /// assert!(result.is_err());
530    /// ```
531    pub fn try_from_fields<F>(fields: F) -> Result<Self, ArrowError>
532    where
533        F: IntoIterator,
534        F::Item: Into<FieldRef>,
535    {
536        let mut out = Vec::with_capacity(i8::MAX as usize + 1);
537
538        for (i, field) in fields.into_iter().enumerate() {
539            let id = i8::try_from(i).map_err(|_| {
540                ArrowError::InvalidArgumentError(
541                    "UnionFields cannot contain more than 128 fields".into(),
542                )
543            })?;
544
545            out.push((id, field.into()));
546        }
547
548        Ok(Self(out.into()))
549    }
550
551    /// Return size of this instance in bytes.
552    pub fn size(&self) -> usize {
553        self.iter()
554            .map(|(_, field)| field.size() + std::mem::size_of::<(i8, FieldRef)>())
555            .sum()
556    }
557
558    /// Returns the number of fields in this [`UnionFields`]
559    pub fn len(&self) -> usize {
560        self.0.len()
561    }
562
563    /// Returns `true` if this is empty
564    pub fn is_empty(&self) -> bool {
565        self.0.is_empty()
566    }
567
568    /// Returns an iterator over the fields and type ids in this [`UnionFields`]
569    pub fn iter(&self) -> impl Iterator<Item = (i8, &FieldRef)> + '_ {
570        self.0.iter().map(|(id, f)| (*id, f))
571    }
572
573    /// Returns a reference to the field at the given index, or `None` if out of bounds.
574    ///
575    /// This is a safe alternative to direct indexing via `[]`.
576    ///
577    /// # Example
578    ///
579    /// ```
580    /// use arrow_schema::{DataType, Field, UnionFields};
581    ///
582    /// let fields = UnionFields::try_new(
583    ///     vec![1, 3],
584    ///     vec![
585    ///         Field::new("field1", DataType::UInt8, false),
586    ///         Field::new("field3", DataType::Utf8, false),
587    ///     ],
588    /// ).unwrap();
589    ///
590    /// assert!(fields.get(0).is_some());
591    /// assert!(fields.get(1).is_some());
592    /// assert!(fields.get(2).is_none());
593    /// ```
594    pub fn get(&self, index: usize) -> Option<&(i8, FieldRef)> {
595        self.0.get(index)
596    }
597
598    /// Searches for a field by its type id, returning the type id and field reference if found.
599    /// Returns `None` if no field with the given type id exists.
600    pub fn find_by_type_id(&self, type_id: i8) -> Option<(i8, &FieldRef)> {
601        self.iter().find(|&(i, _)| i == type_id)
602    }
603
604    /// Searches for a field by value equality, returning its type id and reference if found.
605    /// Returns `None` if no matching field exists in this [`UnionFields`].
606    pub fn find_by_field(&self, field: &Field) -> Option<(i8, &FieldRef)> {
607        self.iter().find(|&(_, f)| f.as_ref() == field)
608    }
609
610    /// Merge this field into self if it is compatible.
611    ///
612    /// See [`Field::try_merge`]
613    pub(crate) fn try_merge(&mut self, other: &Self) -> Result<(), ArrowError> {
614        let mut output: Vec<_> = self.iter().map(|(id, f)| (id, f.clone())).collect();
615        for (field_type_id, from_field) in other.iter() {
616            let mut is_new_field = true;
617            for (self_type_id, self_field) in &mut output {
618                if from_field == self_field {
619                    // If the nested fields in two unions are the same, they must have same
620                    // type id.
621                    if *self_type_id != field_type_id {
622                        return Err(ArrowError::SchemaError(format!(
623                            "Fail to merge schema field '{}' because the self_type_id = {} does not equal field_type_id = {}",
624                            self_field.name(),
625                            self_type_id,
626                            field_type_id
627                        )));
628                    }
629
630                    is_new_field = false;
631                    break;
632                }
633            }
634
635            if is_new_field {
636                if let Some((_, self_field)) = output.iter().find(|(id, _)| *id == field_type_id) {
637                    return Err(ArrowError::SchemaError(format!(
638                        "Fail to merge schema field '{}' because type id {} is already used by field '{}'",
639                        from_field.name(),
640                        field_type_id,
641                        self_field.name()
642                    )));
643                }
644                output.push((field_type_id, from_field.clone()))
645            }
646        }
647        *self = Self(output.into());
648        Ok(())
649    }
650}
651
652impl FromIterator<(i8, FieldRef)> for UnionFields {
653    /// # Panics
654    ///
655    /// This will panic on invalid inputs:
656    /// - Duplicate type IDs
657    /// - Negative type IDs
658    ///
659    /// Prefer using [`UnionFields::try_new`] for a fallible construction method.
660    fn from_iter<T: IntoIterator<Item = (i8, FieldRef)>>(iter: T) -> Self {
661        let iter = iter.into_iter();
662        let (type_ids, fields): (Vec<i8>, Vec<FieldRef>) = iter.unzip();
663        Self::try_new(type_ids, fields).expect("invalid arguments passed")
664    }
665}
666
667#[cfg(test)]
668mod tests {
669    use super::*;
670    use crate::UnionMode;
671
672    #[test]
673    fn test_filter() {
674        let floats = Fields::from(vec![
675            Field::new("a", DataType::Float32, false),
676            Field::new("b", DataType::Float32, false),
677        ]);
678        let fields = Fields::from(vec![
679            Field::new("a", DataType::Int32, true),
680            Field::new("floats", DataType::Struct(floats.clone()), true),
681            Field::new("b", DataType::Int16, true),
682            Field::new(
683                "c",
684                DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
685                false,
686            ),
687            Field::new(
688                "d",
689                DataType::Dictionary(
690                    Box::new(DataType::Int32),
691                    Box::new(DataType::Struct(floats.clone())),
692                ),
693                false,
694            ),
695            Field::new_list(
696                "e",
697                Field::new("floats", DataType::Struct(floats.clone()), true),
698                true,
699            ),
700            Field::new_fixed_size_list(
701                "f",
702                Field::new_list_field(DataType::Int32, false),
703                3,
704                false,
705            ),
706            Field::new_map(
707                "g",
708                Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
709                Field::new(
710                    Field::MAP_KEY_FIELD_DEFAULT_NAME,
711                    DataType::LargeUtf8,
712                    false,
713                ),
714                Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, true),
715                false,
716                false,
717            ),
718            Field::new(
719                "h",
720                DataType::Union(
721                    UnionFields::try_new(
722                        vec![1, 3],
723                        vec![
724                            Field::new("field1", DataType::UInt8, false),
725                            Field::new("field3", DataType::Utf8, false),
726                        ],
727                    )
728                    .unwrap(),
729                    UnionMode::Dense,
730                ),
731                true,
732            ),
733            Field::new(
734                "i",
735                DataType::RunEndEncoded(
736                    Arc::new(Field::new(
737                        Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
738                        DataType::Int32,
739                        false,
740                    )),
741                    Arc::new(Field::new(
742                        Field::REE_VALUES_FIELD_DEFAULT_NAME,
743                        DataType::Struct(floats.clone()),
744                        true,
745                    )),
746                ),
747                false,
748            ),
749        ]);
750
751        let floats_a = DataType::Struct(vec![floats[0].clone()].into());
752
753        let r = fields.filter_leaves(|idx, _| idx == 0 || idx == 1);
754        assert_eq!(r.len(), 2);
755        assert_eq!(r[0], fields[0]);
756        assert_eq!(r[1].data_type(), &floats_a);
757
758        let r = fields.filter_leaves(|_, f| f.name() == "a");
759        assert_eq!(r.len(), 5);
760        assert_eq!(r[0], fields[0]);
761        assert_eq!(r[1].data_type(), &floats_a);
762        assert_eq!(
763            r[2].data_type(),
764            &DataType::Dictionary(Box::new(DataType::Int32), Box::new(floats_a.clone()))
765        );
766        assert_eq!(
767            r[3].as_ref(),
768            &Field::new_list("e", Field::new("floats", floats_a.clone(), true), true)
769        );
770        assert_eq!(
771            r[4].as_ref(),
772            &Field::new(
773                "i",
774                DataType::RunEndEncoded(
775                    Arc::new(Field::new(
776                        Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
777                        DataType::Int32,
778                        false
779                    )),
780                    Arc::new(Field::new(
781                        Field::REE_VALUES_FIELD_DEFAULT_NAME,
782                        floats_a.clone(),
783                        true
784                    )),
785                ),
786                false,
787            )
788        );
789
790        let r = fields.filter_leaves(|_, f| f.name() == "floats");
791        assert_eq!(r.len(), 0);
792
793        let r = fields.filter_leaves(|idx, _| idx == 9);
794        assert_eq!(r.len(), 1);
795        assert_eq!(r[0], fields[6]);
796
797        let r = fields.filter_leaves(|idx, _| idx == 10 || idx == 11);
798        assert_eq!(r.len(), 1);
799        assert_eq!(r[0], fields[7]);
800
801        let union = DataType::Union(
802            UnionFields::try_new(vec![1], vec![Field::new("field1", DataType::UInt8, false)])
803                .unwrap(),
804            UnionMode::Dense,
805        );
806
807        let r = fields.filter_leaves(|idx, _| idx == 12);
808        assert_eq!(r.len(), 1);
809        assert_eq!(r[0].data_type(), &union);
810
811        let r = fields.filter_leaves(|idx, _| idx == 14 || idx == 15);
812        assert_eq!(r.len(), 1);
813        assert_eq!(r[0], fields[9]);
814
815        // Propagate error
816        let r = fields.try_filter_leaves(|_, _| Err(ArrowError::SchemaError("error".to_string())));
817        assert!(r.is_err());
818    }
819
820    #[test]
821    fn test_filter_list_view() {
822        let floats = Fields::from(vec![
823            Field::new("a", DataType::Float32, false),
824            Field::new("b", DataType::Float32, false),
825        ]);
826        let fields = Fields::from(vec![
827            Field::new(
828                "lv_struct",
829                DataType::ListView(Arc::new(Field::new(
830                    "floats",
831                    DataType::Struct(floats.clone()),
832                    true,
833                ))),
834                true,
835            ),
836            Field::new(
837                "llv_struct",
838                DataType::LargeListView(Arc::new(Field::new(
839                    "floats",
840                    DataType::Struct(floats.clone()),
841                    true,
842                ))),
843                false,
844            ),
845            Field::new(
846                "lv_prim",
847                DataType::ListView(Arc::new(Field::new_list_field(DataType::Int32, false))),
848                true,
849            ),
850            Field::new(
851                "llv_prim",
852                DataType::LargeListView(Arc::new(Field::new_list_field(DataType::Int64, true))),
853                false,
854            ),
855        ]);
856
857        let floats_a = DataType::Struct(vec![floats[0].clone()].into());
858
859        // Filter by leaf field name "a" inside nested structs of ListView and LargeListView
860        let r = fields.filter_leaves(|_, f| f.name() == "a");
861        assert_eq!(r.len(), 2);
862        assert_eq!(
863            r[0].as_ref(),
864            &Field::new(
865                "lv_struct",
866                DataType::ListView(Arc::new(Field::new("floats", floats_a.clone(), true))),
867                true,
868            )
869        );
870        assert_eq!(
871            r[1].as_ref(),
872            &Field::new(
873                "llv_struct",
874                DataType::LargeListView(Arc::new(Field::new("floats", floats_a, true))),
875                false,
876            )
877        );
878
879        // Parent ListView/LargeListView fields are not leaves themselves
880        let r = fields.filter_leaves(|_, f| f.name() == "lv_struct" || f.name() == "llv_struct");
881        assert_eq!(r.len(), 0);
882
883        // Leaf indices: 0,1 (lv_struct.floats.a,b), 2,3 (llv_struct.floats.a,b), 4 (lv_prim.item), 5 (llv_prim.item)
884        let r = fields.filter_leaves(|idx, _| idx == 4 || idx == 5);
885        assert_eq!(r.len(), 2);
886        assert_eq!(r[0], fields[2]);
887        assert_eq!(r[1], fields[3]);
888    }
889
890    #[test]
891    fn test_union_fields_try_new_valid() {
892        let res = UnionFields::try_new(
893            vec![1, 6, 7],
894            vec![
895                Field::new("f1", DataType::UInt8, false),
896                Field::new("f6", DataType::Utf8, false),
897                Field::new("f7", DataType::Int32, true),
898            ],
899        );
900        assert!(res.is_ok());
901        let union_fields = res.unwrap();
902        assert_eq!(union_fields.len(), 3);
903        assert_eq!(
904            union_fields.iter().map(|(id, _)| id).collect::<Vec<_>>(),
905            vec![1, 6, 7]
906        );
907    }
908
909    #[test]
910    fn test_union_fields_try_new_empty() {
911        let res = UnionFields::try_new(Vec::<i8>::new(), Vec::<Field>::new());
912        assert!(res.is_ok());
913        assert!(res.unwrap().is_empty());
914    }
915
916    #[test]
917    fn test_union_fields_try_new_duplicate_type_id() {
918        let res = UnionFields::try_new(
919            vec![1, 1],
920            vec![
921                Field::new("f1", DataType::UInt8, false),
922                Field::new("f2", DataType::Utf8, false),
923            ],
924        );
925        assert!(res.is_err());
926        assert!(
927            res.unwrap_err()
928                .to_string()
929                .contains("duplicate type id: 1")
930        );
931    }
932
933    #[test]
934    fn test_union_fields_try_new_duplicate_field() {
935        let field = Field::new("field", DataType::UInt8, false);
936        let res = UnionFields::try_new(vec![1, 2], vec![field.clone(), field]);
937        assert!(res.is_ok());
938    }
939
940    #[test]
941    fn test_union_fields_try_new_more_type_ids() {
942        let res = UnionFields::try_new(
943            vec![1, 2, 3],
944            vec![
945                Field::new("f1", DataType::UInt8, false),
946                Field::new("f2", DataType::Utf8, false),
947            ],
948        );
949        assert!(res.is_err());
950        assert!(
951            res.unwrap_err()
952                .to_string()
953                .contains("type_ids iterator has more elements")
954        );
955    }
956
957    #[test]
958    fn test_union_fields_try_new_more_fields() {
959        let res = UnionFields::try_new(
960            vec![1, 2],
961            vec![
962                Field::new("f1", DataType::UInt8, false),
963                Field::new("f2", DataType::Utf8, false),
964                Field::new("f3", DataType::Int32, true),
965            ],
966        );
967        assert!(res.is_err());
968        assert!(
969            res.unwrap_err()
970                .to_string()
971                .contains("fields iterator has more elements")
972        );
973    }
974
975    #[test]
976    fn test_union_fields_try_new_negative_type_ids() {
977        let res = UnionFields::try_new(
978            vec![-128, -1, 0, 127],
979            vec![
980                Field::new("field_min", DataType::UInt8, false),
981                Field::new("field_neg", DataType::Utf8, false),
982                Field::new("field_zero", DataType::Int32, true),
983                Field::new("field_max", DataType::Boolean, false),
984            ],
985        );
986        assert!(res.is_err());
987        assert!(
988            res.unwrap_err()
989                .to_string()
990                .contains("type ids must be non-negative")
991        )
992    }
993
994    #[test]
995    fn test_union_fields_try_new_complex_types() {
996        let res = UnionFields::try_new(
997            vec![0, 1, 2],
998            vec![
999                Field::new(
1000                    "struct_field",
1001                    DataType::Struct(Fields::from(vec![
1002                        Field::new("a", DataType::Int32, false),
1003                        Field::new("b", DataType::Utf8, true),
1004                    ])),
1005                    false,
1006                ),
1007                Field::new_list(
1008                    "list_field",
1009                    Field::new("item", DataType::Float64, true),
1010                    true,
1011                ),
1012                Field::new(
1013                    "dict_field",
1014                    DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
1015                    false,
1016                ),
1017            ],
1018        );
1019        assert!(res.is_ok());
1020        assert_eq!(res.unwrap().len(), 3);
1021    }
1022
1023    #[test]
1024    fn test_union_fields_try_new_single_field() {
1025        let res = UnionFields::try_new(
1026            vec![42],
1027            vec![Field::new("only_field", DataType::Int64, false)],
1028        );
1029        assert!(res.is_ok());
1030        let union_fields = res.unwrap();
1031        assert_eq!(union_fields.len(), 1);
1032        assert_eq!(union_fields.iter().next().unwrap().0, 42);
1033    }
1034
1035    #[test]
1036    fn test_union_fields_try_from_fields_empty() {
1037        let res = UnionFields::try_from_fields(Vec::<Field>::new());
1038        assert!(res.is_ok());
1039        assert!(res.unwrap().is_empty());
1040    }
1041
1042    #[test]
1043    fn test_union_fields_try_from_fields_single() {
1044        let res = UnionFields::try_from_fields(vec![Field::new("only", DataType::Int64, false)]);
1045        assert!(res.is_ok());
1046        let union_fields = res.unwrap();
1047        assert_eq!(union_fields.len(), 1);
1048        assert_eq!(union_fields.iter().next().unwrap().0, 0);
1049    }
1050
1051    #[test]
1052    fn test_union_fields_try_from_fields_too_many() {
1053        let many_fields: Vec<_> = (0..200)
1054            .map(|i| Field::new(format!("field{i}"), DataType::Int32, false))
1055            .collect();
1056        let res = UnionFields::try_from_fields(many_fields);
1057        assert!(res.is_err());
1058        assert!(
1059            res.unwrap_err()
1060                .to_string()
1061                .contains("UnionFields cannot contain more than 128 fields")
1062        );
1063    }
1064
1065    #[test]
1066    fn test_union_fields_try_from_fields_max_valid() {
1067        let fields: Vec<_> = (0..=i8::MAX)
1068            .map(|i| Field::new(format!("field{i}"), DataType::Int32, false))
1069            .collect();
1070        let res = UnionFields::try_from_fields(fields);
1071        assert!(res.is_ok());
1072        let union_fields = res.unwrap();
1073        assert_eq!(union_fields.len(), 128);
1074        assert_eq!(union_fields.iter().map(|(id, _)| id).min().unwrap(), 0);
1075        assert_eq!(union_fields.iter().map(|(id, _)| id).max().unwrap(), 127);
1076    }
1077
1078    #[test]
1079    fn test_union_fields_try_from_fields_over_max() {
1080        // 129 fields should fail
1081        let fields: Vec<_> = (0..129)
1082            .map(|i| Field::new(format!("field{i}"), DataType::Int32, false))
1083            .collect();
1084        let res = UnionFields::try_from_fields(fields);
1085        assert!(res.is_err());
1086    }
1087
1088    #[test]
1089    fn test_union_fields_try_from_fields_complex_types() {
1090        let res = UnionFields::try_from_fields(vec![
1091            Field::new(
1092                "struct_field",
1093                DataType::Struct(Fields::from(vec![
1094                    Field::new("a", DataType::Int32, false),
1095                    Field::new("b", DataType::Utf8, true),
1096                ])),
1097                false,
1098            ),
1099            Field::new_list(
1100                "list_field",
1101                Field::new("item", DataType::Float64, true),
1102                true,
1103            ),
1104            Field::new(
1105                "dict_field",
1106                DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
1107                false,
1108            ),
1109        ]);
1110        assert!(res.is_ok());
1111        assert_eq!(res.unwrap().len(), 3);
1112    }
1113
1114    #[test]
1115    #[should_panic(
1116        expected = "invalid arguments passed: InvalidArgumentError(\"type ids must be non-negative: -1\")"
1117    )]
1118    fn test_union_fields_from_iter_should_validate() {
1119        let fields: Vec<(i8, FieldRef)> = vec![
1120            (0, Arc::new(Field::new("a", DataType::Int32, true))),
1121            (-1, Arc::new(Field::new("a", DataType::Int32, true))),
1122        ];
1123        UnionFields::from_iter(fields);
1124    }
1125
1126    #[cfg(feature = "serde")]
1127    #[test]
1128    fn test_union_fields_serde_round_trip() {
1129        let fields = UnionFields::try_new(
1130            vec![1, 3],
1131            vec![
1132                Field::new("a", DataType::Int32, true),
1133                Field::new("b", DataType::Utf8, true),
1134            ],
1135        )
1136        .unwrap();
1137        let serialized = postcard::to_stdvec(&fields).unwrap();
1138        let deserialized: UnionFields = postcard::from_bytes(&serialized).unwrap();
1139        assert_eq!(fields, deserialized);
1140    }
1141
1142    #[cfg(feature = "serde")]
1143    #[test]
1144    fn test_union_fields_deserialize_should_validate() {
1145        let field = Arc::new(Field::new("a", DataType::Int32, true));
1146
1147        let duplicate: Vec<(i8, FieldRef)> = vec![(1, Arc::clone(&field)), (1, Arc::clone(&field))];
1148        let serialized = serde_json::to_string(&duplicate).unwrap();
1149        let err = serde_json::from_str::<UnionFields>(&serialized).unwrap_err();
1150        assert!(err.to_string().contains("duplicate type id: 1"), "{err}");
1151
1152        let negative: Vec<(i8, FieldRef)> = vec![(0, Arc::clone(&field)), (-1, field)];
1153        let serialized = serde_json::to_string(&negative).unwrap();
1154        let err = serde_json::from_str::<UnionFields>(&serialized).unwrap_err();
1155        assert!(
1156            err.to_string()
1157                .contains("type ids must be non-negative: -1"),
1158            "{err}"
1159        );
1160    }
1161}
arrow_schema/fields.rs

arrow_schema/
fields.rs
