polars_core/series/arithmetic/
borrowed.rs

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

polars_core/series/arithmetic/
borrowed.rs
