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>(
420 lhs: &'a Series,
421 rhs: &'a Series,
422) -> PolarsResult<(Cow<'a, Series>, Cow<'a, Series>)> {
423 let (l, r) = (lhs.dtype(), rhs.dtype());
425 let is_numeric = |dt: &DataType| dt.is_primitive_numeric() || dt.is_decimal();
426 polars_ensure!(
427 !((l.is_string() && is_numeric(r)) || (is_numeric(l) && r.is_string())),
428 InvalidOperation: "arithmetic on dtypes {l} and {r} is not allowed; try an explicit cast first"
429 );
430
431 #[cfg(feature = "dtype-decimal")]
432 if let Some(out) = decimal_op_operands(lhs, rhs) {
433 return out;
434 }
435 coerce_lhs_rhs(lhs, rhs)
436}
437
438fn coerce_time_units<'a>(
443 lhs: &'a Series,
444 rhs: &'a Series,
445) -> Option<(Cow<'a, Series>, Cow<'a, Series>)> {
446 match (lhs.dtype(), rhs.dtype()) {
447 (DataType::Datetime(lu, t), DataType::Duration(ru)) => {
448 let units = get_time_units(lu, ru);
449 let left = if *lu == units {
450 Cow::Borrowed(lhs)
451 } else {
452 Cow::Owned(lhs.cast(&DataType::Datetime(units, t.clone())).ok()?)
453 };
454 let right = if *ru == units {
455 Cow::Borrowed(rhs)
456 } else {
457 Cow::Owned(rhs.cast(&DataType::Duration(units)).ok()?)
458 };
459 Some((left, right))
460 },
461 (DataType::Date, DataType::Duration(_)) => Some((Cow::Borrowed(lhs), Cow::Borrowed(rhs))),
463 (DataType::Duration(lu), DataType::Duration(ru)) => {
464 let units = get_time_units(lu, ru);
465 let left = if *lu == units {
466 Cow::Borrowed(lhs)
467 } else {
468 Cow::Owned(lhs.cast(&DataType::Duration(units)).ok()?)
469 };
470 let right = if *ru == units {
471 Cow::Borrowed(rhs)
472 } else {
473 Cow::Owned(rhs.cast(&DataType::Duration(units)).ok()?)
474 };
475 Some((left, right))
476 },
477 (DataType::Duration(_), DataType::Datetime(_, _))
479 | (DataType::Duration(_), DataType::Date) => {
480 let (right, left) = coerce_time_units(rhs, lhs)?;
481 Some((left, right))
482 },
483 _ => None,
484 }
485}
486
487#[cfg(feature = "dtype-struct")]
488pub fn _struct_arithmetic<F: FnMut(&Series, &Series) -> PolarsResult<Series>>(
489 s: &Series,
490 rhs: &Series,
491 mut func: F,
492) -> PolarsResult<Series> {
493 let s = s.struct_().unwrap();
494 let rhs = rhs.struct_().unwrap();
495
496 let s_fields = s.fields_as_series();
497 let rhs_fields = rhs.fields_as_series();
498
499 match (s_fields.len(), rhs_fields.len()) {
500 (_, 1) => {
501 let rhs = &rhs.fields_as_series()[0];
502 Ok(s.try_apply_fields(|s| func(s, rhs))?.into_series())
503 },
504 (1, _) => {
505 let s = &s.fields_as_series()[0];
506 Ok(rhs.try_apply_fields(|rhs| func(s, rhs))?.into_series())
507 },
508 _ => {
509 let len = broadcast_len([s, rhs]).context("struct arithmetic")?;
510 let s = s.broadcast_to(len)?;
511 let rhs = rhs.broadcast_to(len)?;
512
513 let (s, rhs) = align_chunks_binary(&s, &rhs);
514 let mut s = s.into_owned();
515
516 s.zip_outer_validity(rhs.as_ref());
518
519 let mut rhs_iter = rhs.fields_as_series().into_iter();
520
521 Ok(s.try_apply_fields(|s| match rhs_iter.next() {
522 Some(rhs) => func(s, &rhs),
523 None => Ok(s.clone()),
524 })?
525 .into_series())
526 },
527 }
528}
529
530fn check_lengths(a: &Series, b: &Series) -> PolarsResult<()> {
531 match (a.len(), b.len()) {
532 (1, _) | (_, 1) => Ok(()),
534 (a, b) if a == b => Ok(()),
536 (a, b) => {
538 polars_bail!(InvalidOperation: "cannot do arithmetic operation on series of different lengths: got {} and {}", a, b)
539 },
540 }
541}
542
543impl Add for &Series {
544 type Output = PolarsResult<Series>;
545
546 fn add(self, rhs: Self) -> Self::Output {
547 check_lengths(self, rhs)?;
548 match (self.dtype(), rhs.dtype()) {
549 #[cfg(feature = "dtype-struct")]
550 (DataType::Struct(_), DataType::Struct(_)) => {
551 _struct_arithmetic(self, rhs, |a, b| a.add(b))
552 },
553 (DataType::List(_), _) | (_, DataType::List(_)) => {
554 list::NumericListOp::add().execute(self, rhs)
555 },
556 #[cfg(feature = "dtype-array")]
557 (DataType::Array(..), _) | (_, DataType::Array(..)) => {
558 fixed_size_list::NumericFixedSizeListOp::add().execute(self, rhs)
559 },
560 (l_dtype, r_dtype) if l_dtype.is_temporal() != r_dtype.is_temporal() => {
561 polars_bail!(opq = add, l_dtype, r_dtype)
562 },
563 _ => {
564 let (lhs, rhs) = coerce_lhs_rhs_numeric_op(self, rhs)?;
565 lhs.add_to(rhs.as_ref())
566 },
567 }
568 }
569}
570
571impl Sub for &Series {
572 type Output = PolarsResult<Series>;
573
574 fn sub(self, rhs: Self) -> Self::Output {
575 check_lengths(self, rhs)?;
576 match (self.dtype(), rhs.dtype()) {
577 #[cfg(feature = "dtype-struct")]
578 (DataType::Struct(_), DataType::Struct(_)) => {
579 _struct_arithmetic(self, rhs, |a, b| a.sub(b))
580 },
581 (DataType::List(_), _) | (_, DataType::List(_)) => {
582 list::NumericListOp::sub().execute(self, rhs)
583 },
584 #[cfg(feature = "dtype-array")]
585 (DataType::Array(..), _) | (_, DataType::Array(..)) => {
586 fixed_size_list::NumericFixedSizeListOp::sub().execute(self, rhs)
587 },
588 (l_dtype, r_dtype) if l_dtype.is_temporal() != r_dtype.is_temporal() => {
589 polars_bail!(opq = sub, l_dtype, r_dtype)
590 },
591 _ => {
592 let (lhs, rhs) = coerce_lhs_rhs_numeric_op(self, rhs)?;
593 lhs.subtract(rhs.as_ref())
594 },
595 }
596 }
597}
598
599impl Mul for &Series {
600 type Output = PolarsResult<Series>;
601
602 fn mul(self, rhs: Self) -> Self::Output {
608 check_lengths(self, rhs)?;
609
610 use DataType::*;
611 match (self.dtype(), rhs.dtype()) {
612 #[cfg(feature = "dtype-struct")]
613 (Struct(_), Struct(_)) => _struct_arithmetic(self, rhs, |a, b| a.mul(b)),
614 (Duration(_), _) | (Date, _) | (Datetime(_, _), _) | (Time, _) => self.multiply(rhs),
616 (_, Date) | (_, Datetime(_, _)) | (_, Time) => {
618 polars_bail!(opq = mul, self.dtype(), rhs.dtype())
619 },
620 (_, Duration(_)) => {
621 let out = rhs.multiply(self)?;
623 Ok(out.with_name(self.name().clone()))
624 },
625 (DataType::List(_), _) | (_, DataType::List(_)) => {
626 list::NumericListOp::mul().execute(self, rhs)
627 },
628 #[cfg(feature = "dtype-array")]
629 (DataType::Array(..), _) | (_, DataType::Array(..)) => {
630 fixed_size_list::NumericFixedSizeListOp::mul().execute(self, rhs)
631 },
632 _ => {
633 let (lhs, rhs) = coerce_lhs_rhs_numeric_op(self, rhs)?;
634 lhs.multiply(rhs.as_ref())
635 },
636 }
637 }
638}
639
640impl Div for &Series {
641 type Output = PolarsResult<Series>;
642
643 fn div(self, rhs: Self) -> Self::Output {
649 check_lengths(self, rhs)?;
650 use DataType::*;
651 match (self.dtype(), rhs.dtype()) {
652 #[cfg(feature = "dtype-struct")]
653 (Struct(_), Struct(_)) => _struct_arithmetic(self, rhs, |a, b| a.div(b)),
654 (Duration(_), _) => self.divide(rhs),
655 (Date, _)
656 | (Datetime(_, _), _)
657 | (Time, _)
658 | (_, Duration(_))
659 | (_, Time)
660 | (_, Date)
661 | (_, Datetime(_, _)) => polars_bail!(opq = div, self.dtype(), rhs.dtype()),
662 (DataType::List(_), _) | (_, DataType::List(_)) => {
663 list::NumericListOp::div().execute(self, rhs)
664 },
665 #[cfg(feature = "dtype-array")]
666 (DataType::Array(..), _) | (_, DataType::Array(..)) => {
667 fixed_size_list::NumericFixedSizeListOp::div().execute(self, rhs)
668 },
669 _ => {
670 let (lhs, rhs) = coerce_lhs_rhs_numeric_op(self, rhs)?;
671 lhs.divide(rhs.as_ref())
672 },
673 }
674 }
675}
676
677impl Rem for &Series {
678 type Output = PolarsResult<Series>;
679
680 fn rem(self, rhs: Self) -> Self::Output {
686 check_lengths(self, rhs)?;
687 match (self.dtype(), rhs.dtype()) {
688 #[cfg(feature = "dtype-struct")]
689 (DataType::Struct(_), DataType::Struct(_)) => {
690 _struct_arithmetic(self, rhs, |a, b| a.rem(b))
691 },
692 (DataType::List(_), _) | (_, DataType::List(_)) => {
693 list::NumericListOp::rem().execute(self, rhs)
694 },
695 #[cfg(feature = "dtype-array")]
696 (DataType::Array(..), _) | (_, DataType::Array(..)) => {
697 fixed_size_list::NumericFixedSizeListOp::rem().execute(self, rhs)
698 },
699 _ => {
700 let (lhs, rhs) = coerce_lhs_rhs_numeric_op(self, rhs)?;
701 lhs.remainder(rhs.as_ref())
702 },
703 }
704 }
705}
706
707fn finish_cast(inp: &Series, out: Series) -> Series {
710 match inp.dtype() {
711 #[cfg(feature = "dtype-date")]
712 DataType::Date => out.into_date(),
713 #[cfg(feature = "dtype-datetime")]
714 DataType::Datetime(tu, tz) => out.into_datetime(*tu, tz.clone()),
715 #[cfg(feature = "dtype-duration")]
716 DataType::Duration(tu) => out.into_duration(*tu),
717 #[cfg(feature = "dtype-time")]
718 DataType::Time => out.into_time(),
719 _ => out,
720 }
721}
722
723impl<T> Sub<T> for &Series
724where
725 T: Num + NumCast,
726{
727 type Output = Series;
728
729 fn sub(self, rhs: T) -> Self::Output {
730 let s = self.to_physical_repr();
731 macro_rules! sub {
732 ($ca:expr) => {{ $ca.sub(rhs).into_series() }};
733 }
734
735 let out = downcast_as_macro_arg_physical!(s, sub);
736 finish_cast(self, out)
737 }
738}
739
740impl<T> Sub<T> for Series
741where
742 T: Num + NumCast,
743{
744 type Output = Self;
745
746 fn sub(self, rhs: T) -> Self::Output {
747 (&self).sub(rhs)
748 }
749}
750
751impl<T> Add<T> for &Series
752where
753 T: Num + NumCast,
754{
755 type Output = Series;
756
757 fn add(self, rhs: T) -> Self::Output {
758 let s = self.to_physical_repr();
759 macro_rules! add {
760 ($ca:expr) => {{ $ca.add(rhs).into_series() }};
761 }
762 let out = downcast_as_macro_arg_physical!(s, add);
763 finish_cast(self, out)
764 }
765}
766
767impl<T> Add<T> for Series
768where
769 T: Num + NumCast,
770{
771 type Output = Self;
772
773 fn add(self, rhs: T) -> Self::Output {
774 (&self).add(rhs)
775 }
776}
777
778impl<T> Div<T> for &Series
779where
780 T: Num + NumCast,
781{
782 type Output = Series;
783
784 fn div(self, rhs: T) -> Self::Output {
785 let s = self.to_physical_repr();
786 macro_rules! div {
787 ($ca:expr) => {{ $ca.div(rhs).into_series() }};
788 }
789
790 let out = downcast_as_macro_arg_physical!(s, div);
791 finish_cast(self, out)
792 }
793}
794
795impl<T> Div<T> for Series
796where
797 T: Num + NumCast,
798{
799 type Output = Self;
800
801 fn div(self, rhs: T) -> Self::Output {
802 (&self).div(rhs)
803 }
804}
805
806impl Series {
808 pub fn wrapping_trunc_div_scalar<T: Num + NumCast>(&self, rhs: T) -> Self {
809 let s = self.to_physical_repr();
810 macro_rules! div {
811 ($ca:expr) => {{
812 let rhs = NumCast::from(rhs).unwrap();
813 $ca.wrapping_trunc_div_scalar(rhs).into_series()
814 }};
815 }
816
817 let out = downcast_as_macro_arg_physical!(s, div);
818 finish_cast(self, out)
819 }
820}
821
822impl<T> Mul<T> for &Series
823where
824 T: Num + NumCast,
825{
826 type Output = Series;
827
828 fn mul(self, rhs: T) -> Self::Output {
829 let s = self.to_physical_repr();
830 macro_rules! mul {
831 ($ca:expr) => {{ $ca.mul(rhs).into_series() }};
832 }
833 let out = downcast_as_macro_arg_physical!(s, mul);
834 finish_cast(self, out)
835 }
836}
837
838impl<T> Mul<T> for Series
839where
840 T: Num + NumCast,
841{
842 type Output = Self;
843
844 fn mul(self, rhs: T) -> Self::Output {
845 (&self).mul(rhs)
846 }
847}
848
849impl<T> Rem<T> for &Series
850where
851 T: Num + NumCast,
852{
853 type Output = Series;
854
855 fn rem(self, rhs: T) -> Self::Output {
856 let s = self.to_physical_repr();
857 macro_rules! rem {
858 ($ca:expr) => {{ $ca.rem(rhs).into_series() }};
859 }
860 let out = downcast_as_macro_arg_physical!(s, rem);
861 finish_cast(self, out)
862 }
863}
864
865impl<T> Rem<T> for Series
866where
867 T: Num + NumCast,
868{
869 type Output = Self;
870
871 fn rem(self, rhs: T) -> Self::Output {
872 (&self).rem(rhs)
873 }
874}
875
876impl<T: PolarsNumericType> ChunkedArray<T> {
880 #[must_use]
882 pub fn lhs_sub<N: Num + NumCast>(&self, lhs: N) -> Self {
883 let lhs: T::Native = NumCast::from(lhs).expect("could not cast");
884 ArithmeticChunked::wrapping_sub_scalar_lhs(lhs, self)
885 }
886
887 #[must_use]
889 pub fn lhs_div<N: Num + NumCast>(&self, lhs: N) -> Self {
890 let lhs: T::Native = NumCast::from(lhs).expect("could not cast");
891 ArithmeticChunked::legacy_div_scalar_lhs(lhs, self)
892 }
893
894 #[must_use]
896 pub fn lhs_rem<N: Num + NumCast>(&self, lhs: N) -> Self {
897 let lhs: T::Native = NumCast::from(lhs).expect("could not cast");
898 ArithmeticChunked::wrapping_mod_scalar_lhs(lhs, self)
899 }
900}
901
902pub trait LhsNumOps {
903 type Output;
904
905 fn add(self, rhs: &Series) -> Self::Output;
906 fn sub(self, rhs: &Series) -> Self::Output;
907 fn div(self, rhs: &Series) -> Self::Output;
908 fn mul(self, rhs: &Series) -> Self::Output;
909 fn rem(self, rem: &Series) -> Self::Output;
910}
911
912impl<T> LhsNumOps for T
913where
914 T: Num + NumCast,
915{
916 type Output = Series;
917
918 fn add(self, rhs: &Series) -> Self::Output {
919 rhs + self
921 }
922 fn sub(self, rhs: &Series) -> Self::Output {
923 let s = rhs.to_physical_repr();
924 macro_rules! sub {
925 ($rhs:expr) => {{ $rhs.lhs_sub(self).into_series() }};
926 }
927 let out = downcast_as_macro_arg_physical!(s, sub);
928
929 finish_cast(rhs, out)
930 }
931 fn div(self, rhs: &Series) -> Self::Output {
932 let s = rhs.to_physical_repr();
933 macro_rules! div {
934 ($rhs:expr) => {{ $rhs.lhs_div(self).into_series() }};
935 }
936 let out = downcast_as_macro_arg_physical!(s, div);
937
938 finish_cast(rhs, out)
939 }
940 fn mul(self, rhs: &Series) -> Self::Output {
941 rhs * self
943 }
944 fn rem(self, rhs: &Series) -> Self::Output {
945 let s = rhs.to_physical_repr();
946 macro_rules! rem {
947 ($rhs:expr) => {{ $rhs.lhs_rem(self).into_series() }};
948 }
949
950 let out = downcast_as_macro_arg_physical!(s, rem);
951
952 finish_cast(rhs, out)
953 }
954}
955
956#[cfg(test)]
957mod test {
958 use crate::prelude::*;
959
960 #[test]
961 #[allow(clippy::eq_op)]
962 fn test_arithmetic_series() -> PolarsResult<()> {
963 let s = Series::new("foo".into(), [1, 2, 3]);
965 assert_eq!(
966 Vec::from((&s * &s)?.i32().unwrap()),
967 [Some(1), Some(4), Some(9)]
968 );
969 assert_eq!(
970 Vec::from((&s / &s)?.i32().unwrap()),
971 [Some(1), Some(1), Some(1)]
972 );
973 assert_eq!(
974 Vec::from((&s - &s)?.i32().unwrap()),
975 [Some(0), Some(0), Some(0)]
976 );
977 assert_eq!(
978 Vec::from((&s + &s)?.i32().unwrap()),
979 [Some(2), Some(4), Some(6)]
980 );
981 assert_eq!(
983 Vec::from((&s + 1).i32().unwrap()),
984 [Some(2), Some(3), Some(4)]
985 );
986 assert_eq!(
987 Vec::from((&s - 1).i32().unwrap()),
988 [Some(0), Some(1), Some(2)]
989 );
990 assert_eq!(
991 Vec::from((&s * 2).i32().unwrap()),
992 [Some(2), Some(4), Some(6)]
993 );
994 assert_eq!(
995 Vec::from((&s / 2).i32().unwrap()),
996 [Some(0), Some(1), Some(1)]
997 );
998
999 assert_eq!(
1001 Vec::from((1.add(&s)).i32().unwrap()),
1002 [Some(2), Some(3), Some(4)]
1003 );
1004 assert_eq!(
1005 Vec::from((1.sub(&s)).i32().unwrap()),
1006 [Some(0), Some(-1), Some(-2)]
1007 );
1008 assert_eq!(
1009 Vec::from((1.div(&s)).i32().unwrap()),
1010 [Some(1), Some(0), Some(0)]
1011 );
1012 assert_eq!(
1013 Vec::from((1.mul(&s)).i32().unwrap()),
1014 [Some(1), Some(2), Some(3)]
1015 );
1016 assert_eq!(
1017 Vec::from((1.rem(&s)).i32().unwrap()),
1018 [Some(0), Some(1), Some(1)]
1019 );
1020
1021 assert_eq!((&s * &s)?.name().as_str(), "foo");
1022 assert_eq!((&s * 1).name().as_str(), "foo");
1023 assert_eq!((1.div(&s)).name().as_str(), "foo");
1024
1025 Ok(())
1026 }
1027
1028 #[test]
1029 #[cfg(feature = "checked_arithmetic")]
1030 fn test_checked_div() {
1031 let s = Series::new("foo".into(), [1i32, 0, 1]);
1032 let out = s.checked_div(&s).unwrap();
1033 assert_eq!(Vec::from(out.i32().unwrap()), &[Some(1), None, Some(1)]);
1034 let out = s.checked_div_num(0).unwrap();
1035 assert_eq!(Vec::from(out.i32().unwrap()), &[None, None, None]);
1036
1037 let s_f32 = Series::new("float32".into(), [1.0f32, 0.0, 1.0]);
1038 let out = s_f32.checked_div(&s_f32).unwrap();
1039 assert_eq!(
1040 Vec::from(out.f32().unwrap()),
1041 &[Some(1.0f32), None, Some(1.0f32)]
1042 );
1043 let out = s_f32.checked_div_num(0.0f32).unwrap();
1044 assert_eq!(Vec::from(out.f32().unwrap()), &[None, None, None]);
1045
1046 let s_f64 = Series::new("float64".into(), [1.0f64, 0.0, 1.0]);
1047 let out = s_f64.checked_div(&s_f64).unwrap();
1048 assert_eq!(
1049 Vec::from(out.f64().unwrap()),
1050 &[Some(1.0f64), None, Some(1.0f64)]
1051 );
1052 let out = s_f64.checked_div_num(0.0f64).unwrap();
1053 assert_eq!(Vec::from(out.f64().unwrap()), &[None, None, None]);
1054 }
1055}