polars_core/schema/
mod.rs

1use std::fmt::Debug;
2
3use polars_utils::pl_str::PlSmallStr;
4
5use crate::prelude::*;
6use crate::utils::try_get_supertype;
7
8pub mod iceberg;
9
10pub type SchemaRef = Arc<Schema>;
11pub type Schema = polars_schema::Schema<DataType, ()>;
12
13pub trait SchemaExt {
14    fn from_arrow_schema(value: &ArrowSchema) -> Self;
15
16    fn get_field(&self, name: &str) -> Option<Field>;
17
18    fn try_get_field(&self, name: &str) -> PolarsResult<Field>;
19
20    fn to_arrow(&self, compat_level: CompatLevel) -> ArrowSchema;
21
22    fn iter_fields(&self) -> impl ExactSizeIterator<Item = Field> + '_;
23
24    fn to_supertype(&mut self, other: &Schema) -> PolarsResult<bool>;
25
26    fn contains_dtype(&self, dtype: &DataType, recursive: bool) -> bool;
27}
28
29impl SchemaExt for Schema {
30    fn from_arrow_schema(value: &ArrowSchema) -> Self {
31        value
32            .iter_values()
33            .map(|x| (x.name.clone(), DataType::from_arrow_field(x)))
34            .collect()
35    }
36
37    /// Look up the name in the schema and return an owned [`Field`] by cloning the data.
38    ///
39    /// Returns `None` if the field does not exist.
40    ///
41    /// This method constructs the `Field` by cloning the name and dtype. For a version that returns references, see
42    /// [`get`][Self::get] or [`get_full`][Self::get_full].
43    fn get_field(&self, name: &str) -> Option<Field> {
44        self.get_full(name)
45            .map(|(_, name, dtype)| Field::new(name.clone(), dtype.clone()))
46    }
47
48    /// Look up the name in the schema and return an owned [`Field`] by cloning the data.
49    ///
50    /// Returns `Err(PolarsErr)` if the field does not exist.
51    ///
52    /// This method constructs the `Field` by cloning the name and dtype. For a version that returns references, see
53    /// [`get`][Self::get] or [`get_full`][Self::get_full].
54    fn try_get_field(&self, name: &str) -> PolarsResult<Field> {
55        self.get_full(name)
56            .ok_or_else(|| polars_err!(SchemaFieldNotFound: "{name}"))
57            .map(|(_, name, dtype)| Field::new(name.clone(), dtype.clone()))
58    }
59
60    /// Convert self to `ArrowSchema` by cloning the fields.
61    fn to_arrow(&self, compat_level: CompatLevel) -> ArrowSchema {
62        self.iter()
63            .map(|(name, dtype)| {
64                (
65                    name.clone(),
66                    dtype.to_arrow_field(name.clone(), compat_level),
67                )
68            })
69            .collect()
70    }
71
72    /// Iterates the [`Field`]s in this schema, constructing them anew by cloning each `(&name, &dtype)` pair.
73    ///
74    /// Note that this clones each name and dtype in order to form an owned [`Field`]. For a clone-free version, use
75    /// [`iter`][Self::iter], which returns `(&name, &dtype)`.
76    fn iter_fields(&self) -> impl ExactSizeIterator<Item = Field> + '_ {
77        self.iter()
78            .map(|(name, dtype)| Field::new(name.clone(), dtype.clone()))
79    }
80
81    /// Take another [`Schema`] and try to find the supertypes between them.
82    fn to_supertype(&mut self, other: &Schema) -> PolarsResult<bool> {
83        polars_ensure!(self.len() == other.len(), ComputeError: "schema lengths differ");
84
85        let mut changed = false;
86        for ((k, dt), (other_k, other_dt)) in self.iter_mut().zip(other.iter()) {
87            polars_ensure!(k == other_k, ComputeError: "schema names differ: got {}, expected {}", k, other_k);
88
89            let st = try_get_supertype(dt, other_dt)?;
90            changed |= (&st != dt) || (&st != other_dt);
91            *dt = st
92        }
93        Ok(changed)
94    }
95
96    fn contains_dtype(&self, dtype: &DataType, recursive: bool) -> bool {
97        if !recursive {
98            self.iter_values().any(|dt| dt == dtype)
99        } else {
100            self.iter_values()
101                .any(|dt| dt.contains_dtype_recursive(dtype))
102        }
103    }
104}
105
106pub trait SchemaNamesAndDtypes {
107    const IS_ARROW: bool;
108    type DataType: Debug + Clone + PartialEq;
109
110    fn iter_names_and_dtypes(
111        &self,
112    ) -> impl ExactSizeIterator<Item = (&PlSmallStr, &Self::DataType)>;
113}
114
115impl SchemaNamesAndDtypes for ArrowSchema {
116    const IS_ARROW: bool = true;
117    type DataType = ArrowDataType;
118
119    fn iter_names_and_dtypes(
120        &self,
121    ) -> impl ExactSizeIterator<Item = (&PlSmallStr, &Self::DataType)> {
122        self.iter_values().map(|x| (&x.name, &x.dtype))
123    }
124}
125
126impl SchemaNamesAndDtypes for Schema {
127    const IS_ARROW: bool = false;
128    type DataType = DataType;
129
130    fn iter_names_and_dtypes(
131        &self,
132    ) -> impl ExactSizeIterator<Item = (&PlSmallStr, &Self::DataType)> {
133        self.iter()
134    }
135}
polars_core/schema/mod.rs

polars_core/schema/
mod.rs
