1use polars_utils::broadcast::broadcast_len;
2
3use super::*;
4use crate::utils::align_chunks_binary;
5
6pub trait NumOpsDispatchInner: PolarsDataType + Sized {
7 fn subtract(lhs: &ChunkedArray<Self>, rhs: &Series) -> PolarsResult<Series> {
8 polars_bail!(opq = sub, lhs.dtype(), rhs.dtype());
9 }
10 fn add_to(lhs: &ChunkedArray<Self>, rhs: &Series) -> PolarsResult<Series> {
11 polars_bail!(opq = add, lhs.dtype(), rhs.dtype());
12 }
13 fn multiply(lhs: &ChunkedArray<Self>, rhs: &Series) -> PolarsResult<Series> {
14 polars_bail!(opq = mul, lhs.dtype(), rhs.dtype());
15 }
16 fn divide(lhs: &ChunkedArray<Self>, rhs: &Series) -> PolarsResult<Series> {
17 polars_bail!(opq = div, lhs.dtype(), rhs.dtype());
18 }
19 fn remainder(lhs: &ChunkedArray<Self>, rhs: &Series) -> PolarsResult<Series> {
20 polars_bail!(opq = rem, lhs.dtype(), rhs.dtype());
21 }
22}
23
24pub trait NumOpsDispatch {
25 fn subtract(&self, rhs: &Series) -> PolarsResult<Series>;
26 fn add_to(&self, rhs: &Series) -> PolarsResult<Series>;
27 fn multiply(&self, rhs: &Series) -> PolarsResult<Series>;
28 fn divide(&self, rhs: &Series) -> PolarsResult<Series>;
29 fn remainder(&self, rhs: &Series) -> PolarsResult<Series>;
30}
31
32impl<T: NumOpsDispatchInner> NumOpsDispatch for ChunkedArray<T> {
33 fn subtract(&self, rhs: &Series) -> PolarsResult<Series> {
34 T::subtract(self, rhs)
35 }
36 fn add_to(&self, rhs: &Series) -> PolarsResult<Series> {
37 T::add_to(self, rhs)
38 }
39 fn multiply(&self, rhs: &Series) -> PolarsResult<Series> {
40 T::multiply(self, rhs)
41 }
42 fn divide(&self, rhs: &Series) -> PolarsResult<Series> {
43 T::divide(self, rhs)
44 }
45 fn remainder(&self, rhs: &Series) -> PolarsResult<Series> {
46 T::remainder(self, rhs)
47 }
48}
49
50impl<T: PolarsNumericType> NumOpsDispatchInner for T {
51 fn subtract(lhs: &ChunkedArray<T>, rhs: &Series) -> PolarsResult<Series> {
52 polars_ensure!(
53 lhs.dtype() == rhs.dtype(),
54 opq = add,
55 rhs.dtype(),
56 rhs.dtype()
57 );
58
59 let rhs = unsafe { lhs.unpack_series_matching_physical_type(rhs) };
65 let out = lhs - rhs;
66 Ok(out.into_series())
67 }
68 fn add_to(lhs: &ChunkedArray<T>, rhs: &Series) -> PolarsResult<Series> {
69 polars_ensure!(
70 lhs.dtype() == rhs.dtype(),
71 opq = add,
72 rhs.dtype(),
73 rhs.dtype()
74 );
75
76 let rhs = unsafe { lhs.unpack_series_matching_physical_type(rhs) };
79 let out = lhs + rhs;
80 Ok(out.into_series())
81 }
82 fn multiply(lhs: &ChunkedArray<T>, rhs: &Series) -> PolarsResult<Series> {
83 polars_ensure!(
84 lhs.dtype() == rhs.dtype(),
85 opq = add,
86 rhs.dtype(),
87 rhs.dtype()
88 );
89
90 let rhs = unsafe { lhs.unpack_series_matching_physical_type(rhs) };
93 let out = lhs * rhs;
94 Ok(out.into_series())
95 }
96 fn divide(lhs: &ChunkedArray<T>, rhs: &Series) -> PolarsResult<Series> {
97 polars_ensure!(
98 lhs.dtype() == rhs.dtype(),
99 opq = add,
100 rhs.dtype(),
101 rhs.dtype()
102 );
103
104 let rhs = unsafe { lhs.unpack_series_matching_physical_type(rhs) };
107 let out = lhs / rhs;
108 Ok(out.into_series())
109 }
110 fn remainder(lhs: &ChunkedArray<T>, rhs: &Series) -> PolarsResult<Series> {
111 polars_ensure!(
112 lhs.dtype() == rhs.dtype(),
113 opq = add,
114 rhs.dtype(),
115 rhs.dtype()
116 );
117
118 let rhs = unsafe { lhs.unpack_series_matching_physical_type(rhs) };
121 let out = lhs % rhs;
122 Ok(out.into_series())
123 }
124}
125
126impl NumOpsDispatchInner for StringType {
127 fn add_to(lhs: &StringChunked, rhs: &Series) -> PolarsResult<Series> {
128 let rhs = lhs.unpack_series_matching_type(rhs)?;
129 let out = lhs + rhs;
130 Ok(out.into_series())
131 }
132}
133
134impl NumOpsDispatchInner for BinaryType {
135 fn add_to(lhs: &BinaryChunked, rhs: &Series) -> PolarsResult<Series> {
136 let rhs = lhs.unpack_series_matching_type(rhs)?;
137 let out = lhs + rhs;
138 Ok(out.into_series())
139 }
140}
141
142impl NumOpsDispatchInner for BooleanType {
143 fn add_to(lhs: &BooleanChunked, rhs: &Series) -> PolarsResult<Series> {
144 let rhs = lhs.unpack_series_matching_type(rhs)?;
145 let out = lhs + rhs;
146 Ok(out.into_series())
147 }
148}
149
150#[cfg(feature = "checked_arithmetic")]
151pub mod checked {
152 use num_traits::{CheckedDiv, One, ToPrimitive, Zero};
153
154 use super::*;
155
156 pub trait NumOpsDispatchCheckedInner: PolarsDataType + Sized {
157 fn checked_div(lhs: &ChunkedArray<Self>, rhs: &Series) -> PolarsResult<Series> {
159 polars_bail!(opq = checked_div, lhs.dtype(), rhs.dtype());
160 }
161 fn checked_div_num<T: ToPrimitive>(
162 lhs: &ChunkedArray<Self>,
163 _rhs: T,
164 ) -> PolarsResult<Series> {
165 polars_bail!(opq = checked_div_num, lhs.dtype(), Self::get_static_dtype());
166 }
167 }
168
169 pub trait NumOpsDispatchChecked {
170 fn checked_div(&self, rhs: &Series) -> PolarsResult<Series>;
172 fn checked_div_num<T: ToPrimitive>(&self, _rhs: T) -> PolarsResult<Series>;
173 }
174
175 impl<S: NumOpsDispatchCheckedInner> NumOpsDispatchChecked for ChunkedArray<S> {
176 fn checked_div(&self, rhs: &Series) -> PolarsResult<Series> {
177 S::checked_div(self, rhs)
178 }
179 fn checked_div_num<T: ToPrimitive>(&self, rhs: T) -> PolarsResult<Series> {
180 S::checked_div_num(self, rhs)
181 }
182 }
183
184 impl<T> NumOpsDispatchCheckedInner for T
185 where
186 T: PolarsIntegerType,
187 T::Native: CheckedDiv<Output = T::Native> + CheckedDiv<Output = T::Native> + Zero + One,
188 {
189 fn checked_div(lhs: &ChunkedArray<T>, rhs: &Series) -> PolarsResult<Series> {
190 let rhs = unsafe { lhs.unpack_series_matching_physical_type(rhs) };
196
197 let ca: ChunkedArray<T> =
198 arity::binary_elementwise(lhs, rhs, |opt_l, opt_r| match (opt_l, opt_r) {
199 (Some(l), Some(r)) => l.checked_div(&r),
200 _ => None,
201 });
202 Ok(ca.into_series())
203 }
204 }
205
206 #[cfg(feature = "dtype-f16")]
207 impl NumOpsDispatchCheckedInner for Float16Type {
208 fn checked_div(lhs: &Float16Chunked, rhs: &Series) -> PolarsResult<Series> {
209 let rhs = unsafe { lhs.unpack_series_matching_physical_type(rhs) };
212
213 let ca: Float16Chunked =
214 arity::binary_elementwise(lhs, rhs, |opt_l, opt_r| match (opt_l, opt_r) {
215 (Some(l), Some(r)) => {
216 if r.is_zero() {
217 None
218 } else {
219 Some(l / r)
220 }
221 },
222 _ => None,
223 });
224 Ok(ca.into_series())
225 }
226 }
227
228 impl NumOpsDispatchCheckedInner for Float32Type {
229 fn checked_div(lhs: &Float32Chunked, rhs: &Series) -> PolarsResult<Series> {
230 let rhs = unsafe { lhs.unpack_series_matching_physical_type(rhs) };
233
234 let ca: Float32Chunked =
235 arity::binary_elementwise(lhs, rhs, |opt_l, opt_r| match (opt_l, opt_r) {
236 (Some(l), Some(r)) => {
237 if r.is_zero() {
238 None
239 } else {
240 Some(l / r)
241 }
242 },
243 _ => None,
244 });
245 Ok(ca.into_series())
246 }
247 }
248
249 impl NumOpsDispatchCheckedInner for Float64Type {
250 fn checked_div(lhs: &Float64Chunked, rhs: &Series) -> PolarsResult<Series> {
251 let rhs = unsafe { lhs.unpack_series_matching_physical_type(rhs) };
254
255 let ca: Float64Chunked =
256 arity::binary_elementwise(lhs, rhs, |opt_l, opt_r| match (opt_l, opt_r) {
257 (Some(l), Some(r)) => {
258 if r.is_zero() {
259 None
260 } else {
261 Some(l / r)
262 }
263 },
264 _ => None,
265 });
266 Ok(ca.into_series())
267 }
268 }
269
270 impl NumOpsDispatchChecked for Series {
271 fn checked_div(&self, rhs: &Series) -> PolarsResult<Series> {
272 let (lhs, rhs) = coerce_lhs_rhs(self, rhs).expect("cannot coerce datatypes");
273 lhs.as_ref().as_ref().checked_div(rhs.as_ref())
274 }
275
276 fn checked_div_num<T: ToPrimitive>(&self, rhs: T) -> PolarsResult<Series> {
277 use DataType::*;
278 let s = self.to_physical_repr();
279
280 let out = match s.dtype() {
281 #[cfg(feature = "dtype-u8")]
282 UInt8 => s
283 .u8()
284 .unwrap()
285 .apply(|opt_v| opt_v.and_then(|v| v.checked_div(rhs.to_u8().unwrap())))
286 .into_series(),
287 #[cfg(feature = "dtype-i8")]
288 Int8 => s
289 .i8()
290 .unwrap()
291 .apply(|opt_v| opt_v.and_then(|v| v.checked_div(rhs.to_i8().unwrap())))
292 .into_series(),
293 #[cfg(feature = "dtype-i16")]
294 Int16 => s
295 .i16()
296 .unwrap()
297 .apply(|opt_v| opt_v.and_then(|v| v.checked_div(rhs.to_i16().unwrap())))
298 .into_series(),
299 #[cfg(feature = "dtype-u16")]
300 UInt16 => s
301 .u16()
302 .unwrap()
303 .apply(|opt_v| opt_v.and_then(|v| v.checked_div(rhs.to_u16().unwrap())))
304 .into_series(),
305 UInt32 => s
306 .u32()
307 .unwrap()
308 .apply(|opt_v| opt_v.and_then(|v| v.checked_div(rhs.to_u32().unwrap())))
309 .into_series(),
310 Int32 => s
311 .i32()
312 .unwrap()
313 .apply(|opt_v| opt_v.and_then(|v| v.checked_div(rhs.to_i32().unwrap())))
314 .into_series(),
315 UInt64 => s
316 .u64()
317 .unwrap()
318 .apply(|opt_v| opt_v.and_then(|v| v.checked_div(rhs.to_u64().unwrap())))
319 .into_series(),
320 Int64 => s
321 .i64()
322 .unwrap()
323 .apply(|opt_v| opt_v.and_then(|v| v.checked_div(rhs.to_i64().unwrap())))
324 .into_series(),
325 Float32 => s
326 .f32()
327 .unwrap()
328 .apply(|opt_v| {
329 opt_v.and_then(|v| {
330 let res = rhs.to_f32().unwrap();
331 if res.is_zero() { None } else { Some(v / res) }
332 })
333 })
334 .into_series(),
335 Float64 => s
336 .f64()
337 .unwrap()
338 .apply(|opt_v| {
339 opt_v.and_then(|v| {
340 let res = rhs.to_f64().unwrap();
341 if res.is_zero() { None } else { Some(v / res) }
342 })
343 })
344 .into_series(),
345 _ => panic!("dtype not yet supported in checked div"),
346 };
347 out.cast(self.dtype())
348 }
349 }
350}
351
352pub fn coerce_lhs_rhs<'a>(
353 lhs: &'a Series,
354 rhs: &'a Series,
355) -> PolarsResult<(Cow<'a, Series>, Cow<'a, Series>)> {
356 if let Some(result) = coerce_time_units(lhs, rhs) {
357 return Ok(result);
358 }
359 let (left_dtype, right_dtype) = (lhs.dtype(), rhs.dtype());
360 let leaf_super_dtype = try_get_supertype(left_dtype.leaf_dtype(), right_dtype.leaf_dtype())?;
361
362 let mut new_left_dtype = left_dtype.cast_leaf(leaf_super_dtype.clone());
363 let mut new_right_dtype = right_dtype.cast_leaf(leaf_super_dtype);
364
365 if left_dtype.is_list()
369 || right_dtype.is_list()
370 || left_dtype.is_array()
371 || right_dtype.is_array()
372 {
373 new_left_dtype = try_get_supertype(&new_left_dtype, &new_right_dtype)?;
374 new_right_dtype = new_left_dtype.clone();
375 }
376
377 let cast = |s: &'a Series, dtype: &DataType| -> PolarsResult<Cow<'a, Series>> {
380 Ok(if s.dtype() == dtype {
381 Cow::Borrowed(s)
382 } else if dtype.leaf_dtype().is_decimal() {
383 Cow::Owned(s.strict_cast(dtype)?)
384 } else {
385 Cow::Owned(s.cast(dtype)?)
386 })
387 };
388 Ok((cast(lhs, &new_left_dtype)?, cast(rhs, &new_right_dtype)?))
389}
390
391#[cfg(feature = "dtype-decimal")]
395pub(crate) fn decimal_op_operands<'a>(
396 lhs: &'a Series,
397 rhs: &'a Series,
398) -> Option<PolarsResult<(Cow<'a, Series>, Cow<'a, Series>)>> {
399 let (l, r) = (lhs.dtype(), rhs.dtype());
400 if !((l.is_decimal() && (r.is_decimal() || r.is_integer()))
401 || (l.is_integer() && r.is_decimal()))
402 {
403 return None;
404 }
405 let to_decimal = |s: &'a Series| -> PolarsResult<Cow<'a, Series>> {
406 if s.dtype().is_decimal() {
407 Ok(Cow::Borrowed(s))
408 } else {
409 let dtype = DataType::Decimal(polars_compute::decimal::DEC128_MAX_PREC, 0);
410 Ok(Cow::Owned(s.strict_cast(&dtype)?))
411 }
412 };
413 Some(to_decimal(lhs).and_then(|l| Ok((l, to_decimal(rhs)?))))
414}
415
416pub fn coerce_lhs_rhs_numeric_op<'a>(
419 lhs: &'a Series,
420 rhs: &'a Series,
421) -> PolarsResult<(Cow<'a, Series>, Cow<'a, Series>)> {
422 #[cfg(feature = "dtype-decimal")]
423 if let Some(out) = decimal_op_operands(lhs, rhs) {
424 return out;
425 }
426 coerce_lhs_rhs(lhs, rhs)
427}
428
429fn coerce_time_units<'a>(
434 lhs: &'a Series,
435 rhs: &'a Series,
436) -> Option<(Cow<'a, Series>, Cow<'a, Series>)> {
437 match (lhs.dtype(), rhs.dtype()) {
438 (DataType::Datetime(lu, t), DataType::Duration(ru)) => {
439 let units = get_time_units(lu, ru);
440 let left = if *lu == units {
441 Cow::Borrowed(lhs)
442 } else {
443 Cow::Owned(lhs.cast(&DataType::Datetime(units, t.clone())).ok()?)
444 };
445 let right = if *ru == units {
446 Cow::Borrowed(rhs)
447 } else {
448 Cow::Owned(rhs.cast(&DataType::Duration(units)).ok()?)
449 };
450 Some((left, right))
451 },
452 (DataType::Date, DataType::Duration(_)) => Some((Cow::Borrowed(lhs), Cow::Borrowed(rhs))),
454 (DataType::Duration(lu), DataType::Duration(ru)) => {
455 let units = get_time_units(lu, ru);
456 let left = if *lu == units {
457 Cow::Borrowed(lhs)
458 } else {
459 Cow::Owned(lhs.cast(&DataType::Duration(units)).ok()?)
460 };
461 let right = if *ru == units {
462 Cow::Borrowed(rhs)
463 } else {
464 Cow::Owned(rhs.cast(&DataType::Duration(units)).ok()?)
465 };
466 Some((left, right))
467 },
468 (DataType::Duration(_), DataType::Datetime(_, _))
470 | (DataType::Duration(_), DataType::Date) => {
471 let (right, left) = coerce_time_units(rhs, lhs)?;
472 Some((left, right))
473 },
474 _ => None,
475 }
476}
477
478#[cfg(feature = "dtype-struct")]
479pub fn _struct_arithmetic<F: FnMut(&Series, &Series) -> PolarsResult<Series>>(
480 s: &Series,
481 rhs: &Series,
482 mut func: F,
483) -> PolarsResult<Series> {
484 let s = s.struct_().unwrap();
485 let rhs = rhs.struct_().unwrap();
486
487 let s_fields = s.fields_as_series();
488 let rhs_fields = rhs.fields_as_series();
489
490 match (s_fields.len(), rhs_fields.len()) {
491 (_, 1) => {
492 let rhs = &rhs.fields_as_series()[0];
493 Ok(s.try_apply_fields(|s| func(s, rhs))?.into_series())
494 },
495 (1, _) => {
496 let s = &s.fields_as_series()[0];
497 Ok(rhs.try_apply_fields(|rhs| func(s, rhs))?.into_series())
498 },
499 _ => {
500 let len = broadcast_len([s, rhs]).context("struct arithmetic")?;
501 let s = s.broadcast_to(len)?;
502 let rhs = rhs.broadcast_to(len)?;
503
504 let (s, rhs) = align_chunks_binary(&s, &rhs);
505 let mut s = s.into_owned();
506
507 s.zip_outer_validity(rhs.as_ref());
509
510 let mut rhs_iter = rhs.fields_as_series().into_iter();
511
512 Ok(s.try_apply_fields(|s| match rhs_iter.next() {
513 Some(rhs) => func(s, &rhs),
514 None => Ok(s.clone()),
515 })?
516 .into_series())
517 },
518 }
519}
520
521fn check_lengths(a: &Series, b: &Series) -> PolarsResult<()> {
522 match (a.len(), b.len()) {
523 (1, _) | (_, 1) => Ok(()),
525 (a, b) if a == b => Ok(()),
527 (a, b) => {
529 polars_bail!(InvalidOperation: "cannot do arithmetic operation on series of different lengths: got {} and {}", a, b)
530 },
531 }
532}
533
534impl Add for &Series {
535 type Output = PolarsResult<Series>;
536
537 fn add(self, rhs: Self) -> Self::Output {
538 check_lengths(self, rhs)?;
539 match (self.dtype(), rhs.dtype()) {
540 #[cfg(feature = "dtype-struct")]
541 (DataType::Struct(_), DataType::Struct(_)) => {
542 _struct_arithmetic(self, rhs, |a, b| a.add(b))
543 },
544 (DataType::List(_), _) | (_, DataType::List(_)) => {
545 list::NumericListOp::add().execute(self, rhs)
546 },
547 #[cfg(feature = "dtype-array")]
548 (DataType::Array(..), _) | (_, DataType::Array(..)) => {
549 fixed_size_list::NumericFixedSizeListOp::add().execute(self, rhs)
550 },
551 (l_dtype, r_dtype) if l_dtype.is_temporal() != r_dtype.is_temporal() => {
552 polars_bail!(opq = add, l_dtype, r_dtype)
553 },
554 _ => {
555 let (lhs, rhs) = coerce_lhs_rhs_numeric_op(self, rhs)?;
556 lhs.add_to(rhs.as_ref())
557 },
558 }
559 }
560}
561
562impl Sub for &Series {
563 type Output = PolarsResult<Series>;
564
565 fn sub(self, rhs: Self) -> Self::Output {
566 check_lengths(self, rhs)?;
567 match (self.dtype(), rhs.dtype()) {
568 #[cfg(feature = "dtype-struct")]
569 (DataType::Struct(_), DataType::Struct(_)) => {
570 _struct_arithmetic(self, rhs, |a, b| a.sub(b))
571 },
572 (DataType::List(_), _) | (_, DataType::List(_)) => {
573 list::NumericListOp::sub().execute(self, rhs)
574 },
575 #[cfg(feature = "dtype-array")]
576 (DataType::Array(..), _) | (_, DataType::Array(..)) => {
577 fixed_size_list::NumericFixedSizeListOp::sub().execute(self, rhs)
578 },
579 (l_dtype, r_dtype) if l_dtype.is_temporal() != r_dtype.is_temporal() => {
580 polars_bail!(opq = sub, l_dtype, r_dtype)
581 },
582 _ => {
583 let (lhs, rhs) = coerce_lhs_rhs_numeric_op(self, rhs)?;
584 lhs.subtract(rhs.as_ref())
585 },
586 }
587 }
588}
589
590impl Mul for &Series {
591 type Output = PolarsResult<Series>;
592
593 fn mul(self, rhs: Self) -> Self::Output {
599 check_lengths(self, rhs)?;
600
601 use DataType::*;
602 match (self.dtype(), rhs.dtype()) {
603 #[cfg(feature = "dtype-struct")]
604 (Struct(_), Struct(_)) => _struct_arithmetic(self, rhs, |a, b| a.mul(b)),
605 (Duration(_), _) | (Date, _) | (Datetime(_, _), _) | (Time, _) => self.multiply(rhs),
607 (_, Date) | (_, Datetime(_, _)) | (_, Time) => {
609 polars_bail!(opq = mul, self.dtype(), rhs.dtype())
610 },
611 (_, Duration(_)) => {
612 let out = rhs.multiply(self)?;
614 Ok(out.with_name(self.name().clone()))
615 },
616 (DataType::List(_), _) | (_, DataType::List(_)) => {
617 list::NumericListOp::mul().execute(self, rhs)
618 },
619 #[cfg(feature = "dtype-array")]
620 (DataType::Array(..), _) | (_, DataType::Array(..)) => {
621 fixed_size_list::NumericFixedSizeListOp::mul().execute(self, rhs)
622 },
623 _ => {
624 let (lhs, rhs) = coerce_lhs_rhs_numeric_op(self, rhs)?;
625 lhs.multiply(rhs.as_ref())
626 },
627 }
628 }
629}
630
631impl Div for &Series {
632 type Output = PolarsResult<Series>;
633
634 fn div(self, rhs: Self) -> Self::Output {
640 check_lengths(self, rhs)?;
641 use DataType::*;
642 match (self.dtype(), rhs.dtype()) {
643 #[cfg(feature = "dtype-struct")]
644 (Struct(_), Struct(_)) => _struct_arithmetic(self, rhs, |a, b| a.div(b)),
645 (Duration(_), _) => self.divide(rhs),
646 (Date, _)
647 | (Datetime(_, _), _)
648 | (Time, _)
649 | (_, Duration(_))
650 | (_, Time)
651 | (_, Date)
652 | (_, Datetime(_, _)) => polars_bail!(opq = div, self.dtype(), rhs.dtype()),
653 (DataType::List(_), _) | (_, DataType::List(_)) => {
654 list::NumericListOp::div().execute(self, rhs)
655 },
656 #[cfg(feature = "dtype-array")]
657 (DataType::Array(..), _) | (_, DataType::Array(..)) => {
658 fixed_size_list::NumericFixedSizeListOp::div().execute(self, rhs)
659 },
660 _ => {
661 let (lhs, rhs) = coerce_lhs_rhs_numeric_op(self, rhs)?;
662 lhs.divide(rhs.as_ref())
663 },
664 }
665 }
666}
667
668impl Rem for &Series {
669 type Output = PolarsResult<Series>;
670
671 fn rem(self, rhs: Self) -> Self::Output {
677 check_lengths(self, rhs)?;
678 match (self.dtype(), rhs.dtype()) {
679 #[cfg(feature = "dtype-struct")]
680 (DataType::Struct(_), DataType::Struct(_)) => {
681 _struct_arithmetic(self, rhs, |a, b| a.rem(b))
682 },
683 (DataType::List(_), _) | (_, DataType::List(_)) => {
684 list::NumericListOp::rem().execute(self, rhs)
685 },
686 #[cfg(feature = "dtype-array")]
687 (DataType::Array(..), _) | (_, DataType::Array(..)) => {
688 fixed_size_list::NumericFixedSizeListOp::rem().execute(self, rhs)
689 },
690 _ => {
691 let (lhs, rhs) = coerce_lhs_rhs_numeric_op(self, rhs)?;
692 lhs.remainder(rhs.as_ref())
693 },
694 }
695 }
696}
697
698fn finish_cast(inp: &Series, out: Series) -> Series {
701 match inp.dtype() {
702 #[cfg(feature = "dtype-date")]
703 DataType::Date => out.into_date(),
704 #[cfg(feature = "dtype-datetime")]
705 DataType::Datetime(tu, tz) => out.into_datetime(*tu, tz.clone()),
706 #[cfg(feature = "dtype-duration")]
707 DataType::Duration(tu) => out.into_duration(*tu),
708 #[cfg(feature = "dtype-time")]
709 DataType::Time => out.into_time(),
710 _ => out,
711 }
712}
713
714impl<T> Sub<T> for &Series
715where
716 T: Num + NumCast,
717{
718 type Output = Series;
719
720 fn sub(self, rhs: T) -> Self::Output {
721 let s = self.to_physical_repr();
722 macro_rules! sub {
723 ($ca:expr) => {{ $ca.sub(rhs).into_series() }};
724 }
725
726 let out = downcast_as_macro_arg_physical!(s, sub);
727 finish_cast(self, out)
728 }
729}
730
731impl<T> Sub<T> for Series
732where
733 T: Num + NumCast,
734{
735 type Output = Self;
736
737 fn sub(self, rhs: T) -> Self::Output {
738 (&self).sub(rhs)
739 }
740}
741
742impl<T> Add<T> for &Series
743where
744 T: Num + NumCast,
745{
746 type Output = Series;
747
748 fn add(self, rhs: T) -> Self::Output {
749 let s = self.to_physical_repr();
750 macro_rules! add {
751 ($ca:expr) => {{ $ca.add(rhs).into_series() }};
752 }
753 let out = downcast_as_macro_arg_physical!(s, add);
754 finish_cast(self, out)
755 }
756}
757
758impl<T> Add<T> for Series
759where
760 T: Num + NumCast,
761{
762 type Output = Self;
763
764 fn add(self, rhs: T) -> Self::Output {
765 (&self).add(rhs)
766 }
767}
768
769impl<T> Div<T> for &Series
770where
771 T: Num + NumCast,
772{
773 type Output = Series;
774
775 fn div(self, rhs: T) -> Self::Output {
776 let s = self.to_physical_repr();
777 macro_rules! div {
778 ($ca:expr) => {{ $ca.div(rhs).into_series() }};
779 }
780
781 let out = downcast_as_macro_arg_physical!(s, div);
782 finish_cast(self, out)
783 }
784}
785
786impl<T> Div<T> for Series
787where
788 T: Num + NumCast,
789{
790 type Output = Self;
791
792 fn div(self, rhs: T) -> Self::Output {
793 (&self).div(rhs)
794 }
795}
796
797impl Series {
799 pub fn wrapping_trunc_div_scalar<T: Num + NumCast>(&self, rhs: T) -> Self {
800 let s = self.to_physical_repr();
801 macro_rules! div {
802 ($ca:expr) => {{
803 let rhs = NumCast::from(rhs).unwrap();
804 $ca.wrapping_trunc_div_scalar(rhs).into_series()
805 }};
806 }
807
808 let out = downcast_as_macro_arg_physical!(s, div);
809 finish_cast(self, out)
810 }
811}
812
813impl<T> Mul<T> for &Series
814where
815 T: Num + NumCast,
816{
817 type Output = Series;
818
819 fn mul(self, rhs: T) -> Self::Output {
820 let s = self.to_physical_repr();
821 macro_rules! mul {
822 ($ca:expr) => {{ $ca.mul(rhs).into_series() }};
823 }
824 let out = downcast_as_macro_arg_physical!(s, mul);
825 finish_cast(self, out)
826 }
827}
828
829impl<T> Mul<T> for Series
830where
831 T: Num + NumCast,
832{
833 type Output = Self;
834
835 fn mul(self, rhs: T) -> Self::Output {
836 (&self).mul(rhs)
837 }
838}
839
840impl<T> Rem<T> for &Series
841where
842 T: Num + NumCast,
843{
844 type Output = Series;
845
846 fn rem(self, rhs: T) -> Self::Output {
847 let s = self.to_physical_repr();
848 macro_rules! rem {
849 ($ca:expr) => {{ $ca.rem(rhs).into_series() }};
850 }
851 let out = downcast_as_macro_arg_physical!(s, rem);
852 finish_cast(self, out)
853 }
854}
855
856impl<T> Rem<T> for Series
857where
858 T: Num + NumCast,
859{
860 type Output = Self;
861
862 fn rem(self, rhs: T) -> Self::Output {
863 (&self).rem(rhs)
864 }
865}
866
867impl<T: PolarsNumericType> ChunkedArray<T> {
871 #[must_use]
873 pub fn lhs_sub<N: Num + NumCast>(&self, lhs: N) -> Self {
874 let lhs: T::Native = NumCast::from(lhs).expect("could not cast");
875 ArithmeticChunked::wrapping_sub_scalar_lhs(lhs, self)
876 }
877
878 #[must_use]
880 pub fn lhs_div<N: Num + NumCast>(&self, lhs: N) -> Self {
881 let lhs: T::Native = NumCast::from(lhs).expect("could not cast");
882 ArithmeticChunked::legacy_div_scalar_lhs(lhs, self)
883 }
884
885 #[must_use]
887 pub fn lhs_rem<N: Num + NumCast>(&self, lhs: N) -> Self {
888 let lhs: T::Native = NumCast::from(lhs).expect("could not cast");
889 ArithmeticChunked::wrapping_mod_scalar_lhs(lhs, self)
890 }
891}
892
893pub trait LhsNumOps {
894 type Output;
895
896 fn add(self, rhs: &Series) -> Self::Output;
897 fn sub(self, rhs: &Series) -> Self::Output;
898 fn div(self, rhs: &Series) -> Self::Output;
899 fn mul(self, rhs: &Series) -> Self::Output;
900 fn rem(self, rem: &Series) -> Self::Output;
901}
902
903impl<T> LhsNumOps for T
904where
905 T: Num + NumCast,
906{
907 type Output = Series;
908
909 fn add(self, rhs: &Series) -> Self::Output {
910 rhs + self
912 }
913 fn sub(self, rhs: &Series) -> Self::Output {
914 let s = rhs.to_physical_repr();
915 macro_rules! sub {
916 ($rhs:expr) => {{ $rhs.lhs_sub(self).into_series() }};
917 }
918 let out = downcast_as_macro_arg_physical!(s, sub);
919
920 finish_cast(rhs, out)
921 }
922 fn div(self, rhs: &Series) -> Self::Output {
923 let s = rhs.to_physical_repr();
924 macro_rules! div {
925 ($rhs:expr) => {{ $rhs.lhs_div(self).into_series() }};
926 }
927 let out = downcast_as_macro_arg_physical!(s, div);
928
929 finish_cast(rhs, out)
930 }
931 fn mul(self, rhs: &Series) -> Self::Output {
932 rhs * self
934 }
935 fn rem(self, rhs: &Series) -> Self::Output {
936 let s = rhs.to_physical_repr();
937 macro_rules! rem {
938 ($rhs:expr) => {{ $rhs.lhs_rem(self).into_series() }};
939 }
940
941 let out = downcast_as_macro_arg_physical!(s, rem);
942
943 finish_cast(rhs, out)
944 }
945}
946
947#[cfg(test)]
948mod test {
949 use crate::prelude::*;
950
951 #[test]
952 #[allow(clippy::eq_op)]
953 fn test_arithmetic_series() -> PolarsResult<()> {
954 let s = Series::new("foo".into(), [1, 2, 3]);
956 assert_eq!(
957 Vec::from((&s * &s)?.i32().unwrap()),
958 [Some(1), Some(4), Some(9)]
959 );
960 assert_eq!(
961 Vec::from((&s / &s)?.i32().unwrap()),
962 [Some(1), Some(1), Some(1)]
963 );
964 assert_eq!(
965 Vec::from((&s - &s)?.i32().unwrap()),
966 [Some(0), Some(0), Some(0)]
967 );
968 assert_eq!(
969 Vec::from((&s + &s)?.i32().unwrap()),
970 [Some(2), Some(4), Some(6)]
971 );
972 assert_eq!(
974 Vec::from((&s + 1).i32().unwrap()),
975 [Some(2), Some(3), Some(4)]
976 );
977 assert_eq!(
978 Vec::from((&s - 1).i32().unwrap()),
979 [Some(0), Some(1), Some(2)]
980 );
981 assert_eq!(
982 Vec::from((&s * 2).i32().unwrap()),
983 [Some(2), Some(4), Some(6)]
984 );
985 assert_eq!(
986 Vec::from((&s / 2).i32().unwrap()),
987 [Some(0), Some(1), Some(1)]
988 );
989
990 assert_eq!(
992 Vec::from((1.add(&s)).i32().unwrap()),
993 [Some(2), Some(3), Some(4)]
994 );
995 assert_eq!(
996 Vec::from((1.sub(&s)).i32().unwrap()),
997 [Some(0), Some(-1), Some(-2)]
998 );
999 assert_eq!(
1000 Vec::from((1.div(&s)).i32().unwrap()),
1001 [Some(1), Some(0), Some(0)]
1002 );
1003 assert_eq!(
1004 Vec::from((1.mul(&s)).i32().unwrap()),
1005 [Some(1), Some(2), Some(3)]
1006 );
1007 assert_eq!(
1008 Vec::from((1.rem(&s)).i32().unwrap()),
1009 [Some(0), Some(1), Some(1)]
1010 );
1011
1012 assert_eq!((&s * &s)?.name().as_str(), "foo");
1013 assert_eq!((&s * 1).name().as_str(), "foo");
1014 assert_eq!((1.div(&s)).name().as_str(), "foo");
1015
1016 Ok(())
1017 }
1018
1019 #[test]
1020 #[cfg(feature = "checked_arithmetic")]
1021 fn test_checked_div() {
1022 let s = Series::new("foo".into(), [1i32, 0, 1]);
1023 let out = s.checked_div(&s).unwrap();
1024 assert_eq!(Vec::from(out.i32().unwrap()), &[Some(1), None, Some(1)]);
1025 let out = s.checked_div_num(0).unwrap();
1026 assert_eq!(Vec::from(out.i32().unwrap()), &[None, None, None]);
1027
1028 let s_f32 = Series::new("float32".into(), [1.0f32, 0.0, 1.0]);
1029 let out = s_f32.checked_div(&s_f32).unwrap();
1030 assert_eq!(
1031 Vec::from(out.f32().unwrap()),
1032 &[Some(1.0f32), None, Some(1.0f32)]
1033 );
1034 let out = s_f32.checked_div_num(0.0f32).unwrap();
1035 assert_eq!(Vec::from(out.f32().unwrap()), &[None, None, None]);
1036
1037 let s_f64 = Series::new("float64".into(), [1.0f64, 0.0, 1.0]);
1038 let out = s_f64.checked_div(&s_f64).unwrap();
1039 assert_eq!(
1040 Vec::from(out.f64().unwrap()),
1041 &[Some(1.0f64), None, Some(1.0f64)]
1042 );
1043 let out = s_f64.checked_div_num(0.0f64).unwrap();
1044 assert_eq!(Vec::from(out.f64().unwrap()), &[None, None, None]);
1045 }
1046}