polars_ops/series/ops/
bitwise.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
use polars_core::chunked_array::ops::arity::unary_mut_values;
use polars_core::chunked_array::ChunkedArray;
use polars_core::prelude::DataType;
use polars_core::series::Series;
use polars_core::{with_match_physical_float_polars_type, with_match_physical_integer_polars_type};
use polars_error::{polars_bail, PolarsResult};

use super::*;

macro_rules! apply_bitwise_op {
    ($($op:ident),+ $(,)?) => {
        $(
        pub fn $op(s: &Series) -> PolarsResult<Series> {
            match s.dtype() {
                DataType::Boolean => {
                    let ca: &ChunkedArray<BooleanType> = s.as_any().downcast_ref().unwrap();
                    Ok(unary_mut_values::<BooleanType, UInt32Type, _, _>(
                        ca,
                        |a| polars_compute::bitwise::BitwiseKernel::$op(a),
                    ).into_series())
                },
                dt if dt.is_integer() => {
                    with_match_physical_integer_polars_type!(dt, |$T| {
                        let ca: &ChunkedArray<$T> = s.as_any().downcast_ref().unwrap();
                        Ok(unary_mut_values::<$T, UInt32Type, _, _>(
                            ca,
                            |a| polars_compute::bitwise::BitwiseKernel::$op(a),
                        ).into_series())
                    })
                },
                dt if dt.is_float() => {
                    with_match_physical_float_polars_type!(dt, |$T| {
                        let ca: &ChunkedArray<$T> = s.as_any().downcast_ref().unwrap();
                        Ok(unary_mut_values::<$T, UInt32Type, _, _>(
                            ca,
                            |a| polars_compute::bitwise::BitwiseKernel::$op(a),
                        ).into_series())
                    })
                },
                dt => {
                    polars_bail!(InvalidOperation: "dtype {:?} not supported in '{}' operation", dt, stringify!($op))
                },
            }
        }
        )+

    };
}

apply_bitwise_op! {
    count_ones,
    count_zeros,
    leading_ones,
    leading_zeros,
    trailing_ones,
    trailing_zeros,
}