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
469macro_rules! binary_eq_ineq_impl {
470    ($($ca:ident),+) => {
471        $(
472        impl ChunkCompareEq<&$ca> for $ca {
473            type Item = BooleanChunked;
474
475            fn equal(&self, rhs: &$ca) -> BooleanChunked {
476                // Broadcast.
477                match (self.len(), rhs.len()) {
478                    (_, 1) => {
479                        if let Some(value) = rhs.get(0) {
480                            self.equal(value)
481                        } else {
482                            BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
483                        }
484                    },
485                    (1, _) => {
486                        if let Some(value) = self.get(0) {
487                            rhs.equal(value)
488                        } else {
489                            BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
490                        }
491                    },
492                    _ => arity::binary_mut_values(
493                        self,
494                        rhs,
495                        |a, b| a.tot_eq_kernel(b).into(),
496                        PlSmallStr::EMPTY,
497                    ),
498                }
499            }
500
501            fn equal_missing(&self, rhs: &$ca) -> BooleanChunked {
502                // Broadcast.
503                match (self.len(), rhs.len()) {
504                    (_, 1) => {
505                        if let Some(value) = rhs.get(0) {
506                            self.equal_missing(value)
507                        } else {
508                            self.is_null()
509                        }
510                    },
511                    (1, _) => {
512                        if let Some(value) = self.get(0) {
513                            rhs.equal_missing(value)
514                        } else {
515                            rhs.is_null()
516                        }
517                    },
518                    _ => arity::binary_mut_with_options(
519                        self,
520                        rhs,
521                        |a, b| a.tot_eq_missing_kernel(b).into(),
522                        PlSmallStr::EMPTY,
523                    ),
524                }
525            }
526
527            fn not_equal(&self, rhs: &$ca) -> BooleanChunked {
528                // Broadcast.
529                match (self.len(), rhs.len()) {
530                    (_, 1) => {
531                        if let Some(value) = rhs.get(0) {
532                            self.not_equal(value)
533                        } else {
534                            BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
535                        }
536                    },
537                    (1, _) => {
538                        if let Some(value) = self.get(0) {
539                            rhs.not_equal(value)
540                        } else {
541                            BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
542                        }
543                    },
544                    _ => arity::binary_mut_values(
545                        self,
546                        rhs,
547                        |a, b| a.tot_ne_kernel(b).into(),
548                        PlSmallStr::EMPTY,
549                    ),
550                }
551            }
552
553            fn not_equal_missing(&self, rhs: &$ca) -> BooleanChunked {
554                // Broadcast.
555                match (self.len(), rhs.len()) {
556                    (_, 1) => {
557                        if let Some(value) = rhs.get(0) {
558                            self.not_equal_missing(value)
559                        } else {
560                            self.is_not_null()
561                        }
562                    },
563                    (1, _) => {
564                        if let Some(value) = self.get(0) {
565                            rhs.not_equal_missing(value)
566                        } else {
567                            rhs.is_not_null()
568                        }
569                    },
570                    _ => arity::binary_mut_with_options(
571                        self,
572                        rhs,
573                        |a, b| a.tot_ne_missing_kernel(b).into(),
574                        PlSmallStr::EMPTY,
575                    ),
576                }
577            }
578        }
579
580        impl ChunkCompareIneq<&$ca> for $ca {
581            type Item = BooleanChunked;
582
583            fn lt(&self, rhs: &$ca) -> BooleanChunked {
584                // Broadcast.
585                match (self.len(), rhs.len()) {
586                    (_, 1) => {
587                        if let Some(value) = rhs.get(0) {
588                            self.lt(value)
589                        } else {
590                            BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
591                        }
592                    },
593                    (1, _) => {
594                        if let Some(value) = self.get(0) {
595                            rhs.gt(value)
596                        } else {
597                            BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
598                        }
599                    },
600                    _ => arity::binary_mut_values(
601                        self,
602                        rhs,
603                        |a, b| a.tot_lt_kernel(b).into(),
604                        PlSmallStr::EMPTY,
605                    ),
606                }
607            }
608
609            fn lt_eq(&self, rhs: &$ca) -> BooleanChunked {
610                // Broadcast.
611                match (self.len(), rhs.len()) {
612                    (_, 1) => {
613                        if let Some(value) = rhs.get(0) {
614                            self.lt_eq(value)
615                        } else {
616                            BooleanChunked::full_null(PlSmallStr::EMPTY, self.len())
617                        }
618                    },
619                    (1, _) => {
620                        if let Some(value) = self.get(0) {
621                            rhs.gt_eq(value)
622                        } else {
623                            BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len())
624                        }
625                    },
626                    _ => arity::binary_mut_values(
627                        self,
628                        rhs,
629                        |a, b| a.tot_le_kernel(b).into(),
630                        PlSmallStr::EMPTY,
631                    ),
632                }
633            }
634
635            fn gt(&self, rhs: &Self) -> BooleanChunked {
636                rhs.lt(self)
637            }
638
639            fn gt_eq(&self, rhs: &Self) -> BooleanChunked {
640                rhs.lt_eq(self)
641            }
642        }
643        )+
644    };
645}
646
647binary_eq_ineq_impl!(BinaryChunked, BinaryOffsetChunked);
648
649fn _list_comparison_helper<F, B>(
650    lhs: &ListChunked,
651    rhs: &ListChunked,
652    op: F,
653    broadcast_op: B,
654    missing: bool,
655    is_ne: bool,
656) -> BooleanChunked
657where
658    F: Fn(&ListArray<i64>, &ListArray<i64>) -> Bitmap,
659    B: Fn(&ListArray<i64>, &Box<dyn Array>) -> Bitmap,
660{
661    match (lhs.len(), rhs.len()) {
662        (_, 1) => {
663            let right = rhs
664                .downcast_iter()
665                .find(|x| !x.is_empty())
666                .unwrap()
667                .as_any()
668                .downcast_ref::<ListArray<i64>>()
669                .unwrap();
670
671            if !right.validity().is_none_or(|v| v.get(0).unwrap()) {
672                if missing {
673                    if is_ne {
674                        return lhs.is_not_null();
675                    } else {
676                        return lhs.is_null();
677                    }
678                } else {
679                    return BooleanChunked::full_null(PlSmallStr::EMPTY, lhs.len());
680                }
681            }
682
683            let values = right.values().sliced(
684                (*right.offsets().first()).try_into().unwrap(),
685                right.offsets().range().try_into().unwrap(),
686            );
687
688            if missing {
689                arity::unary_mut_with_options(lhs, |a| broadcast_op(a, &values).into())
690            } else {
691                arity::unary_mut_values(lhs, |a| broadcast_op(a, &values).into())
692            }
693        },
694        (1, _) => {
695            let left = lhs
696                .downcast_iter()
697                .find(|x| !x.is_empty())
698                .unwrap()
699                .as_any()
700                .downcast_ref::<ListArray<i64>>()
701                .unwrap();
702
703            if !left.validity().is_none_or(|v| v.get(0).unwrap()) {
704                if missing {
705                    if is_ne {
706                        return rhs.is_not_null();
707                    } else {
708                        return rhs.is_null();
709                    }
710                } else {
711                    return BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len());
712                }
713            }
714
715            let values = left.values().sliced(
716                (*left.offsets().first()).try_into().unwrap(),
717                left.offsets().range().try_into().unwrap(),
718            );
719
720            if missing {
721                arity::unary_mut_with_options(rhs, |a| broadcast_op(a, &values).into())
722            } else {
723                arity::unary_mut_values(rhs, |a| broadcast_op(a, &values).into())
724            }
725        },
726        _ => {
727            if missing {
728                arity::binary_mut_with_options(lhs, rhs, |a, b| op(a, b).into(), PlSmallStr::EMPTY)
729            } else {
730                arity::binary_mut_values(lhs, rhs, |a, b| op(a, b).into(), PlSmallStr::EMPTY)
731            }
732        },
733    }
734}
735
736impl ChunkCompareEq<&ListChunked> for ListChunked {
737    type Item = BooleanChunked;
738    fn equal(&self, rhs: &ListChunked) -> BooleanChunked {
739        _list_comparison_helper(
740            self,
741            rhs,
742            TotalEqKernel::tot_eq_kernel,
743            TotalEqKernel::tot_eq_kernel_broadcast,
744            false,
745            false,
746        )
747    }
748
749    fn equal_missing(&self, rhs: &ListChunked) -> BooleanChunked {
750        _list_comparison_helper(
751            self,
752            rhs,
753            TotalEqKernel::tot_eq_missing_kernel,
754            TotalEqKernel::tot_eq_missing_kernel_broadcast,
755            true,
756            false,
757        )
758    }
759
760    fn not_equal(&self, rhs: &ListChunked) -> BooleanChunked {
761        _list_comparison_helper(
762            self,
763            rhs,
764            TotalEqKernel::tot_ne_kernel,
765            TotalEqKernel::tot_ne_kernel_broadcast,
766            false,
767            true,
768        )
769    }
770
771    fn not_equal_missing(&self, rhs: &ListChunked) -> BooleanChunked {
772        _list_comparison_helper(
773            self,
774            rhs,
775            TotalEqKernel::tot_ne_missing_kernel,
776            TotalEqKernel::tot_ne_missing_kernel_broadcast,
777            true,
778            true,
779        )
780    }
781}
782
783#[cfg(feature = "dtype-struct")]
784fn struct_helper<F, R>(
785    a: &StructChunked,
786    b: &StructChunked,
787    op: F,
788    reduce: R,
789    op_is_ne: bool,
790    is_missing: bool,
791) -> BooleanChunked
792where
793    F: Fn(&Series, &Series) -> BooleanChunked,
794    R: Fn(BooleanChunked, BooleanChunked) -> BooleanChunked,
795{
796    let len_a = a.len();
797    let len_b = b.len();
798    let broadcasts = len_a == 1 || len_b == 1;
799    if (a.len() != b.len() && !broadcasts) || a.struct_fields().len() != b.struct_fields().len() {
800        BooleanChunked::full(PlSmallStr::EMPTY, op_is_ne, a.len())
801    } else {
802        let (a, b) = align_chunks_binary(a, b);
803
804        let mut out = a
805            .fields_as_series()
806            .iter()
807            .zip(b.fields_as_series().iter())
808            .map(|(l, r)| op(l, r))
809            .reduce(&reduce)
810            .unwrap_or_else(|| BooleanChunked::full(PlSmallStr::EMPTY, !op_is_ne, a.len()));
811
812        if is_missing && (a.has_nulls() || b.has_nulls()) {
813            // Do some allocations so that we can use the Series dispatch, it otherwise
814            // gets complicated dealing with combinations of ==, != and broadcasting.
815            let default = || {
816                BooleanChunked::with_chunk(PlSmallStr::EMPTY, BooleanArray::from_slice([true]))
817                    .into_series()
818            };
819            let validity_to_series = |x| unsafe {
820                BooleanChunked::with_chunk(
821                    PlSmallStr::EMPTY,
822                    BooleanArray::from_inner_unchecked(ArrowDataType::Boolean, x, None),
823                )
824                .into_series()
825            };
826
827            out = reduce(
828                out,
829                op(
830                    &a.rechunk_validity()
831                        .map_or_else(default, validity_to_series),
832                    &b.rechunk_validity()
833                        .map_or_else(default, validity_to_series),
834                ),
835            )
836        }
837
838        if !is_missing && (a.null_count() > 0 || b.null_count() > 0) {
839            let mut a = a;
840            let mut b = b;
841
842            if broadcasts {
843                if a.len() == 1 {
844                    a = std::borrow::Cow::Owned(a.new_from_index(0, b.len()));
845                }
846                if b.len() == 1 {
847                    b = std::borrow::Cow::Owned(b.new_from_index(0, a.len()));
848                }
849            }
850
851            let mut a = a.into_owned();
852            a.zip_outer_validity(&b);
853            unsafe {
854                let mut new_null_count = 0;
855                for (arr, a) in out.downcast_iter_mut().zip(a.downcast_iter()) {
856                    arr.set_validity(a.validity().cloned());
857                    new_null_count += arr.null_count();
858                }
859                out.set_null_count(new_null_count);
860            }
861        }
862
863        out
864    }
865}
866
867#[cfg(feature = "dtype-struct")]
868impl ChunkCompareEq<&StructChunked> for StructChunked {
869    type Item = BooleanChunked;
870    fn equal(&self, rhs: &StructChunked) -> BooleanChunked {
871        struct_helper(
872            self,
873            rhs,
874            |l, r| l.equal_missing(r).unwrap(),
875            |a, b| a.bitand(b),
876            false,
877            false,
878        )
879    }
880
881    fn equal_missing(&self, rhs: &StructChunked) -> BooleanChunked {
882        struct_helper(
883            self,
884            rhs,
885            |l, r| l.equal_missing(r).unwrap(),
886            |a, b| a.bitand(b),
887            false,
888            true,
889        )
890    }
891
892    fn not_equal(&self, rhs: &StructChunked) -> BooleanChunked {
893        struct_helper(
894            self,
895            rhs,
896            |l, r| l.not_equal_missing(r).unwrap(),
897            |a, b| a | b,
898            true,
899            false,
900        )
901    }
902
903    fn not_equal_missing(&self, rhs: &StructChunked) -> BooleanChunked {
904        struct_helper(
905            self,
906            rhs,
907            |l, r| l.not_equal_missing(r).unwrap(),
908            |a, b| a | b,
909            true,
910            true,
911        )
912    }
913}
914
915#[cfg(feature = "dtype-array")]
916fn _array_comparison_helper<F, B>(
917    lhs: &ArrayChunked,
918    rhs: &ArrayChunked,
919    op: F,
920    broadcast_op: B,
921    missing: bool,
922    is_ne: bool,
923) -> BooleanChunked
924where
925    F: Fn(&FixedSizeListArray, &FixedSizeListArray) -> Bitmap,
926    B: Fn(&FixedSizeListArray, &Box<dyn Array>) -> Bitmap,
927{
928    match (lhs.len(), rhs.len()) {
929        (_, 1) => {
930            let right = rhs
931                .downcast_iter()
932                .find(|x| !x.is_empty())
933                .unwrap()
934                .as_any()
935                .downcast_ref::<FixedSizeListArray>()
936                .unwrap();
937
938            if !right.validity().is_none_or(|v| v.get(0).unwrap()) {
939                if missing {
940                    if is_ne {
941                        return lhs.is_not_null();
942                    } else {
943                        return lhs.is_null();
944                    }
945                } else {
946                    return BooleanChunked::full_null(PlSmallStr::EMPTY, lhs.len());
947                }
948            }
949
950            if missing {
951                arity::unary_mut_with_options(lhs, |a| broadcast_op(a, right.values()).into())
952            } else {
953                arity::unary_mut_values(lhs, |a| broadcast_op(a, right.values()).into())
954            }
955        },
956        (1, _) => {
957            let left = lhs
958                .downcast_iter()
959                .find(|x| !x.is_empty())
960                .unwrap()
961                .as_any()
962                .downcast_ref::<FixedSizeListArray>()
963                .unwrap();
964
965            if !left.validity().is_none_or(|v| v.get(0).unwrap()) {
966                if missing {
967                    if is_ne {
968                        return rhs.is_not_null();
969                    } else {
970                        return rhs.is_null();
971                    }
972                } else {
973                    return BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len());
974                }
975            }
976
977            if missing {
978                arity::unary_mut_with_options(rhs, |a| broadcast_op(a, left.values()).into())
979            } else {
980                arity::unary_mut_values(rhs, |a| broadcast_op(a, left.values()).into())
981            }
982        },
983        _ => {
984            if missing {
985                arity::binary_mut_with_options(lhs, rhs, |a, b| op(a, b).into(), PlSmallStr::EMPTY)
986            } else {
987                arity::binary_mut_values(lhs, rhs, |a, b| op(a, b).into(), PlSmallStr::EMPTY)
988            }
989        },
990    }
991}
992
993#[cfg(feature = "dtype-array")]
994impl ChunkCompareEq<&ArrayChunked> for ArrayChunked {
995    type Item = BooleanChunked;
996    fn equal(&self, rhs: &ArrayChunked) -> BooleanChunked {
997        _array_comparison_helper(
998            self,
999            rhs,
1000            TotalEqKernel::tot_eq_kernel,
1001            TotalEqKernel::tot_eq_kernel_broadcast,
1002            false,
1003            false,
1004        )
1005    }
1006
1007    fn equal_missing(&self, rhs: &ArrayChunked) -> BooleanChunked {
1008        _array_comparison_helper(
1009            self,
1010            rhs,
1011            TotalEqKernel::tot_eq_missing_kernel,
1012            TotalEqKernel::tot_eq_missing_kernel_broadcast,
1013            true,
1014            false,
1015        )
1016    }
1017
1018    fn not_equal(&self, rhs: &ArrayChunked) -> BooleanChunked {
1019        _array_comparison_helper(
1020            self,
1021            rhs,
1022            TotalEqKernel::tot_ne_kernel,
1023            TotalEqKernel::tot_ne_kernel_broadcast,
1024            false,
1025            true,
1026        )
1027    }
1028
1029    fn not_equal_missing(&self, rhs: &ArrayChunked) -> Self::Item {
1030        _array_comparison_helper(
1031            self,
1032            rhs,
1033            TotalEqKernel::tot_ne_missing_kernel,
1034            TotalEqKernel::tot_ne_missing_kernel_broadcast,
1035            true,
1036            true,
1037        )
1038    }
1039}
1040
1041impl Not for &BooleanChunked {
1042    type Output = BooleanChunked;
1043
1044    fn not(self) -> Self::Output {
1045        let chunks = self.downcast_iter().map(compute::boolean::not);
1046        ChunkedArray::from_chunk_iter(self.name().clone(), chunks)
1047    }
1048}
1049
1050impl Not for BooleanChunked {
1051    type Output = BooleanChunked;
1052
1053    fn not(self) -> Self::Output {
1054        (&self).not()
1055    }
1056}
1057
1058impl BooleanChunked {
1059    /// Returns whether any of the values in the column are `true`.
1060    ///
1061    /// Null values are ignored.
1062    pub fn any(&self) -> bool {
1063        self.downcast_iter().any(compute::boolean::any)
1064    }
1065
1066    /// Returns whether all values in the array are `true`.
1067    ///
1068    /// Null values are ignored.
1069    pub fn all(&self) -> bool {
1070        self.downcast_iter().all(compute::boolean::all)
1071    }
1072
1073    /// Returns whether any of the values in the column are `true`.
1074    ///
1075    /// The output is unknown (`None`) if the array contains any null values and
1076    /// no `true` values.
1077    pub fn any_kleene(&self) -> Option<bool> {
1078        let mut result = Some(false);
1079        for arr in self.downcast_iter() {
1080            match compute::boolean_kleene::any(arr) {
1081                Some(true) => return Some(true),
1082                None => result = None,
1083                _ => (),
1084            };
1085        }
1086        result
1087    }
1088
1089    /// Returns whether all values in the column are `true`.
1090    ///
1091    /// The output is unknown (`None`) if the array contains any null values and
1092    /// no `false` values.
1093    pub fn all_kleene(&self) -> Option<bool> {
1094        let mut result = Some(true);
1095        for arr in self.downcast_iter() {
1096            match compute::boolean_kleene::all(arr) {
1097                Some(false) => return Some(false),
1098                None => result = None,
1099                _ => (),
1100            };
1101        }
1102        result
1103    }
1104}
1105
1106// private
1107pub(crate) trait ChunkEqualElement {
1108    /// Only meant for physical types.
1109    /// Check if element in self is equal to element in other, assumes same dtypes
1110    ///
1111    /// # Safety
1112    ///
1113    /// No type checks.
1114    unsafe fn equal_element(&self, _idx_self: usize, _idx_other: usize, _other: &Series) -> bool {
1115        unimplemented!()
1116    }
1117}
1118
1119impl<T> ChunkEqualElement for ChunkedArray<T>
1120where
1121    T: PolarsNumericType,
1122{
1123    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
1124        let ca_other = other.as_ref().as_ref();
1125        debug_assert!(self.dtype() == other.dtype());
1126        let ca_other = &*(ca_other as *const ChunkedArray<T>);
1127        // Should be get and not get_unchecked, because there could be nulls
1128        self.get_unchecked(idx_self)
1129            .tot_eq(&ca_other.get_unchecked(idx_other))
1130    }
1131}
1132
1133impl ChunkEqualElement for BooleanChunked {
1134    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
1135        let ca_other = other.as_ref().as_ref();
1136        debug_assert!(self.dtype() == other.dtype());
1137        let ca_other = &*(ca_other as *const BooleanChunked);
1138        self.get_unchecked(idx_self) == ca_other.get_unchecked(idx_other)
1139    }
1140}
1141
1142impl ChunkEqualElement for StringChunked {
1143    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
1144        let ca_other = other.as_ref().as_ref();
1145        debug_assert!(self.dtype() == other.dtype());
1146        let ca_other = &*(ca_other as *const StringChunked);
1147        self.get_unchecked(idx_self) == ca_other.get_unchecked(idx_other)
1148    }
1149}
1150
1151impl ChunkEqualElement for BinaryChunked {
1152    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
1153        let ca_other = other.as_ref().as_ref();
1154        debug_assert!(self.dtype() == other.dtype());
1155        let ca_other = &*(ca_other as *const BinaryChunked);
1156        self.get_unchecked(idx_self) == ca_other.get_unchecked(idx_other)
1157    }
1158}
1159
1160impl ChunkEqualElement for BinaryOffsetChunked {
1161    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
1162        let ca_other = other.as_ref().as_ref();
1163        debug_assert!(self.dtype() == other.dtype());
1164        let ca_other = &*(ca_other as *const BinaryOffsetChunked);
1165        self.get_unchecked(idx_self) == ca_other.get_unchecked(idx_other)
1166    }
1167}
1168
1169impl ChunkEqualElement for ListChunked {}
1170#[cfg(feature = "dtype-array")]
1171impl ChunkEqualElement for ArrayChunked {}
1172
1173#[cfg(test)]
1174#[cfg_attr(feature = "nightly", allow(clippy::manual_repeat_n))] // remove once stable
1175mod test {
1176    use std::iter::repeat_n;
1177
1178    use super::super::test::get_chunked_array;
1179    use crate::prelude::*;
1180
1181    pub(crate) fn create_two_chunked() -> (Int32Chunked, Int32Chunked) {
1182        let mut a1 = Int32Chunked::new(PlSmallStr::from_static("a"), &[1, 2, 3]);
1183        let a2 = Int32Chunked::new(PlSmallStr::from_static("a"), &[4, 5, 6]);
1184        let a3 = Int32Chunked::new(PlSmallStr::from_static("a"), &[1, 2, 3, 4, 5, 6]);
1185        a1.append(&a2).unwrap();
1186        (a1, a3)
1187    }
1188
1189    #[test]
1190    fn test_bitwise_ops() {
1191        let a = BooleanChunked::new(PlSmallStr::from_static("a"), &[true, false, false]);
1192        let b = BooleanChunked::new(
1193            PlSmallStr::from_static("b"),
1194            &[Some(true), Some(true), None],
1195        );
1196        assert_eq!(Vec::from(&a | &b), &[Some(true), Some(true), None]);
1197        assert_eq!(Vec::from(&a & &b), &[Some(true), Some(false), Some(false)]);
1198        assert_eq!(Vec::from(!b), &[Some(false), Some(false), None]);
1199    }
1200
1201    #[test]
1202    fn test_compare_chunk_diff() {
1203        let (a1, a2) = create_two_chunked();
1204
1205        assert_eq!(
1206            a1.equal(&a2).into_iter().collect::<Vec<_>>(),
1207            repeat_n(Some(true), 6).collect::<Vec<_>>()
1208        );
1209        assert_eq!(
1210            a2.equal(&a1).into_iter().collect::<Vec<_>>(),
1211            repeat_n(Some(true), 6).collect::<Vec<_>>()
1212        );
1213        assert_eq!(
1214            a1.not_equal(&a2).into_iter().collect::<Vec<_>>(),
1215            repeat_n(Some(false), 6).collect::<Vec<_>>()
1216        );
1217        assert_eq!(
1218            a2.not_equal(&a1).into_iter().collect::<Vec<_>>(),
1219            repeat_n(Some(false), 6).collect::<Vec<_>>()
1220        );
1221        assert_eq!(
1222            a1.gt(&a2).into_iter().collect::<Vec<_>>(),
1223            repeat_n(Some(false), 6).collect::<Vec<_>>()
1224        );
1225        assert_eq!(
1226            a2.gt(&a1).into_iter().collect::<Vec<_>>(),
1227            repeat_n(Some(false), 6).collect::<Vec<_>>()
1228        );
1229        assert_eq!(
1230            a1.gt_eq(&a2).into_iter().collect::<Vec<_>>(),
1231            repeat_n(Some(true), 6).collect::<Vec<_>>()
1232        );
1233        assert_eq!(
1234            a2.gt_eq(&a1).into_iter().collect::<Vec<_>>(),
1235            repeat_n(Some(true), 6).collect::<Vec<_>>()
1236        );
1237        assert_eq!(
1238            a1.lt_eq(&a2).into_iter().collect::<Vec<_>>(),
1239            repeat_n(Some(true), 6).collect::<Vec<_>>()
1240        );
1241        assert_eq!(
1242            a2.lt_eq(&a1).into_iter().collect::<Vec<_>>(),
1243            repeat_n(Some(true), 6).collect::<Vec<_>>()
1244        );
1245        assert_eq!(
1246            a1.lt(&a2).into_iter().collect::<Vec<_>>(),
1247            repeat_n(Some(false), 6).collect::<Vec<_>>()
1248        );
1249        assert_eq!(
1250            a2.lt(&a1).into_iter().collect::<Vec<_>>(),
1251            repeat_n(Some(false), 6).collect::<Vec<_>>()
1252        );
1253    }
1254
1255    #[test]
1256    fn test_equal_chunks() {
1257        let a1 = get_chunked_array();
1258        let a2 = get_chunked_array();
1259
1260        assert_eq!(
1261            a1.equal(&a2).into_iter().collect::<Vec<_>>(),
1262            repeat_n(Some(true), 3).collect::<Vec<_>>()
1263        );
1264        assert_eq!(
1265            a2.equal(&a1).into_iter().collect::<Vec<_>>(),
1266            repeat_n(Some(true), 3).collect::<Vec<_>>()
1267        );
1268        assert_eq!(
1269            a1.not_equal(&a2).into_iter().collect::<Vec<_>>(),
1270            repeat_n(Some(false), 3).collect::<Vec<_>>()
1271        );
1272        assert_eq!(
1273            a2.not_equal(&a1).into_iter().collect::<Vec<_>>(),
1274            repeat_n(Some(false), 3).collect::<Vec<_>>()
1275        );
1276        assert_eq!(
1277            a1.gt(&a2).into_iter().collect::<Vec<_>>(),
1278            repeat_n(Some(false), 3).collect::<Vec<_>>()
1279        );
1280        assert_eq!(
1281            a2.gt(&a1).into_iter().collect::<Vec<_>>(),
1282            repeat_n(Some(false), 3).collect::<Vec<_>>()
1283        );
1284        assert_eq!(
1285            a1.gt_eq(&a2).into_iter().collect::<Vec<_>>(),
1286            repeat_n(Some(true), 3).collect::<Vec<_>>()
1287        );
1288        assert_eq!(
1289            a2.gt_eq(&a1).into_iter().collect::<Vec<_>>(),
1290            repeat_n(Some(true), 3).collect::<Vec<_>>()
1291        );
1292        assert_eq!(
1293            a1.lt_eq(&a2).into_iter().collect::<Vec<_>>(),
1294            repeat_n(Some(true), 3).collect::<Vec<_>>()
1295        );
1296        assert_eq!(
1297            a2.lt_eq(&a1).into_iter().collect::<Vec<_>>(),
1298            repeat_n(Some(true), 3).collect::<Vec<_>>()
1299        );
1300        assert_eq!(
1301            a1.lt(&a2).into_iter().collect::<Vec<_>>(),
1302            repeat_n(Some(false), 3).collect::<Vec<_>>()
1303        );
1304        assert_eq!(
1305            a2.lt(&a1).into_iter().collect::<Vec<_>>(),
1306            repeat_n(Some(false), 3).collect::<Vec<_>>()
1307        );
1308    }
1309
1310    #[test]
1311    fn test_null_handling() {
1312        // assert we comply with arrows way of handling null data
1313        // we check comparison on two arrays with one chunk and verify it is equal to a differently
1314        // chunked array comparison.
1315
1316        // two same chunked arrays
1317        let a1: Int32Chunked = [Some(1), None, Some(3)].iter().copied().collect();
1318        let a2: Int32Chunked = [Some(1), Some(2), Some(3)].iter().copied().collect();
1319
1320        let mut a2_2chunks: Int32Chunked = [Some(1), Some(2)].iter().copied().collect();
1321        a2_2chunks
1322            .append(&[Some(3)].iter().copied().collect())
1323            .unwrap();
1324
1325        assert_eq!(
1326            a1.equal(&a2).into_iter().collect::<Vec<_>>(),
1327            a1.equal(&a2_2chunks).into_iter().collect::<Vec<_>>()
1328        );
1329
1330        assert_eq!(
1331            a1.not_equal(&a2).into_iter().collect::<Vec<_>>(),
1332            a1.not_equal(&a2_2chunks).into_iter().collect::<Vec<_>>()
1333        );
1334        assert_eq!(
1335            a1.not_equal(&a2).into_iter().collect::<Vec<_>>(),
1336            a2_2chunks.not_equal(&a1).into_iter().collect::<Vec<_>>()
1337        );
1338
1339        assert_eq!(
1340            a1.gt(&a2).into_iter().collect::<Vec<_>>(),
1341            a1.gt(&a2_2chunks).into_iter().collect::<Vec<_>>()
1342        );
1343        assert_eq!(
1344            a1.gt(&a2).into_iter().collect::<Vec<_>>(),
1345            a2_2chunks.gt(&a1).into_iter().collect::<Vec<_>>()
1346        );
1347
1348        assert_eq!(
1349            a1.gt_eq(&a2).into_iter().collect::<Vec<_>>(),
1350            a1.gt_eq(&a2_2chunks).into_iter().collect::<Vec<_>>()
1351        );
1352        assert_eq!(
1353            a1.gt_eq(&a2).into_iter().collect::<Vec<_>>(),
1354            a2_2chunks.gt_eq(&a1).into_iter().collect::<Vec<_>>()
1355        );
1356
1357        assert_eq!(
1358            a1.lt_eq(&a2).into_iter().collect::<Vec<_>>(),
1359            a1.lt_eq(&a2_2chunks).into_iter().collect::<Vec<_>>()
1360        );
1361        assert_eq!(
1362            a1.lt_eq(&a2).into_iter().collect::<Vec<_>>(),
1363            a2_2chunks.lt_eq(&a1).into_iter().collect::<Vec<_>>()
1364        );
1365
1366        assert_eq!(
1367            a1.lt(&a2).into_iter().collect::<Vec<_>>(),
1368            a1.lt(&a2_2chunks).into_iter().collect::<Vec<_>>()
1369        );
1370        assert_eq!(
1371            a1.lt(&a2).into_iter().collect::<Vec<_>>(),
1372            a2_2chunks.lt(&a1).into_iter().collect::<Vec<_>>()
1373        );
1374    }
1375
1376    #[test]
1377    fn test_left_right() {
1378        // This failed with arrow comparisons.
1379        // sliced
1380        let a1: Int32Chunked = [Some(1), Some(2)].iter().copied().collect();
1381        let a1 = a1.slice(1, 1);
1382        let a2: Int32Chunked = [Some(2)].iter().copied().collect();
1383        assert_eq!(a1.equal(&a2).sum(), a2.equal(&a1).sum());
1384        assert_eq!(a1.not_equal(&a2).sum(), a2.not_equal(&a1).sum());
1385        assert_eq!(a1.gt(&a2).sum(), a2.gt(&a1).sum());
1386        assert_eq!(a1.lt(&a2).sum(), a2.lt(&a1).sum());
1387        assert_eq!(a1.lt_eq(&a2).sum(), a2.lt_eq(&a1).sum());
1388        assert_eq!(a1.gt_eq(&a2).sum(), a2.gt_eq(&a1).sum());
1389
1390        let a1: StringChunked = ["a", "b"].iter().copied().collect();
1391        let a1 = a1.slice(1, 1);
1392        let a2: StringChunked = ["b"].iter().copied().collect();
1393        assert_eq!(a1.equal(&a2).sum(), a2.equal(&a1).sum());
1394        assert_eq!(a1.not_equal(&a2).sum(), a2.not_equal(&a1).sum());
1395        assert_eq!(a1.gt(&a2).sum(), a2.gt(&a1).sum());
1396        assert_eq!(a1.lt(&a2).sum(), a2.lt(&a1).sum());
1397        assert_eq!(a1.lt_eq(&a2).sum(), a2.lt_eq(&a1).sum());
1398        assert_eq!(a1.gt_eq(&a2).sum(), a2.gt_eq(&a1).sum());
1399    }
1400
1401    #[test]
1402    fn test_kleene() {
1403        let a = BooleanChunked::new(PlSmallStr::EMPTY, &[Some(true), Some(false), None]);
1404        let trues = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[true, true, true]);
1405        let falses = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[false, false, false]);
1406
1407        let c = &a | &trues;
1408        assert_eq!(Vec::from(&c), &[Some(true), Some(true), Some(true)]);
1409
1410        let c = &a | &falses;
1411        assert_eq!(Vec::from(&c), &[Some(true), Some(false), None])
1412    }
1413
1414    #[test]
1415    fn list_broadcasting_lists() {
1416        let s_el = Series::new(PlSmallStr::EMPTY, &[1, 2, 3]);
1417        let s_lhs = Series::new(PlSmallStr::EMPTY, &[s_el.clone(), s_el.clone()]);
1418        let s_rhs = Series::new(PlSmallStr::EMPTY, std::slice::from_ref(&s_el));
1419
1420        let result = s_lhs.list().unwrap().equal(s_rhs.list().unwrap());
1421        assert_eq!(result.len(), 2);
1422        assert!(result.all());
1423    }
1424
1425    #[test]
1426    fn test_broadcasting_bools() {
1427        let a = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[true, false, true]);
1428        let true_ = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[true]);
1429        let false_ = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[false]);
1430
1431        let out = a.equal(&true_);
1432        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1433        let out = true_.equal(&a);
1434        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1435        let out = a.equal(&false_);
1436        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1437        let out = false_.equal(&a);
1438        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1439
1440        let out = a.not_equal(&true_);
1441        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1442        let out = true_.not_equal(&a);
1443        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1444        let out = a.not_equal(&false_);
1445        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1446        let out = false_.not_equal(&a);
1447        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1448
1449        let out = a.gt(&true_);
1450        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1451        let out = true_.gt(&a);
1452        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1453        let out = a.gt(&false_);
1454        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1455        let out = false_.gt(&a);
1456        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1457
1458        let out = a.gt_eq(&true_);
1459        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1460        let out = true_.gt_eq(&a);
1461        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1462        let out = a.gt_eq(&false_);
1463        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1464        let out = false_.gt_eq(&a);
1465        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1466
1467        let out = a.lt(&true_);
1468        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1469        let out = true_.lt(&a);
1470        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1471        let out = a.lt(&false_);
1472        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1473        let out = false_.lt(&a);
1474        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1475
1476        let out = a.lt_eq(&true_);
1477        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1478        let out = true_.lt_eq(&a);
1479        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1480        let out = a.lt_eq(&false_);
1481        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1482        let out = false_.lt_eq(&a);
1483        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1484
1485        let a =
1486            BooleanChunked::from_slice_options(PlSmallStr::EMPTY, &[Some(true), Some(false), None]);
1487        let all_true = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[true, true, true]);
1488        let all_false = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[false, false, false]);
1489        let out = a.equal(&true_);
1490        assert_eq!(Vec::from(&out), &[Some(true), Some(false), None]);
1491        let out = a.not_equal(&true_);
1492        assert_eq!(Vec::from(&out), &[Some(false), Some(true), None]);
1493
1494        let out = a.equal(&all_true);
1495        assert_eq!(Vec::from(&out), &[Some(true), Some(false), None]);
1496        let out = a.not_equal(&all_true);
1497        assert_eq!(Vec::from(&out), &[Some(false), Some(true), None]);
1498        let out = a.equal(&false_);
1499        assert_eq!(Vec::from(&out), &[Some(false), Some(true), None]);
1500        let out = a.not_equal(&false_);
1501        assert_eq!(Vec::from(&out), &[Some(true), Some(false), None]);
1502        let out = a.equal(&all_false);
1503        assert_eq!(Vec::from(&out), &[Some(false), Some(true), None]);
1504        let out = a.not_equal(&all_false);
1505        assert_eq!(Vec::from(&out), &[Some(true), Some(false), None]);
1506    }
1507
1508    #[test]
1509    fn test_broadcasting_numeric() {
1510        let a = Int32Chunked::from_slice(PlSmallStr::EMPTY, &[1, 2, 3]);
1511        let one = Int32Chunked::from_slice(PlSmallStr::EMPTY, &[1]);
1512        let three = Int32Chunked::from_slice(PlSmallStr::EMPTY, &[3]);
1513
1514        let out = a.equal(&one);
1515        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(false)]);
1516        let out = one.equal(&a);
1517        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(false)]);
1518        let out = a.equal(&three);
1519        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(true)]);
1520        let out = three.equal(&a);
1521        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(true)]);
1522
1523        let out = a.not_equal(&one);
1524        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(true)]);
1525        let out = one.not_equal(&a);
1526        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(true)]);
1527        let out = a.not_equal(&three);
1528        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(false)]);
1529        let out = three.not_equal(&a);
1530        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(false)]);
1531
1532        let out = a.gt(&one);
1533        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(true)]);
1534        let out = one.gt(&a);
1535        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1536        let out = a.gt(&three);
1537        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1538        let out = three.gt(&a);
1539        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(false)]);
1540
1541        let out = a.lt(&one);
1542        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1543        let out = one.lt(&a);
1544        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(true)]);
1545        let out = a.lt(&three);
1546        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(false)]);
1547        let out = three.lt(&a);
1548        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1549
1550        let out = a.gt_eq(&one);
1551        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1552        let out = one.gt_eq(&a);
1553        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(false)]);
1554        let out = a.gt_eq(&three);
1555        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(true)]);
1556        let out = three.gt_eq(&a);
1557        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1558
1559        let out = a.lt_eq(&one);
1560        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(false)]);
1561        let out = one.lt_eq(&a);
1562        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1563        let out = a.lt_eq(&three);
1564        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1565        let out = three.lt_eq(&a);
1566        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(true)]);
1567    }
1568}