polars_core/schema/
iceberg.rs

1//! TODO
2//!
3//! This should ideally be moved to `polars-schema`, currently it cannot due to dependency on
4//! `polars_core::DataType`.
5use std::borrow::Cow;
6use std::sync::Arc;
7
8use polars_arrow::datatypes::{ArrowDataType, ArrowSchema, Field as ArrowField};
9use polars_error::{PolarsResult, feature_gated, polars_bail, polars_err};
10use polars_utils::aliases::InitHashMaps;
11use polars_utils::pl_str::PlSmallStr;
12
13use crate::prelude::{DataType, Field, PlIndexMap};
14
15pub const LIST_ELEMENT_DEFAULT_ID: u32 = u32::MAX;
16
17/// Maps Iceberg physical IDs to columns.
18///
19/// Note: This doesn't use `Schema<F>` as the keys are u32's.
20#[derive(Debug, Clone, Eq, PartialEq)]
21#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
22#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
23pub struct IcebergSchema(PlIndexMap<u32, IcebergColumn>);
24pub type IcebergSchemaRef = Arc<IcebergSchema>;
25
26impl IcebergSchema {
27    /// Constructs a schema keyed by the physical ID stored in the arrow field metadata.
28    pub fn from_arrow_schema(schema: &ArrowSchema) -> PolarsResult<Self> {
29        Self::try_from_arrow_fields_iter(schema.iter_values())
30    }
31
32    pub fn try_from_arrow_fields_iter<'a, I>(iter: I) -> PolarsResult<Self>
33    where
34        I: IntoIterator<Item = &'a ArrowField>,
35    {
36        let iter = iter.into_iter();
37
38        let mut out = PlIndexMap::with_capacity(iter.size_hint().0);
39
40        for arrow_field in iter {
41            let col: IcebergColumn = arrow_field_to_iceberg_column_rec(arrow_field, None)?;
42            let existing = out.insert(col.physical_id, col);
43
44            if let Some(existing) = existing {
45                polars_bail!(
46                    Duplicate:
47                    "IcebergSchema: duplicate physical ID {:?}",
48                    existing,
49                )
50            }
51        }
52
53        Ok(Self(out))
54    }
55}
56
57#[derive(Debug, Clone, Eq, Hash, PartialEq)]
58#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
59#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
60pub struct IcebergColumn {
61    /// Output name
62    pub name: PlSmallStr,
63    /// This is expected to map from 'PARQUET:field_id'
64    pub physical_id: u32,
65    pub type_: IcebergColumnType,
66}
67
68#[derive(Debug, Clone, Eq, Hash, PartialEq)]
69#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
70#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
71pub enum IcebergColumnType {
72    Primitive {
73        /// This must not be a nested data type.
74        dtype: DataType,
75    },
76    List(Box<IcebergColumn>),
77    /// (values, width)
78    FixedSizeList(Box<IcebergColumn>, usize),
79    /// (keys, values)
80    Map(Box<IcebergColumn>, Box<IcebergColumn>),
81    Struct(IcebergSchema),
82}
83
84impl IcebergColumnType {
85    pub fn to_polars_dtype(&self) -> DataType {
86        use IcebergColumnType::*;
87
88        match self {
89            Primitive { dtype } => dtype.clone(),
90            List(inner) => DataType::List(Box::new(inner.type_.to_polars_dtype())),
91            FixedSizeList(inner, width) => {
92                feature_gated!("dtype-array", {
93                    DataType::Array(Box::new(inner.type_.to_polars_dtype()), *width)
94                })
95            },
96            Map(key, value) => feature_gated!("dtype-map", {
97                DataType::Map(
98                    Box::new(key.type_.to_polars_dtype()),
99                    Box::new(value.type_.to_polars_dtype()),
100                )
101            }),
102            Struct(fields) => feature_gated!("dtype-struct", {
103                DataType::Struct(
104                    fields
105                        .values()
106                        .map(|col| Field::new(col.name.clone(), col.type_.to_polars_dtype()))
107                        .collect(),
108                )
109            }),
110        }
111    }
112
113    pub fn is_nested(&self) -> bool {
114        use IcebergColumnType::*;
115
116        match self {
117            List(_) | FixedSizeList(..) | Map(..) | Struct(_) => true,
118            Primitive { .. } => false,
119        }
120    }
121}
122
123fn arrow_field_to_iceberg_column_rec(
124    field: &ArrowField,
125    field_id_override: Option<u32>,
126) -> PolarsResult<IcebergColumn> {
127    const PARQUET_FIELD_ID_KEY: &str = "PARQUET:field_id";
128
129    let physical_id: u32 = field_id_override.ok_or(Cow::Borrowed("")).or_else(|_| {
130        field
131            .metadata
132            .as_deref()
133            .ok_or(Cow::Borrowed("metadata was None"))
134            .and_then(|md| {
135                md.get(PARQUET_FIELD_ID_KEY)
136                    .ok_or(Cow::Borrowed("key not found in metadata"))
137            })
138            .and_then(|x| {
139                x.parse()
140                    .map_err(|_| Cow::Owned(format!("could not parse value as u32: '{x}'")))
141            })
142            .map_err(|failed_reason: Cow<'_, str>| {
143                polars_err!(
144                    SchemaFieldNotFound:
145                    "IcebergSchema: failed to load '{PARQUET_FIELD_ID_KEY}' for field {}: {}",
146                    &field.name,
147                    failed_reason,
148                )
149            })
150    })?;
151
152    // Prevent accidental re-use.
153    #[expect(unused)]
154    let field_id_override: ();
155
156    use ArrowDataType as ADT;
157
158    let name = field.name.clone();
159
160    let list_column_type = |field: &ArrowField| -> PolarsResult<IcebergColumnType> {
161        let field_id_override = field
162            .metadata
163            .as_ref()
164            .is_none_or(|x| !x.contains_key(PARQUET_FIELD_ID_KEY))
165            .then_some(LIST_ELEMENT_DEFAULT_ID);
166
167        Ok(IcebergColumnType::List(Box::new(
168            arrow_field_to_iceberg_column_rec(field, field_id_override)?,
169        )))
170    };
171
172    let type_ = match &field.dtype {
173        ADT::List(field) | ADT::LargeList(field) => list_column_type(field)?,
174
175        ADT::Map(entries, _) => {
176            #[cfg(feature = "dtype-map")]
177            {
178                // The `entries` field itself carries no physical ID - the IDs sit on the `key` and
179                // `value` fields underneath it.
180                let ADT::Struct(entry_fields) = &entries.dtype else {
181                    polars_bail!(
182                        ComputeError:
183                        "IcebergSchema: expected struct under arrow map type, got: {:?}",
184                        &entries.dtype,
185                    )
186                };
187
188                let [key, value] = entry_fields.as_slice() else {
189                    polars_bail!(
190                        ComputeError:
191                        "IcebergSchema: expected 2 fields under arrow map type, got: {}",
192                        entry_fields.len(),
193                    )
194                };
195
196                IcebergColumnType::Map(
197                    Box::new(arrow_field_to_iceberg_column_rec(key, None)?),
198                    Box::new(arrow_field_to_iceberg_column_rec(value, None)?),
199                )
200            }
201            #[cfg(not(feature = "dtype-map"))]
202            {
203                // Without the `Map` dtype these columns load as `List(Struct { key, value })`.
204                list_column_type(entries)?
205            }
206        },
207
208        #[cfg(feature = "dtype-array")]
209        ADT::FixedSizeList(field, width) => IcebergColumnType::FixedSizeList(
210            Box::new(arrow_field_to_iceberg_column_rec(field, None)?),
211            *width,
212        ),
213
214        #[cfg(feature = "dtype-struct")]
215        ADT::Struct(fields) => {
216            IcebergColumnType::Struct(IcebergSchema::try_from_arrow_fields_iter(fields)?)
217        },
218
219        dtype => {
220            if let ADT::Dictionary(_key_type, value_type, _is_ordered) = dtype
221                && !value_type.is_nested()
222            {
223                let mut new_field = ArrowField::new(name.clone(), dtype.clone(), true);
224                if let Some(metadata) = field.metadata.as_ref() {
225                    new_field = new_field.with_metadata((**metadata).clone());
226                }
227                let dtype = DataType::from_arrow_field(&new_field);
228
229                IcebergColumnType::Primitive { dtype }
230            } else if let ADT::Extension(ext_type) = dtype
231                && let DataType::Binary = DataType::from_arrow_dtype(&ext_type.inner)
232            {
233                // Iceberg UUID type will hit this branch.
234                IcebergColumnType::Primitive {
235                    dtype: DataType::Binary,
236                }
237            } else if dtype.is_nested() {
238                polars_bail!(
239                    ComputeError:
240                    "IcebergSchema: unsupported arrow type: {:?}",
241                    dtype,
242                )
243            } else {
244                let dtype =
245                    DataType::from_arrow_field(&ArrowField::new(name.clone(), dtype.clone(), true));
246
247                IcebergColumnType::Primitive { dtype }
248            }
249        },
250    };
251
252    let out = IcebergColumn {
253        name,
254        physical_id,
255        type_,
256    };
257
258    Ok(out)
259}
260
261impl<T> FromIterator<T> for IcebergSchema
262where
263    PlIndexMap<u32, IcebergColumn>: FromIterator<T>,
264{
265    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
266        Self(PlIndexMap::<u32, IcebergColumn>::from_iter(iter))
267    }
268}
269
270impl std::hash::Hash for IcebergSchema {
271    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
272        for col in self.values() {
273            col.hash(state);
274        }
275    }
276}
277
278impl std::ops::Deref for IcebergSchema {
279    type Target = PlIndexMap<u32, IcebergColumn>;
280
281    fn deref(&self) -> &Self::Target {
282        &self.0
283    }
284}
285
286impl std::ops::DerefMut for IcebergSchema {
287    fn deref_mut(&mut self) -> &mut Self::Target {
288        &mut self.0
289    }
290}
polars_core/schema/iceberg.rs

polars_core/schema/
iceberg.rs
