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

polars_core/series/arithmetic/
borrowed.rs
