polars_core/chunked_array/ops/
decimal.rs

1use polars_compute::decimal::dec128_verify_prec_scale;
2
3use crate::chunked_array::cast::CastOptions;
4use crate::prelude::*;
5
6impl StringChunked {
7    /// Convert an [`StringChunked`] to a [`Series`] of [`DataType::Decimal`].
8    /// Scale needed for the decimal type are inferred.  Parsing is not strict.  
9    /// Scale inference assumes that all tested strings are well-formed numbers,
10    /// and may produce unexpected results for scale if this is not the case.
11    ///
12    /// If the decimal `precision` and `scale` are already known, consider
13    /// using the `cast` method.
14    pub fn to_decimal_infer(&self, infer_length: usize) -> PolarsResult<Series> {
15        let mut scale = 0;
16        let mut prec_past_scale = 0;
17        let mut iter = self.into_iter();
18        let mut valid_count = 0;
19        while let Some(Some(v)) = iter.next() {
20            let mut bytes = v.as_bytes();
21            if bytes.first() == Some(&b'-') || bytes.first() == Some(&b'+') {
22                bytes = &bytes[1..];
23            }
24            while bytes.first() == Some(&b'0') {
25                bytes = &bytes[1..];
26            }
27            if let Some(separator) = bytes.iter().position(|b| *b == b'.') {
28                scale = scale.max(bytes.len() - 1 - separator);
29                prec_past_scale = prec_past_scale.max(separator);
30            } else {
31                prec_past_scale = prec_past_scale.max(bytes.len());
32            }
33
34            valid_count += 1;
35            if valid_count == infer_length {
36                break;
37            }
38        }
39
40        self.to_decimal(prec_past_scale + scale, scale)
41    }
42
43    pub fn to_decimal(&self, prec: usize, scale: usize) -> PolarsResult<Series> {
44        dec128_verify_prec_scale(prec, scale)?;
45        self.cast_with_options(&DataType::Decimal(prec, scale), CastOptions::NonStrict)
46    }
47}
48
49#[cfg(test)]
50mod test {
51    #[test]
52    fn test_inferred_length() {
53        use super::*;
54        let vals = [
55            "1.0",
56            "bad",
57            "225.0",
58            "3.00045",
59            "-4.0",
60            "5.104",
61            "5.25251525353",
62        ];
63        let s = StringChunked::from_slice(PlSmallStr::from_str("test"), &vals);
64        let s = s.to_decimal_infer(6).unwrap();
65        assert_eq!(s.dtype(), &DataType::Decimal(8, 5));
66        assert_eq!(s.len(), 7);
67        assert_eq!(s.get(0).unwrap(), AnyValue::Decimal(100000, 6, 5));
68        assert_eq!(s.get(1).unwrap(), AnyValue::Null);
69        assert_eq!(s.get(3).unwrap(), AnyValue::Decimal(300045, 6, 5));
70        assert_eq!(s.get(4).unwrap(), AnyValue::Decimal(-400000, 6, 5));
71        assert_eq!(s.get(6).unwrap(), AnyValue::Decimal(525252, 6, 5));
72    }
73}