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_private(FunctionExpr::StructExpr(StructFunction::FieldByIndex(
11                index,
12            )))
13            .with_function_options(|mut options| {
14                options.flags |= FunctionFlags::ALLOW_RENAME;
15                options
16            })
17    }
18
19    /// Retrieve one or multiple of the fields of this [`StructChunked`] as a new Series.
20    /// This expression also expands the `"*"` wildcard column.
21    pub fn field_by_names<I, S>(self, names: I) -> Expr
22    where
23        I: IntoIterator<Item = S>,
24        S: Into<PlSmallStr>,
25    {
26        self.field_by_names_impl(names.into_iter().map(|x| x.into()).collect())
27    }
28
29    fn field_by_names_impl(self, names: Arc<[PlSmallStr]>) -> Expr {
30        self.0
31            .map_private(FunctionExpr::StructExpr(StructFunction::MultipleFields(
32                names,
33            )))
34            .with_function_options(|mut options| {
35                options.flags |= FunctionFlags::ALLOW_RENAME;
36                options
37            })
38    }
39
40    /// Retrieve one of the fields of this [`StructChunked`] as a new Series.
41    /// This expression also supports wildcard "*" and regex expansion.
42    pub fn field_by_name(self, name: &str) -> Expr {
43        if name == "*" || is_regex_projection(name) {
44            return self.field_by_names([name]);
45        }
46        self.0
47            .map_private(FunctionExpr::StructExpr(StructFunction::FieldByName(
48                name.into(),
49            )))
50            .with_function_options(|mut options| {
51                options.flags |= FunctionFlags::ALLOW_RENAME;
52                options
53            })
54    }
55
56    /// Rename the fields of the [`StructChunked`].
57    pub fn rename_fields<I, S>(self, names: I) -> Expr
58    where
59        I: IntoIterator<Item = S>,
60        S: Into<PlSmallStr>,
61    {
62        self._rename_fields_impl(names.into_iter().map(|x| x.into()).collect())
63    }
64
65    pub fn _rename_fields_impl(self, names: Arc<[PlSmallStr]>) -> Expr {
66        self.0
67            .map_private(FunctionExpr::StructExpr(StructFunction::RenameFields(
68                names,
69            )))
70    }
71
72    #[cfg(feature = "json")]
73    pub fn json_encode(self) -> Expr {
74        self.0
75            .map_private(FunctionExpr::StructExpr(StructFunction::JsonEncode))
76    }
77
78    pub fn with_fields(self, fields: Vec<Expr>) -> PolarsResult<Expr> {
79        fn materialize_field(this: &Expr, field: Expr) -> PolarsResult<Expr> {
80            field.try_map_expr(|e| match e {
81                Expr::Field(names) => {
82                    let this = this.clone().struct_();
83                    Ok(if names.len() == 1 {
84                        this.field_by_name(names[0].as_ref())
85                    } else {
86                        this.field_by_names_impl(names)
87                    })
88                },
89                Expr::Exclude(_, _) => {
90                    polars_bail!(InvalidOperation: "'exclude' not allowed in 'field'")
91                },
92                _ => Ok(e),
93            })
94        }
95
96        let mut new_fields = Vec::with_capacity(fields.len());
97        new_fields.push(Default::default());
98
99        for e in fields.into_iter().map(|e| materialize_field(&self.0, e)) {
100            new_fields.push(e?)
101        }
102        new_fields[0] = self.0;
103        Ok(Expr::Function {
104            input: new_fields,
105            function: FunctionExpr::StructExpr(StructFunction::WithFields),
106            options: FunctionOptions {
107                collect_groups: ApplyOptions::ElementWise,
108                flags: FunctionFlags::default()
109                    | FunctionFlags::PASS_NAME_TO_APPLY
110                    | FunctionFlags::INPUT_WILDCARD_EXPANSION,
111                ..Default::default()
112            },
113        })
114    }
115}
polars_plan/dsl/struct_.rs

polars_plan/dsl/
struct_.rs
