polars_plan/dsl/functions/
arity.rs

1use super::*;
2
3macro_rules! prepare_binary_function {
4    ($f:ident) => {
5        move |c: &mut [Column]| {
6            let s0 = std::mem::take(&mut c[0]);
7            let s1 = std::mem::take(&mut c[1]);
8
9            $f(s0, s1)
10        }
11    };
12}
13
14/// Apply a closure on the two columns that are evaluated from [`Expr`] a and [`Expr`] b.
15///
16/// The closure takes two arguments, each a [`Series`]. `output_type` must be the output dtype of the resulting [`Series`].
17pub fn map_binary<F>(a: Expr, b: Expr, f: F, output_type: GetOutput) -> Expr
18where
19    F: 'static + Fn(Column, Column) -> PolarsResult<Option<Column>> + Send + Sync,
20{
21    let function = prepare_binary_function!(f);
22    a.map_many(function, &[b], output_type)
23}
24
25/// Like [`map_binary`], but used in a group_by-aggregation context.
26///
27/// See [`Expr::apply`] for the difference between [`map`](Expr::map) and [`apply`](Expr::apply).
28pub fn apply_binary<F>(a: Expr, b: Expr, f: F, output_type: GetOutput) -> Expr
29where
30    F: 'static + Fn(Column, Column) -> PolarsResult<Option<Column>> + Send + Sync,
31{
32    let function = prepare_binary_function!(f);
33    a.apply_many(function, &[b], output_type)
34}