arrow_schema/
fields.rs

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

arrow_schema/
fields.rs
