polars_compute/bitwise/
mod.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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
use std::convert::identity;

use arrow::array::{Array, BooleanArray, PrimitiveArray};
use arrow::datatypes::ArrowDataType;
use arrow::legacy::utils::CustomIterTools;

pub trait BitwiseKernel {
    type Scalar;

    fn count_ones(&self) -> PrimitiveArray<u32>;
    fn count_zeros(&self) -> PrimitiveArray<u32>;

    fn leading_ones(&self) -> PrimitiveArray<u32>;
    fn leading_zeros(&self) -> PrimitiveArray<u32>;

    fn trailing_ones(&self) -> PrimitiveArray<u32>;
    fn trailing_zeros(&self) -> PrimitiveArray<u32>;

    fn reduce_and(&self) -> Option<Self::Scalar>;
    fn reduce_or(&self) -> Option<Self::Scalar>;
    fn reduce_xor(&self) -> Option<Self::Scalar>;

    fn bit_and(lhs: Self::Scalar, rhs: Self::Scalar) -> Self::Scalar;
    fn bit_or(lhs: Self::Scalar, rhs: Self::Scalar) -> Self::Scalar;
    fn bit_xor(lhs: Self::Scalar, rhs: Self::Scalar) -> Self::Scalar;
}

macro_rules! impl_bitwise_kernel {
    ($(($T:ty, $to_bits:expr, $from_bits:expr)),+ $(,)?) => {
        $(
        impl BitwiseKernel for PrimitiveArray<$T> {
            type Scalar = $T;

            #[inline(never)]
            fn count_ones(&self) -> PrimitiveArray<u32> {
                PrimitiveArray::new(
                    ArrowDataType::UInt32,
                    self.values_iter()
                        .map(|&v| $to_bits(v).count_ones())
                        .collect_trusted::<Vec<_>>()
                        .into(),
                    self.validity().cloned(),
                )
            }

            #[inline(never)]
            fn count_zeros(&self) -> PrimitiveArray<u32> {
                PrimitiveArray::new(
                    ArrowDataType::UInt32,
                    self.values_iter()
                        .map(|&v| $to_bits(v).count_zeros())
                        .collect_trusted::<Vec<_>>()
                        .into(),
                    self.validity().cloned(),
                )
            }

            #[inline(never)]
            fn leading_ones(&self) -> PrimitiveArray<u32> {
                PrimitiveArray::new(
                    ArrowDataType::UInt32,
                    self.values_iter()
                        .map(|&v| $to_bits(v).leading_ones())
                        .collect_trusted::<Vec<_>>()
                        .into(),
                    self.validity().cloned(),
                )
            }

            #[inline(never)]
            fn leading_zeros(&self) -> PrimitiveArray<u32> {
                PrimitiveArray::new(
                    ArrowDataType::UInt32,
                    self.values_iter()
                        .map(|&v| $to_bits(v).leading_zeros())
                        .collect_trusted::<Vec<_>>()
                        .into(),
                    self.validity().cloned(),
                )
            }

            #[inline(never)]
            fn trailing_ones(&self) -> PrimitiveArray<u32> {
                PrimitiveArray::new(
                    ArrowDataType::UInt32,
                    self.values_iter()
                        .map(|&v| $to_bits(v).trailing_ones())
                        .collect_trusted::<Vec<_>>()
                        .into(),
                    self.validity().cloned(),
                )
            }

            #[inline(never)]
            fn trailing_zeros(&self) -> PrimitiveArray<u32> {
                PrimitiveArray::new(
                    ArrowDataType::UInt32,
                    self.values().iter()
                        .map(|&v| $to_bits(v).trailing_zeros())
                        .collect_trusted::<Vec<_>>()
                        .into(),
                    self.validity().cloned(),
                )
            }

            #[inline(never)]
            fn reduce_and(&self) -> Option<Self::Scalar> {
                if !self.has_nulls() {
                    self.values_iter().copied().map($to_bits).reduce(|a, b| a & b).map($from_bits)
                } else {
                    self.non_null_values_iter().map($to_bits).reduce(|a, b| a & b).map($from_bits)
                }
            }

            #[inline(never)]
            fn reduce_or(&self) -> Option<Self::Scalar> {
                if !self.has_nulls() {
                    self.values_iter().copied().map($to_bits).reduce(|a, b| a | b).map($from_bits)
                } else {
                    self.non_null_values_iter().map($to_bits).reduce(|a, b| a | b).map($from_bits)
                }
            }

            #[inline(never)]
            fn reduce_xor(&self) -> Option<Self::Scalar> {
                if !self.has_nulls() {
                    self.values_iter().copied().map($to_bits).reduce(|a, b| a ^ b).map($from_bits)
                } else {
                    self.non_null_values_iter().map($to_bits).reduce(|a, b| a ^ b).map($from_bits)
                }
            }

            fn bit_and(lhs: Self::Scalar, rhs: Self::Scalar) -> Self::Scalar {
                $from_bits($to_bits(lhs) & $to_bits(rhs))
            }
            fn bit_or(lhs: Self::Scalar, rhs: Self::Scalar) -> Self::Scalar {
                $from_bits($to_bits(lhs) | $to_bits(rhs))
            }
            fn bit_xor(lhs: Self::Scalar, rhs: Self::Scalar) -> Self::Scalar {
                $from_bits($to_bits(lhs) ^ $to_bits(rhs))
            }
        }
        )+
    };
}

impl_bitwise_kernel! {
    (i8, identity, identity),
    (i16, identity, identity),
    (i32, identity, identity),
    (i64, identity, identity),
    (u8, identity, identity),
    (u16, identity, identity),
    (u32, identity, identity),
    (u64, identity, identity),
    (f32, f32::to_bits, f32::from_bits),
    (f64, f64::to_bits, f64::from_bits),
}

impl BitwiseKernel for BooleanArray {
    type Scalar = bool;

    #[inline(never)]
    fn count_ones(&self) -> PrimitiveArray<u32> {
        PrimitiveArray::new(
            ArrowDataType::UInt32,
            self.values_iter()
                .map(u32::from)
                .collect_trusted::<Vec<_>>()
                .into(),
            self.validity().cloned(),
        )
    }

    #[inline(never)]
    fn count_zeros(&self) -> PrimitiveArray<u32> {
        PrimitiveArray::new(
            ArrowDataType::UInt32,
            self.values_iter()
                .map(|v| u32::from(!v))
                .collect_trusted::<Vec<_>>()
                .into(),
            self.validity().cloned(),
        )
    }

    #[inline(always)]
    fn leading_ones(&self) -> PrimitiveArray<u32> {
        self.count_ones()
    }

    #[inline(always)]
    fn leading_zeros(&self) -> PrimitiveArray<u32> {
        self.count_zeros()
    }

    #[inline(always)]
    fn trailing_ones(&self) -> PrimitiveArray<u32> {
        self.count_ones()
    }

    #[inline(always)]
    fn trailing_zeros(&self) -> PrimitiveArray<u32> {
        self.count_zeros()
    }

    fn reduce_and(&self) -> Option<Self::Scalar> {
        if self.len() == self.null_count() {
            None
        } else if !self.has_nulls() {
            Some(self.values().unset_bits() == 0)
        } else {
            Some((self.values() & self.validity().unwrap()).unset_bits() == 0)
        }
    }

    fn reduce_or(&self) -> Option<Self::Scalar> {
        if self.len() == self.null_count() {
            None
        } else if !self.has_nulls() {
            Some(self.values().set_bits() > 0)
        } else {
            Some((self.values() & self.validity().unwrap()).set_bits() > 0)
        }
    }

    fn reduce_xor(&self) -> Option<Self::Scalar> {
        if self.len() == self.null_count() {
            None
        } else if !self.has_nulls() {
            Some(self.values().set_bits() % 2 == 1)
        } else {
            Some((self.values() & self.validity().unwrap()).set_bits() % 2 == 1)
        }
    }

    fn bit_and(lhs: Self::Scalar, rhs: Self::Scalar) -> Self::Scalar {
        lhs & rhs
    }
    fn bit_or(lhs: Self::Scalar, rhs: Self::Scalar) -> Self::Scalar {
        lhs | rhs
    }
    fn bit_xor(lhs: Self::Scalar, rhs: Self::Scalar) -> Self::Scalar {
        lhs ^ rhs
    }
}