/Users/andrewlamb/Software/arrow-rs/arrow-schema/src/fields.rs

Source
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use std::ops::Deref;
use std::sync::Arc;

use crate::{ArrowError, DataType, Field, FieldRef};

/// A cheaply cloneable, owned slice of [`FieldRef`]
///
/// Similar to `Arc<Vec<FieldRef>>` or `Arc<[FieldRef]>`
///
/// Can be constructed in a number of ways
///
/// ```
/// # use std::sync::Arc;
/// # use arrow_schema::{DataType, Field, Fields, SchemaBuilder};
/// // Can be constructed from Vec<Field>
/// Fields::from(vec![Field::new("a", DataType::Boolean, false)]);
/// // Can be constructed from Vec<FieldRef>
/// Fields::from(vec![Arc::new(Field::new("a", DataType::Boolean, false))]);
/// // Can be constructed from an iterator of Field
/// std::iter::once(Field::new("a", DataType::Boolean, false)).collect::<Fields>();
/// // Can be constructed from an iterator of FieldRef
/// std::iter::once(Arc::new(Field::new("a", DataType::Boolean, false))).collect::<Fields>();
/// ```
///
/// See [`SchemaBuilder`] for mutating or updating [`Fields`]
///
/// ```
/// # use arrow_schema::{DataType, Field, SchemaBuilder};
/// let mut builder = SchemaBuilder::new();
/// builder.push(Field::new("a", DataType::Boolean, false));
/// builder.push(Field::new("b", DataType::Boolean, false));
/// let fields = builder.finish().fields;
///
/// let mut builder = SchemaBuilder::from(&fields);
/// builder.remove(0);
/// let new = builder.finish().fields;
/// ```
///
/// [`SchemaBuilder`]: crate::SchemaBuilder
#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "serde", serde(transparent))]
pub struct Fields(Arc<[FieldRef]>);

impl std::fmt::Debug for Fields {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.0.as_ref().fmt(f)
    }
}

impl Fields {
    /// Returns a new empty [`Fields`]
    pub fn empty() -> Self {
        Self(Arc::new([]))
    }

    /// Return size of this instance in bytes.
    pub fn size(&self) -> usize {
        self.iter()
            .map(|field| field.size() + std::mem::size_of::<FieldRef>())
            .sum()
    }

    /// Searches for a field by name, returning it along with its index if found
    pub fn find(&self, name: &str) -> Option<(usize, &FieldRef)> {
        self.0.iter().enumerate().find(|(_, b)| b.name() == name)
    }

    /// Check to see if `self` is a superset of `other`
    ///
    /// In particular returns true if both have the same number of fields, and [`Field::contains`]
    /// for each field across self and other
    ///
    /// In other words, any record that conforms to `other` should also conform to `self`
    pub fn contains(&self, other: &Fields) -> bool {
        if Arc::ptr_eq(&self.0, &other.0) {
            return true;
        }
        self.len() == other.len()
            && self
                .iter()
                .zip(other.iter())
                .all(|(a, b)| Arc::ptr_eq(a, b) || a.contains(b))
    }

    /// Returns a copy of this [`Fields`] containing only those [`FieldRef`] passing a predicate
    ///
    /// Performs a depth-first scan of [`Fields`] invoking `filter` for each [`FieldRef`]
    /// containing no child [`FieldRef`], a leaf field, along with a count of the number
    /// of such leaves encountered so far. Only [`FieldRef`] for which `filter`
    /// returned `true` will be included in the result.
    ///
    /// This can therefore be used to select a subset of fields from nested types
    /// such as [`DataType::Struct`] or [`DataType::List`].
    ///
    /// ```
    /// # use arrow_schema::{DataType, Field, Fields};
    /// let fields = Fields::from(vec![
    ///     Field::new("a", DataType::Int32, true), // Leaf 0
    ///     Field::new("b", DataType::Struct(Fields::from(vec![
    ///         Field::new("c", DataType::Float32, false), // Leaf 1
    ///         Field::new("d", DataType::Float64, false), // Leaf 2
    ///         Field::new("e", DataType::Struct(Fields::from(vec![
    ///             Field::new("f", DataType::Int32, false),   // Leaf 3
    ///             Field::new("g", DataType::Float16, false), // Leaf 4
    ///         ])), true),
    ///     ])), false)
    /// ]);
    /// let filtered = fields.filter_leaves(|idx, _| [0, 2, 3, 4].contains(&idx));
    /// let expected = Fields::from(vec![
    ///     Field::new("a", DataType::Int32, true),
    ///     Field::new("b", DataType::Struct(Fields::from(vec![
    ///         Field::new("d", DataType::Float64, false),
    ///         Field::new("e", DataType::Struct(Fields::from(vec![
    ///             Field::new("f", DataType::Int32, false),
    ///             Field::new("g", DataType::Float16, false),
    ///         ])), true),
    ///     ])), false)
    /// ]);
    /// assert_eq!(filtered, expected);
    /// ```
    pub fn filter_leaves<F: FnMut(usize, &FieldRef) -> bool>(&self, mut filter: F) -> Self {
        self.try_filter_leaves(|idx, field| Ok(filter(idx, field)))
            .unwrap()
    }

    /// Returns a copy of this [`Fields`] containing only those [`FieldRef`] passing a predicate
    /// or an error if the predicate fails.
    ///
    /// See [`Fields::filter_leaves`] for more information.
    pub fn try_filter_leaves<F: FnMut(usize, &FieldRef) -> Result<bool, ArrowError>>(
        &self,
        mut filter: F,
    ) -> Result<Self, ArrowError> {
        fn filter_field<F: FnMut(&FieldRef) -> Result<bool, ArrowError>>(
            f: &FieldRef,
            filter: &mut F,
        ) -> Result<Option<FieldRef>, ArrowError> {
            use DataType::*;

            let v = match f.data_type() {
                Dictionary(_, v) => v.as_ref(),       // Key must be integer
                RunEndEncoded(_, v) => v.data_type(), // Run-ends must be integer
                d => d,
            };
            let d = match v {
                List(child) => {
                    let fields = filter_field(child, filter)?;
                    if let Some(fields) = fields {
                        List(fields)
                    } else {
                        return Ok(None);
                    }
                }
                LargeList(child) => {
                    let fields = filter_field(child, filter)?;
                    if let Some(fields) = fields {
                        LargeList(fields)
                    } else {
                        return Ok(None);
                    }
                }
                Map(child, ordered) => {
                    let fields = filter_field(child, filter)?;
                    if let Some(fields) = fields {
                        Map(fields, *ordered)
                    } else {
                        return Ok(None);
                    }
                }
                FixedSizeList(child, size) => {
                    let fields = filter_field(child, filter)?;
                    if let Some(fields) = fields {
                        FixedSizeList(fields, *size)
                    } else {
                        return Ok(None);
                    }
                }
                Struct(fields) => {
                    let filtered: Result<Vec<_>, _> =
                        fields.iter().map(|f| filter_field(f, filter)).collect();
                    let filtered: Fields = filtered?
                        .iter()
                        .filter_map(|f| f.as_ref().cloned())
                        .collect();

                    if filtered.is_empty() {
                        return Ok(None);
                    }

                    Struct(filtered)
                }
                Union(fields, mode) => {
                    let filtered: Result<Vec<_>, _> = fields
                        .iter()
                        .map(|(id, f)| filter_field(f, filter).map(|f| f.map(|f| (id, f))))
                        .collect();
                    let filtered: UnionFields = filtered?
                        .iter()
                        .filter_map(|f| f.as_ref().cloned())
                        .collect();

                    if filtered.is_empty() {
                        return Ok(None);
                    }

                    Union(filtered, *mode)
                }
                _ => {
                    let filtered = filter(f)?;
                    return Ok(filtered.then(|| f.clone()));
                }
            };
            let d = match f.data_type() {
                Dictionary(k, _) => Dictionary(k.clone(), Box::new(d)),
                RunEndEncoded(v, f) => {
                    RunEndEncoded(v.clone(), Arc::new(f.as_ref().clone().with_data_type(d)))
                }
                _ => d,
            };
            Ok(Some(Arc::new(f.as_ref().clone().with_data_type(d))))
        }

        let mut leaf_idx = 0;
        let mut filter = |f: &FieldRef| {
            let t = filter(leaf_idx, f)?;
            leaf_idx += 1;
            Ok(t)
        };

        let filtered: Result<Vec<_>, _> = self
            .0
            .iter()
            .map(|f| filter_field(f, &mut filter))
            .collect();
        let filtered = filtered?
            .iter()
            .filter_map(|f| f.as_ref().cloned())
            .collect();
        Ok(filtered)
    }
}

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

impl FromIterator<Field> for Fields {
    fn from_iter<T: IntoIterator<Item = Field>>(iter: T) -> Self {
        iter.into_iter().map(Arc::new).collect()
    }
}

impl FromIterator<FieldRef> for Fields {
    fn from_iter<T: IntoIterator<Item = FieldRef>>(iter: T) -> Self {
        Self(iter.into_iter().collect())
    }
}

impl From<Vec<Field>> for Fields {
    fn from(value: Vec<Field>) -> Self {
        value.into_iter().collect()
    }
}

impl From<Vec<FieldRef>> for Fields {
    fn from(value: Vec<FieldRef>) -> Self {
        Self(value.into())
    }
}

impl From<&[FieldRef]> for Fields {
    fn from(value: &[FieldRef]) -> Self {
        Self(value.into())
    }
}

impl<const N: usize> From<[FieldRef; N]> for Fields {
    fn from(value: [FieldRef; N]) -> Self {
        Self(Arc::new(value))
    }
}

impl Deref for Fields {
    type Target = [FieldRef];

    fn deref(&self) -> &Self::Target {
        self.0.as_ref()
    }
}

impl<'a> IntoIterator for &'a Fields {
    type Item = &'a FieldRef;
    type IntoIter = std::slice::Iter<'a, FieldRef>;

    fn into_iter(self) -> Self::IntoIter {
        self.0.iter()
    }
}

/// A cheaply cloneable, owned collection of [`FieldRef`] and their corresponding type ids
#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "serde", serde(transparent))]
pub struct UnionFields(Arc<[(i8, FieldRef)]>);

impl std::fmt::Debug for UnionFields {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.0.as_ref().fmt(f)
    }
}

impl UnionFields {
    /// Create a new [`UnionFields`] with no fields
    pub fn empty() -> Self {
        Self(Arc::from([]))
    }

    /// Create a new [`UnionFields`] from a [`Fields`] and array of type_ids
    ///
    /// See <https://arrow.apache.org/docs/format/Columnar.html#union-layout>
    ///
    /// ```
    /// use arrow_schema::{DataType, Field, UnionFields};
    /// // Create a new UnionFields with type id mapping
    /// // 1 -> DataType::UInt8
    /// // 3 -> DataType::Utf8
    /// UnionFields::new(
    ///     vec![1, 3],
    ///     vec![
    ///         Field::new("field1", DataType::UInt8, false),
    ///         Field::new("field3", DataType::Utf8, false),
    ///     ],
    /// );
    /// ```
    pub fn new<F, T>(type_ids: T, fields: F) -> Self
    where
        F: IntoIterator,
        F::Item: Into<FieldRef>,
        T: IntoIterator<Item = i8>,
    {
        let fields = fields.into_iter().map(Into::into);
        let mut set = 0_u128;
        type_ids
            .into_iter()
            .inspect(|&idx| {
                let mask = 1_u128 << idx;
                if (set & mask) != 0 {
                    panic!("duplicate type id: {idx}");
                } else {
                    set |= mask;
                }
            })
            .zip(fields)
            .collect()
    }

    /// Return size of this instance in bytes.
    pub fn size(&self) -> usize {
        self.iter()
            .map(|(_, field)| field.size() + std::mem::size_of::<(i8, FieldRef)>())
            .sum()
    }

    /// Returns the number of fields in this [`UnionFields`]
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Returns `true` if this is empty
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    /// Returns an iterator over the fields and type ids in this [`UnionFields`]
    pub fn iter(&self) -> impl Iterator<Item = (i8, &FieldRef)> + '_ {
        self.0.iter().map(|(id, f)| (*id, f))
    }

    /// Merge this field into self if it is compatible.
    ///
    /// See [`Field::try_merge`]
    pub(crate) fn try_merge(&mut self, other: &Self) -> Result<(), ArrowError> {
        // TODO: This currently may produce duplicate type IDs (#3982)
        let mut output: Vec<_> = self.iter().map(|(id, f)| (id, f.clone())).collect();
        for (field_type_id, from_field) in other.iter() {
            let mut is_new_field = true;
            for (self_type_id, self_field) in output.iter_mut() {
                if from_field == self_field {
                    // If the nested fields in two unions are the same, they must have same
                    // type id.
                    if *self_type_id != field_type_id {
                        return Err(ArrowError::SchemaError(
                            format!("Fail to merge schema field '{}' because the self_type_id = {} does not equal field_type_id = {}",
                                    self_field.name(), self_type_id, field_type_id)
                        ));
                    }

                    is_new_field = false;
                    break;
                }
            }

            if is_new_field {
                output.push((field_type_id, from_field.clone()))
            }
        }
        *self = output.into_iter().collect();
        Ok(())
    }
}

impl FromIterator<(i8, FieldRef)> for UnionFields {
    fn from_iter<T: IntoIterator<Item = (i8, FieldRef)>>(iter: T) -> Self {
        // TODO: Should this validate type IDs are unique (#3982)
        Self(iter.into_iter().collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::UnionMode;

    #[test]
    fn test_filter() {
        let floats = Fields::from(vec![
            Field::new("a", DataType::Float32, false),
            Field::new("b", DataType::Float32, false),
        ]);
        let fields = Fields::from(vec![
            Field::new("a", DataType::Int32, true),
            Field::new("floats", DataType::Struct(floats.clone()), true),
            Field::new("b", DataType::Int16, true),
            Field::new(
                "c",
                DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
                false,
            ),
            Field::new(
                "d",
                DataType::Dictionary(
                    Box::new(DataType::Int32),
                    Box::new(DataType::Struct(floats.clone())),
                ),
                false,
            ),
            Field::new_list(
                "e",
                Field::new("floats", DataType::Struct(floats.clone()), true),
                true,
            ),
            Field::new_fixed_size_list(
                "f",
                Field::new_list_field(DataType::Int32, false),
                3,
                false,
            ),
            Field::new_map(
                "g",
                "entries",
                Field::new("keys", DataType::LargeUtf8, false),
                Field::new("values", DataType::Int32, true),
                false,
                false,
            ),
            Field::new(
                "h",
                DataType::Union(
                    UnionFields::new(
                        vec![1, 3],
                        vec![
                            Field::new("field1", DataType::UInt8, false),
                            Field::new("field3", DataType::Utf8, false),
                        ],
                    ),
                    UnionMode::Dense,
                ),
                true,
            ),
            Field::new(
                "i",
                DataType::RunEndEncoded(
                    Arc::new(Field::new("run_ends", DataType::Int32, false)),
                    Arc::new(Field::new("values", DataType::Struct(floats.clone()), true)),
                ),
                false,
            ),
        ]);

        let floats_a = DataType::Struct(vec![floats[0].clone()].into());

        let r = fields.filter_leaves(|idx, _| idx == 0 || idx == 1);
        assert_eq!(r.len(), 2);
        assert_eq!(r[0], fields[0]);
        assert_eq!(r[1].data_type(), &floats_a);

        let r = fields.filter_leaves(|_, f| f.name() == "a");
        assert_eq!(r.len(), 5);
        assert_eq!(r[0], fields[0]);
        assert_eq!(r[1].data_type(), &floats_a);
        assert_eq!(
            r[2].data_type(),
            &DataType::Dictionary(Box::new(DataType::Int32), Box::new(floats_a.clone()))
        );
        assert_eq!(
            r[3].as_ref(),
            &Field::new_list("e", Field::new("floats", floats_a.clone(), true), true)
        );
        assert_eq!(
            r[4].as_ref(),
            &Field::new(
                "i",
                DataType::RunEndEncoded(
                    Arc::new(Field::new("run_ends", DataType::Int32, false)),
                    Arc::new(Field::new("values", floats_a.clone(), true)),
                ),
                false,
            )
        );

        let r = fields.filter_leaves(|_, f| f.name() == "floats");
        assert_eq!(r.len(), 0);

        let r = fields.filter_leaves(|idx, _| idx == 9);
        assert_eq!(r.len(), 1);
        assert_eq!(r[0], fields[6]);

        let r = fields.filter_leaves(|idx, _| idx == 10 || idx == 11);
        assert_eq!(r.len(), 1);
        assert_eq!(r[0], fields[7]);

        let union = DataType::Union(
            UnionFields::new(vec![1], vec![Field::new("field1", DataType::UInt8, false)]),
            UnionMode::Dense,
        );

        let r = fields.filter_leaves(|idx, _| idx == 12);
        assert_eq!(r.len(), 1);
        assert_eq!(r[0].data_type(), &union);

        let r = fields.filter_leaves(|idx, _| idx == 14 || idx == 15);
        assert_eq!(r.len(), 1);
        assert_eq!(r[0], fields[9]);

        // Propagate error
        let r = fields.try_filter_leaves(|_, _| Err(ArrowError::SchemaError("error".to_string())));
        assert!(r.is_err());
    }
}

Coverage Report

Created: 2025-08-26 07:03

Line	Count	Source
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
18		use std::ops::Deref;
19		use std::sync::Arc;
20
21		use 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))]
60		pub struct Fields(Arc<[FieldRef]>);
61
62		impl std::fmt::Debug for Fields {
63	0	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64	0	self.0.as_ref().fmt(f)
65	0	}
66		}
67
68		impl Fields {
69		/// Returns a new empty [`Fields`]
70	0	pub fn empty() -> Self {
71	0	Self(Arc::new([]))
72	0	}
73
74		/// Return size of this instance in bytes.
75	0	pub fn size(&self) -> usize {
76	0	self.iter()
77	0	.map(\|field\| field.size() + std::mem::size_of::<FieldRef>())
78	0	.sum()
79	0	}
80
81		/// Searches for a field by name, returning it along with its index if found
82	0	pub fn find(&self, name: &str) -> Option<(usize, &FieldRef)> {
83	0	self.0.iter().enumerate().find(\|(_, b)\| b.name() == name)
84	0	}
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	0	pub fn contains(&self, other: &Fields) -> bool {
93	0	if Arc::ptr_eq(&self.0, &other.0) {
94	0	return true;
95	0	}
96	0	self.len() == other.len()
97	0	&& self
98	0	.iter()
99	0	.zip(other.iter())
100	0	.all(\|(a, b)\| Arc::ptr_eq(a, b) \|\| a.contains(b))
101	0	}
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	0	pub fn filter_leaves<F: FnMut(usize, &FieldRef) -> bool>(&self, mut filter: F) -> Self {
140	0	self.try_filter_leaves(\|idx, field\| Ok(filter(idx, field)))
141	0	.unwrap()
142	0	}
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	0	pub fn try_filter_leaves<F: FnMut(usize, &FieldRef) -> Result<bool, ArrowError>>(
149	0	&self,
150	0	mut filter: F,
151	0	) -> Result<Self, ArrowError> {
152	0	fn filter_field<F: FnMut(&FieldRef) -> Result<bool, ArrowError>>(
153	0	f: &FieldRef,
154	0	filter: &mut F,
155	0	) -> Result<Option<FieldRef>, ArrowError> {
156		use DataType::*;
157
158	0	let v = match f.data_type() {
159	0	Dictionary(_, v) => v.as_ref(), // Key must be integer
160	0	RunEndEncoded(_, v) => v.data_type(), // Run-ends must be integer
161	0	d => d,
162		};
163	0	let d = match v {
164	0	List(child) => {
165	0	let fields = filter_field(child, filter)?;
166	0	if let Some(fields) = fields {
167	0	List(fields)
168		} else {
169	0	return Ok(None);
170		}
171		}
172	0	LargeList(child) => {
173	0	let fields = filter_field(child, filter)?;
174	0	if let Some(fields) = fields {
175	0	LargeList(fields)
176		} else {
177	0	return Ok(None);
178		}
179		}
180	0	Map(child, ordered) => {
181	0	let fields = filter_field(child, filter)?;
182	0	if let Some(fields) = fields {
183	0	Map(fields, *ordered)
184		} else {
185	0	return Ok(None);
186		}
187		}
188	0	FixedSizeList(child, size) => {
189	0	let fields = filter_field(child, filter)?;
190	0	if let Some(fields) = fields {
191	0	FixedSizeList(fields, *size)
192		} else {
193	0	return Ok(None);
194		}
195		}
196	0	Struct(fields) => {
197	0	let filtered: Result<Vec<_>, _> =
198	0	fields.iter().map(\|f\| filter_field(f, filter)).collect();
199	0	let filtered: Fields = filtered?
200	0	.iter()
201	0	.filter_map(\|f\| f.as_ref().cloned())
202	0	.collect();
203
204	0	if filtered.is_empty() {
205	0	return Ok(None);
206	0	}
207
208	0	Struct(filtered)
209		}
210	0	Union(fields, mode) => {
211	0	let filtered: Result<Vec<_>, _> = fields
212	0	.iter()
213	0	.map(\|(id, f)\| filter_field(f, filter).map(\|f\| f.map(\|f\| (id, f))))
214	0	.collect();
215	0	let filtered: UnionFields = filtered?
216	0	.iter()
217	0	.filter_map(\|f\| f.as_ref().cloned())
218	0	.collect();
219
220	0	if filtered.is_empty() {
221	0	return Ok(None);
222	0	}
223
224	0	Union(filtered, *mode)
225		}
226		_ => {
227	0	let filtered = filter(f)?;
228	0	return Ok(filtered.then(\|\| f.clone()));
229		}
230		};
231	0	let d = match f.data_type() {
232	0	Dictionary(k, _) => Dictionary(k.clone(), Box::new(d)),
233	0	RunEndEncoded(v, f) => {
234	0	RunEndEncoded(v.clone(), Arc::new(f.as_ref().clone().with_data_type(d)))
235		}
236	0	_ => d,
237		};
238	0	Ok(Some(Arc::new(f.as_ref().clone().with_data_type(d))))
239	0	}
240
241	0	let mut leaf_idx = 0;
242	0	let mut filter = \|f: &FieldRef\| {
243	0	let t = filter(leaf_idx, f)?;
244	0	leaf_idx += 1;
245	0	Ok(t)
246	0	};
247
248	0	let filtered: Result<Vec<_>, _> = self
249	0	.0
250	0	.iter()
251	0	.map(\|f\| filter_field(f, &mut filter))
252	0	.collect();
253	0	let filtered = filtered?
254	0	.iter()
255	0	.filter_map(\|f\| f.as_ref().cloned())
256	0	.collect();
257	0	Ok(filtered)
258	0	}
259		}
260
261		impl Default for Fields {
262	0	fn default() -> Self {
263	0	Self::empty()
264	0	}
265		}
266
267		impl FromIterator<Field> for Fields {
268	291	fn from_iter<T: IntoIterator<Item = Field>>(iter: T) -> Self {
269	291	iter.into_iter().map(Arc::new).collect()
270	291	}
271		}
272
273		impl FromIterator<FieldRef> for Fields {
274	291	fn from_iter<T: IntoIterator<Item = FieldRef>>(iter: T) -> Self {
275	291	Self(iter.into_iter().collect())
276	291	}
277		}
278
279		impl From<Vec<Field>> for Fields {
280	195	fn from(value: Vec<Field>) -> Self {
281	195	value.into_iter().collect()
282	195	}
283		}
284
285		impl From<Vec<FieldRef>> for Fields {
286	102	fn from(value: Vec<FieldRef>) -> Self {
287	102	Self(value.into())
288	102	}
289		}
290
291		impl From<&[FieldRef]> for Fields {
292	0	fn from(value: &[FieldRef]) -> Self {
293	0	Self(value.into())
294	0	}
295		}
296
297		impl<const N: usize> From<[FieldRef; N]> for Fields {
298	0	fn from(value: [FieldRef; N]) -> Self {
299	0	Self(Arc::new(value))
300	0	}
301		}
302
303		impl Deref for Fields {
304		type Target = [FieldRef];
305
306	993	fn deref(&self) -> &Self::Target {
307	993	self.0.as_ref()
308	993	}
309		}
310
311		impl<'a> IntoIterator for &'a Fields {
312		type Item = &'a FieldRef;
313		type IntoIter = std::slice::Iter<'a, FieldRef>;
314
315	248	fn into_iter(self) -> Self::IntoIter {
316	248	self.0.iter()
317	248	}
318		}
319
320		/// A cheaply cloneable, owned collection of [`FieldRef`] and their corresponding type ids
321		#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
322		#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
323		#[cfg_attr(feature = "serde", serde(transparent))]
324		pub struct UnionFields(Arc<[(i8, FieldRef)]>);
325
326		impl std::fmt::Debug for UnionFields {
327	0	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
328	0	self.0.as_ref().fmt(f)
329	0	}
330		}
331
332		impl UnionFields {
333		/// Create a new [`UnionFields`] with no fields
334	0	pub fn empty() -> Self {
335	0	Self(Arc::from([]))
336	0	}
337
338		/// Create a new [`UnionFields`] from a [`Fields`] and array of type_ids
339		///
340		/// See <https://arrow.apache.org/docs/format/Columnar.html#union-layout>
341		///
342		/// ```
343		/// use arrow_schema::{DataType, Field, UnionFields};
344		/// // Create a new UnionFields with type id mapping
345		/// // 1 -> DataType::UInt8
346		/// // 3 -> DataType::Utf8
347		/// UnionFields::new(
348		/// vec![1, 3],
349		/// vec![
350		/// Field::new("field1", DataType::UInt8, false),
351		/// Field::new("field3", DataType::Utf8, false),
352		/// ],
353		/// );
354		/// ```
355	0	pub fn new<F, T>(type_ids: T, fields: F) -> Self
356	0	where
357	0	F: IntoIterator,
358	0	F::Item: Into<FieldRef>,
359	0	T: IntoIterator<Item = i8>,
360		{
361	0	let fields = fields.into_iter().map(Into::into);
362	0	let mut set = 0_u128;
363	0	type_ids
364	0	.into_iter()
365	0	.inspect(\|&idx\| {
366	0	let mask = 1_u128 << idx;
367	0	if (set & mask) != 0 {
368	0	panic!("duplicate type id: {idx}");
369	0	} else {
370	0	set \|= mask;
371	0	}
372	0	})
373	0	.zip(fields)
374	0	.collect()
375	0	}
376
377		/// Return size of this instance in bytes.
378	0	pub fn size(&self) -> usize {
379	0	self.iter()
380	0	.map(\|(_, field)\| field.size() + std::mem::size_of::<(i8, FieldRef)>())
381	0	.sum()
382	0	}
383
384		/// Returns the number of fields in this [`UnionFields`]
385	0	pub fn len(&self) -> usize {
386	0	self.0.len()
387	0	}
388
389		/// Returns `true` if this is empty
390	0	pub fn is_empty(&self) -> bool {
391	0	self.0.is_empty()
392	0	}
393
394		/// Returns an iterator over the fields and type ids in this [`UnionFields`]
395	0	pub fn iter(&self) -> impl Iterator<Item = (i8, &FieldRef)> + '_ {
396	0	self.0.iter().map(\|(id, f)\| (*id, f))
397	0	}
398
399		/// Merge this field into self if it is compatible.
400		///
401		/// See [`Field::try_merge`]
402	0	pub(crate) fn try_merge(&mut self, other: &Self) -> Result<(), ArrowError> {
403		// TODO: This currently may produce duplicate type IDs (#3982)
404	0	let mut output: Vec<_> = self.iter().map(\|(id, f)\| (id, f.clone())).collect();
405	0	for (field_type_id, from_field) in other.iter() {
406	0	let mut is_new_field = true;
407	0	for (self_type_id, self_field) in output.iter_mut() {
408	0	if from_field == self_field {
409		// If the nested fields in two unions are the same, they must have same
410		// type id.
411	0	if *self_type_id != field_type_id {
412	0	return Err(ArrowError::SchemaError(
413	0	format!("Fail to merge schema field '{}' because the self_type_id = {} does not equal field_type_id = {}",
414	0	self_field.name(), self_type_id, field_type_id)
415	0	));
416	0	}
417
418	0	is_new_field = false;
419	0	break;
420	0	}
421		}
422
423	0	if is_new_field {
424	0	output.push((field_type_id, from_field.clone()))
425	0	}
426		}
427	0	*self = output.into_iter().collect();
428	0	Ok(())
429	0	}
430		}
431
432		impl FromIterator<(i8, FieldRef)> for UnionFields {
433	1	fn from_iter<T: IntoIterator<Item = (i8, FieldRef)>>(iter: T) -> Self {
434		// TODO: Should this validate type IDs are unique (#3982)
435	1	Self(iter.into_iter().collect())
436	1	}
437		}
438
439		#[cfg(test)]
440		mod tests {
441		use super::*;
442		use crate::UnionMode;
443
444		#[test]
445		fn test_filter() {
446		let floats = Fields::from(vec![
447		Field::new("a", DataType::Float32, false),
448		Field::new("b", DataType::Float32, false),
449		]);
450		let fields = Fields::from(vec![
451		Field::new("a", DataType::Int32, true),
452		Field::new("floats", DataType::Struct(floats.clone()), true),
453		Field::new("b", DataType::Int16, true),
454		Field::new(
455		"c",
456		DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
457		false,
458		),
459		Field::new(
460		"d",
461		DataType::Dictionary(
462		Box::new(DataType::Int32),
463		Box::new(DataType::Struct(floats.clone())),
464		),
465		false,
466		),
467		Field::new_list(
468		"e",
469		Field::new("floats", DataType::Struct(floats.clone()), true),
470		true,
471		),
472		Field::new_fixed_size_list(
473		"f",
474		Field::new_list_field(DataType::Int32, false),
475		3,
476		false,
477		),
478		Field::new_map(
479		"g",
480		"entries",
481		Field::new("keys", DataType::LargeUtf8, false),
482		Field::new("values", DataType::Int32, true),
483		false,
484		false,
485		),
486		Field::new(
487		"h",
488		DataType::Union(
489		UnionFields::new(
490		vec![1, 3],
491		vec![
492		Field::new("field1", DataType::UInt8, false),
493		Field::new("field3", DataType::Utf8, false),
494		],
495		),
496		UnionMode::Dense,
497		),
498		true,
499		),
500		Field::new(
501		"i",
502		DataType::RunEndEncoded(
503		Arc::new(Field::new("run_ends", DataType::Int32, false)),
504		Arc::new(Field::new("values", DataType::Struct(floats.clone()), true)),
505		),
506		false,
507		),
508		]);
509
510		let floats_a = DataType::Struct(vec![floats[0].clone()].into());
511
512		let r = fields.filter_leaves(\|idx, _\| idx == 0 \|\| idx == 1);
513		assert_eq!(r.len(), 2);
514		assert_eq!(r[0], fields[0]);
515		assert_eq!(r[1].data_type(), &floats_a);
516
517		let r = fields.filter_leaves(\|_, f\| f.name() == "a");
518		assert_eq!(r.len(), 5);
519		assert_eq!(r[0], fields[0]);
520		assert_eq!(r[1].data_type(), &floats_a);
521		assert_eq!(
522		r[2].data_type(),
523		&DataType::Dictionary(Box::new(DataType::Int32), Box::new(floats_a.clone()))
524		);
525		assert_eq!(
526		r[3].as_ref(),
527		&Field::new_list("e", Field::new("floats", floats_a.clone(), true), true)
528		);
529		assert_eq!(
530		r[4].as_ref(),
531		&Field::new(
532		"i",
533		DataType::RunEndEncoded(
534		Arc::new(Field::new("run_ends", DataType::Int32, false)),
535		Arc::new(Field::new("values", floats_a.clone(), true)),
536		),
537		false,
538		)
539		);
540
541		let r = fields.filter_leaves(\|_, f\| f.name() == "floats");
542		assert_eq!(r.len(), 0);
543
544		let r = fields.filter_leaves(\|idx, _\| idx == 9);
545		assert_eq!(r.len(), 1);
546		assert_eq!(r[0], fields[6]);
547
548		let r = fields.filter_leaves(\|idx, _\| idx == 10 \|\| idx == 11);
549		assert_eq!(r.len(), 1);
550		assert_eq!(r[0], fields[7]);
551
552		let union = DataType::Union(
553		UnionFields::new(vec![1], vec![Field::new("field1", DataType::UInt8, false)]),
554		UnionMode::Dense,
555		);
556
557		let r = fields.filter_leaves(\|idx, _\| idx == 12);
558		assert_eq!(r.len(), 1);
559		assert_eq!(r[0].data_type(), &union);
560
561		let r = fields.filter_leaves(\|idx, _\| idx == 14 \|\| idx == 15);
562		assert_eq!(r.len(), 1);
563		assert_eq!(r[0], fields[9]);
564
565		// Propagate error
566		let r = fields.try_filter_leaves(\|_, _\| Err(ArrowError::SchemaError("error".to_string())));
567		assert!(r.is_err());
568		}
569		}