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    args: UnionArgs,
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 options = HConcatOptions {
64        parallel: args.parallel,
65        strict: args.strict,
66    };
67    let lp = DslPlan::HConcat {
68        inputs: lfs.iter().map(|lf| lf.logical_plan.clone()).collect(),
69        options,
70    };
71    Ok(LazyFrame::from_inner(lp, opt_state, cached_arena))
72}
73
74/// Concat multiple [`LazyFrame`]s vertically.
75pub fn concat<L: AsRef<[LazyFrame]>>(inputs: L, args: UnionArgs) -> PolarsResult<LazyFrame> {
76    concat_impl(inputs, args)
77}
78
79#[cfg(test)]
80mod test {
81    // used only if feature="diagonal_concat"
82    #[allow(unused_imports)]
83    use super::*;
84
85    #[test]
86    #[cfg(feature = "diagonal_concat")]
87    fn test_diag_concat_lf() -> PolarsResult<()> {
88        let a = df![
89            "a" => [1, 2],
90            "b" => ["a", "b"]
91        ]?;
92
93        let b = df![
94            "b" => ["a", "b"],
95            "c" => [1, 2]
96        ]?;
97
98        let c = df![
99            "a" => [5, 7],
100            "c" => [1, 2],
101            "d" => [1, 2]
102        ]?;
103
104        let out = concat_lf_diagonal(
105            &[a.lazy(), b.lazy(), c.lazy()],
106            UnionArgs {
107                rechunk: false,
108                parallel: false,
109                ..Default::default()
110            },
111        )?
112        .collect()?;
113
114        let expected = df![
115            "a" => [Some(1), Some(2), None, None, Some(5), Some(7)],
116            "b" => [Some("a"), Some("b"), Some("a"), Some("b"), None, None],
117            "c" => [None, None, Some(1), Some(2), Some(1), Some(2)],
118            "d" => [None, None, None, None, Some(1), Some(2)]
119        ]?;
120
121        assert!(out.equals_missing(&expected));
122
123        Ok(())
124    }
125}
polars_lazy/dsl/functions.rs

polars_lazy/dsl/
functions.rs
