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::implementations::null::NullChunked;
16use crate::series::IsSorted;
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.into_owned();
832            a.zip_outer_validity(&b);
833            unsafe {
834                for (arr, a) in out.downcast_iter_mut().zip(a.downcast_iter()) {
835                    arr.set_validity(a.validity().cloned())
836                }
837            }
838        }
839
840        out
841    }
842}
843
844#[cfg(feature = "dtype-struct")]
845impl ChunkCompareEq<&StructChunked> for StructChunked {
846    type Item = BooleanChunked;
847    fn equal(&self, rhs: &StructChunked) -> BooleanChunked {
848        struct_helper(
849            self,
850            rhs,
851            |l, r| l.equal_missing(r).unwrap(),
852            |a, b| a.bitand(b),
853            false,
854            false,
855        )
856    }
857
858    fn equal_missing(&self, rhs: &StructChunked) -> BooleanChunked {
859        struct_helper(
860            self,
861            rhs,
862            |l, r| l.equal_missing(r).unwrap(),
863            |a, b| a.bitand(b),
864            false,
865            true,
866        )
867    }
868
869    fn not_equal(&self, rhs: &StructChunked) -> BooleanChunked {
870        struct_helper(
871            self,
872            rhs,
873            |l, r| l.not_equal_missing(r).unwrap(),
874            |a, b| a | b,
875            true,
876            false,
877        )
878    }
879
880    fn not_equal_missing(&self, rhs: &StructChunked) -> BooleanChunked {
881        struct_helper(
882            self,
883            rhs,
884            |l, r| l.not_equal_missing(r).unwrap(),
885            |a, b| a | b,
886            true,
887            true,
888        )
889    }
890}
891
892#[cfg(feature = "dtype-array")]
893fn _array_comparison_helper<F, B>(
894    lhs: &ArrayChunked,
895    rhs: &ArrayChunked,
896    op: F,
897    broadcast_op: B,
898    missing: bool,
899    is_ne: bool,
900) -> BooleanChunked
901where
902    F: Fn(&FixedSizeListArray, &FixedSizeListArray) -> Bitmap,
903    B: Fn(&FixedSizeListArray, &Box<dyn Array>) -> Bitmap,
904{
905    match (lhs.len(), rhs.len()) {
906        (_, 1) => {
907            let right = rhs
908                .downcast_iter()
909                .find(|x| !x.is_empty())
910                .unwrap()
911                .as_any()
912                .downcast_ref::<FixedSizeListArray>()
913                .unwrap();
914
915            if !right.validity().is_none_or(|v| v.get(0).unwrap()) {
916                if missing {
917                    if is_ne {
918                        return lhs.is_not_null();
919                    } else {
920                        return lhs.is_null();
921                    }
922                } else {
923                    return BooleanChunked::full_null(PlSmallStr::EMPTY, lhs.len());
924                }
925            }
926
927            if missing {
928                arity::unary_mut_with_options(lhs, |a| broadcast_op(a, right.values()).into())
929            } else {
930                arity::unary_mut_values(lhs, |a| broadcast_op(a, right.values()).into())
931            }
932        },
933        (1, _) => {
934            let left = lhs
935                .downcast_iter()
936                .find(|x| !x.is_empty())
937                .unwrap()
938                .as_any()
939                .downcast_ref::<FixedSizeListArray>()
940                .unwrap();
941
942            if !left.validity().is_none_or(|v| v.get(0).unwrap()) {
943                if missing {
944                    if is_ne {
945                        return rhs.is_not_null();
946                    } else {
947                        return rhs.is_null();
948                    }
949                } else {
950                    return BooleanChunked::full_null(PlSmallStr::EMPTY, rhs.len());
951                }
952            }
953
954            if missing {
955                arity::unary_mut_with_options(rhs, |a| broadcast_op(a, left.values()).into())
956            } else {
957                arity::unary_mut_values(rhs, |a| broadcast_op(a, left.values()).into())
958            }
959        },
960        _ => {
961            if missing {
962                arity::binary_mut_with_options(lhs, rhs, |a, b| op(a, b).into(), PlSmallStr::EMPTY)
963            } else {
964                arity::binary_mut_values(lhs, rhs, |a, b| op(a, b).into(), PlSmallStr::EMPTY)
965            }
966        },
967    }
968}
969
970#[cfg(feature = "dtype-array")]
971impl ChunkCompareEq<&ArrayChunked> for ArrayChunked {
972    type Item = BooleanChunked;
973    fn equal(&self, rhs: &ArrayChunked) -> BooleanChunked {
974        _array_comparison_helper(
975            self,
976            rhs,
977            TotalEqKernel::tot_eq_kernel,
978            TotalEqKernel::tot_eq_kernel_broadcast,
979            false,
980            false,
981        )
982    }
983
984    fn equal_missing(&self, rhs: &ArrayChunked) -> BooleanChunked {
985        _array_comparison_helper(
986            self,
987            rhs,
988            TotalEqKernel::tot_eq_missing_kernel,
989            TotalEqKernel::tot_eq_missing_kernel_broadcast,
990            true,
991            false,
992        )
993    }
994
995    fn not_equal(&self, rhs: &ArrayChunked) -> BooleanChunked {
996        _array_comparison_helper(
997            self,
998            rhs,
999            TotalEqKernel::tot_ne_kernel,
1000            TotalEqKernel::tot_ne_kernel_broadcast,
1001            false,
1002            true,
1003        )
1004    }
1005
1006    fn not_equal_missing(&self, rhs: &ArrayChunked) -> Self::Item {
1007        _array_comparison_helper(
1008            self,
1009            rhs,
1010            TotalEqKernel::tot_ne_missing_kernel,
1011            TotalEqKernel::tot_ne_missing_kernel_broadcast,
1012            true,
1013            true,
1014        )
1015    }
1016}
1017
1018impl Not for &BooleanChunked {
1019    type Output = BooleanChunked;
1020
1021    fn not(self) -> Self::Output {
1022        let chunks = self.downcast_iter().map(compute::boolean::not);
1023        ChunkedArray::from_chunk_iter(self.name().clone(), chunks)
1024    }
1025}
1026
1027impl Not for BooleanChunked {
1028    type Output = BooleanChunked;
1029
1030    fn not(self) -> Self::Output {
1031        (&self).not()
1032    }
1033}
1034
1035impl BooleanChunked {
1036    /// Returns whether any of the values in the column are `true`.
1037    ///
1038    /// Null values are ignored.
1039    pub fn any(&self) -> bool {
1040        self.downcast_iter().any(compute::boolean::any)
1041    }
1042
1043    /// Returns whether all values in the array are `true`.
1044    ///
1045    /// Null values are ignored.
1046    pub fn all(&self) -> bool {
1047        self.downcast_iter().all(compute::boolean::all)
1048    }
1049
1050    /// Returns whether any of the values in the column are `true`.
1051    ///
1052    /// The output is unknown (`None`) if the array contains any null values and
1053    /// no `true` values.
1054    pub fn any_kleene(&self) -> Option<bool> {
1055        let mut result = Some(false);
1056        for arr in self.downcast_iter() {
1057            match compute::boolean_kleene::any(arr) {
1058                Some(true) => return Some(true),
1059                None => result = None,
1060                _ => (),
1061            };
1062        }
1063        result
1064    }
1065
1066    /// Returns whether all values in the column are `true`.
1067    ///
1068    /// The output is unknown (`None`) if the array contains any null values and
1069    /// no `false` values.
1070    pub fn all_kleene(&self) -> Option<bool> {
1071        let mut result = Some(true);
1072        for arr in self.downcast_iter() {
1073            match compute::boolean_kleene::all(arr) {
1074                Some(false) => return Some(false),
1075                None => result = None,
1076                _ => (),
1077            };
1078        }
1079        result
1080    }
1081}
1082
1083// private
1084pub(crate) trait ChunkEqualElement {
1085    /// Only meant for physical types.
1086    /// Check if element in self is equal to element in other, assumes same dtypes
1087    ///
1088    /// # Safety
1089    ///
1090    /// No type checks.
1091    unsafe fn equal_element(&self, _idx_self: usize, _idx_other: usize, _other: &Series) -> bool {
1092        unimplemented!()
1093    }
1094}
1095
1096impl<T> ChunkEqualElement for ChunkedArray<T>
1097where
1098    T: PolarsNumericType,
1099{
1100    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
1101        let ca_other = other.as_ref().as_ref();
1102        debug_assert!(self.dtype() == other.dtype());
1103        let ca_other = &*(ca_other as *const ChunkedArray<T>);
1104        // Should be get and not get_unchecked, because there could be nulls
1105        self.get_unchecked(idx_self)
1106            .tot_eq(&ca_other.get_unchecked(idx_other))
1107    }
1108}
1109
1110impl ChunkEqualElement for BooleanChunked {
1111    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
1112        let ca_other = other.as_ref().as_ref();
1113        debug_assert!(self.dtype() == other.dtype());
1114        let ca_other = &*(ca_other as *const BooleanChunked);
1115        self.get_unchecked(idx_self) == ca_other.get_unchecked(idx_other)
1116    }
1117}
1118
1119impl ChunkEqualElement for StringChunked {
1120    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
1121        let ca_other = other.as_ref().as_ref();
1122        debug_assert!(self.dtype() == other.dtype());
1123        let ca_other = &*(ca_other as *const StringChunked);
1124        self.get_unchecked(idx_self) == ca_other.get_unchecked(idx_other)
1125    }
1126}
1127
1128impl ChunkEqualElement for BinaryChunked {
1129    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
1130        let ca_other = other.as_ref().as_ref();
1131        debug_assert!(self.dtype() == other.dtype());
1132        let ca_other = &*(ca_other as *const BinaryChunked);
1133        self.get_unchecked(idx_self) == ca_other.get_unchecked(idx_other)
1134    }
1135}
1136
1137impl ChunkEqualElement for BinaryOffsetChunked {
1138    unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
1139        let ca_other = other.as_ref().as_ref();
1140        debug_assert!(self.dtype() == other.dtype());
1141        let ca_other = &*(ca_other as *const BinaryOffsetChunked);
1142        self.get_unchecked(idx_self) == ca_other.get_unchecked(idx_other)
1143    }
1144}
1145
1146impl ChunkEqualElement for ListChunked {}
1147#[cfg(feature = "dtype-array")]
1148impl ChunkEqualElement for ArrayChunked {}
1149
1150#[cfg(test)]
1151mod test {
1152    use std::iter::repeat;
1153
1154    use super::super::test::get_chunked_array;
1155    use crate::prelude::*;
1156
1157    pub(crate) fn create_two_chunked() -> (Int32Chunked, Int32Chunked) {
1158        let mut a1 = Int32Chunked::new(PlSmallStr::from_static("a"), &[1, 2, 3]);
1159        let a2 = Int32Chunked::new(PlSmallStr::from_static("a"), &[4, 5, 6]);
1160        let a3 = Int32Chunked::new(PlSmallStr::from_static("a"), &[1, 2, 3, 4, 5, 6]);
1161        a1.append(&a2).unwrap();
1162        (a1, a3)
1163    }
1164
1165    #[test]
1166    fn test_bitwise_ops() {
1167        let a = BooleanChunked::new(PlSmallStr::from_static("a"), &[true, false, false]);
1168        let b = BooleanChunked::new(
1169            PlSmallStr::from_static("b"),
1170            &[Some(true), Some(true), None],
1171        );
1172        assert_eq!(Vec::from(&a | &b), &[Some(true), Some(true), None]);
1173        assert_eq!(Vec::from(&a & &b), &[Some(true), Some(false), Some(false)]);
1174        assert_eq!(Vec::from(!b), &[Some(false), Some(false), None]);
1175    }
1176
1177    #[test]
1178    fn test_compare_chunk_diff() {
1179        let (a1, a2) = create_two_chunked();
1180
1181        assert_eq!(
1182            a1.equal(&a2).into_iter().collect::<Vec<_>>(),
1183            repeat(Some(true)).take(6).collect::<Vec<_>>()
1184        );
1185        assert_eq!(
1186            a2.equal(&a1).into_iter().collect::<Vec<_>>(),
1187            repeat(Some(true)).take(6).collect::<Vec<_>>()
1188        );
1189        assert_eq!(
1190            a1.not_equal(&a2).into_iter().collect::<Vec<_>>(),
1191            repeat(Some(false)).take(6).collect::<Vec<_>>()
1192        );
1193        assert_eq!(
1194            a2.not_equal(&a1).into_iter().collect::<Vec<_>>(),
1195            repeat(Some(false)).take(6).collect::<Vec<_>>()
1196        );
1197        assert_eq!(
1198            a1.gt(&a2).into_iter().collect::<Vec<_>>(),
1199            repeat(Some(false)).take(6).collect::<Vec<_>>()
1200        );
1201        assert_eq!(
1202            a2.gt(&a1).into_iter().collect::<Vec<_>>(),
1203            repeat(Some(false)).take(6).collect::<Vec<_>>()
1204        );
1205        assert_eq!(
1206            a1.gt_eq(&a2).into_iter().collect::<Vec<_>>(),
1207            repeat(Some(true)).take(6).collect::<Vec<_>>()
1208        );
1209        assert_eq!(
1210            a2.gt_eq(&a1).into_iter().collect::<Vec<_>>(),
1211            repeat(Some(true)).take(6).collect::<Vec<_>>()
1212        );
1213        assert_eq!(
1214            a1.lt_eq(&a2).into_iter().collect::<Vec<_>>(),
1215            repeat(Some(true)).take(6).collect::<Vec<_>>()
1216        );
1217        assert_eq!(
1218            a2.lt_eq(&a1).into_iter().collect::<Vec<_>>(),
1219            repeat(Some(true)).take(6).collect::<Vec<_>>()
1220        );
1221        assert_eq!(
1222            a1.lt(&a2).into_iter().collect::<Vec<_>>(),
1223            repeat(Some(false)).take(6).collect::<Vec<_>>()
1224        );
1225        assert_eq!(
1226            a2.lt(&a1).into_iter().collect::<Vec<_>>(),
1227            repeat(Some(false)).take(6).collect::<Vec<_>>()
1228        );
1229    }
1230
1231    #[test]
1232    fn test_equal_chunks() {
1233        let a1 = get_chunked_array();
1234        let a2 = get_chunked_array();
1235
1236        assert_eq!(
1237            a1.equal(&a2).into_iter().collect::<Vec<_>>(),
1238            repeat(Some(true)).take(3).collect::<Vec<_>>()
1239        );
1240        assert_eq!(
1241            a2.equal(&a1).into_iter().collect::<Vec<_>>(),
1242            repeat(Some(true)).take(3).collect::<Vec<_>>()
1243        );
1244        assert_eq!(
1245            a1.not_equal(&a2).into_iter().collect::<Vec<_>>(),
1246            repeat(Some(false)).take(3).collect::<Vec<_>>()
1247        );
1248        assert_eq!(
1249            a2.not_equal(&a1).into_iter().collect::<Vec<_>>(),
1250            repeat(Some(false)).take(3).collect::<Vec<_>>()
1251        );
1252        assert_eq!(
1253            a1.gt(&a2).into_iter().collect::<Vec<_>>(),
1254            repeat(Some(false)).take(3).collect::<Vec<_>>()
1255        );
1256        assert_eq!(
1257            a2.gt(&a1).into_iter().collect::<Vec<_>>(),
1258            repeat(Some(false)).take(3).collect::<Vec<_>>()
1259        );
1260        assert_eq!(
1261            a1.gt_eq(&a2).into_iter().collect::<Vec<_>>(),
1262            repeat(Some(true)).take(3).collect::<Vec<_>>()
1263        );
1264        assert_eq!(
1265            a2.gt_eq(&a1).into_iter().collect::<Vec<_>>(),
1266            repeat(Some(true)).take(3).collect::<Vec<_>>()
1267        );
1268        assert_eq!(
1269            a1.lt_eq(&a2).into_iter().collect::<Vec<_>>(),
1270            repeat(Some(true)).take(3).collect::<Vec<_>>()
1271        );
1272        assert_eq!(
1273            a2.lt_eq(&a1).into_iter().collect::<Vec<_>>(),
1274            repeat(Some(true)).take(3).collect::<Vec<_>>()
1275        );
1276        assert_eq!(
1277            a1.lt(&a2).into_iter().collect::<Vec<_>>(),
1278            repeat(Some(false)).take(3).collect::<Vec<_>>()
1279        );
1280        assert_eq!(
1281            a2.lt(&a1).into_iter().collect::<Vec<_>>(),
1282            repeat(Some(false)).take(3).collect::<Vec<_>>()
1283        );
1284    }
1285
1286    #[test]
1287    fn test_null_handling() {
1288        // assert we comply with arrows way of handling null data
1289        // we check comparison on two arrays with one chunk and verify it is equal to a differently
1290        // chunked array comparison.
1291
1292        // two same chunked arrays
1293        let a1: Int32Chunked = [Some(1), None, Some(3)].iter().copied().collect();
1294        let a2: Int32Chunked = [Some(1), Some(2), Some(3)].iter().copied().collect();
1295
1296        let mut a2_2chunks: Int32Chunked = [Some(1), Some(2)].iter().copied().collect();
1297        a2_2chunks
1298            .append(&[Some(3)].iter().copied().collect())
1299            .unwrap();
1300
1301        assert_eq!(
1302            a1.equal(&a2).into_iter().collect::<Vec<_>>(),
1303            a1.equal(&a2_2chunks).into_iter().collect::<Vec<_>>()
1304        );
1305
1306        assert_eq!(
1307            a1.not_equal(&a2).into_iter().collect::<Vec<_>>(),
1308            a1.not_equal(&a2_2chunks).into_iter().collect::<Vec<_>>()
1309        );
1310        assert_eq!(
1311            a1.not_equal(&a2).into_iter().collect::<Vec<_>>(),
1312            a2_2chunks.not_equal(&a1).into_iter().collect::<Vec<_>>()
1313        );
1314
1315        assert_eq!(
1316            a1.gt(&a2).into_iter().collect::<Vec<_>>(),
1317            a1.gt(&a2_2chunks).into_iter().collect::<Vec<_>>()
1318        );
1319        assert_eq!(
1320            a1.gt(&a2).into_iter().collect::<Vec<_>>(),
1321            a2_2chunks.gt(&a1).into_iter().collect::<Vec<_>>()
1322        );
1323
1324        assert_eq!(
1325            a1.gt_eq(&a2).into_iter().collect::<Vec<_>>(),
1326            a1.gt_eq(&a2_2chunks).into_iter().collect::<Vec<_>>()
1327        );
1328        assert_eq!(
1329            a1.gt_eq(&a2).into_iter().collect::<Vec<_>>(),
1330            a2_2chunks.gt_eq(&a1).into_iter().collect::<Vec<_>>()
1331        );
1332
1333        assert_eq!(
1334            a1.lt_eq(&a2).into_iter().collect::<Vec<_>>(),
1335            a1.lt_eq(&a2_2chunks).into_iter().collect::<Vec<_>>()
1336        );
1337        assert_eq!(
1338            a1.lt_eq(&a2).into_iter().collect::<Vec<_>>(),
1339            a2_2chunks.lt_eq(&a1).into_iter().collect::<Vec<_>>()
1340        );
1341
1342        assert_eq!(
1343            a1.lt(&a2).into_iter().collect::<Vec<_>>(),
1344            a1.lt(&a2_2chunks).into_iter().collect::<Vec<_>>()
1345        );
1346        assert_eq!(
1347            a1.lt(&a2).into_iter().collect::<Vec<_>>(),
1348            a2_2chunks.lt(&a1).into_iter().collect::<Vec<_>>()
1349        );
1350    }
1351
1352    #[test]
1353    fn test_left_right() {
1354        // This failed with arrow comparisons.
1355        // sliced
1356        let a1: Int32Chunked = [Some(1), Some(2)].iter().copied().collect();
1357        let a1 = a1.slice(1, 1);
1358        let a2: Int32Chunked = [Some(2)].iter().copied().collect();
1359        assert_eq!(a1.equal(&a2).sum(), a2.equal(&a1).sum());
1360        assert_eq!(a1.not_equal(&a2).sum(), a2.not_equal(&a1).sum());
1361        assert_eq!(a1.gt(&a2).sum(), a2.gt(&a1).sum());
1362        assert_eq!(a1.lt(&a2).sum(), a2.lt(&a1).sum());
1363        assert_eq!(a1.lt_eq(&a2).sum(), a2.lt_eq(&a1).sum());
1364        assert_eq!(a1.gt_eq(&a2).sum(), a2.gt_eq(&a1).sum());
1365
1366        let a1: StringChunked = ["a", "b"].iter().copied().collect();
1367        let a1 = a1.slice(1, 1);
1368        let a2: StringChunked = ["b"].iter().copied().collect();
1369        assert_eq!(a1.equal(&a2).sum(), a2.equal(&a1).sum());
1370        assert_eq!(a1.not_equal(&a2).sum(), a2.not_equal(&a1).sum());
1371        assert_eq!(a1.gt(&a2).sum(), a2.gt(&a1).sum());
1372        assert_eq!(a1.lt(&a2).sum(), a2.lt(&a1).sum());
1373        assert_eq!(a1.lt_eq(&a2).sum(), a2.lt_eq(&a1).sum());
1374        assert_eq!(a1.gt_eq(&a2).sum(), a2.gt_eq(&a1).sum());
1375    }
1376
1377    #[test]
1378    fn test_kleene() {
1379        let a = BooleanChunked::new(PlSmallStr::EMPTY, &[Some(true), Some(false), None]);
1380        let trues = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[true, true, true]);
1381        let falses = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[false, false, false]);
1382
1383        let c = &a | &trues;
1384        assert_eq!(Vec::from(&c), &[Some(true), Some(true), Some(true)]);
1385
1386        let c = &a | &falses;
1387        assert_eq!(Vec::from(&c), &[Some(true), Some(false), None])
1388    }
1389
1390    #[test]
1391    fn list_broadcasting_lists() {
1392        let s_el = Series::new(PlSmallStr::EMPTY, &[1, 2, 3]);
1393        let s_lhs = Series::new(PlSmallStr::EMPTY, &[s_el.clone(), s_el.clone()]);
1394        let s_rhs = Series::new(PlSmallStr::EMPTY, &[s_el.clone()]);
1395
1396        let result = s_lhs.list().unwrap().equal(s_rhs.list().unwrap());
1397        assert_eq!(result.len(), 2);
1398        assert!(result.all());
1399    }
1400
1401    #[test]
1402    fn test_broadcasting_bools() {
1403        let a = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[true, false, true]);
1404        let true_ = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[true]);
1405        let false_ = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[false]);
1406
1407        let out = a.equal(&true_);
1408        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1409        let out = true_.equal(&a);
1410        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1411        let out = a.equal(&false_);
1412        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1413        let out = false_.equal(&a);
1414        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1415
1416        let out = a.not_equal(&true_);
1417        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1418        let out = true_.not_equal(&a);
1419        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1420        let out = a.not_equal(&false_);
1421        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1422        let out = false_.not_equal(&a);
1423        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1424
1425        let out = a.gt(&true_);
1426        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1427        let out = true_.gt(&a);
1428        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1429        let out = a.gt(&false_);
1430        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1431        let out = false_.gt(&a);
1432        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1433
1434        let out = a.gt_eq(&true_);
1435        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1436        let out = true_.gt_eq(&a);
1437        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1438        let out = a.gt_eq(&false_);
1439        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1440        let out = false_.gt_eq(&a);
1441        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1442
1443        let out = a.lt(&true_);
1444        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1445        let out = true_.lt(&a);
1446        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1447        let out = a.lt(&false_);
1448        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1449        let out = false_.lt(&a);
1450        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1451
1452        let out = a.lt_eq(&true_);
1453        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1454        let out = true_.lt_eq(&a);
1455        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(true)]);
1456        let out = a.lt_eq(&false_);
1457        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(false)]);
1458        let out = false_.lt_eq(&a);
1459        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1460
1461        let a =
1462            BooleanChunked::from_slice_options(PlSmallStr::EMPTY, &[Some(true), Some(false), None]);
1463        let all_true = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[true, true, true]);
1464        let all_false = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[false, false, false]);
1465        let out = a.equal(&true_);
1466        assert_eq!(Vec::from(&out), &[Some(true), Some(false), None]);
1467        let out = a.not_equal(&true_);
1468        assert_eq!(Vec::from(&out), &[Some(false), Some(true), None]);
1469
1470        let out = a.equal(&all_true);
1471        assert_eq!(Vec::from(&out), &[Some(true), Some(false), None]);
1472        let out = a.not_equal(&all_true);
1473        assert_eq!(Vec::from(&out), &[Some(false), Some(true), None]);
1474        let out = a.equal(&false_);
1475        assert_eq!(Vec::from(&out), &[Some(false), Some(true), None]);
1476        let out = a.not_equal(&false_);
1477        assert_eq!(Vec::from(&out), &[Some(true), Some(false), None]);
1478        let out = a.equal(&all_false);
1479        assert_eq!(Vec::from(&out), &[Some(false), Some(true), None]);
1480        let out = a.not_equal(&all_false);
1481        assert_eq!(Vec::from(&out), &[Some(true), Some(false), None]);
1482    }
1483
1484    #[test]
1485    fn test_broadcasting_numeric() {
1486        let a = Int32Chunked::from_slice(PlSmallStr::EMPTY, &[1, 2, 3]);
1487        let one = Int32Chunked::from_slice(PlSmallStr::EMPTY, &[1]);
1488        let three = Int32Chunked::from_slice(PlSmallStr::EMPTY, &[3]);
1489
1490        let out = a.equal(&one);
1491        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(false)]);
1492        let out = one.equal(&a);
1493        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(false)]);
1494        let out = a.equal(&three);
1495        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(true)]);
1496        let out = three.equal(&a);
1497        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(true)]);
1498
1499        let out = a.not_equal(&one);
1500        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(true)]);
1501        let out = one.not_equal(&a);
1502        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(true)]);
1503        let out = a.not_equal(&three);
1504        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(false)]);
1505        let out = three.not_equal(&a);
1506        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(false)]);
1507
1508        let out = a.gt(&one);
1509        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(true)]);
1510        let out = one.gt(&a);
1511        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1512        let out = a.gt(&three);
1513        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1514        let out = three.gt(&a);
1515        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(false)]);
1516
1517        let out = a.lt(&one);
1518        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1519        let out = one.lt(&a);
1520        assert_eq!(Vec::from(&out), &[Some(false), Some(true), Some(true)]);
1521        let out = a.lt(&three);
1522        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(false)]);
1523        let out = three.lt(&a);
1524        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(false)]);
1525
1526        let out = a.gt_eq(&one);
1527        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1528        let out = one.gt_eq(&a);
1529        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(false)]);
1530        let out = a.gt_eq(&three);
1531        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(true)]);
1532        let out = three.gt_eq(&a);
1533        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1534
1535        let out = a.lt_eq(&one);
1536        assert_eq!(Vec::from(&out), &[Some(true), Some(false), Some(false)]);
1537        let out = one.lt_eq(&a);
1538        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1539        let out = a.lt_eq(&three);
1540        assert_eq!(Vec::from(&out), &[Some(true), Some(true), Some(true)]);
1541        let out = three.lt_eq(&a);
1542        assert_eq!(Vec::from(&out), &[Some(false), Some(false), Some(true)]);
1543    }
1544}