polars_plan/dsl/functions/
concat.rs

1use super::*;
2
3#[cfg(all(feature = "concat_str", feature = "strings"))]
4/// Horizontally concat string columns in linear time
5pub fn concat_str<E: AsRef<[Expr]>>(s: E, separator: &str, ignore_nulls: bool) -> Expr {
6    let input = s.as_ref().to_vec();
7    let separator = separator.into();
8
9    Expr::Function {
10        input,
11        function: StringFunction::ConcatHorizontal {
12            delimiter: separator,
13            ignore_nulls,
14        }
15        .into(),
16        options: FunctionOptions::elementwise()
17            .with_flags(|f| f | FunctionFlags::INPUT_WILDCARD_EXPANSION),
18    }
19}
20
21#[cfg(all(feature = "concat_str", feature = "strings"))]
22/// Format the results of an array of expressions using a format string
23pub fn format_str<E: AsRef<[Expr]>>(format: &str, args: E) -> PolarsResult<Expr> {
24    let mut args: std::collections::VecDeque<Expr> = args.as_ref().to_vec().into();
25
26    // Parse the format string, and separate substrings between placeholders
27    let segments: Vec<&str> = format.split("{}").collect();
28
29    polars_ensure!(
30        segments.len() - 1 == args.len(),
31        ShapeMismatch: "number of placeholders should equal the number of arguments"
32    );
33
34    let mut exprs: Vec<Expr> = Vec::new();
35
36    for (i, s) in segments.iter().enumerate() {
37        if i > 0 {
38            if let Some(arg) = args.pop_front() {
39                exprs.push(arg);
40            }
41        }
42
43        if !s.is_empty() {
44            exprs.push(lit(s.to_string()))
45        }
46    }
47
48    Ok(concat_str(exprs, "", false))
49}
50
51/// Concat lists entries.
52pub fn concat_list<E: AsRef<[IE]>, IE: Into<Expr> + Clone>(s: E) -> PolarsResult<Expr> {
53    let s: Vec<_> = s.as_ref().iter().map(|e| e.clone().into()).collect();
54
55    polars_ensure!(!s.is_empty(), ComputeError: "`concat_list` needs one or more expressions");
56
57    Ok(Expr::Function {
58        input: s,
59        function: FunctionExpr::ListExpr(ListFunction::Concat),
60        options: FunctionOptions::elementwise()
61            .with_flags(|f| f | FunctionFlags::INPUT_WILDCARD_EXPANSION),
62    })
63}
64
65/// Horizontally concatenate columns into a single array-type column.
66pub fn concat_arr(input: Vec<Expr>) -> PolarsResult<Expr> {
67    feature_gated!("dtype-array", {
68        polars_ensure!(!input.is_empty(), ComputeError: "`concat_arr` needs one or more expressions");
69
70        Ok(Expr::Function {
71            input,
72            function: FunctionExpr::ArrayExpr(ArrayFunction::Concat),
73            options: FunctionOptions::elementwise()
74                .with_flags(|f| f | FunctionFlags::INPUT_WILDCARD_EXPANSION),
75        })
76    })
77}
78
79pub fn concat_expr<E: AsRef<[IE]>, IE: Into<Expr> + Clone>(
80    s: E,
81    rechunk: bool,
82) -> PolarsResult<Expr> {
83    let s: Vec<_> = s.as_ref().iter().map(|e| e.clone().into()).collect();
84    polars_ensure!(!s.is_empty(), ComputeError: "`concat_expr` needs one or more expressions");
85    Ok(Expr::n_ary(FunctionExpr::ConcatExpr(rechunk), s))
86}