polars_plan/dsl/
struct_.rs

1use super::*;
2use crate::plans::conversion::is_regex_projection;
3
4/// Specialized expressions for Struct dtypes.
5pub struct StructNameSpace(pub(crate) Expr);
6
7impl StructNameSpace {
8    pub fn field_by_index(self, index: i64) -> Expr {
9        self.0
10            .map_unary(FunctionExpr::StructExpr(StructFunction::FieldByIndex(
11                index,
12            )))
13    }
14
15    /// Retrieve one or multiple of the fields of this [`StructChunked`] as a new Series.
16    /// This expression also expands the `"*"` wildcard column.
17    pub fn field_by_names<I, S>(self, names: I) -> Expr
18    where
19        I: IntoIterator<Item = S>,
20        S: Into<PlSmallStr>,
21    {
22        self.field_by_names_impl(names.into_iter().map(|x| x.into()).collect())
23    }
24
25    fn field_by_names_impl(self, names: Arc<[PlSmallStr]>) -> Expr {
26        self.0
27            .map_unary(FunctionExpr::StructExpr(StructFunction::MultipleFields(
28                names,
29            )))
30    }
31
32    /// Retrieve one of the fields of this [`StructChunked`] as a new Series.
33    /// This expression also supports wildcard "*" and regex expansion.
34    pub fn field_by_name(self, name: &str) -> Expr {
35        if name == "*" || is_regex_projection(name) {
36            return self.field_by_names([name]);
37        }
38        self.0
39            .map_unary(FunctionExpr::StructExpr(StructFunction::FieldByName(
40                name.into(),
41            )))
42    }
43
44    /// Rename the fields of the [`StructChunked`].
45    pub fn rename_fields<I, S>(self, names: I) -> Expr
46    where
47        I: IntoIterator<Item = S>,
48        S: Into<PlSmallStr>,
49    {
50        self._rename_fields_impl(names.into_iter().map(|x| x.into()).collect())
51    }
52
53    pub fn _rename_fields_impl(self, names: Arc<[PlSmallStr]>) -> Expr {
54        self.0
55            .map_unary(FunctionExpr::StructExpr(StructFunction::RenameFields(
56                names,
57            )))
58    }
59
60    #[cfg(feature = "json")]
61    pub fn json_encode(self) -> Expr {
62        self.0
63            .map_unary(FunctionExpr::StructExpr(StructFunction::JsonEncode))
64    }
65
66    pub fn with_fields(self, fields: Vec<Expr>) -> PolarsResult<Expr> {
67        fn materialize_field(this: &Expr, field: Expr) -> PolarsResult<Expr> {
68            field.try_map_expr(|e| match e {
69                Expr::Field(names) => {
70                    let this = this.clone().struct_();
71                    Ok(if names.len() == 1 {
72                        this.field_by_name(names[0].as_ref())
73                    } else {
74                        this.field_by_names_impl(names)
75                    })
76                },
77                Expr::Exclude(_, _) => {
78                    polars_bail!(InvalidOperation: "'exclude' not allowed in 'field'")
79                },
80                _ => Ok(e),
81            })
82        }
83
84        let s = self.0.clone();
85        self.0.try_map_n_ary(
86            FunctionExpr::StructExpr(StructFunction::WithFields),
87            fields.into_iter().map(|e| materialize_field(&s, e)),
88        )
89    }
90}
polars_plan/dsl/struct_.rs

polars_plan/dsl/
struct_.rs
