comparison.rs

use polars_compute::decimal::{DEC128_MAX_PREC, dec128_fits};
use pyo3::prelude::*;

use crate::PySeries;
use crate::error::PyPolarsErr;
use crate::prelude::*;
use crate::utils::EnterPolarsExt;

#[pymethods]
impl PySeries {
    fn eq(&self, py: Python<'_>, rhs: &PySeries) -> PyResult<Self> {
        py.enter_polars_series(|| self.series.read().equal(&*rhs.series.read()))
    }

    fn neq(&self, py: Python<'_>, rhs: &PySeries) -> PyResult<Self> {
        py.enter_polars_series(|| self.series.read().not_equal(&*rhs.series.read()))
    }

    fn gt(&self, py: Python<'_>, rhs: &PySeries) -> PyResult<Self> {
        py.enter_polars_series(|| self.series.read().gt(&*rhs.series.read()))
    }

    fn gt_eq(&self, py: Python<'_>, rhs: &PySeries) -> PyResult<Self> {
        py.enter_polars_series(|| self.series.read().gt_eq(&*rhs.series.read()))
    }

    fn lt(&self, py: Python<'_>, rhs: &PySeries) -> PyResult<Self> {
        py.enter_polars_series(|| self.series.read().lt(&*rhs.series.read()))
    }

    fn lt_eq(&self, py: Python<'_>, rhs: &PySeries) -> PyResult<Self> {
        py.enter_polars_series(|| self.series.read().lt_eq(&*rhs.series.read()))
    }
}

macro_rules! impl_op {
    ($op:ident, $name:ident, $type:ty) => {
        #[pymethods]
        impl PySeries {
            fn $name(&self, py: Python<'_>, rhs: $type) -> PyResult<Self> {
                py.enter_polars_series(|| self.series.read().$op(rhs))
            }
        }
    };
}

impl_op!(equal, eq_u8, u8);
impl_op!(equal, eq_u16, u16);
impl_op!(equal, eq_u32, u32);
impl_op!(equal, eq_u64, u64);
impl_op!(equal, eq_i8, i8);
impl_op!(equal, eq_i16, i16);
impl_op!(equal, eq_i32, i32);
impl_op!(equal, eq_i64, i64);
impl_op!(equal, eq_i128, i128);
impl_op!(equal, eq_f16, pf16);
impl_op!(equal, eq_f32, f32);
impl_op!(equal, eq_f64, f64);
impl_op!(equal, eq_str, &str);

impl_op!(not_equal, neq_u8, u8);
impl_op!(not_equal, neq_u16, u16);
impl_op!(not_equal, neq_u32, u32);
impl_op!(not_equal, neq_u64, u64);
impl_op!(not_equal, neq_i8, i8);
impl_op!(not_equal, neq_i16, i16);
impl_op!(not_equal, neq_i32, i32);
impl_op!(not_equal, neq_i64, i64);
impl_op!(not_equal, neq_i128, i128);
impl_op!(not_equal, neq_f16, pf16);
impl_op!(not_equal, neq_f32, f32);
impl_op!(not_equal, neq_f64, f64);
impl_op!(not_equal, neq_str, &str);

impl_op!(gt, gt_u8, u8);
impl_op!(gt, gt_u16, u16);
impl_op!(gt, gt_u32, u32);
impl_op!(gt, gt_u64, u64);
impl_op!(gt, gt_u128, u128);
impl_op!(gt, gt_i8, i8);
impl_op!(gt, gt_i16, i16);
impl_op!(gt, gt_i32, i32);
impl_op!(gt, gt_i64, i64);
impl_op!(gt, gt_i128, i128);
impl_op!(gt, gt_f16, pf16);
impl_op!(gt, gt_f32, f32);
impl_op!(gt, gt_f64, f64);
impl_op!(gt, gt_str, &str);

impl_op!(gt_eq, gt_eq_u8, u8);
impl_op!(gt_eq, gt_eq_u16, u16);
impl_op!(gt_eq, gt_eq_u32, u32);
impl_op!(gt_eq, gt_eq_u64, u64);
impl_op!(gt_eq, gt_eq_i8, i8);
impl_op!(gt_eq, gt_eq_i16, i16);
impl_op!(gt_eq, gt_eq_i32, i32);
impl_op!(gt_eq, gt_eq_i64, i64);
impl_op!(gt_eq, gt_eq_i128, i128);
impl_op!(gt_eq, gt_eq_f16, pf16);
impl_op!(gt_eq, gt_eq_f32, f32);
impl_op!(gt_eq, gt_eq_f64, f64);
impl_op!(gt_eq, gt_eq_str, &str);

impl_op!(lt, lt_u8, u8);
impl_op!(lt, lt_u16, u16);
impl_op!(lt, lt_u32, u32);
impl_op!(lt, lt_u64, u64);
impl_op!(lt, lt_u128, u128);
impl_op!(lt, lt_i8, i8);
impl_op!(lt, lt_i16, i16);
impl_op!(lt, lt_i32, i32);
impl_op!(lt, lt_i64, i64);
impl_op!(lt, lt_i128, i128);
impl_op!(lt, lt_f16, pf16);
impl_op!(lt, lt_f32, f32);
impl_op!(lt, lt_f64, f64);
impl_op!(lt, lt_str, &str);

impl_op!(lt_eq, lt_eq_u8, u8);
impl_op!(lt_eq, lt_eq_u16, u16);
impl_op!(lt_eq, lt_eq_u32, u32);
impl_op!(lt_eq, lt_eq_u64, u64);
impl_op!(lt_eq, lt_eq_i8, i8);
impl_op!(lt_eq, lt_eq_i16, i16);
impl_op!(lt_eq, lt_eq_i32, i32);
impl_op!(lt_eq, lt_eq_i64, i64);
impl_op!(lt_eq, lt_eq_i128, i128);
impl_op!(lt_eq, lt_eq_f16, pf16);
impl_op!(lt_eq, lt_eq_f32, f32);
impl_op!(lt_eq, lt_eq_f64, f64);
impl_op!(lt_eq, lt_eq_str, &str);

struct PyDecimal(i128, usize, usize);

impl<'a, 'py> FromPyObject<'a, 'py> for PyDecimal {
    type Error = PyErr;

    fn extract(obj: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
        if let Ok(val) = obj.extract() {
            return Ok(PyDecimal(val, DEC128_MAX_PREC, 0));
        }

        let err = || {
            Err(PyPolarsErr::from(polars_err!(ComputeError: "overflow in Python Decimal to Polars Decimal conversion")).into())
        };

        let (sign, digits, exponent) = obj
            .call_method0("as_tuple")?
            .extract::<(i8, Vec<u8>, i8)>()?;
        let mut val = 0_i128;
        for d in digits {
            if let Some(v) = val.checked_mul(10).and_then(|val| val.checked_add(d as _)) {
                val = v;
            } else {
                return err();
            }
        }
        let scale = if exponent > 0 {
            if let Some(v) = val.checked_mul(10_i128.pow((-exponent) as u32)) {
                val = v;
            } else {
                return err();
            };
            0_usize
        } else {
            (-exponent) as usize
        };
        if sign == 1 {
            val = -val
        };
        if dec128_fits(val, DEC128_MAX_PREC) {
            Ok(PyDecimal(val, DEC128_MAX_PREC, scale))
        } else {
            err()
        }
    }
}

macro_rules! impl_decimal {
    ($name:ident, $method:ident) => {
        #[pymethods]
        impl PySeries {
            fn $name(&self, py: Python<'_>, rhs: PyDecimal) -> PyResult<Self> {
                let rhs = Series::new(
                    PlSmallStr::from_static("decimal"),
                    &[AnyValue::Decimal(rhs.0, rhs.1, rhs.2)],
                );
                py.enter_polars_series(|| self.series.read().$method(&rhs))
            }
        }
    };
}

impl_decimal!(eq_decimal, equal);
impl_decimal!(neq_decimal, not_equal);
impl_decimal!(gt_decimal, gt);
impl_decimal!(gt_eq_decimal, gt_eq);
impl_decimal!(lt_decimal, lt);
impl_decimal!(lt_eq_decimal, lt_eq);