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

polars_lazy/dsl/
functions.rs
