Skip to main content

polars_core/series/arithmetic/
borrowed.rs

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        // SAFETY:
60        // There will be UB if a ChunkedArray is alive with the wrong datatype.
61        // we now only create the potentially wrong dtype for a short time.
62        // Note that the physical type correctness is checked!
63        // The ChunkedArray with the wrong dtype is dropped after this operation
64        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        // SAFETY:
77        // see subtract
78        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        // SAFETY:
91        // see subtract
92        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        // SAFETY:
105        // see subtract
106        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        // SAFETY:
119        // see subtract
120        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        /// Checked integer division. Computes self / rhs, returning None if rhs == 0 or the division results in overflow.
158        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        /// Checked integer division. Computes self / rhs, returning None if rhs == 0 or the division results in overflow.
171        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            // SAFETY:
191            // There will be UB if a ChunkedArray is alive with the wrong datatype.
192            // we now only create the potentially wrong dtype for a short time.
193            // Note that the physical type correctness is checked!
194            // The ChunkedArray with the wrong dtype is dropped after this operation
195            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            // SAFETY:
210            // see check_div for chunkedarray<T>
211            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            // SAFETY:
231            // see check_div for chunkedarray<T>
232            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            // SAFETY:
252            // see check_div
253            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    // Correct the list and array types
366    //
367    // This also casts Lists <-> Array.
368    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    // A decimal supertype can be narrower than an input, which must raise rather than
378    // produce nulls.
379    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/// Returns both operands as decimals if one is a decimal and the other a decimal or
392/// integer. Decimals keep their own scale, since the decimal kernels handle mixed scales
393/// without a lossy common cast, and integers are cast to `Decimal(38, 0)`.
394#[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
416/// Like [`coerce_lhs_rhs`], but decimal operands keep their own scales, and an integer
417/// combined with a decimal is cast to `Decimal(38, 0)`. Also disallows mixed
418/// string <=> numeric upcasts.
419pub fn coerce_lhs_rhs_numeric_op<'a>(
420    lhs: &'a Series,
421    rhs: &'a Series,
422) -> PolarsResult<(Cow<'a, Series>, Cow<'a, Series>)> {
423    // String is a supertype of the numeric types, but they must not be mixed in arithmetic.
424    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
438// Handle (Date | Datetime) +/- (Duration) | (Duration) +/- (Date | Datetime) | (Duration) +-
439// (Duration)
440// Time arithmetic is only implemented on the date / datetime so ensure that's on left
441
442fn 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        // make sure to return Some here, so we don't cast to supertype.
462        (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        // swap the order
478        (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            // Expects lengths to be equal.
517            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        // broadcasting
533        (1, _) | (_, 1) => Ok(()),
534        // equal
535        (a, b) if a == b => Ok(()),
536        // unequal
537        (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    /// ```
603    /// # use polars_core::prelude::*;
604    /// let s: Series = [1, 2, 3].iter().collect();
605    /// let out = (&s * &s).unwrap();
606    /// ```
607    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            // temporal lh
615            (Duration(_), _) | (Date, _) | (Datetime(_, _), _) | (Time, _) => self.multiply(rhs),
616            // temporal rhs
617            (_, Date) | (_, Datetime(_, _)) | (_, Time) => {
618                polars_bail!(opq = mul, self.dtype(), rhs.dtype())
619            },
620            (_, Duration(_)) => {
621                // swap order
622                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    /// ```
644    /// # use polars_core::prelude::*;
645    /// let s: Series = [1, 2, 3].iter().collect();
646    /// let out = (&s / &s).unwrap();
647    /// ```
648    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    /// ```
681    /// # use polars_core::prelude::*;
682    /// let s: Series = [1, 2, 3].iter().collect();
683    /// let out = (&s / &s).unwrap();
684    /// ```
685    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
707// Series +-/* numbers instead of Series
708
709fn 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
806// TODO: remove this, temporary band-aid.
807impl 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
876/// We cannot override the left hand side behaviour. So we create a trait LhsNumOps.
877/// This allows for 1.add(&Series)
878///
879impl<T: PolarsNumericType> ChunkedArray<T> {
880    /// Apply lhs - self
881    #[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    /// Apply lhs / self
888    #[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    /// Apply lhs % self
895    #[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        // order doesn't matter, dispatch to rhs + lhs
920        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        // order doesn't matter, dispatch to rhs * lhs
942        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        // Series +-/* Series
964        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        // Series +-/* Number
982        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        // Lhs operations
1000        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}