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() {
314                                None
315                            } else {
316                                Some(v / res)
317                            }
318                        })
319                    })
320                    .into_series(),
321                Float64 => s
322                    .f64()
323                    .unwrap()
324                    .apply(|opt_v| {
325                        opt_v.and_then(|v| {
326                            let res = rhs.to_f64().unwrap();
327                            if res.is_zero() {
328                                None
329                            } else {
330                                Some(v / res)
331                            }
332                        })
333                    })
334                    .into_series(),
335                _ => panic!("dtype not yet supported in checked div"),
336            };
337            out.cast(self.dtype())
338        }
339    }
340}
341
342pub fn coerce_lhs_rhs<'a>(
343    lhs: &'a Series,
344    rhs: &'a Series,
345) -> PolarsResult<(Cow<'a, Series>, Cow<'a, Series>)> {
346    if let Some(result) = coerce_time_units(lhs, rhs) {
347        return Ok(result);
348    }
349    let (left_dtype, right_dtype) = (lhs.dtype(), rhs.dtype());
350    let leaf_super_dtype = try_get_supertype(left_dtype.leaf_dtype(), right_dtype.leaf_dtype())?;
351
352    let mut new_left_dtype = left_dtype.cast_leaf(leaf_super_dtype.clone());
353    let mut new_right_dtype = right_dtype.cast_leaf(leaf_super_dtype);
354
355    // Correct the list and array types
356    //
357    // This also casts Lists <-> Array.
358    if left_dtype.is_list()
359        || right_dtype.is_list()
360        || left_dtype.is_array()
361        || right_dtype.is_array()
362    {
363        new_left_dtype = try_get_supertype(&new_left_dtype, &new_right_dtype)?;
364        new_right_dtype = new_left_dtype.clone();
365    }
366
367    let left = if lhs.dtype() == &new_left_dtype {
368        Cow::Borrowed(lhs)
369    } else {
370        Cow::Owned(lhs.cast(&new_left_dtype)?)
371    };
372    let right = if rhs.dtype() == &new_right_dtype {
373        Cow::Borrowed(rhs)
374    } else {
375        Cow::Owned(rhs.cast(&new_right_dtype)?)
376    };
377    Ok((left, right))
378}
379
380// Handle (Date | Datetime) +/- (Duration) | (Duration) +/- (Date | Datetime) | (Duration) +-
381// (Duration)
382// Time arithmetic is only implemented on the date / datetime so ensure that's on left
383
384fn coerce_time_units<'a>(
385    lhs: &'a Series,
386    rhs: &'a Series,
387) -> Option<(Cow<'a, Series>, Cow<'a, Series>)> {
388    match (lhs.dtype(), rhs.dtype()) {
389        (DataType::Datetime(lu, t), DataType::Duration(ru)) => {
390            let units = get_time_units(lu, ru);
391            let left = if *lu == units {
392                Cow::Borrowed(lhs)
393            } else {
394                Cow::Owned(lhs.cast(&DataType::Datetime(units, t.clone())).ok()?)
395            };
396            let right = if *ru == units {
397                Cow::Borrowed(rhs)
398            } else {
399                Cow::Owned(rhs.cast(&DataType::Duration(units)).ok()?)
400            };
401            Some((left, right))
402        },
403        // make sure to return Some here, so we don't cast to supertype.
404        (DataType::Date, DataType::Duration(_)) => Some((Cow::Borrowed(lhs), Cow::Borrowed(rhs))),
405        (DataType::Duration(lu), DataType::Duration(ru)) => {
406            let units = get_time_units(lu, ru);
407            let left = if *lu == units {
408                Cow::Borrowed(lhs)
409            } else {
410                Cow::Owned(lhs.cast(&DataType::Duration(units)).ok()?)
411            };
412            let right = if *ru == units {
413                Cow::Borrowed(rhs)
414            } else {
415                Cow::Owned(rhs.cast(&DataType::Duration(units)).ok()?)
416            };
417            Some((left, right))
418        },
419        // swap the order
420        (DataType::Duration(_), DataType::Datetime(_, _))
421        | (DataType::Duration(_), DataType::Date) => {
422            let (right, left) = coerce_time_units(rhs, lhs)?;
423            Some((left, right))
424        },
425        _ => None,
426    }
427}
428
429#[cfg(feature = "dtype-struct")]
430pub fn _struct_arithmetic<F: FnMut(&Series, &Series) -> PolarsResult<Series>>(
431    s: &Series,
432    rhs: &Series,
433    mut func: F,
434) -> PolarsResult<Series> {
435    let s = s.struct_().unwrap();
436    let rhs = rhs.struct_().unwrap();
437
438    let s_fields = s.fields_as_series();
439    let rhs_fields = rhs.fields_as_series();
440
441    match (s_fields.len(), rhs_fields.len()) {
442        (_, 1) => {
443            let rhs = &rhs.fields_as_series()[0];
444            Ok(s.try_apply_fields(|s| func(s, rhs))?.into_series())
445        },
446        (1, _) => {
447            let s = &s.fields_as_series()[0];
448            Ok(rhs.try_apply_fields(|rhs| func(s, rhs))?.into_series())
449        },
450        _ => {
451            let (s, rhs) = align_chunks_binary(s, rhs);
452            let mut s = s.into_owned();
453            s.zip_outer_validity(rhs.as_ref());
454
455            let mut rhs_iter = rhs.fields_as_series().into_iter();
456
457            Ok(s.try_apply_fields(|s| match rhs_iter.next() {
458                Some(rhs) => func(s, &rhs),
459                None => Ok(s.clone()),
460            })?
461            .into_series())
462        },
463    }
464}
465
466fn check_lengths(a: &Series, b: &Series) -> PolarsResult<()> {
467    match (a.len(), b.len()) {
468        // broadcasting
469        (1, _) | (_, 1) => Ok(()),
470        // equal
471        (a, b) if a == b => Ok(()),
472        // unequal
473        (a, b) => {
474            polars_bail!(InvalidOperation: "cannot do arithmetic operation on series of different lengths: got {} and {}", a, b)
475        },
476    }
477}
478
479impl Add for &Series {
480    type Output = PolarsResult<Series>;
481
482    fn add(self, rhs: Self) -> Self::Output {
483        check_lengths(self, rhs)?;
484        match (self.dtype(), rhs.dtype()) {
485            #[cfg(feature = "dtype-struct")]
486            (DataType::Struct(_), DataType::Struct(_)) => {
487                _struct_arithmetic(self, rhs, |a, b| a.add(b))
488            },
489            (DataType::List(_), _) | (_, DataType::List(_)) => {
490                list::NumericListOp::add().execute(self, rhs)
491            },
492            #[cfg(feature = "dtype-array")]
493            (DataType::Array(..), _) | (_, DataType::Array(..)) => {
494                fixed_size_list::NumericFixedSizeListOp::add().execute(self, rhs)
495            },
496            _ => {
497                let (lhs, rhs) = coerce_lhs_rhs(self, rhs)?;
498                lhs.add_to(rhs.as_ref())
499            },
500        }
501    }
502}
503
504impl Sub for &Series {
505    type Output = PolarsResult<Series>;
506
507    fn sub(self, rhs: Self) -> Self::Output {
508        check_lengths(self, rhs)?;
509        match (self.dtype(), rhs.dtype()) {
510            #[cfg(feature = "dtype-struct")]
511            (DataType::Struct(_), DataType::Struct(_)) => {
512                _struct_arithmetic(self, rhs, |a, b| a.sub(b))
513            },
514            (DataType::List(_), _) | (_, DataType::List(_)) => {
515                list::NumericListOp::sub().execute(self, rhs)
516            },
517            #[cfg(feature = "dtype-array")]
518            (DataType::Array(..), _) | (_, DataType::Array(..)) => {
519                fixed_size_list::NumericFixedSizeListOp::sub().execute(self, rhs)
520            },
521            _ => {
522                let (lhs, rhs) = coerce_lhs_rhs(self, rhs)?;
523                lhs.subtract(rhs.as_ref())
524            },
525        }
526    }
527}
528
529impl Mul for &Series {
530    type Output = PolarsResult<Series>;
531
532    /// ```
533    /// # use polars_core::prelude::*;
534    /// let s: Series = [1, 2, 3].iter().collect();
535    /// let out = (&s * &s).unwrap();
536    /// ```
537    fn mul(self, rhs: Self) -> Self::Output {
538        check_lengths(self, rhs)?;
539
540        use DataType::*;
541        match (self.dtype(), rhs.dtype()) {
542            #[cfg(feature = "dtype-struct")]
543            (Struct(_), Struct(_)) => _struct_arithmetic(self, rhs, |a, b| a.mul(b)),
544            // temporal lh
545            (Duration(_), _) | (Date, _) | (Datetime(_, _), _) | (Time, _) => self.multiply(rhs),
546            // temporal rhs
547            (_, Date) | (_, Datetime(_, _)) | (_, Time) => {
548                polars_bail!(opq = mul, self.dtype(), rhs.dtype())
549            },
550            (_, Duration(_)) => {
551                // swap order
552                let out = rhs.multiply(self)?;
553                Ok(out.with_name(self.name().clone()))
554            },
555            (DataType::List(_), _) | (_, DataType::List(_)) => {
556                list::NumericListOp::mul().execute(self, rhs)
557            },
558            #[cfg(feature = "dtype-array")]
559            (DataType::Array(..), _) | (_, DataType::Array(..)) => {
560                fixed_size_list::NumericFixedSizeListOp::mul().execute(self, rhs)
561            },
562            _ => {
563                let (lhs, rhs) = coerce_lhs_rhs(self, rhs)?;
564                lhs.multiply(rhs.as_ref())
565            },
566        }
567    }
568}
569
570impl Div for &Series {
571    type Output = PolarsResult<Series>;
572
573    /// ```
574    /// # use polars_core::prelude::*;
575    /// let s: Series = [1, 2, 3].iter().collect();
576    /// let out = (&s / &s).unwrap();
577    /// ```
578    fn div(self, rhs: Self) -> Self::Output {
579        check_lengths(self, rhs)?;
580        use DataType::*;
581        match (self.dtype(), rhs.dtype()) {
582            #[cfg(feature = "dtype-struct")]
583            (Struct(_), Struct(_)) => _struct_arithmetic(self, rhs, |a, b| a.div(b)),
584            (Duration(_), _) => self.divide(rhs),
585            (Date, _)
586            | (Datetime(_, _), _)
587            | (Time, _)
588            | (_, Duration(_))
589            | (_, Time)
590            | (_, Date)
591            | (_, Datetime(_, _)) => polars_bail!(opq = div, self.dtype(), rhs.dtype()),
592            (DataType::List(_), _) | (_, DataType::List(_)) => {
593                list::NumericListOp::div().execute(self, rhs)
594            },
595            #[cfg(feature = "dtype-array")]
596            (DataType::Array(..), _) | (_, DataType::Array(..)) => {
597                fixed_size_list::NumericFixedSizeListOp::div().execute(self, rhs)
598            },
599            _ => {
600                let (lhs, rhs) = coerce_lhs_rhs(self, rhs)?;
601                lhs.divide(rhs.as_ref())
602            },
603        }
604    }
605}
606
607impl Rem for &Series {
608    type Output = PolarsResult<Series>;
609
610    /// ```
611    /// # use polars_core::prelude::*;
612    /// let s: Series = [1, 2, 3].iter().collect();
613    /// let out = (&s / &s).unwrap();
614    /// ```
615    fn rem(self, rhs: Self) -> Self::Output {
616        check_lengths(self, rhs)?;
617        match (self.dtype(), rhs.dtype()) {
618            #[cfg(feature = "dtype-struct")]
619            (DataType::Struct(_), DataType::Struct(_)) => {
620                _struct_arithmetic(self, rhs, |a, b| a.rem(b))
621            },
622            (DataType::List(_), _) | (_, DataType::List(_)) => {
623                list::NumericListOp::rem().execute(self, rhs)
624            },
625            #[cfg(feature = "dtype-array")]
626            (DataType::Array(..), _) | (_, DataType::Array(..)) => {
627                fixed_size_list::NumericFixedSizeListOp::rem().execute(self, rhs)
628            },
629            _ => {
630                let (lhs, rhs) = coerce_lhs_rhs(self, rhs)?;
631                lhs.remainder(rhs.as_ref())
632            },
633        }
634    }
635}
636
637// Series +-/* numbers instead of Series
638
639fn finish_cast(inp: &Series, out: Series) -> Series {
640    match inp.dtype() {
641        #[cfg(feature = "dtype-date")]
642        DataType::Date => out.into_date(),
643        #[cfg(feature = "dtype-datetime")]
644        DataType::Datetime(tu, tz) => out.into_datetime(*tu, tz.clone()),
645        #[cfg(feature = "dtype-duration")]
646        DataType::Duration(tu) => out.into_duration(*tu),
647        #[cfg(feature = "dtype-time")]
648        DataType::Time => out.into_time(),
649        _ => out,
650    }
651}
652
653impl<T> Sub<T> for &Series
654where
655    T: Num + NumCast,
656{
657    type Output = Series;
658
659    fn sub(self, rhs: T) -> Self::Output {
660        let s = self.to_physical_repr();
661        macro_rules! sub {
662            ($ca:expr) => {{
663                $ca.sub(rhs).into_series()
664            }};
665        }
666
667        let out = downcast_as_macro_arg_physical!(s, sub);
668        finish_cast(self, out)
669    }
670}
671
672impl<T> Sub<T> for Series
673where
674    T: Num + NumCast,
675{
676    type Output = Self;
677
678    fn sub(self, rhs: T) -> Self::Output {
679        (&self).sub(rhs)
680    }
681}
682
683impl<T> Add<T> for &Series
684where
685    T: Num + NumCast,
686{
687    type Output = Series;
688
689    fn add(self, rhs: T) -> Self::Output {
690        let s = self.to_physical_repr();
691        macro_rules! add {
692            ($ca:expr) => {{
693                $ca.add(rhs).into_series()
694            }};
695        }
696        let out = downcast_as_macro_arg_physical!(s, add);
697        finish_cast(self, out)
698    }
699}
700
701impl<T> Add<T> for Series
702where
703    T: Num + NumCast,
704{
705    type Output = Self;
706
707    fn add(self, rhs: T) -> Self::Output {
708        (&self).add(rhs)
709    }
710}
711
712impl<T> Div<T> for &Series
713where
714    T: Num + NumCast,
715{
716    type Output = Series;
717
718    fn div(self, rhs: T) -> Self::Output {
719        let s = self.to_physical_repr();
720        macro_rules! div {
721            ($ca:expr) => {{
722                $ca.div(rhs).into_series()
723            }};
724        }
725
726        let out = downcast_as_macro_arg_physical!(s, div);
727        finish_cast(self, out)
728    }
729}
730
731impl<T> Div<T> for Series
732where
733    T: Num + NumCast,
734{
735    type Output = Self;
736
737    fn div(self, rhs: T) -> Self::Output {
738        (&self).div(rhs)
739    }
740}
741
742// TODO: remove this, temporary band-aid.
743impl Series {
744    pub fn wrapping_trunc_div_scalar<T: Num + NumCast>(&self, rhs: T) -> Self {
745        let s = self.to_physical_repr();
746        macro_rules! div {
747            ($ca:expr) => {{
748                let rhs = NumCast::from(rhs).unwrap();
749                $ca.wrapping_trunc_div_scalar(rhs).into_series()
750            }};
751        }
752
753        let out = downcast_as_macro_arg_physical!(s, div);
754        finish_cast(self, out)
755    }
756}
757
758impl<T> Mul<T> for &Series
759where
760    T: Num + NumCast,
761{
762    type Output = Series;
763
764    fn mul(self, rhs: T) -> Self::Output {
765        let s = self.to_physical_repr();
766        macro_rules! mul {
767            ($ca:expr) => {{
768                $ca.mul(rhs).into_series()
769            }};
770        }
771        let out = downcast_as_macro_arg_physical!(s, mul);
772        finish_cast(self, out)
773    }
774}
775
776impl<T> Mul<T> for Series
777where
778    T: Num + NumCast,
779{
780    type Output = Self;
781
782    fn mul(self, rhs: T) -> Self::Output {
783        (&self).mul(rhs)
784    }
785}
786
787impl<T> Rem<T> for &Series
788where
789    T: Num + NumCast,
790{
791    type Output = Series;
792
793    fn rem(self, rhs: T) -> Self::Output {
794        let s = self.to_physical_repr();
795        macro_rules! rem {
796            ($ca:expr) => {{
797                $ca.rem(rhs).into_series()
798            }};
799        }
800        let out = downcast_as_macro_arg_physical!(s, rem);
801        finish_cast(self, out)
802    }
803}
804
805impl<T> Rem<T> for Series
806where
807    T: Num + NumCast,
808{
809    type Output = Self;
810
811    fn rem(self, rhs: T) -> Self::Output {
812        (&self).rem(rhs)
813    }
814}
815
816/// We cannot override the left hand side behaviour. So we create a trait LhsNumOps.
817/// This allows for 1.add(&Series)
818///
819impl<T> ChunkedArray<T>
820where
821    T: PolarsNumericType,
822    ChunkedArray<T>: IntoSeries,
823{
824    /// Apply lhs - self
825    #[must_use]
826    pub fn lhs_sub<N: Num + NumCast>(&self, lhs: N) -> Self {
827        let lhs: T::Native = NumCast::from(lhs).expect("could not cast");
828        ArithmeticChunked::wrapping_sub_scalar_lhs(lhs, self)
829    }
830
831    /// Apply lhs / self
832    #[must_use]
833    pub fn lhs_div<N: Num + NumCast>(&self, lhs: N) -> Self {
834        let lhs: T::Native = NumCast::from(lhs).expect("could not cast");
835        ArithmeticChunked::legacy_div_scalar_lhs(lhs, self)
836    }
837
838    /// Apply lhs % self
839    #[must_use]
840    pub fn lhs_rem<N: Num + NumCast>(&self, lhs: N) -> Self {
841        let lhs: T::Native = NumCast::from(lhs).expect("could not cast");
842        ArithmeticChunked::wrapping_mod_scalar_lhs(lhs, self)
843    }
844}
845
846pub trait LhsNumOps {
847    type Output;
848
849    fn add(self, rhs: &Series) -> Self::Output;
850    fn sub(self, rhs: &Series) -> Self::Output;
851    fn div(self, rhs: &Series) -> Self::Output;
852    fn mul(self, rhs: &Series) -> Self::Output;
853    fn rem(self, rem: &Series) -> Self::Output;
854}
855
856impl<T> LhsNumOps for T
857where
858    T: Num + NumCast,
859{
860    type Output = Series;
861
862    fn add(self, rhs: &Series) -> Self::Output {
863        // order doesn't matter, dispatch to rhs + lhs
864        rhs + self
865    }
866    fn sub(self, rhs: &Series) -> Self::Output {
867        let s = rhs.to_physical_repr();
868        macro_rules! sub {
869            ($rhs:expr) => {{
870                $rhs.lhs_sub(self).into_series()
871            }};
872        }
873        let out = downcast_as_macro_arg_physical!(s, sub);
874
875        finish_cast(rhs, out)
876    }
877    fn div(self, rhs: &Series) -> Self::Output {
878        let s = rhs.to_physical_repr();
879        macro_rules! div {
880            ($rhs:expr) => {{
881                $rhs.lhs_div(self).into_series()
882            }};
883        }
884        let out = downcast_as_macro_arg_physical!(s, div);
885
886        finish_cast(rhs, out)
887    }
888    fn mul(self, rhs: &Series) -> Self::Output {
889        // order doesn't matter, dispatch to rhs * lhs
890        rhs * self
891    }
892    fn rem(self, rhs: &Series) -> Self::Output {
893        let s = rhs.to_physical_repr();
894        macro_rules! rem {
895            ($rhs:expr) => {{
896                $rhs.lhs_rem(self).into_series()
897            }};
898        }
899
900        let out = downcast_as_macro_arg_physical!(s, rem);
901
902        finish_cast(rhs, out)
903    }
904}
905
906#[cfg(test)]
907mod test {
908    use crate::prelude::*;
909
910    #[test]
911    #[allow(clippy::eq_op)]
912    fn test_arithmetic_series() -> PolarsResult<()> {
913        // Series +-/* Series
914        let s = Series::new("foo".into(), [1, 2, 3]);
915        assert_eq!(
916            Vec::from((&s * &s)?.i32().unwrap()),
917            [Some(1), Some(4), Some(9)]
918        );
919        assert_eq!(
920            Vec::from((&s / &s)?.i32().unwrap()),
921            [Some(1), Some(1), Some(1)]
922        );
923        assert_eq!(
924            Vec::from((&s - &s)?.i32().unwrap()),
925            [Some(0), Some(0), Some(0)]
926        );
927        assert_eq!(
928            Vec::from((&s + &s)?.i32().unwrap()),
929            [Some(2), Some(4), Some(6)]
930        );
931        // Series +-/* Number
932        assert_eq!(
933            Vec::from((&s + 1).i32().unwrap()),
934            [Some(2), Some(3), Some(4)]
935        );
936        assert_eq!(
937            Vec::from((&s - 1).i32().unwrap()),
938            [Some(0), Some(1), Some(2)]
939        );
940        assert_eq!(
941            Vec::from((&s * 2).i32().unwrap()),
942            [Some(2), Some(4), Some(6)]
943        );
944        assert_eq!(
945            Vec::from((&s / 2).i32().unwrap()),
946            [Some(0), Some(1), Some(1)]
947        );
948
949        // Lhs operations
950        assert_eq!(
951            Vec::from((1.add(&s)).i32().unwrap()),
952            [Some(2), Some(3), Some(4)]
953        );
954        assert_eq!(
955            Vec::from((1.sub(&s)).i32().unwrap()),
956            [Some(0), Some(-1), Some(-2)]
957        );
958        assert_eq!(
959            Vec::from((1.div(&s)).i32().unwrap()),
960            [Some(1), Some(0), Some(0)]
961        );
962        assert_eq!(
963            Vec::from((1.mul(&s)).i32().unwrap()),
964            [Some(1), Some(2), Some(3)]
965        );
966        assert_eq!(
967            Vec::from((1.rem(&s)).i32().unwrap()),
968            [Some(0), Some(1), Some(1)]
969        );
970
971        assert_eq!((&s * &s)?.name().as_str(), "foo");
972        assert_eq!((&s * 1).name().as_str(), "foo");
973        assert_eq!((1.div(&s)).name().as_str(), "foo");
974
975        Ok(())
976    }
977
978    #[test]
979    #[cfg(feature = "checked_arithmetic")]
980    fn test_checked_div() {
981        let s = Series::new("foo".into(), [1i32, 0, 1]);
982        let out = s.checked_div(&s).unwrap();
983        assert_eq!(Vec::from(out.i32().unwrap()), &[Some(1), None, Some(1)]);
984        let out = s.checked_div_num(0).unwrap();
985        assert_eq!(Vec::from(out.i32().unwrap()), &[None, None, None]);
986
987        let s_f32 = Series::new("float32".into(), [1.0f32, 0.0, 1.0]);
988        let out = s_f32.checked_div(&s_f32).unwrap();
989        assert_eq!(
990            Vec::from(out.f32().unwrap()),
991            &[Some(1.0f32), None, Some(1.0f32)]
992        );
993        let out = s_f32.checked_div_num(0.0f32).unwrap();
994        assert_eq!(Vec::from(out.f32().unwrap()), &[None, None, None]);
995
996        let s_f64 = Series::new("float64".into(), [1.0f64, 0.0, 1.0]);
997        let out = s_f64.checked_div(&s_f64).unwrap();
998        assert_eq!(
999            Vec::from(out.f64().unwrap()),
1000            &[Some(1.0f64), None, Some(1.0f64)]
1001        );
1002        let out = s_f64.checked_div_num(0.0f64).unwrap();
1003        assert_eq!(Vec::from(out.f64().unwrap()), &[None, None, None]);
1004    }
1005}
polars_core/series/arithmetic/borrowed.rs

polars_core/series/arithmetic/
borrowed.rs
