polars_rows_iter/iter_from_column/
iter_from_column_str.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
use super::*;
use iter_from_column::IterFromColumn;
use polars::prelude::*;
use value_functions::*;

impl<'a> IterFromColumn<'a> for &'a str {
    fn create_iter(
        dataframe: &'a DataFrame,
        column_name: &'a str,
    ) -> PolarsResult<Box<dyn Iterator<Item = PolarsResult<Self>> + 'a>> {
        let iter = create_iter(dataframe, column_name)?;
        let iter = Box::new(iter.map(mandatory_value));
        Ok(iter)
    }
}

impl<'a> IterFromColumn<'a> for Option<&'a str> {
    fn create_iter(
        dataframe: &'a DataFrame,
        column_name: &'a str,
    ) -> PolarsResult<Box<dyn Iterator<Item = PolarsResult<Self>> + 'a>> {
        let iter = create_iter(dataframe, column_name)?;
        let iter = Box::new(iter.map(optional_value));
        Ok(iter)
    }
}

struct PolarsIteratorStrWrapper<'a> {
    inner: Box<dyn PolarsIterator<Item = Option<&'a str>> + 'a>,
}

impl<'a> Iterator for PolarsIteratorStrWrapper<'a> {
    type Item = Option<&'a str>;

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next()
    }
}

fn create_str_iter<'a>(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<&'a str>> + 'a>> {
    let inner = column.str()?.into_iter();
    Ok(Box::new(PolarsIteratorStrWrapper { inner }))
}

#[cfg(feature = "dtype-categorical")]
struct CategoricalIteratorWrapper<'a> {
    inner: CatIter<'a>,
}

#[cfg(feature = "dtype-categorical")]
impl<'a> Iterator for CategoricalIteratorWrapper<'a> {
    type Item = Option<&'a str>;

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next()
    }
}

#[cfg(feature = "dtype-categorical")]
fn create_cat_iter<'a>(column: &'a Column) -> PolarsResult<Box<dyn Iterator<Item = Option<&'a str>> + 'a>> {
    let inner = column.categorical()?.iter_str();
    Ok(Box::new(CategoricalIteratorWrapper { inner }))
}

fn create_iter<'a>(
    dataframe: &'a DataFrame,
    column_name: &'a str,
) -> PolarsResult<Box<dyn Iterator<Item = Option<&'a str>> + 'a>> {
    let column = dataframe.column(column_name)?;

    let iter = match column.dtype() {
        DataType::String => create_str_iter(column)?,
        #[cfg(feature = "dtype-categorical")]
        DataType::Categorical(_, _) => create_cat_iter(column)?,
        dtype => {
            return Err(polars_err!(SchemaMismatch: "Cannot get &str from column '{column_name}' with dtype : {dtype}"))
        }
    };

    Ok(iter)
}