polars_plan/plans/
apply.rs

1use std::fmt::{Debug, Formatter};
2
3use polars_core::prelude::*;
4
5pub trait DataFrameUdf: Send + Sync {
6    fn call_udf(&self, df: DataFrame) -> PolarsResult<DataFrame>;
7}
8
9impl<F> DataFrameUdf for F
10where
11    F: Fn(DataFrame) -> PolarsResult<DataFrame> + Send + Sync,
12{
13    fn call_udf(&self, df: DataFrame) -> PolarsResult<DataFrame> {
14        self(df)
15    }
16}
17
18pub trait DataFrameUdfMut: Send + Sync {
19    fn call_udf(&mut self, df: DataFrame) -> PolarsResult<DataFrame>;
20}
21
22impl<F> DataFrameUdfMut for F
23where
24    F: FnMut(DataFrame) -> PolarsResult<DataFrame> + Send + Sync,
25{
26    fn call_udf(&mut self, df: DataFrame) -> PolarsResult<DataFrame> {
27        self(df)
28    }
29}
30
31impl Debug for dyn DataFrameUdf {
32    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
33        write!(f, "dyn DataFrameUdf")
34    }
35}
36impl Debug for dyn DataFrameUdfMut {
37    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
38        write!(f, "dyn DataFrameUdfMut")
39    }
40}
41
42pub trait UdfSchema: Send + Sync {
43    fn get_schema(&self, input_schema: &Schema) -> PolarsResult<SchemaRef>;
44}
45
46impl<F> UdfSchema for F
47where
48    F: Fn(&Schema) -> PolarsResult<SchemaRef> + Send + Sync,
49{
50    fn get_schema(&self, input_schema: &Schema) -> PolarsResult<SchemaRef> {
51        self(input_schema)
52    }
53}
54
55impl Debug for dyn UdfSchema {
56    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
57        write!(f, "dyn UdfSchema")
58    }
59}