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)`.
418pub 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
429// Handle (Date | Datetime) +/- (Duration) | (Duration) +/- (Date | Datetime) | (Duration) +-
430// (Duration)
431// Time arithmetic is only implemented on the date / datetime so ensure that's on left
432
433fn 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        // make sure to return Some here, so we don't cast to supertype.
453        (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        // swap the order
469        (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            // Expects lengths to be equal.
508            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        // broadcasting
524        (1, _) | (_, 1) => Ok(()),
525        // equal
526        (a, b) if a == b => Ok(()),
527        // unequal
528        (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    /// ```
594    /// # use polars_core::prelude::*;
595    /// let s: Series = [1, 2, 3].iter().collect();
596    /// let out = (&s * &s).unwrap();
597    /// ```
598    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            // temporal lh
606            (Duration(_), _) | (Date, _) | (Datetime(_, _), _) | (Time, _) => self.multiply(rhs),
607            // temporal rhs
608            (_, Date) | (_, Datetime(_, _)) | (_, Time) => {
609                polars_bail!(opq = mul, self.dtype(), rhs.dtype())
610            },
611            (_, Duration(_)) => {
612                // swap order
613                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    /// ```
635    /// # use polars_core::prelude::*;
636    /// let s: Series = [1, 2, 3].iter().collect();
637    /// let out = (&s / &s).unwrap();
638    /// ```
639    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    /// ```
672    /// # use polars_core::prelude::*;
673    /// let s: Series = [1, 2, 3].iter().collect();
674    /// let out = (&s / &s).unwrap();
675    /// ```
676    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
698// Series +-/* numbers instead of Series
699
700fn 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
797// TODO: remove this, temporary band-aid.
798impl 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
867/// We cannot override the left hand side behaviour. So we create a trait LhsNumOps.
868/// This allows for 1.add(&Series)
869///
870impl<T: PolarsNumericType> ChunkedArray<T> {
871    /// Apply lhs - self
872    #[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    /// Apply lhs / self
879    #[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    /// Apply lhs % self
886    #[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        // order doesn't matter, dispatch to rhs + lhs
911        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        // order doesn't matter, dispatch to rhs * lhs
933        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        // Series +-/* Series
955        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        // Series +-/* Number
973        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        // Lhs operations
991        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}