bookdata/arrow/
dfext.rs

1//! Book data extensions to Polars.
2use log::*;
3use polars::prelude::*;
4
5use crate::cleaning::names::clean_name;
6
7pub fn udf_clean_name(col: Series) -> PolarsResult<Option<Series>> {
8    let col = col.str()?;
9    let res: StringChunked = col
10        .into_iter()
11        .map(|n| {
12            if let Some(s) = n {
13                Some(clean_name(s))
14            } else {
15                None
16            }
17        })
18        .collect();
19    Ok(Some(res.into_series()))
20}
21
22/// Get a schema from a data frame with maximal nullability.
23pub fn nonnull_schema(df: &DataFrame) -> ArrowSchema {
24    let schema = df.schema().to_arrow(false);
25    debug!("de-nullifying schema: {:?}", schema);
26    let clean = ArrowSchema {
27        fields: schema
28            .fields
29            .into_iter()
30            .map(|f| ArrowField {
31                is_nullable: df.column(&f.name).expect("missing column").null_count() > 0,
32                ..f
33            })
34            .collect(),
35        metadata: schema.metadata,
36    };
37    debug!("converted schema: {:?}", clean);
38    clean
39}