polars_core/chunked_array/comparison/
mod.rs

1mod scalar;
2
3#[cfg(feature = "dtype-categorical")]
4mod categorical;
5
6use std::ops::{BitAnd, Not};
7
8use arrow::array::BooleanArray;
9use arrow::bitmap::{Bitmap, BitmapBuilder};
10use arrow::compute;
11use num_traits::{NumCast, ToPrimitive};
12use polars_compute::comparisons::{TotalEqKernel, TotalOrdKernel};
13
14use crate::prelude::*;
15use crate::series::IsSorted;
16use crate::series::implementations::null::NullChunked;
17use crate::utils::align_chunks_binary;
18
19impl<T> ChunkCompareEq<&ChunkedArray<T>> for ChunkedArray<T>
20where
21    T: PolarsNumericType,
22    T::Array: TotalOrdKernel<Scalar = T::Native> + TotalEqKernel<Scalar = T::Native>,
23{
24    type Item = BooleanChunked;
25
26    fn equal(&self, rhs: &ChunkedArray<T>) -> BooleanChunked {
27        // Broadcast.
28        match (self.len(), rhs.len()) {
29            (_, 1) => {
30                if let Some(value) = rhs.get(0) {
31                    self.equal(value)
32                } else {
33                    BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
34                }
35            },
36            (1, _) => {
37                if let Some(value) = self.get(0) {
38                    rhs.equal(value)
39                } else {
40                    BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
41                }
42            },
43            _ => arity::binary_mut_values(
44                self,
45                rhs,
46                |a, b| a.tot_eq_kernel(b).into(),
47                PlSmallStr::EMPTY,
48            ),
49        }
50    }
51
52    fn equal_missing(&self, rhs: &ChunkedArray<T>) -> BooleanChunked {
53        // Broadcast.
54        match (self.len(), rhs.len()) {
55            (_, 1) => {
56                if let Some(value) = rhs.get(0) {
57                    self.equal_missing(value)
58                } else {
59                    self.is_null()
60                }
61            },
62            (1, _) => {
63                if let Some(value) = self.get(0) {
64                    rhs.equal_missing(value)
65                } else {
66                    rhs.is_null()
67                }
68            },
69            _ => arity::binary_mut_with_options(
70                self,
71                rhs,
72                |a, b| a.tot_eq_missing_kernel(b).into(),
73                PlSmallStr::EMPTY,
74            ),
75        }
76    }
77
78    fn not_equal(&self, rhs: &ChunkedArray<T>) -> BooleanChunked {
79        // Broadcast.
80        match (self.len(), rhs.len()) {
81            (_, 1) => {
82                if let Some(value) = rhs.get(0) {
83                    self.not_equal(value)
84                } else {
85                    BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
86                }
87            },
88            (1, _) => {
89                if let Some(value) = self.get(0) {
90                    rhs.not_equal(value)
91                } else {
92                    BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
93                }
94            },
95            _ => arity::binary_mut_values(
96                self,
97                rhs,
98                |a, b| a.tot_ne_kernel(b).into(),
99                PlSmallStr::EMPTY,
100            ),
101        }
102    }
103
104    fn not_equal_missing(&self, rhs: &ChunkedArray<T>) -> BooleanChunked {
105        // Broadcast.
106        match (self.len(), rhs.len()) {
107            (_, 1) => {
108                if let Some(value) = rhs.get(0) {
109                    self.not_equal_missing(value)
110                } else {
111                    self.is_not_null()
112                }
113            },
114            (1, _) => {
115                if let Some(value) = self.get(0) {
116                    rhs.not_equal_missing(value)
117                } else {
118                    rhs.is_not_null()
119                }
120            },
121            _ => arity::binary_mut_with_options(
122                self,
123                rhs,
124                |a, b| a.tot_ne_missing_kernel(b).into(),
125                PlSmallStr::EMPTY,
126            ),
127        }
128    }
129}
130
131impl<T> ChunkCompareIneq<&ChunkedArray<T>> for ChunkedArray<T>
132where
133    T: PolarsNumericType,
134    T::Array: TotalOrdKernel<Scalar = T::Native> + TotalEqKernel<Scalar = T::Native>,
135{
136    type Item = BooleanChunked;
137
138    fn lt(&self, rhs: &ChunkedArray<T>) -> BooleanChunked {
139        // Broadcast.
140        match (self.len(), rhs.len()) {
141            (_, 1) => {
142                if let Some(value) = rhs.get(0) {
143                    self.lt(value)
144                } else {
145                    BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
146                }
147            },
148            (1, _) => {
149                if let Some(value) = self.get(0) {
150                    rhs.gt(value)
151                } else {
152                    BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
153                }
154            },
155            _ => arity::binary_mut_values(
156                self,
157                rhs,
158                |a, b| a.tot_lt_kernel(b).into(),
159                PlSmallStr::EMPTY,
160            ),
161        }
162    }
163
164    fn lt_eq(&self, rhs: &ChunkedArray<T>) -> BooleanChunked {
165        // Broadcast.
166        match (self.len(), rhs.len()) {
167            (_, 1) => {
168                if let Some(value) = rhs.get(0) {
169                    self.lt_eq(value)
170                } else {
171                    BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
172                }
173            },
174            (1, _) => {
175                if let Some(value) = self.get(0) {
176                    rhs.gt_eq(value)
177                } else {
178                    BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
179                }
180            },
181            _ => arity::binary_mut_values(
182                self,
183                rhs,
184                |a, b| a.tot_le_kernel(b).into(),
185                PlSmallStr::EMPTY,
186            ),
187        }
188    }
189
190    fn gt(&self, rhs: &Self) -> BooleanChunked {
191        rhs.lt(self)
192    }
193
194    fn gt_eq(&self, rhs: &Self) -> BooleanChunked {
195        rhs.lt_eq(self)
196    }
197}
198
199impl ChunkCompareEq<&NullChunked> for NullChunked {
200    type Item = BooleanChunked;
201
202    fn equal(&self, rhs: &NullChunked) -> Self::Item {
203        BooleanChunked::full_null(self.name().clone(), get_broadcast_length(self, rhs))
204    }
205
206    fn equal_missing(&self, rhs: &NullChunked) -> Self::Item {
207        BooleanChunked::full(self.name().clone(), true, get_broadcast_length(self, rhs))
208    }
209
210    fn not_equal(&self, rhs: &NullChunked) -> Self::Item {
211        BooleanChunked::full_null(self.name().clone(), get_broadcast_length(self, rhs))
212    }
213
214    fn not_equal_missing(&self, rhs: &NullChunked) -> Self::Item {
215        BooleanChunked::full(self.name().clone(), false, get_broadcast_length(self, rhs))
216    }
217}
218
219impl ChunkCompareIneq<&NullChunked> for NullChunked {
220    type Item = BooleanChunked;
221
222    fn gt(&self, rhs: &NullChunked) -> Self::Item {
223        BooleanChunked::full_null(self.name().clone(), get_broadcast_length(self, rhs))
224    }
225
226    fn gt_eq(&self, rhs: &NullChunked) -> Self::Item {
227        BooleanChunked::full_null(self.name().clone(), get_broadcast_length(self, rhs))
228    }
229
230    fn lt(&self, rhs: &NullChunked) -> Self::Item {
231        BooleanChunked::full_null(self.name().clone(), get_broadcast_length(self, rhs))
232    }
233
234    fn lt_eq(&self, rhs: &NullChunked) -> Self::Item {
235        BooleanChunked::full_null(self.name().clone(), get_broadcast_length(self, rhs))
236    }
237}
238
239#[inline]
240fn get_broadcast_length(lhs: &NullChunked, rhs: &NullChunked) -> usize {
241    match (lhs.len(), rhs.len()) {
242        (1, len_r) => len_r,
243        (len_l, 1) => len_l,
244        (len_l, len_r) if len_l == len_r => len_l,
245        _ => panic!("Cannot compare two series of different lengths."),
246    }
247}
248
249impl ChunkCompareEq<&BooleanChunked> for BooleanChunked {
250    type Item = BooleanChunked;
251
252    fn equal(&self, rhs: &BooleanChunked) -> BooleanChunked {
253        // Broadcast.
254        match (self.len(), rhs.len()) {
255            (_, 1) => {
256                if let Some(value) = rhs.get(0) {
257                    arity::unary_mut_values(self, |arr| arr.tot_eq_kernel_broadcast(&value).into())
258                } else {
259                    BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
260                }
261            },
262            (1, _) => {
263                if let Some(value) = self.get(0) {
264                    arity::unary_mut_values(rhs, |arr| arr.tot_eq_kernel_broadcast(&value).into())
265                } else {
266                    BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
267                }
268            },
269            _ => arity::binary_mut_values(
270                self,
271                rhs,
272                |a, b| a.tot_eq_kernel(b).into(),
273                PlSmallStr::EMPTY,
274            ),
275        }
276    }
277
278    fn equal_missing(&self, rhs: &BooleanChunked) -> BooleanChunked {
279        // Broadcast.
280        match (self.len(), rhs.len()) {
281            (_, 1) => {
282                if let Some(value) = rhs.get(0) {
283                    arity::unary_mut_with_options(self, |arr| {
284                        arr.tot_eq_missing_kernel_broadcast(&value).into()
285                    })
286                } else {
287                    self.is_null()
288                }
289            },
290            (1, _) => {
291                if let Some(value) = self.get(0) {
292                    arity::unary_mut_with_options(rhs, |arr| {
293                        arr.tot_eq_missing_kernel_broadcast(&value).into()
294                    })
295                } else {
296                    rhs.is_null()
297                }
298            },
299            _ => arity::binary_mut_with_options(
300                self,
301                rhs,
302                |a, b| a.tot_eq_missing_kernel(b).into(),
303                PlSmallStr::EMPTY,
304            ),
305        }
306    }
307
308    fn not_equal(&self, rhs: &BooleanChunked) -> BooleanChunked {
309        // Broadcast.
310        match (self.len(), rhs.len()) {
311            (_, 1) => {
312                if let Some(value) = rhs.get(0) {
313                    arity::unary_mut_values(self, |arr| arr.tot_ne_kernel_broadcast(&value).into())
314                } else {
315                    BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
316                }
317            },
318            (1, _) => {
319                if let Some(value) = self.get(0) {
320                    arity::unary_mut_values(rhs, |arr| arr.tot_ne_kernel_broadcast(&value).into())
321                } else {
322                    BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
323                }
324            },
325            _ => arity::binary_mut_values(
326                self,
327                rhs,
328                |a, b| a.tot_ne_kernel(b).into(),
329                PlSmallStr::EMPTY,
330            ),
331        }
332    }
333
334    fn not_equal_missing(&self, rhs: &BooleanChunked) -> BooleanChunked {
335        // Broadcast.
336        match (self.len(), rhs.len()) {
337            (_, 1) => {
338                if let Some(value) = rhs.get(0) {
339                    arity::unary_mut_with_options(self, |arr| {
340                        arr.tot_ne_missing_kernel_broadcast(&value).into()
341                    })
342                } else {
343                    self.is_not_null()
344                }
345            },
346            (1, _) => {
347                if let Some(value) = self.get(0) {
348                    arity::unary_mut_with_options(rhs, |arr| {
349                        arr.tot_ne_missing_kernel_broadcast(&value).into()
350                    })
351                } else {
352                    rhs.is_not_null()
353                }
354            },
355            _ => arity::binary_mut_with_options(
356                self,
357                rhs,
358                |a, b| a.tot_ne_missing_kernel(b).into(),
359                PlSmallStr::EMPTY,
360            ),
361        }
362    }
363}
364
365impl ChunkCompareIneq<&BooleanChunked> for BooleanChunked {
366    type Item = BooleanChunked;
367
368    fn lt(&self, rhs: &BooleanChunked) -> BooleanChunked {
369        // Broadcast.
370        match (self.len(), rhs.len()) {
371            (_, 1) => {
372                if let Some(value) = rhs.get(0) {
373                    arity::unary_mut_values(self, |arr| arr.tot_lt_kernel_broadcast(&value).into())
374                } else {
375                    BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
376                }
377            },
378            (1, _) => {
379                if let Some(value) = self.get(0) {
380                    arity::unary_mut_values(rhs, |arr| arr.tot_gt_kernel_broadcast(&value).into())
381                } else {
382                    BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
383                }
384            },
385            _ => arity::binary_mut_values(
386                self,
387                rhs,
388                |a, b| a.tot_lt_kernel(b).into(),
389                PlSmallStr::EMPTY,
390            ),
391        }
392    }
393
394    fn lt_eq(&self, rhs: &BooleanChunked) -> BooleanChunked {
395        // Broadcast.
396        match (self.len(), rhs.len()) {
397            (_, 1) => {
398                if let Some(value) = rhs.get(0) {
399                    arity::unary_mut_values(self, |arr| arr.tot_le_kernel_broadcast(&value).into())
400                } else {
401                    BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
402                }
403            },
404            (1, _) => {
405                if let Some(value) = self.get(0) {
406                    arity::unary_mut_values(rhs, |arr| arr.tot_ge_kernel_broadcast(&value).into())
407                } else {
408                    BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
409                }
410            },
411            _ => arity::binary_mut_values(
412                self,
413                rhs,
414                |a, b| a.tot_le_kernel(b).into(),
415                PlSmallStr::EMPTY,
416            ),
417        }
418    }
419
420    fn gt(&self, rhs: &Self) -> BooleanChunked {
421        rhs.lt(self)
422    }
423
424    fn gt_eq(&self, rhs: &Self) -> BooleanChunked {
425        rhs.lt_eq(self)
426    }
427}
428
429impl ChunkCompareEq<&StringChunked> for StringChunked {
430    type Item = BooleanChunked;
431
432    fn equal(&self, rhs: &StringChunked) -> BooleanChunked {
433        self.as_binary().equal(&rhs.as_binary())
434    }
435
436    fn equal_missing(&self, rhs: &StringChunked) -> BooleanChunked {
437        self.as_binary().equal_missing(&rhs.as_binary())
438    }
439
440    fn not_equal(&self, rhs: &StringChunked) -> BooleanChunked {
441        self.as_binary().not_equal(&rhs.as_binary())
442    }
443
444    fn not_equal_missing(&self, rhs: &StringChunked) -> BooleanChunked {
445        self.as_binary().not_equal_missing(&rhs.as_binary())
446    }
447}
448
449impl ChunkCompareIneq<&StringChunked> for StringChunked {
450    type Item = BooleanChunked;
451
452    fn gt(&self, rhs: &StringChunked) -> BooleanChunked {
453        self.as_binary().gt(&rhs.as_binary())
454    }
455
456    fn gt_eq(&self, rhs: &StringChunked) -> BooleanChunked {
457        self.as_binary().gt_eq(&rhs.as_binary())
458    }
459
460    fn lt(&self, rhs: &StringChunked) -> BooleanChunked {
461        self.as_binary().lt(&rhs.as_binary())
462    }
463
464    fn lt_eq(&self, rhs: &StringChunked) -> BooleanChunked {
465        self.as_binary().lt_eq(&rhs.as_binary())
466    }
467}
468
469impl ChunkCompareEq<&BinaryChunked> for BinaryChunked {
470    type Item = BooleanChunked;
471
472    fn equal(&self, rhs: &BinaryChunked) -> BooleanChunked {
473        // Broadcast.
474        match (self.len(), rhs.len()) {
475            (_, 1) => {
476                if let Some(value) = rhs.get(0) {
477                    self.equal(value)
478                } else {
479                    BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
480                }
481            },
482            (1, _) => {
483                if let Some(value) = self.get(0) {
484                    rhs.equal(value)
485                } else {
486                    BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
487                }
488            },
489            _ => arity::binary_mut_values(
490                self,
491                rhs,
492                |a, b| a.tot_eq_kernel(b).into(),
493                PlSmallStr::EMPTY,
494            ),
495        }
496    }
497
498    fn equal_missing(&self, rhs: &BinaryChunked) -> BooleanChunked {
499        // Broadcast.
500        match (self.len(), rhs.len()) {
501            (_, 1) => {
502                if let Some(value) = rhs.get(0) {
503                    self.equal_missing(value)
504                } else {
505                    self.is_null()
506                }
507            },
508            (1, _) => {
509                if let Some(value) = self.get(0) {
510                    rhs.equal_missing(value)
511                } else {
512                    rhs.is_null()
513                }
514            },
515            _ => arity::binary_mut_with_options(
516                self,
517                rhs,
518                |a, b| a.tot_eq_missing_kernel(b).into(),
519                PlSmallStr::EMPTY,
520            ),
521        }
522    }
523
524    fn not_equal(&self, rhs: &BinaryChunked) -> BooleanChunked {
525        // Broadcast.
526        match (self.len(), rhs.len()) {
527            (_, 1) => {
528                if let Some(value) = rhs.get(0) {
529                    self.not_equal(value)
530                } else {
531                    BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
532                }
533            },
534            (1, _) => {
535                if let Some(value) = self.get(0) {
536                    rhs.not_equal(value)
537                } else {
538                    BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
539                }
540            },
541            _ => arity::binary_mut_values(
542                self,
543                rhs,
544                |a, b| a.tot_ne_kernel(b).into(),
545                PlSmallStr::EMPTY,
546            ),
547        }
548    }
549
550    fn not_equal_missing(&self, rhs: &BinaryChunked) -> BooleanChunked {
551        // Broadcast.
552        match (self.len(), rhs.len()) {
553            (_, 1) => {
554                if let Some(value) = rhs.get(0) {
555                    self.not_equal_missing(value)
556                } else {
557                    self.is_not_null()
558                }
559            },
560            (1, _) => {
561                if let Some(value) = self.get(0) {
562                    rhs.not_equal_missing(value)
563                } else {
564                    rhs.is_not_null()
565                }
566            },
567            _ => arity::binary_mut_with_options(
568                self,
569                rhs,
570                |a, b| a.tot_ne_missing_kernel(b).into(),
571                PlSmallStr::EMPTY,
572            ),
573        }
574    }
575}
576
577impl ChunkCompareIneq<&BinaryChunked> for BinaryChunked {
578    type Item = BooleanChunked;
579
580    fn lt(&self, rhs: &BinaryChunked) -> BooleanChunked {
581        // Broadcast.
582        match (self.len(), rhs.len()) {
583            (_, 1) => {
584                if let Some(value) = rhs.get(0) {
585                    self.lt(value)
586                } else {
587                    BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
588                }
589            },
590            (1, _) => {
591                if let Some(value) = self.get(0) {
592                    rhs.gt(value)
593                } else {
594                    BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
595                }
596            },
597            _ => arity::binary_mut_values(
598                self,
599                rhs,
600                |a, b| a.tot_lt_kernel(b).into(),
601                PlSmallStr::EMPTY,
602            ),
603        }
604    }
605
606    fn lt_eq(&self, rhs: &BinaryChunked) -> BooleanChunked {
607        // Broadcast.
608        match (self.len(), rhs.len()) {
609            (_, 1) => {
610                if let Some(value) = rhs.get(0) {
611                    self.lt_eq(value)
612                } else {
613                    BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
614                }
615            },
616            (1, _) => {
617                if let Some(value) = self.get(0) {
618                    rhs.gt_eq(value)
619                } else {
620                    BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
621                }
622            },
623            _ => arity::binary_mut_values(
624                self,
625                rhs,
626                |a, b| a.tot_le_kernel(b).into(),
627                PlSmallStr::EMPTY,
628            ),
629        }
630    }
631
632    fn gt(&self, rhs: &Self) -> BooleanChunked {
633        rhs.lt(self)
634    }
635
636    fn gt_eq(&self, rhs: &Self) -> BooleanChunked {
637        rhs.lt_eq(self)
638    }
639}
640
641fn _list_comparison_helper<F, B>(
642    lhs: &ListChunked,
643    rhs: &ListChunked,
644    op: F,
645    broadcast_op: B,
646    missing: bool,
647    is_ne: bool,
648) -> BooleanChunked
649where
650    F: Fn(&ListArray<i64>, &ListArray<i64>) -> Bitmap,
651    B: Fn(&ListArray<i64>, &Box<dyn Array>) -> Bitmap,
652{
653    match (lhs.len(), rhs.len()) {
654        (_, 1) => {
655            let right = rhs
656                .downcast_iter()
657                .find(|x| !x.is_empty())
658                .unwrap()
659                .as_any()
660                .downcast_ref::<ListArray<i64>>()
661                .unwrap();
662
663            if !right.validity().is_none_or(|v| v.get(0).unwrap()) {
664                if missing {
665                    if is_ne {
666                        return lhs.is_not_null();
667                    } else {
668                        return lhs.is_null();
669                    }
670                } else {
671                    return BooleanChunked::full_null(PlSmallStr::EMPTY, lhs.len());
672                }
673            }
674
675            let values = right.values().sliced(
676                (*right.offsets().first()).try_into().unwrap(),
677                right.offsets().range().try_into().unwrap(),
678            );
679
680            if missing {
681                arity::unary_mut_with_options(lhs, |a| broadcast_op(a, &values).into())
682            } else {
683                arity::unary_mut_values(lhs, |a| broadcast_op(a, &values).into())
684            }
685        },
686        (1, _) => {
687            let left = lhs
688                .downcast_iter()
689                .find(|x| !x.is_empty())
690                .unwrap()
691                .as_any()
692                .downcast_ref::<ListArray<i64>>()
693                .unwrap();
694
695            if !left.validity().is_none_or(|v| v.get(0).unwrap()) {
696                if missing {
697                    if is_ne {
698                        return rhs.is_not_null();
699                    } else {
700                        return rhs.is_null();
701                    }
702                } else {
703                    return BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len());
704                }
705            }
706
707            let values = left.values().sliced(
708                (*left.offsets().first()).try_into().unwrap(),
709                left.offsets().range().try_into().unwrap(),
710            );
711
712            if missing {
713                arity::unary_mut_with_options(rhs, |a| broadcast_op(a, &values).into())
714            } else {
715                arity::unary_mut_values(rhs, |a| broadcast_op(a, &values).into())
716            }
717        },
718        _ => {
719            if missing {
720                arity::binary_mut_with_options(lhs, rhs, |a, b| op(a, b).into(), PlSmallStr::EMPTY)
721            } else {
722                arity::binary_mut_values(lhs, rhs, |a, b| op(a, b).into(), PlSmallStr::EMPTY)
723            }
724        },
725    }
726}
727
728impl ChunkCompareEq<&ListChunked> for ListChunked {
729    type Item = BooleanChunked;
730    fn equal(&self, rhs: &ListChunked) -> BooleanChunked {
731        _list_comparison_helper(
732            self,
733            rhs,
734            TotalEqKernel::tot_eq_kernel,
735            TotalEqKernel::tot_eq_kernel_broadcast,
736            false,
737            false,
738        )
739    }
740
741    fn equal_missing(&self, rhs: &ListChunked) -> BooleanChunked {
742        _list_comparison_helper(
743            self,
744            rhs,
745            TotalEqKernel::tot_eq_missing_kernel,
746            TotalEqKernel::tot_eq_missing_kernel_broadcast,
747            true,
748            false,
749        )
750    }
751
752    fn not_equal(&self, rhs: &ListChunked) -> BooleanChunked {
753        _list_comparison_helper(
754            self,
755            rhs,
756            TotalEqKernel::tot_ne_kernel,
757            TotalEqKernel::tot_ne_kernel_broadcast,
758            false,
759            true,
760        )
761    }
762
763    fn not_equal_missing(&self, rhs: &ListChunked) -> BooleanChunked {
764        _list_comparison_helper(
765            self,
766            rhs,
767            TotalEqKernel::tot_ne_missing_kernel,
768            TotalEqKernel::tot_ne_missing_kernel_broadcast,
769            true,
770            true,
771        )
772    }
773}
774
775#[cfg(feature = "dtype-struct")]
776fn struct_helper<F, R>(
777    a: &StructChunked,
778    b: &StructChunked,
779    op: F,
780    reduce: R,
781    op_is_ne: bool,
782    is_missing: bool,
783) -> BooleanChunked
784where
785    F: Fn(&Series, &Series) -> BooleanChunked,
786    R: Fn(BooleanChunked, BooleanChunked) -> BooleanChunked,
787{
788    let len_a = a.len();
789    let len_b = b.len();
790    let broadcasts = len_a == 1 || len_b == 1;
791    if (a.len() != b.len() && !broadcasts) || a.struct_fields().len() != b.struct_fields().len() {
792        BooleanChunked::full(PlSmallStr::EMPTY, op_is_ne, a.len())
793    } else {
794        let (a, b) = align_chunks_binary(a, b);
795
796        let mut out = a
797            .fields_as_series()
798            .iter()
799            .zip(b.fields_as_series().iter())
800            .map(|(l, r)| op(l, r))
801            .reduce(&reduce)
802            .unwrap_or_else(|| BooleanChunked::full(PlSmallStr::EMPTY, !op_is_ne, a.len()));
803
804        if is_missing && (a.has_nulls() || b.has_nulls()) {
805            // Do some allocations so that we can use the Series dispatch, it otherwise
806            // gets complicated dealing with combinations of ==, != and broadcasting.
807            let default = || {
808                BooleanChunked::with_chunk(PlSmallStr::EMPTY, BooleanArray::from_slice([true]))
809                    .into_series()
810            };
811            let validity_to_series = |x| unsafe {
812                BooleanChunked::with_chunk(
813                    PlSmallStr::EMPTY,
814                    BooleanArray::from_inner_unchecked(ArrowDataType::Boolean, x, None),
815                )
816                .into_series()
817            };
818
819            out = reduce(
820                out,
821                op(
822                    &a.rechunk_validity()
823                        .map_or_else(default, validity_to_series),
824                    &b.rechunk_validity()
825                        .map_or_else(default, validity_to_series),
826                ),
827            )
828        }
829
830        if !is_missing && (a.null_count() > 0 || b.null_count() > 0) {
831            let mut a = a;
832            let mut b = b;
833
834            if broadcasts {
835                if a.len() == 1 {
836                    a = std::borrow::Cow::Owned(a.new_from_index(0, b.len()));
837                }
838                if b.len() == 1 {
839                    b = std::borrow::Cow::Owned(b.new_from_index(0, a.len()));
840                }
841            }
842
843            let mut a = a.into_owned();
844            a.zip_outer_validity(&b);
845            unsafe {
846                let mut new_null_count = 0;
847                for (arr, a) in out.downcast_iter_mut().zip(a.downcast_iter()) {
848                    arr.set_validity(a.validity().cloned());
849                    new_null_count += arr.null_count();
850                }
851                out.set_null_count(new_null_count);
852            }
853        }
854
855        out
856    }
857}
858
859#[cfg(feature = "dtype-struct")]
860impl ChunkCompareEq<&StructChunked> for StructChunked {
861    type Item = BooleanChunked;
862    fn equal(&self, rhs: &StructChunked) -> BooleanChunked {
863        struct_helper(
864            self,
865            rhs,
866            |l, r| l.equal_missing(r).unwrap(),
867            |a, b| a.bitand(b),
868            false,
869            false,
870        )
871    }
872
873    fn equal_missing(&self, rhs: &StructChunked) -> BooleanChunked {
874        struct_helper(
875            self,
876            rhs,
877            |l, r| l.equal_missing(r).unwrap(),
878            |a, b| a.bitand(b),
879            false,
880            true,
881        )
882    }
883
884    fn not_equal(&self, rhs: &StructChunked) -> BooleanChunked {
885        struct_helper(
886            self,
887            rhs,
888            |l, r| l.not_equal_missing(r).unwrap(),
889            |a, b| a | b,
890            true,
891            false,
892        )
893    }
894
895    fn not_equal_missing(&self, rhs: &StructChunked) -> BooleanChunked {
896        struct_helper(
897            self,
898            rhs,
899            |l, r| l.not_equal_missing(r).unwrap(),
900            |a, b| a | b,
901            true,
902            true,
903        )
904    }
905}
906
907#[cfg(feature = "dtype-array")]
908fn _array_comparison_helper<F, B>(
909    lhs: &ArrayChunked,
910    rhs: &ArrayChunked,
911    op: F,
912    broadcast_op: B,
913    missing: bool,
914    is_ne: bool,
915) -> BooleanChunked
916where
917    F: Fn(&FixedSizeListArray, &FixedSizeListArray) -> Bitmap,
918    B: Fn(&FixedSizeListArray, &Box<dyn Array>) -> Bitmap,
919{
920    match (lhs.len(), rhs.len()) {
921        (_, 1) => {
922            let right = rhs
923                .downcast_iter()
924                .find(|x| !x.is_empty())
925                .unwrap()
926                .as_any()
927                .downcast_ref::<FixedSizeListArray>()
928                .unwrap();
929
930            if !right.validity().is_none_or(|v| v.get(0).unwrap()) {
931                if missing {
932                    if is_ne {
933                        return lhs.is_not_null();
934                    } else {
935                        return lhs.is_null();
936                    }
937                } else {
938                    return BooleanChunked::full_null(PlSmallStr::EMPTY, lhs.len());
939                }
940            }
941
942            if missing {
943                arity::unary_mut_with_options(lhs, |a| broadcast_op(a, right.values()).into())
944            } else {
945                arity::unary_mut_values(lhs, |a| broadcast_op(a, right.values()).into())
946            }
947        },
948        (1, _) => {
949            let left = lhs
950                .downcast_iter()
951                .find(|x| !x.is_empty())
952                .unwrap()
953                .as_any()
954                .downcast_ref::<FixedSizeListArray>()
955                .unwrap();
956
957            if !left.validity().is_none_or(|v| v.get(0).unwrap()) {
958                if missing {
959                    if is_ne {
960                        return rhs.is_not_null();
961                    } else {
962                        return rhs.is_null();
963                    }
964                } else {
965                    return BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len());
966                }
967            }
968
969            if missing {
970                arity::unary_mut_with_options(rhs, |a| broadcast_op(a, left.values()).into())
971            } else {
972                arity::unary_mut_values(rhs, |a| broadcast_op(a, left.values()).into())
973            }
974        },
975        _ => {
976            if missing {
977                arity::binary_mut_with_options(lhs, rhs, |a, b| op(a, b).into(), PlSmallStr::EMPTY)
978            } else {
979                arity::binary_mut_values(lhs, rhs, |a, b| op(a, b).into(), PlSmallStr::EMPTY)
980            }
981        },
982    }
983}
984
985#[cfg(feature = "dtype-array")]
986impl ChunkCompareEq<&ArrayChunked> for ArrayChunked {
987    type Item = BooleanChunked;
988    fn equal(&self, rhs: &ArrayChunked) -> BooleanChunked {
989        _array_comparison_helper(
990            self,
991            rhs,
992            TotalEqKernel::tot_eq_kernel,
993            TotalEqKernel::tot_eq_kernel_broadcast,
994            false,
995            false,
996        )
997    }
998
999    fn equal_missing(&self, rhs: &ArrayChunked) -> BooleanChunked {
1000        _array_comparison_helper(
1001            self,
1002            rhs,
1003            TotalEqKernel::tot_eq_missing_kernel,
1004            TotalEqKernel::tot_eq_missing_kernel_broadcast,
1005            true,
1006            false,
1007        )
1008    }
1009
1010    fn not_equal(&self, rhs: &ArrayChunked) -> BooleanChunked {
1011        _array_comparison_helper(
1012            self,
1013            rhs,
1014            TotalEqKernel::tot_ne_kernel,
1015            TotalEqKernel::tot_ne_kernel_broadcast,
1016            false,
1017            true,
1018        )
1019    }
1020
1021    fn not_equal_missing(&self, rhs: &ArrayChunked) -> Self::Item {
1022        _array_comparison_helper(
1023            self,
1024            rhs,
1025            TotalEqKernel::tot_ne_missing_kernel,
1026            TotalEqKernel::tot_ne_missing_kernel_broadcast,
1027            true,
1028            true,
1029        )
1030    }
1031}
1032
1033impl Not for &BooleanChunked {
1034    type Output = BooleanChunked;
1035
1036    fn not(self) -> Self::Output {
1037        let chunks = self.downcast_iter().map(compute::boolean::not);
1038        ChunkedArray::from_chunk_iter(self.name().clone(), chunks)
1039    }
1040}
1041
1042impl Not for BooleanChunked {
1043    type Output = BooleanChunked;
1044
1045    fn not(self) -> Self::Output {
1046        (&self).not()
1047    }
1048}
1049
1050impl BooleanChunked {
1051    /// Returns whether any of the values in the column are `true`.
1052    ///
1053    /// Null values are ignored.
1054    pub fn any(&self) -> bool {
1055        self.downcast_iter().any(compute::boolean::any)
1056    }
1057
1058    /// Returns whether all values in the array are `true`.
1059    ///
1060    /// Null values are ignored.
1061    pub fn all(&self) -> bool {
1062        self.downcast_iter().all(compute::boolean::all)
1063    }
1064
1065    /// Returns whether any of the values in the column are `true`.
1066    ///
1067    /// The output is unknown (`None`) if the array contains any null values and
1068    /// no `true` values.
1069    pub fn any_kleene(&self) -> Option<bool> {
1070        let mut result = Some(false);
1071        for arr in self.downcast_iter() {
1072            match compute::boolean_kleene::any(arr) {
1073                Some(true) => return Some(true),
1074                None => result = None,
1075                _ => (),
1076            };
1077        }
1078        result
1079    }
1080
1081    /// Returns whether all values in the column are `true`.
1082    ///
1083    /// The output is unknown (`None`) if the array contains any null values and
1084    /// no `false` values.
1085    pub fn all_kleene(&self) -> Option<bool> {
1086        let mut result = Some(true);
1087        for arr in self.downcast_iter() {
1088            match compute::boolean_kleene::all(arr) {
1089                Some(false) => return Some(false),
1090                None => result = None,
1091                _ => (),
1092            };
1093        }
1094        result
1095    }
1096}
1097
1098// private
1099pub(crate) trait ChunkEqualElement {
1100    /// Only meant for physical types.
1101    /// Check if element in self is equal to element in other, assumes same dtypes
1102    ///
1103    /// # Safety
1104    ///
1105    /// No type checks.
1106    unsafe fn equal_element(&self, _idx_self: usize, _idx_other: usize, _other: &Series) -> bool {
1107        unimplemented!()
1108    }
1109}
1110
1111impl<T> ChunkEqualElement for ChunkedArray<T>
1112where
1113    T: PolarsNumericType,
1114{
1115    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
1116        let ca_other = other.as_ref().as_ref();
1117        debug_assert!(self.dtype() == other.dtype());
1118        let ca_other = &*(ca_other as *const ChunkedArray<T>);
1119        // Should be get and not get_unchecked, because there could be nulls
1120        self.get_unchecked(idx_self)
1121            .tot_eq(&ca_other.get_unchecked(idx_other))
1122    }
1123}
1124
1125impl ChunkEqualElement for BooleanChunked {
1126    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
1127        let ca_other = other.as_ref().as_ref();
1128        debug_assert!(self.dtype() == other.dtype());
1129        let ca_other = &*(ca_other as *const BooleanChunked);
1130        self.get_unchecked(idx_self) == ca_other.get_unchecked(idx_other)
1131    }
1132}
1133
1134impl ChunkEqualElement for StringChunked {
1135    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
1136        let ca_other = other.as_ref().as_ref();
1137        debug_assert!(self.dtype() == other.dtype());
1138        let ca_other = &*(ca_other as *const StringChunked);
1139        self.get_unchecked(idx_self) == ca_other.get_unchecked(idx_other)
1140    }
1141}
1142
1143impl ChunkEqualElement for BinaryChunked {
1144    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
1145        let ca_other = other.as_ref().as_ref();
1146        debug_assert!(self.dtype() == other.dtype());
1147        let ca_other = &*(ca_other as *const BinaryChunked);
1148        self.get_unchecked(idx_self) == ca_other.get_unchecked(idx_other)
1149    }
1150}
1151
1152impl ChunkEqualElement for BinaryOffsetChunked {
1153    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
1154        let ca_other = other.as_ref().as_ref();
1155        debug_assert!(self.dtype() == other.dtype());
1156        let ca_other = &*(ca_other as *const BinaryOffsetChunked);
1157        self.get_unchecked(idx_self) == ca_other.get_unchecked(idx_other)
1158    }
1159}
1160
1161impl ChunkEqualElement for ListChunked {}
1162#[cfg(feature = "dtype-array")]
1163impl ChunkEqualElement for ArrayChunked {}
1164
1165#[cfg(test)]
1166#[cfg_attr(feature = "nightly", allow(clippy::manual_repeat_n))] // remove once stable
1167mod test {
1168    use std::iter::repeat_n;
1169
1170    use super::super::test::get_chunked_array;
1171    use crate::prelude::*;
1172
1173    pub(crate) fn create_two_chunked() -> (Int32Chunked, Int32Chunked) {
1174        let mut a1 = Int32Chunked::new(PlSmallStr::from_static("a"), &[1, 2, 3]);
1175        let a2 = Int32Chunked::new(PlSmallStr::from_static("a"), &[4, 5, 6]);
1176        let a3 = Int32Chunked::new(PlSmallStr::from_static("a"), &[1, 2, 3, 4, 5, 6]);
1177        a1.append(&a2).unwrap();
1178        (a1, a3)
1179    }
1180
1181    #[test]
1182    fn test_bitwise_ops() {
1183        let a = BooleanChunked::new(PlSmallStr::from_static("a"), &[true, false, false]);
1184        let b = BooleanChunked::new(
1185            PlSmallStr::from_static("b"),
1186            &[Some(true), Some(true), None],
1187        );
1188        assert_eq!(Vec::from(&a | &b), &[Some(true), Some(true), None]);
1189        assert_eq!(Vec::from(&a & &b), &[Some(true), Some(false), Some(false)]);
1190        assert_eq!(Vec::from(!b), &[Some(false), Some(false), None]);
1191    }
1192
1193    #[test]
1194    fn test_compare_chunk_diff() {
1195        let (a1, a2) = create_two_chunked();
1196
1197        assert_eq!(
1198            a1.equal(&a2).into_iter().collect::<Vec<_>>(),
1199            repeat_n(Some(true), 6).collect::<Vec<_>>()
1200        );
1201        assert_eq!(
1202            a2.equal(&a1).into_iter().collect::<Vec<_>>(),
1203            repeat_n(Some(true), 6).collect::<Vec<_>>()
1204        );
1205        assert_eq!(
1206            a1.not_equal(&a2).into_iter().collect::<Vec<_>>(),
1207            repeat_n(Some(false), 6).collect::<Vec<_>>()
1208        );
1209        assert_eq!(
1210            a2.not_equal(&a1).into_iter().collect::<Vec<_>>(),
1211            repeat_n(Some(false), 6).collect::<Vec<_>>()
1212        );
1213        assert_eq!(
1214            a1.gt(&a2).into_iter().collect::<Vec<_>>(),
1215            repeat_n(Some(false), 6).collect::<Vec<_>>()
1216        );
1217        assert_eq!(
1218            a2.gt(&a1).into_iter().collect::<Vec<_>>(),
1219            repeat_n(Some(false), 6).collect::<Vec<_>>()
1220        );
1221        assert_eq!(
1222            a1.gt_eq(&a2).into_iter().collect::<Vec<_>>(),
1223            repeat_n(Some(true), 6).collect::<Vec<_>>()
1224        );
1225        assert_eq!(
1226            a2.gt_eq(&a1).into_iter().collect::<Vec<_>>(),
1227            repeat_n(Some(true), 6).collect::<Vec<_>>()
1228        );
1229        assert_eq!(
1230            a1.lt_eq(&a2).into_iter().collect::<Vec<_>>(),
1231            repeat_n(Some(true), 6).collect::<Vec<_>>()
1232        );
1233        assert_eq!(
1234            a2.lt_eq(&a1).into_iter().collect::<Vec<_>>(),
1235            repeat_n(Some(true), 6).collect::<Vec<_>>()
1236        );
1237        assert_eq!(
1238            a1.lt(&a2).into_iter().collect::<Vec<_>>(),
1239            repeat_n(Some(false), 6).collect::<Vec<_>>()
1240        );
1241        assert_eq!(
1242            a2.lt(&a1).into_iter().collect::<Vec<_>>(),
1243            repeat_n(Some(false), 6).collect::<Vec<_>>()
1244        );
1245    }
1246
1247    #[test]
1248    fn test_equal_chunks() {
1249        let a1 = get_chunked_array();
1250        let a2 = get_chunked_array();
1251
1252        assert_eq!(
1253            a1.equal(&a2).into_iter().collect::<Vec<_>>(),
1254            repeat_n(Some(true), 3).collect::<Vec<_>>()
1255        );
1256        assert_eq!(
1257            a2.equal(&a1).into_iter().collect::<Vec<_>>(),
1258            repeat_n(Some(true), 3).collect::<Vec<_>>()
1259        );
1260        assert_eq!(
1261            a1.not_equal(&a2).into_iter().collect::<Vec<_>>(),
1262            repeat_n(Some(false), 3).collect::<Vec<_>>()
1263        );
1264        assert_eq!(
1265            a2.not_equal(&a1).into_iter().collect::<Vec<_>>(),
1266            repeat_n(Some(false), 3).collect::<Vec<_>>()
1267        );
1268        assert_eq!(
1269            a1.gt(&a2).into_iter().collect::<Vec<_>>(),
1270            repeat_n(Some(false), 3).collect::<Vec<_>>()
1271        );
1272        assert_eq!(
1273            a2.gt(&a1).into_iter().collect::<Vec<_>>(),
1274            repeat_n(Some(false), 3).collect::<Vec<_>>()
1275        );
1276        assert_eq!(
1277            a1.gt_eq(&a2).into_iter().collect::<Vec<_>>(),
1278            repeat_n(Some(true), 3).collect::<Vec<_>>()
1279        );
1280        assert_eq!(
1281            a2.gt_eq(&a1).into_iter().collect::<Vec<_>>(),
1282            repeat_n(Some(true), 3).collect::<Vec<_>>()
1283        );
1284        assert_eq!(
1285            a1.lt_eq(&a2).into_iter().collect::<Vec<_>>(),
1286            repeat_n(Some(true), 3).collect::<Vec<_>>()
1287        );
1288        assert_eq!(
1289            a2.lt_eq(&a1).into_iter().collect::<Vec<_>>(),
1290            repeat_n(Some(true), 3).collect::<Vec<_>>()
1291        );
1292        assert_eq!(
1293            a1.lt(&a2).into_iter().collect::<Vec<_>>(),
1294            repeat_n(Some(false), 3).collect::<Vec<_>>()
1295        );
1296        assert_eq!(
1297            a2.lt(&a1).into_iter().collect::<Vec<_>>(),
1298            repeat_n(Some(false), 3).collect::<Vec<_>>()
1299        );
1300    }
1301
1302    #[test]
1303    fn test_null_handling() {
1304        // assert we comply with arrows way of handling null data
1305        // we check comparison on two arrays with one chunk and verify it is equal to a differently
1306        // chunked array comparison.
1307
1308        // two same chunked arrays
1309        let a1: Int32Chunked = [Some(1), None, Some(3)].iter().copied().collect();
1310        let a2: Int32Chunked = [Some(1), Some(2), Some(3)].iter().copied().collect();
1311
1312        let mut a2_2chunks: Int32Chunked = [Some(1), Some(2)].iter().copied().collect();
1313        a2_2chunks
1314            .append(&[Some(3)].iter().copied().collect())
1315            .unwrap();
1316
1317        assert_eq!(
1318            a1.equal(&a2).into_iter().collect::<Vec<_>>(),
1319            a1.equal(&a2_2chunks).into_iter().collect::<Vec<_>>()
1320        );
1321
1322        assert_eq!(
1323            a1.not_equal(&a2).into_iter().collect::<Vec<_>>(),
1324            a1.not_equal(&a2_2chunks).into_iter().collect::<Vec<_>>()
1325        );
1326        assert_eq!(
1327            a1.not_equal(&a2).into_iter().collect::<Vec<_>>(),
1328            a2_2chunks.not_equal(&a1).into_iter().collect::<Vec<_>>()
1329        );
1330
1331        assert_eq!(
1332            a1.gt(&a2).into_iter().collect::<Vec<_>>(),
1333            a1.gt(&a2_2chunks).into_iter().collect::<Vec<_>>()
1334        );
1335        assert_eq!(
1336            a1.gt(&a2).into_iter().collect::<Vec<_>>(),
1337            a2_2chunks.gt(&a1).into_iter().collect::<Vec<_>>()
1338        );
1339
1340        assert_eq!(
1341            a1.gt_eq(&a2).into_iter().collect::<Vec<_>>(),
1342            a1.gt_eq(&a2_2chunks).into_iter().collect::<Vec<_>>()
1343        );
1344        assert_eq!(
1345            a1.gt_eq(&a2).into_iter().collect::<Vec<_>>(),
1346            a2_2chunks.gt_eq(&a1).into_iter().collect::<Vec<_>>()
1347        );
1348
1349        assert_eq!(
1350            a1.lt_eq(&a2).into_iter().collect::<Vec<_>>(),
1351            a1.lt_eq(&a2_2chunks).into_iter().collect::<Vec<_>>()
1352        );
1353        assert_eq!(
1354            a1.lt_eq(&a2).into_iter().collect::<Vec<_>>(),
1355            a2_2chunks.lt_eq(&a1).into_iter().collect::<Vec<_>>()
1356        );
1357
1358        assert_eq!(
1359            a1.lt(&a2).into_iter().collect::<Vec<_>>(),
1360            a1.lt(&a2_2chunks).into_iter().collect::<Vec<_>>()
1361        );
1362        assert_eq!(
1363            a1.lt(&a2).into_iter().collect::<Vec<_>>(),
1364            a2_2chunks.lt(&a1).into_iter().collect::<Vec<_>>()
1365        );
1366    }
1367
1368    #[test]
1369    fn test_left_right() {
1370        // This failed with arrow comparisons.
1371        // sliced
1372        let a1: Int32Chunked = [Some(1), Some(2)].iter().copied().collect();
1373        let a1 = a1.slice(1, 1);
1374        let a2: Int32Chunked = [Some(2)].iter().copied().collect();
1375        assert_eq!(a1.equal(&a2).sum(), a2.equal(&a1).sum());
1376        assert_eq!(a1.not_equal(&a2).sum(), a2.not_equal(&a1).sum());
1377        assert_eq!(a1.gt(&a2).sum(), a2.gt(&a1).sum());
1378        assert_eq!(a1.lt(&a2).sum(), a2.lt(&a1).sum());
1379        assert_eq!(a1.lt_eq(&a2).sum(), a2.lt_eq(&a1).sum());
1380        assert_eq!(a1.gt_eq(&a2).sum(), a2.gt_eq(&a1).sum());
1381
1382        let a1: StringChunked = ["a", "b"].iter().copied().collect();
1383        let a1 = a1.slice(1, 1);
1384        let a2: StringChunked = ["b"].iter().copied().collect();
1385        assert_eq!(a1.equal(&a2).sum(), a2.equal(&a1).sum());
1386        assert_eq!(a1.not_equal(&a2).sum(), a2.not_equal(&a1).sum());
1387        assert_eq!(a1.gt(&a2).sum(), a2.gt(&a1).sum());
1388        assert_eq!(a1.lt(&a2).sum(), a2.lt(&a1).sum());
1389        assert_eq!(a1.lt_eq(&a2).sum(), a2.lt_eq(&a1).sum());
1390        assert_eq!(a1.gt_eq(&a2).sum(), a2.gt_eq(&a1).sum());
1391    }
1392
1393    #[test]
1394    fn test_kleene() {
1395        let a = BooleanChunked::new(PlSmallStr::EMPTY, &[Some(true), Some(false), None]);
1396        let trues = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[true, true, true]);
1397        let falses = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[false, false, false]);
1398
1399        let c = &a | &trues;
1400        assert_eq!(Vec::from(&c), &[Some(true), Some(true), Some(true)]);
1401
1402        let c = &a | &falses;
1403        assert_eq!(Vec::from(&c), &[Some(true), Some(false), None])
1404    }
1405
1406    #[test]
1407    fn list_broadcasting_lists() {
1408        let s_el = Series::new(PlSmallStr::EMPTY, &[1, 2, 3]);
1409        let s_lhs = Series::new(PlSmallStr::EMPTY, &[s_el.clone(), s_el.clone()]);
1410        let s_rhs = Series::new(PlSmallStr::EMPTY, std::slice::from_ref(&s_el));
1411
1412        let result = s_lhs.list().unwrap().equal(s_rhs.list().unwrap());
1413        assert_eq!(result.len(), 2);
1414        assert!(result.all());
1415    }
1416
1417    #[test]
1418    fn test_broadcasting_bools() {
1419        let a = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[true, false, true]);
1420        let true_ = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[true]);
1421        let false_ = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[false]);
1422
1423        let out = a.equal(&true_);
1424        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1425        let out = true_.equal(&a);
1426        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1427        let out = a.equal(&false_);
1428        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1429        let out = false_.equal(&a);
1430        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1431
1432        let out = a.not_equal(&true_);
1433        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1434        let out = true_.not_equal(&a);
1435        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1436        let out = a.not_equal(&false_);
1437        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1438        let out = false_.not_equal(&a);
1439        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1440
1441        let out = a.gt(&true_);
1442        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1443        let out = true_.gt(&a);
1444        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1445        let out = a.gt(&false_);
1446        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1447        let out = false_.gt(&a);
1448        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1449
1450        let out = a.gt_eq(&true_);
1451        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1452        let out = true_.gt_eq(&a);
1453        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1454        let out = a.gt_eq(&false_);
1455        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1456        let out = false_.gt_eq(&a);
1457        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1458
1459        let out = a.lt(&true_);
1460        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1461        let out = true_.lt(&a);
1462        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1463        let out = a.lt(&false_);
1464        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1465        let out = false_.lt(&a);
1466        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1467
1468        let out = a.lt_eq(&true_);
1469        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1470        let out = true_.lt_eq(&a);
1471        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1472        let out = a.lt_eq(&false_);
1473        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1474        let out = false_.lt_eq(&a);
1475        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1476
1477        let a =
1478            BooleanChunked::from_slice_options(PlSmallStr::EMPTY, &[Some(true), Some(false), None]);
1479        let all_true = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[true, true, true]);
1480        let all_false = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[false, false, false]);
1481        let out = a.equal(&true_);
1482        assert_eq!(Vec::from(&out), &[Some(true), Some(false), None]);
1483        let out = a.not_equal(&true_);
1484        assert_eq!(Vec::from(&out), &[Some(false), Some(true), None]);
1485
1486        let out = a.equal(&all_true);
1487        assert_eq!(Vec::from(&out), &[Some(true), Some(false), None]);
1488        let out = a.not_equal(&all_true);
1489        assert_eq!(Vec::from(&out), &[Some(false), Some(true), None]);
1490        let out = a.equal(&false_);
1491        assert_eq!(Vec::from(&out), &[Some(false), Some(true), None]);
1492        let out = a.not_equal(&false_);
1493        assert_eq!(Vec::from(&out), &[Some(true), Some(false), None]);
1494        let out = a.equal(&all_false);
1495        assert_eq!(Vec::from(&out), &[Some(false), Some(true), None]);
1496        let out = a.not_equal(&all_false);
1497        assert_eq!(Vec::from(&out), &[Some(true), Some(false), None]);
1498    }
1499
1500    #[test]
1501    fn test_broadcasting_numeric() {
1502        let a = Int32Chunked::from_slice(PlSmallStr::EMPTY, &[1, 2, 3]);
1503        let one = Int32Chunked::from_slice(PlSmallStr::EMPTY, &[1]);
1504        let three = Int32Chunked::from_slice(PlSmallStr::EMPTY, &[3]);
1505
1506        let out = a.equal(&one);
1507        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(false)]);
1508        let out = one.equal(&a);
1509        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(false)]);
1510        let out = a.equal(&three);
1511        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(true)]);
1512        let out = three.equal(&a);
1513        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(true)]);
1514
1515        let out = a.not_equal(&one);
1516        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(true)]);
1517        let out = one.not_equal(&a);
1518        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(true)]);
1519        let out = a.not_equal(&three);
1520        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(false)]);
1521        let out = three.not_equal(&a);
1522        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(false)]);
1523
1524        let out = a.gt(&one);
1525        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(true)]);
1526        let out = one.gt(&a);
1527        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1528        let out = a.gt(&three);
1529        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1530        let out = three.gt(&a);
1531        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(false)]);
1532
1533        let out = a.lt(&one);
1534        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1535        let out = one.lt(&a);
1536        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(true)]);
1537        let out = a.lt(&three);
1538        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(false)]);
1539        let out = three.lt(&a);
1540        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1541
1542        let out = a.gt_eq(&one);
1543        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1544        let out = one.gt_eq(&a);
1545        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(false)]);
1546        let out = a.gt_eq(&three);
1547        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(true)]);
1548        let out = three.gt_eq(&a);
1549        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1550
1551        let out = a.lt_eq(&one);
1552        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(false)]);
1553        let out = one.lt_eq(&a);
1554        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1555        let out = a.lt_eq(&three);
1556        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1557        let out = three.lt_eq(&a);
1558        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(true)]);
1559    }
1560}