1use polars_error::feature_gated;
4
5use crate::prelude::*;
6use crate::series::arithmetic::coerce_lhs_rhs;
7use crate::series::nulls::replace_non_null;
8
9macro_rules! impl_eq_compare {
10 ($self:expr, $rhs:expr, $method:ident) => {{
11 use DataType::*;
12 let (lhs, rhs) = ($self, $rhs);
13 validate_types(lhs.dtype(), rhs.dtype())?;
14
15 polars_ensure!(
16 lhs.len() == rhs.len() ||
17
18 lhs.len() == 1 ||
20 rhs.len() == 1,
21 ShapeMismatch: "could not compare between two series of different length ({} != {})",
22 lhs.len(),
23 rhs.len()
24 );
25
26 match (lhs.dtype(), rhs.dtype()) {
27 #[cfg(feature = "dtype-categorical")]
28 (Categorical(lcats, _), Categorical(rcats, _)) => {
29 ensure_same_categories(lcats, rcats)?;
30 return with_match_categorical_physical_type!(lcats.physical(), |$C| {
31 lhs.cat::<$C>().unwrap().$method(rhs.cat::<$C>().unwrap())
32 })
33 },
34 #[cfg(feature = "dtype-categorical")]
35 (Enum(lfcats, _), Enum(rfcats, _)) => {
36 ensure_same_frozen_categories(lfcats, rfcats)?;
37 return with_match_categorical_physical_type!(lfcats.physical(), |$C| {
38 lhs.cat::<$C>().unwrap().$method(rhs.cat::<$C>().unwrap())
39 })
40 },
41 #[cfg(feature = "dtype-categorical")]
42 (Categorical(_, _) | Enum(_, _), String) => {
43 return with_match_categorical_physical_type!(lhs.dtype().cat_physical().unwrap(), |$C| {
44 Ok(lhs.cat::<$C>().unwrap().$method(rhs.str().unwrap()))
45 })
46 },
47 #[cfg(feature = "dtype-categorical")]
48 (String, Categorical(_, _) | Enum(_, _)) => {
49 return with_match_categorical_physical_type!(rhs.dtype().cat_physical().unwrap(), |$C| {
50 Ok(rhs.cat::<$C>().unwrap().$method(lhs.str().unwrap()))
51 })
52 },
53 #[cfg(feature = "dtype-map")]
54 (ldt @ Map(_, _), rdt @ Map(_, _)) if ldt == rdt => {
55 let lhs = lhs.map().unwrap();
56 let rhs = rhs.map().unwrap();
57 return lhs.storage().$method(rhs.storage());
58 },
59
60 #[cfg(feature = "dtype-extension")]
61 (le @ Extension(_, _), re @ Extension(_, _)) if le == re => {
62 let lhs = lhs.ext().unwrap();
63 let rhs = rhs.ext().unwrap();
64 return lhs.storage().$method(rhs.storage());
65 },
66
67 #[cfg(feature = "dtype-extension")]
68 (Extension(_, storage), rdt) if **storage == *rdt => {
69 let lhs = lhs.ext().unwrap();
70 return lhs.storage().$method(rhs);
71 },
72
73 #[cfg(feature = "dtype-extension")]
74 (ldt, Extension(_, storage)) if *ldt == **storage => {
75 let rhs = rhs.ext().unwrap();
76 return lhs.$method(rhs.storage());
77 },
78 _ => (),
79 };
80
81 let (lhs, rhs) = coerce_lhs_rhs(lhs, rhs)
82 .map_err(|_| polars_err!(
83 SchemaMismatch: "could not evaluate comparison between series '{}' of dtype: {:?} and series '{}' of dtype: {:?}",
84 lhs.name(), lhs.dtype(), rhs.name(), rhs.dtype()
85 ))?;
86 let lhs = lhs.to_physical_repr();
87 let rhs = rhs.to_physical_repr();
88 let mut out = match lhs.dtype() {
89 Null => lhs.null().unwrap().$method(rhs.null().unwrap()),
90 Boolean => lhs.bool().unwrap().$method(rhs.bool().unwrap()),
91 String => lhs.str().unwrap().$method(rhs.str().unwrap()),
92 Binary => lhs.binary().unwrap().$method(rhs.binary().unwrap()),
93 BinaryOffset => lhs.binary_offset().unwrap().$method(rhs.binary_offset().unwrap()),
94 UInt8 => feature_gated!("dtype-u8", lhs.u8().unwrap().$method(rhs.u8().unwrap())),
95 UInt16 => feature_gated!("dtype-u16", lhs.u16().unwrap().$method(rhs.u16().unwrap())),
96 UInt32 => lhs.u32().unwrap().$method(rhs.u32().unwrap()),
97 UInt64 => lhs.u64().unwrap().$method(rhs.u64().unwrap()),
98 UInt128 => feature_gated!("dtype-u128", lhs.u128().unwrap().$method(rhs.u128().unwrap())),
99 Int8 => feature_gated!("dtype-i8", lhs.i8().unwrap().$method(rhs.i8().unwrap())),
100 Int16 => feature_gated!("dtype-i16", lhs.i16().unwrap().$method(rhs.i16().unwrap())),
101 Int32 => lhs.i32().unwrap().$method(rhs.i32().unwrap()),
102 Int64 => lhs.i64().unwrap().$method(rhs.i64().unwrap()),
103 Int128 => feature_gated!("dtype-i128", lhs.i128().unwrap().$method(rhs.i128().unwrap())),
104 Float16 => feature_gated!("dtype-f16", lhs.f16().unwrap().$method(rhs.f16().unwrap())),
105 Float32 => lhs.f32().unwrap().$method(rhs.f32().unwrap()),
106 Float64 => lhs.f64().unwrap().$method(rhs.f64().unwrap()),
107 List(_) => lhs.list().unwrap().$method(rhs.list().unwrap()),
108 #[cfg(feature = "dtype-array")]
109 Array(_, _) => lhs.array().unwrap().$method(rhs.array().unwrap()),
110 #[cfg(feature = "dtype-struct")]
111 Struct(_) => lhs.struct_().unwrap().$method(rhs.struct_().unwrap()),
112
113 dt => polars_bail!(InvalidOperation: "could not apply comparison on series of dtype '{}; operand names: '{}', '{}'", dt, lhs.name(), rhs.name()),
114 };
115 out.rename(lhs.name().clone());
116 PolarsResult::Ok(out)
117 }};
118}
119
120macro_rules! bail_invalid_ineq {
121 ($lhs:expr, $rhs:expr, $op:literal) => {
122 polars_bail!(
123 InvalidOperation: "cannot perform '{}' comparison between series '{}' of dtype: {} and series '{}' of dtype: {}",
124 $op,
125 $lhs.name(), $lhs.dtype(),
126 $rhs.name(), $rhs.dtype(),
127 )
128 };
129}
130
131macro_rules! impl_ineq_compare {
132 ($self:expr, $rhs:expr, $method:ident, $op:literal, $rev_method:ident) => {{
133 use DataType::*;
134 let (lhs, rhs) = ($self, $rhs);
135 validate_types(lhs.dtype(), rhs.dtype())?;
136
137 polars_ensure!(
138 lhs.len() == rhs.len() ||
139
140 lhs.len() == 1 ||
142 rhs.len() == 1,
143 ShapeMismatch:
144 "could not perform '{}' comparison between series '{}' of length: {} and series '{}' of length: {}, because they have different lengths",
145 $op,
146 lhs.name(), lhs.len(),
147 rhs.name(), rhs.len()
148 );
149
150 match (lhs.dtype(), rhs.dtype()) {
151 #[cfg(feature = "dtype-categorical")]
152 (Categorical(lcats, _), Categorical(rcats, _)) => {
153 ensure_same_categories(lcats, rcats)?;
154 return with_match_categorical_physical_type!(lcats.physical(), |$C| {
155 lhs.cat::<$C>().unwrap().$method(rhs.cat::<$C>().unwrap())
156 })
157 },
158 #[cfg(feature = "dtype-categorical")]
159 (Enum(lfcats, _), Enum(rfcats, _)) => {
160 ensure_same_frozen_categories(lfcats, rfcats)?;
161 return with_match_categorical_physical_type!(lfcats.physical(), |$C| {
162 lhs.cat::<$C>().unwrap().$method(rhs.cat::<$C>().unwrap())
163 })
164 },
165 #[cfg(feature = "dtype-categorical")]
166 (Categorical(_, _) | Enum(_, _), String) => {
167 return with_match_categorical_physical_type!(lhs.dtype().cat_physical().unwrap(), |$C| {
168 lhs.cat::<$C>().unwrap().$method(rhs.str().unwrap())
169 })
170 },
171 #[cfg(feature = "dtype-categorical")]
172 (String, Categorical(_, _) | Enum(_, _)) => {
173 return with_match_categorical_physical_type!(rhs.dtype().cat_physical().unwrap(), |$C| {
174 rhs.cat::<$C>().unwrap().$rev_method(lhs.str().unwrap())
176 })
177 },
178 #[cfg(feature = "dtype-map")]
180 (Map(_, _), _) | (_, Map(_, _)) => bail_invalid_ineq!(lhs, rhs, $op),
181
182 #[cfg(feature = "dtype-extension")]
183 (le @ Extension(_, _), re @ Extension(_, _)) if le == re => {
184 let lhs = lhs.ext().unwrap();
185 let rhs = rhs.ext().unwrap();
186 return lhs.storage().$method(rhs.storage());
187 },
188
189 #[cfg(feature = "dtype-extension")]
190 (Extension(_, storage), rdt) if **storage == *rdt => {
191 let lhs = lhs.ext().unwrap();
192 return lhs.storage().$method(rhs);
193 },
194
195 #[cfg(feature = "dtype-extension")]
196 (ldt, Extension(_, storage)) if *ldt == **storage => {
197 let rhs = rhs.ext().unwrap();
198 return lhs.$method(rhs.storage());
199 },
200 _ => (),
201 };
202
203 let (lhs, rhs) = coerce_lhs_rhs(lhs, rhs).map_err(|_|
204 polars_err!(
205 SchemaMismatch: "could not evaluate '{}' comparison between series '{}' of dtype: {:?} and series '{}' of dtype: {:?}",
206 $op,
207 lhs.name(), lhs.dtype(),
208 rhs.name(), rhs.dtype()
209 )
210 )?;
211 let lhs = lhs.to_physical_repr();
212 let rhs = rhs.to_physical_repr();
213 let mut out = match lhs.dtype() {
214 Null => lhs.null().unwrap().$method(rhs.null().unwrap()),
215 Boolean => lhs.bool().unwrap().$method(rhs.bool().unwrap()),
216 String => lhs.str().unwrap().$method(rhs.str().unwrap()),
217 Binary => lhs.binary().unwrap().$method(rhs.binary().unwrap()),
218 BinaryOffset => lhs.binary_offset().unwrap().$method(rhs.binary_offset().unwrap()),
219 UInt8 => feature_gated!("dtype-u8", lhs.u8().unwrap().$method(rhs.u8().unwrap())),
220 UInt16 => feature_gated!("dtype-u16", lhs.u16().unwrap().$method(rhs.u16().unwrap())),
221 UInt32 => lhs.u32().unwrap().$method(rhs.u32().unwrap()),
222 UInt64 => lhs.u64().unwrap().$method(rhs.u64().unwrap()),
223 UInt128 => feature_gated!("dtype-u128", lhs.u128().unwrap().$method(rhs.u128().unwrap())),
224 Int8 => feature_gated!("dtype-i8", lhs.i8().unwrap().$method(rhs.i8().unwrap())),
225 Int16 => feature_gated!("dtype-i16", lhs.i16().unwrap().$method(rhs.i16().unwrap())),
226 Int32 => lhs.i32().unwrap().$method(rhs.i32().unwrap()),
227 Int64 => lhs.i64().unwrap().$method(rhs.i64().unwrap()),
228 Int128 => feature_gated!("dtype-i128", lhs.i128().unwrap().$method(rhs.i128().unwrap())),
229 Float16 => feature_gated!("dtype-f16", lhs.f16().unwrap().$method(rhs.f16().unwrap())),
230 Float32 => lhs.f32().unwrap().$method(rhs.f32().unwrap()),
231 Float64 => lhs.f64().unwrap().$method(rhs.f64().unwrap()),
232 List(_) => bail_invalid_ineq!(lhs, rhs, $op),
233 #[cfg(feature = "dtype-array")]
234 Array(_, _) => bail_invalid_ineq!(lhs, rhs, $op),
235 #[cfg(feature = "dtype-struct")]
236 Struct(_) => bail_invalid_ineq!(lhs, rhs, $op),
237
238 dt => polars_bail!(InvalidOperation: "could not apply comparison on series of dtype '{}'; operand names: '{}', '{}'", dt, lhs.name(), rhs.name()),
239 };
240 out.rename(lhs.name().clone());
241 PolarsResult::Ok(out)
242 }};
243}
244
245fn validate_types(left: &DataType, right: &DataType) -> PolarsResult<()> {
246 use DataType::*;
247
248 match (left, right) {
249 (String, dt) | (dt, String) if dt.is_primitive_numeric() => {
250 polars_bail!(ComputeError: "cannot compare string with numeric type ({})", dt)
251 },
252 #[cfg(feature = "dtype-categorical")]
253 (Categorical(_, _) | Enum(_, _), dt) | (dt, Categorical(_, _) | Enum(_, _))
254 if !(dt.is_categorical() | dt.is_string() | dt.is_enum()) =>
255 {
256 polars_bail!(ComputeError: "cannot compare categorical with {}", dt)
257 },
258 #[cfg(feature = "dtype-duration")]
259 (Date, Duration(_)) | (Duration(_), Date) => {
260 polars_bail!(ComputeError: "cannot compare date with duration")
261 },
262 _ => (),
263 };
264 Ok(())
265}
266
267impl ChunkCompareEq<&Series> for Series {
268 type Item = PolarsResult<BooleanChunked>;
269
270 fn equal(&self, rhs: &Series) -> Self::Item {
272 impl_eq_compare!(self, rhs, equal)
273 }
274
275 fn equal_missing(&self, rhs: &Series) -> Self::Item {
277 impl_eq_compare!(self, rhs, equal_missing)
278 }
279
280 fn not_equal(&self, rhs: &Series) -> Self::Item {
282 impl_eq_compare!(self, rhs, not_equal)
283 }
284
285 fn not_equal_missing(&self, rhs: &Series) -> Self::Item {
287 impl_eq_compare!(self, rhs, not_equal_missing)
288 }
289}
290
291impl ChunkCompareIneq<&Series> for Series {
292 type Item = PolarsResult<BooleanChunked>;
293
294 fn gt(&self, rhs: &Series) -> Self::Item {
296 impl_ineq_compare!(self, rhs, gt, ">", lt)
297 }
298
299 fn gt_eq(&self, rhs: &Series) -> Self::Item {
301 impl_ineq_compare!(self, rhs, gt_eq, ">=", lt_eq)
302 }
303
304 fn lt(&self, rhs: &Series) -> Self::Item {
306 impl_ineq_compare!(self, rhs, lt, "<", gt)
307 }
308
309 fn lt_eq(&self, rhs: &Series) -> Self::Item {
311 impl_ineq_compare!(self, rhs, lt_eq, "<=", gt_eq)
312 }
313}
314
315impl<Rhs> ChunkCompareEq<Rhs> for Series
316where
317 Rhs: NumericNative,
318{
319 type Item = PolarsResult<BooleanChunked>;
320
321 fn equal(&self, rhs: Rhs) -> Self::Item {
322 validate_types(self.dtype(), &DataType::Int8)?;
323 let s = self.to_physical_repr();
324 Ok(apply_method_physical_numeric!(&s, equal, rhs))
325 }
326
327 fn equal_missing(&self, rhs: Rhs) -> Self::Item {
328 validate_types(self.dtype(), &DataType::Int8)?;
329 let s = self.to_physical_repr();
330 Ok(apply_method_physical_numeric!(&s, equal_missing, rhs))
331 }
332
333 fn not_equal(&self, rhs: Rhs) -> Self::Item {
334 validate_types(self.dtype(), &DataType::Int8)?;
335 let s = self.to_physical_repr();
336 Ok(apply_method_physical_numeric!(&s, not_equal, rhs))
337 }
338
339 fn not_equal_missing(&self, rhs: Rhs) -> Self::Item {
340 validate_types(self.dtype(), &DataType::Int8)?;
341 let s = self.to_physical_repr();
342 Ok(apply_method_physical_numeric!(&s, not_equal_missing, rhs))
343 }
344}
345
346impl<Rhs> ChunkCompareIneq<Rhs> for Series
347where
348 Rhs: NumericNative,
349{
350 type Item = PolarsResult<BooleanChunked>;
351
352 fn gt(&self, rhs: Rhs) -> Self::Item {
353 validate_types(self.dtype(), &DataType::Int8)?;
354 let s = self.to_physical_repr();
355 Ok(apply_method_physical_numeric!(&s, gt, rhs))
356 }
357
358 fn gt_eq(&self, rhs: Rhs) -> Self::Item {
359 validate_types(self.dtype(), &DataType::Int8)?;
360 let s = self.to_physical_repr();
361 Ok(apply_method_physical_numeric!(&s, gt_eq, rhs))
362 }
363
364 fn lt(&self, rhs: Rhs) -> Self::Item {
365 validate_types(self.dtype(), &DataType::Int8)?;
366 let s = self.to_physical_repr();
367 Ok(apply_method_physical_numeric!(&s, lt, rhs))
368 }
369
370 fn lt_eq(&self, rhs: Rhs) -> Self::Item {
371 validate_types(self.dtype(), &DataType::Int8)?;
372 let s = self.to_physical_repr();
373 Ok(apply_method_physical_numeric!(&s, lt_eq, rhs))
374 }
375}
376
377impl ChunkCompareEq<&str> for Series {
378 type Item = PolarsResult<BooleanChunked>;
379
380 fn equal(&self, rhs: &str) -> PolarsResult<BooleanChunked> {
381 validate_types(self.dtype(), &DataType::String)?;
382 match self.dtype() {
383 DataType::String => Ok(self.str().unwrap().equal(rhs)),
384 #[cfg(feature = "dtype-categorical")]
385 DataType::Categorical(_, _) | DataType::Enum(_, _) => Ok(
386 with_match_categorical_physical_type!(self.dtype().cat_physical().unwrap(), |$C| {
387 self.cat::<$C>().unwrap().equal(rhs)
388 }),
389 ),
390 #[cfg(feature = "dtype-extension")]
391 DataType::Extension(_, _) => self.ext().unwrap().storage().equal(rhs),
392 _ => Ok(BooleanChunked::full(self.name().clone(), false, self.len())),
393 }
394 }
395
396 fn equal_missing(&self, rhs: &str) -> Self::Item {
397 validate_types(self.dtype(), &DataType::String)?;
398 match self.dtype() {
399 DataType::String => Ok(self.str().unwrap().equal_missing(rhs)),
400 #[cfg(feature = "dtype-categorical")]
401 DataType::Categorical(_, _) | DataType::Enum(_, _) => Ok(
402 with_match_categorical_physical_type!(self.dtype().cat_physical().unwrap(), |$C| {
403 self.cat::<$C>().unwrap().equal_missing(rhs)
404 }),
405 ),
406 #[cfg(feature = "dtype-extension")]
407 DataType::Extension(_, _) => self.ext().unwrap().storage().equal_missing(rhs),
408 _ => Ok(replace_non_null(
409 self.name().clone(),
410 self.0.chunks(),
411 false,
412 )),
413 }
414 }
415
416 fn not_equal(&self, rhs: &str) -> PolarsResult<BooleanChunked> {
417 validate_types(self.dtype(), &DataType::String)?;
418 match self.dtype() {
419 DataType::String => Ok(self.str().unwrap().not_equal(rhs)),
420 #[cfg(feature = "dtype-categorical")]
421 DataType::Categorical(_, _) | DataType::Enum(_, _) => Ok(
422 with_match_categorical_physical_type!(self.dtype().cat_physical().unwrap(), |$C| {
423 self.cat::<$C>().unwrap().not_equal(rhs)
424 }),
425 ),
426 #[cfg(feature = "dtype-extension")]
427 DataType::Extension(_, _) => self.ext().unwrap().storage().not_equal(rhs),
428 _ => Ok(BooleanChunked::full(self.name().clone(), true, self.len())),
429 }
430 }
431
432 fn not_equal_missing(&self, rhs: &str) -> Self::Item {
433 validate_types(self.dtype(), &DataType::String)?;
434 match self.dtype() {
435 DataType::String => Ok(self.str().unwrap().not_equal_missing(rhs)),
436 #[cfg(feature = "dtype-categorical")]
437 DataType::Categorical(_, _) | DataType::Enum(_, _) => Ok(
438 with_match_categorical_physical_type!(self.dtype().cat_physical().unwrap(), |$C| {
439 self.cat::<$C>().unwrap().not_equal_missing(rhs)
440 }),
441 ),
442 #[cfg(feature = "dtype-extension")]
443 DataType::Extension(_, _) => self.ext().unwrap().storage().not_equal_missing(rhs),
444 _ => Ok(replace_non_null(self.name().clone(), self.0.chunks(), true)),
445 }
446 }
447}
448
449impl ChunkCompareIneq<&str> for Series {
450 type Item = PolarsResult<BooleanChunked>;
451
452 fn gt(&self, rhs: &str) -> Self::Item {
453 validate_types(self.dtype(), &DataType::String)?;
454 match self.dtype() {
455 DataType::String => Ok(self.str().unwrap().gt(rhs)),
456 #[cfg(feature = "dtype-categorical")]
457 DataType::Categorical(_, _) | DataType::Enum(_, _) => Ok(
458 with_match_categorical_physical_type!(self.dtype().cat_physical().unwrap(), |$C| {
459 self.cat::<$C>().unwrap().gt(rhs)
460 }),
461 ),
462 #[cfg(feature = "dtype-extension")]
463 DataType::Extension(_, _) => self.ext().unwrap().storage().gt(rhs),
464 _ => polars_bail!(
465 ComputeError: "cannot compare str value to series of type {}", self.dtype(),
466 ),
467 }
468 }
469
470 fn gt_eq(&self, rhs: &str) -> Self::Item {
471 validate_types(self.dtype(), &DataType::String)?;
472 match self.dtype() {
473 DataType::String => Ok(self.str().unwrap().gt_eq(rhs)),
474 #[cfg(feature = "dtype-categorical")]
475 DataType::Categorical(_, _) | DataType::Enum(_, _) => Ok(
476 with_match_categorical_physical_type!(self.dtype().cat_physical().unwrap(), |$C| {
477 self.cat::<$C>().unwrap().gt_eq(rhs)
478 }),
479 ),
480 #[cfg(feature = "dtype-extension")]
481 DataType::Extension(_, _) => self.ext().unwrap().storage().gt_eq(rhs),
482 _ => polars_bail!(
483 ComputeError: "cannot compare str value to series of type {}", self.dtype(),
484 ),
485 }
486 }
487
488 fn lt(&self, rhs: &str) -> Self::Item {
489 validate_types(self.dtype(), &DataType::String)?;
490 match self.dtype() {
491 DataType::String => Ok(self.str().unwrap().lt(rhs)),
492 #[cfg(feature = "dtype-categorical")]
493 DataType::Categorical(_, _) | DataType::Enum(_, _) => Ok(
494 with_match_categorical_physical_type!(self.dtype().cat_physical().unwrap(), |$C| {
495 self.cat::<$C>().unwrap().lt(rhs)
496 }),
497 ),
498 #[cfg(feature = "dtype-extension")]
499 DataType::Extension(_, _) => self.ext().unwrap().storage().lt(rhs),
500 _ => polars_bail!(
501 ComputeError: "cannot compare str value to series of type {}", self.dtype(),
502 ),
503 }
504 }
505
506 fn lt_eq(&self, rhs: &str) -> Self::Item {
507 validate_types(self.dtype(), &DataType::String)?;
508 match self.dtype() {
509 DataType::String => Ok(self.str().unwrap().lt_eq(rhs)),
510 #[cfg(feature = "dtype-categorical")]
511 DataType::Categorical(_, _) | DataType::Enum(_, _) => Ok(
512 with_match_categorical_physical_type!(self.dtype().cat_physical().unwrap(), |$C| {
513 self.cat::<$C>().unwrap().lt_eq(rhs)
514 }),
515 ),
516 #[cfg(feature = "dtype-extension")]
517 DataType::Extension(_, _) => self.ext().unwrap().storage().lt_eq(rhs),
518 _ => polars_bail!(
519 ComputeError: "cannot compare str value to series of type {}", self.dtype(),
520 ),
521 }
522 }
523}