polars_lazy/dsl/
functions.rs

1//! # Functions
2//!
3//! Function on multiple expressions.
4//!
5
6use polars_core::prelude::*;
7pub use polars_plan::dsl::functions::*;
8use polars_plan::prelude::UnionArgs;
9
10use crate::prelude::*;
11
12pub(crate) fn concat_impl<L: AsRef<[LazyFrame]>>(
13    inputs: L,
14    args: UnionArgs,
15) -> PolarsResult<LazyFrame> {
16    let mut inputs = inputs.as_ref().to_vec();
17
18    let lf = std::mem::take(
19        inputs
20            .get_mut(0)
21            .ok_or_else(|| polars_err!(NoData: "empty container given"))?,
22    );
23
24    let opt_state = lf.opt_state;
25    let cached_arenas = lf.cached_arena.clone();
26
27    let mut lps = Vec::with_capacity(inputs.len());
28    lps.push(lf.logical_plan);
29
30    for lf in &mut inputs[1..] {
31        let lp = std::mem::take(&mut lf.logical_plan);
32        lps.push(lp)
33    }
34
35    let lp = DslPlan::Union { inputs: lps, args };
36    Ok(LazyFrame::from_inner(lp, opt_state, cached_arenas))
37}
38
39#[cfg(feature = "diagonal_concat")]
40/// Concat [LazyFrame]s diagonally.
41/// Calls [`concat`][concat()] internally.
42pub fn concat_lf_diagonal<L: AsRef<[LazyFrame]>>(
43    inputs: L,
44    mut args: UnionArgs,
45) -> PolarsResult<LazyFrame> {
46    args.diagonal = true;
47    concat_impl(inputs, args)
48}
49
50/// Concat [LazyFrame]s horizontally.
51pub fn concat_lf_horizontal<L: AsRef<[LazyFrame]>>(
52    inputs: L,
53    options: HConcatOptions,
54) -> PolarsResult<LazyFrame> {
55    let lfs = inputs.as_ref();
56    let (opt_state, cached_arena) = lfs
57        .first()
58        .map(|lf| (lf.opt_state, lf.cached_arena.clone()))
59        .ok_or_else(
60            || polars_err!(NoData: "Require at least one LazyFrame for horizontal concatenation"),
61        )?;
62
63    let lp = DslPlan::HConcat {
64        inputs: lfs.iter().map(|lf| lf.logical_plan.clone()).collect(),
65        options,
66    };
67    Ok(LazyFrame::from_inner(lp, opt_state, cached_arena))
68}
69
70/// Concat multiple [`LazyFrame`]s vertically.
71pub fn concat<L: AsRef<[LazyFrame]>>(inputs: L, args: UnionArgs) -> PolarsResult<LazyFrame> {
72    concat_impl(inputs, args)
73}
74
75#[cfg(test)]
76mod test {
77    // used only if feature="diagonal_concat"
78    #[allow(unused_imports)]
79    use super::*;
80
81    #[test]
82    #[cfg(feature = "diagonal_concat")]
83    fn test_diag_concat_lf() -> PolarsResult<()> {
84        let a = df![
85            "a" => [1, 2],
86            "b" => ["a", "b"]
87        ]?;
88
89        let b = df![
90            "b" => ["a", "b"],
91            "c" => [1, 2]
92        ]?;
93
94        let c = df![
95            "a" => [5, 7],
96            "c" => [1, 2],
97            "d" => [1, 2]
98        ]?;
99
100        let out = concat_lf_diagonal(
101            &[a.lazy(), b.lazy(), c.lazy()],
102            UnionArgs {
103                rechunk: false,
104                parallel: false,
105                ..Default::default()
106            },
107        )?
108        .collect()?;
109
110        let expected = df![
111            "a" => [Some(1), Some(2), None, None, Some(5), Some(7)],
112            "b" => [Some("a"), Some("b"), Some("a"), Some("b"), None, None],
113            "c" => [None, None, Some(1), Some(2), Some(1), Some(2)],
114            "d" => [None, None, None, None, Some(1), Some(2)]
115        ]?;
116
117        assert!(out.equals_missing(&expected));
118
119        Ok(())
120    }
121}
polars_lazy/dsl/functions.rs

polars_lazy/dsl/
functions.rs
