polars_core/chunked_array/logical/
decimal.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
use std::borrow::Cow;

use super::*;
use crate::chunked_array::cast::cast_chunks;
use crate::prelude::*;

pub type DecimalChunked = Logical<DecimalType, Int128Type>;

impl Int128Chunked {
    #[inline]
    pub fn into_decimal_unchecked(self, precision: Option<usize>, scale: usize) -> DecimalChunked {
        let mut dt = DecimalChunked::new_logical(self);
        dt.2 = Some(DataType::Decimal(precision, Some(scale)));
        dt
    }

    pub fn into_decimal(
        self,
        precision: Option<usize>,
        scale: usize,
    ) -> PolarsResult<DecimalChunked> {
        // TODO: if precision is None, do we check that the value fits within precision of 38?...
        if let Some(precision) = precision {
            let precision_max = 10_i128.pow(precision as u32);
            if let Some((min, max)) = self.min_max() {
                let max_abs = max.abs().max(min.abs());
                polars_ensure!(
                    max_abs < precision_max,
                    ComputeError: "decimal precision {} can't fit values with {} digits",
                    precision,
                    max_abs.to_string().len()
                );
            }
        }
        Ok(self.into_decimal_unchecked(precision, scale))
    }
}

impl LogicalType for DecimalChunked {
    fn dtype(&self) -> &DataType {
        self.2.as_ref().unwrap()
    }

    #[inline]
    fn get_any_value(&self, i: usize) -> PolarsResult<AnyValue<'_>> {
        polars_ensure!(i < self.len(), oob = i, self.len());
        Ok(unsafe { self.get_any_value_unchecked(i) })
    }

    #[inline]
    unsafe fn get_any_value_unchecked(&self, i: usize) -> AnyValue<'_> {
        match self.0.get_unchecked(i) {
            Some(v) => AnyValue::Decimal(v, self.scale()),
            None => AnyValue::Null,
        }
    }

    fn cast_with_options(
        &self,
        dtype: &DataType,
        cast_options: CastOptions,
    ) -> PolarsResult<Series> {
        let mut dtype = Cow::Borrowed(dtype);
        if let DataType::Decimal(to_precision, to_scale) = dtype.as_ref() {
            let from_precision = self.precision();
            let from_scale = self.scale();

            let to_precision = to_precision.or(from_precision);
            let to_scale = to_scale.unwrap_or(from_scale);

            if to_precision == from_precision && to_scale == from_scale {
                return Ok(self.clone().into_series());
            }

            dtype = Cow::Owned(DataType::Decimal(to_precision, Some(to_scale)));
        }

        let arrow_dtype = self.dtype().to_arrow(CompatLevel::newest());
        let chunks = self
            .chunks
            .iter()
            .map(|arr| {
                arr.as_any()
                    .downcast_ref::<PrimitiveArray<i128>>()
                    .unwrap()
                    .clone()
                    .to(arrow_dtype.clone())
                    .to_boxed()
            })
            .collect::<Vec<_>>();
        let chunks = cast_chunks(&chunks, dtype.as_ref(), cast_options)?;
        Series::try_from((self.name().clone(), chunks))
    }
}

impl DecimalChunked {
    pub fn precision(&self) -> Option<usize> {
        match self.2.as_ref().unwrap() {
            DataType::Decimal(precision, _) => *precision,
            _ => unreachable!(),
        }
    }

    pub fn scale(&self) -> usize {
        match self.2.as_ref().unwrap() {
            DataType::Decimal(_, scale) => scale.unwrap_or_else(|| unreachable!()),
            _ => unreachable!(),
        }
    }

    pub fn to_scale(&self, scale: usize) -> PolarsResult<Cow<'_, Self>> {
        if self.scale() == scale {
            return Ok(Cow::Borrowed(self));
        }

        let mut precision = self.precision();
        if let Some(ref mut precision) = precision {
            if self.scale() < scale {
                *precision += scale;
                *precision = (*precision).min(38);
            }
        }

        let s = self.cast_with_options(
            &DataType::Decimal(precision, Some(scale)),
            CastOptions::NonStrict,
        )?;
        Ok(Cow::Owned(s.decimal().unwrap().clone()))
    }
}