polars_plan/dsl/function_expr/
boolean.rs

1use std::ops::{BitAnd, BitOr};
2
3use polars_core::POOL;
4use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
5
6use super::*;
7#[cfg(feature = "is_in")]
8use crate::wrap;
9use crate::{map, map_as_slice};
10
11#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
12#[derive(Clone, PartialEq, Debug, Eq, Hash)]
13pub enum BooleanFunction {
14    Any {
15        ignore_nulls: bool,
16    },
17    All {
18        ignore_nulls: bool,
19    },
20    IsNull,
21    IsNotNull,
22    IsFinite,
23    IsInfinite,
24    IsNan,
25    IsNotNan,
26    #[cfg(feature = "is_first_distinct")]
27    IsFirstDistinct,
28    #[cfg(feature = "is_last_distinct")]
29    IsLastDistinct,
30    #[cfg(feature = "is_unique")]
31    IsUnique,
32    #[cfg(feature = "is_unique")]
33    IsDuplicated,
34    #[cfg(feature = "is_between")]
35    IsBetween {
36        closed: ClosedInterval,
37    },
38    #[cfg(feature = "is_in")]
39    IsIn {
40        nulls_equal: bool,
41    },
42    AllHorizontal,
43    AnyHorizontal,
44    // Also bitwise negate
45    Not,
46}
47
48impl BooleanFunction {
49    pub(super) fn get_field(&self, mapper: FieldsMapper) -> PolarsResult<Field> {
50        match self {
51            BooleanFunction::Not => {
52                mapper.try_map_dtype(|dtype| {
53                    match dtype {
54                        DataType::Boolean => Ok(DataType::Boolean),
55                        dt if dt.is_integer() => Ok(dt.clone()),
56                        dt => polars_bail!(InvalidOperation: "dtype {:?} not supported in 'not' operation", dt) 
57                    }
58                })
59
60            },
61            _ => mapper.with_dtype(DataType::Boolean),
62        }
63    }
64
65    pub fn function_options(&self) -> FunctionOptions {
66        use BooleanFunction as B;
67        match self {
68            B::Any { .. } | B::All { .. } => FunctionOptions::aggregation(),
69            B::IsNull | B::IsNotNull | B::IsFinite | B::IsInfinite | B::IsNan | B::IsNotNan => {
70                FunctionOptions::elementwise()
71            },
72            #[cfg(feature = "is_first_distinct")]
73            B::IsFirstDistinct => FunctionOptions::length_preserving(),
74            #[cfg(feature = "is_last_distinct")]
75            B::IsLastDistinct => FunctionOptions::length_preserving(),
76            #[cfg(feature = "is_unique")]
77            B::IsUnique => FunctionOptions::length_preserving(),
78            #[cfg(feature = "is_unique")]
79            B::IsDuplicated => FunctionOptions::length_preserving(),
80            #[cfg(feature = "is_between")]
81            B::IsBetween { .. } => FunctionOptions::elementwise().with_supertyping(
82                (SuperTypeFlags::default() & !SuperTypeFlags::ALLOW_PRIMITIVE_TO_STRING).into(),
83            ),
84            #[cfg(feature = "is_in")]
85            B::IsIn { .. } => FunctionOptions::elementwise().with_supertyping(Default::default()),
86            B::AllHorizontal | B::AnyHorizontal => FunctionOptions::elementwise().with_flags(|f| {
87                f | FunctionFlags::INPUT_WILDCARD_EXPANSION | FunctionFlags::ALLOW_EMPTY_INPUTS
88            }),
89            B::Not => FunctionOptions::elementwise(),
90        }
91    }
92}
93
94impl Display for BooleanFunction {
95    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
96        use BooleanFunction::*;
97        let s = match self {
98            All { .. } => "all",
99            Any { .. } => "any",
100            IsNull => "is_null",
101            IsNotNull => "is_not_null",
102            IsFinite => "is_finite",
103            IsInfinite => "is_infinite",
104            IsNan => "is_nan",
105            IsNotNan => "is_not_nan",
106            #[cfg(feature = "is_first_distinct")]
107            IsFirstDistinct => "is_first_distinct",
108            #[cfg(feature = "is_last_distinct")]
109            IsLastDistinct => "is_last_distinct",
110            #[cfg(feature = "is_unique")]
111            IsUnique => "is_unique",
112            #[cfg(feature = "is_unique")]
113            IsDuplicated => "is_duplicated",
114            #[cfg(feature = "is_between")]
115            IsBetween { .. } => "is_between",
116            #[cfg(feature = "is_in")]
117            IsIn { .. } => "is_in",
118            AnyHorizontal => "any_horizontal",
119            AllHorizontal => "all_horizontal",
120            Not => "not",
121        };
122        write!(f, "{s}")
123    }
124}
125
126impl From<BooleanFunction> for SpecialEq<Arc<dyn ColumnsUdf>> {
127    fn from(func: BooleanFunction) -> Self {
128        use BooleanFunction::*;
129        match func {
130            Any { ignore_nulls } => map!(any, ignore_nulls),
131            All { ignore_nulls } => map!(all, ignore_nulls),
132            IsNull => map!(is_null),
133            IsNotNull => map!(is_not_null),
134            IsFinite => map!(is_finite),
135            IsInfinite => map!(is_infinite),
136            IsNan => map!(is_nan),
137            IsNotNan => map!(is_not_nan),
138            #[cfg(feature = "is_first_distinct")]
139            IsFirstDistinct => map!(is_first_distinct),
140            #[cfg(feature = "is_last_distinct")]
141            IsLastDistinct => map!(is_last_distinct),
142            #[cfg(feature = "is_unique")]
143            IsUnique => map!(is_unique),
144            #[cfg(feature = "is_unique")]
145            IsDuplicated => map!(is_duplicated),
146            #[cfg(feature = "is_between")]
147            IsBetween { closed } => map_as_slice!(is_between, closed),
148            #[cfg(feature = "is_in")]
149            IsIn { nulls_equal } => wrap!(is_in, nulls_equal),
150            Not => map!(not),
151            AllHorizontal => map_as_slice!(all_horizontal),
152            AnyHorizontal => map_as_slice!(any_horizontal),
153        }
154    }
155}
156
157impl From<BooleanFunction> for FunctionExpr {
158    fn from(func: BooleanFunction) -> Self {
159        FunctionExpr::Boolean(func)
160    }
161}
162
163fn any(s: &Column, ignore_nulls: bool) -> PolarsResult<Column> {
164    let ca = s.bool()?;
165    if ignore_nulls {
166        Ok(Column::new(s.name().clone(), [ca.any()]))
167    } else {
168        Ok(Column::new(s.name().clone(), [ca.any_kleene()]))
169    }
170}
171
172fn all(s: &Column, ignore_nulls: bool) -> PolarsResult<Column> {
173    let ca = s.bool()?;
174    if ignore_nulls {
175        Ok(Column::new(s.name().clone(), [ca.all()]))
176    } else {
177        Ok(Column::new(s.name().clone(), [ca.all_kleene()]))
178    }
179}
180
181fn is_null(s: &Column) -> PolarsResult<Column> {
182    Ok(s.is_null().into_column())
183}
184
185fn is_not_null(s: &Column) -> PolarsResult<Column> {
186    Ok(s.is_not_null().into_column())
187}
188
189fn is_finite(s: &Column) -> PolarsResult<Column> {
190    s.is_finite().map(|ca| ca.into_column())
191}
192
193fn is_infinite(s: &Column) -> PolarsResult<Column> {
194    s.is_infinite().map(|ca| ca.into_column())
195}
196
197pub(super) fn is_nan(s: &Column) -> PolarsResult<Column> {
198    s.is_nan().map(|ca| ca.into_column())
199}
200
201pub(super) fn is_not_nan(s: &Column) -> PolarsResult<Column> {
202    s.is_not_nan().map(|ca| ca.into_column())
203}
204
205#[cfg(feature = "is_first_distinct")]
206fn is_first_distinct(s: &Column) -> PolarsResult<Column> {
207    polars_ops::prelude::is_first_distinct(s.as_materialized_series()).map(|ca| ca.into_column())
208}
209
210#[cfg(feature = "is_last_distinct")]
211fn is_last_distinct(s: &Column) -> PolarsResult<Column> {
212    polars_ops::prelude::is_last_distinct(s.as_materialized_series()).map(|ca| ca.into_column())
213}
214
215#[cfg(feature = "is_unique")]
216fn is_unique(s: &Column) -> PolarsResult<Column> {
217    polars_ops::prelude::is_unique(s.as_materialized_series()).map(|ca| ca.into_column())
218}
219
220#[cfg(feature = "is_unique")]
221fn is_duplicated(s: &Column) -> PolarsResult<Column> {
222    polars_ops::prelude::is_duplicated(s.as_materialized_series()).map(|ca| ca.into_column())
223}
224
225#[cfg(feature = "is_between")]
226fn is_between(s: &[Column], closed: ClosedInterval) -> PolarsResult<Column> {
227    let ser = &s[0];
228    let lower = &s[1];
229    let upper = &s[2];
230    polars_ops::prelude::is_between(
231        ser.as_materialized_series(),
232        lower.as_materialized_series(),
233        upper.as_materialized_series(),
234        closed,
235    )
236    .map(|ca| ca.into_column())
237}
238
239#[cfg(feature = "is_in")]
240fn is_in(s: &mut [Column], nulls_equal: bool) -> PolarsResult<Option<Column>> {
241    let left = &s[0];
242    let other = &s[1];
243    polars_ops::prelude::is_in(
244        left.as_materialized_series(),
245        other.as_materialized_series(),
246        nulls_equal,
247    )
248    .map(|ca| Some(ca.into_column()))
249}
250
251fn not(s: &Column) -> PolarsResult<Column> {
252    polars_ops::series::negate_bitwise(s.as_materialized_series()).map(Column::from)
253}
254
255// We shouldn't hit these often only on very wide dataframes where we don't reduce to & expressions.
256fn any_horizontal(s: &[Column]) -> PolarsResult<Column> {
257    let out = POOL
258        .install(|| {
259            s.par_iter()
260                .try_fold(
261                    || BooleanChunked::new(PlSmallStr::EMPTY, &[false]),
262                    |acc, b| {
263                        let b = b.cast(&DataType::Boolean)?;
264                        let b = b.bool()?;
265                        PolarsResult::Ok((&acc).bitor(b))
266                    },
267                )
268                .try_reduce(
269                    || BooleanChunked::new(PlSmallStr::EMPTY, [false]),
270                    |a, b| Ok(a.bitor(b)),
271                )
272        })?
273        .with_name(s[0].name().clone());
274    Ok(out.into_column())
275}
276
277// We shouldn't hit these often only on very wide dataframes where we don't reduce to & expressions.
278fn all_horizontal(s: &[Column]) -> PolarsResult<Column> {
279    let out = POOL
280        .install(|| {
281            s.par_iter()
282                .try_fold(
283                    || BooleanChunked::new(PlSmallStr::EMPTY, &[true]),
284                    |acc, b| {
285                        let b = b.cast(&DataType::Boolean)?;
286                        let b = b.bool()?;
287                        PolarsResult::Ok((&acc).bitand(b))
288                    },
289                )
290                .try_reduce(
291                    || BooleanChunked::new(PlSmallStr::EMPTY, [true]),
292                    |a, b| Ok(a.bitand(b)),
293                )
294        })?
295        .with_name(s[0].name().clone());
296    Ok(out.into_column())
297}
polars_plan/dsl/function_expr/boolean.rs

polars_plan/dsl/function_expr/
boolean.rs
