datatypes/vectors/
decimal.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::any::Any;
16use std::fmt::Debug;
17use std::sync::Arc;
18
19use arrow_array::builder::{ArrayBuilder, Decimal128Builder};
20use arrow_array::iterator::ArrayIter;
21use arrow_array::{Array, ArrayRef, Decimal128Array};
22use common_decimal::Decimal128;
23use common_decimal::decimal128::{DECIMAL128_DEFAULT_SCALE, DECIMAL128_MAX_PRECISION};
24use snafu::{OptionExt, ResultExt};
25
26use crate::arrow::datatypes::DataType as ArrowDataType;
27use crate::data_type::ConcreteDataType;
28use crate::error::{
29    self, CastTypeSnafu, InvalidPrecisionOrScaleSnafu, Result, ValueExceedsPrecisionSnafu,
30};
31use crate::prelude::{ScalarVector, ScalarVectorBuilder};
32use crate::serialize::Serializable;
33use crate::value::{Value, ValueRef};
34use crate::vectors;
35use crate::vectors::{MutableVector, Validity, Vector, VectorRef};
36
37/// Decimal128Vector is a vector keep i128 values with precision and scale.
38#[derive(Debug, PartialEq)]
39pub struct Decimal128Vector {
40    array: Decimal128Array,
41}
42
43impl Decimal128Vector {
44    /// New a Decimal128Vector from Arrow Decimal128Array
45    pub fn new(array: Decimal128Array) -> Self {
46        Self { array }
47    }
48
49    /// Construct Vector from i128 values
50    pub fn from_values<I: IntoIterator<Item = i128>>(iter: I) -> Self {
51        Self {
52            array: Decimal128Array::from_iter_values(iter),
53        }
54    }
55
56    /// Construct Vector from i128 values slice
57    pub fn from_slice<P: AsRef<[i128]>>(slice: P) -> Self {
58        let iter = slice.as_ref().iter().copied();
59        Self {
60            array: Decimal128Array::from_iter_values(iter),
61        }
62    }
63
64    /// Construct Vector from Wrapper(Decimal128) values slice
65    pub fn from_wrapper_slice<P: AsRef<[Decimal128]>>(slice: P) -> Self {
66        let iter = slice.as_ref().iter().copied().map(|v| v.val());
67        Self {
68            array: Decimal128Array::from_iter_values(iter),
69        }
70    }
71
72    /// Get decimal128 value from vector by offset and length.
73    pub fn get_slice(&self, offset: usize, length: usize) -> Self {
74        let array = self.array.slice(offset, length);
75        Self { array }
76    }
77
78    /// Returns a Decimal vector with the same data as self, with the
79    /// specified precision and scale(should in Decimal128 range),
80    /// and return error if value is out of precision bound.
81    ///
82    ///
83    /// For example:
84    /// value = 12345, precision = 3, return error.
85    pub fn with_precision_and_scale(self, precision: u8, scale: i8) -> Result<Self> {
86        let array = self
87            .array
88            .with_precision_and_scale(precision, scale)
89            .context(InvalidPrecisionOrScaleSnafu { precision, scale })?;
90        Ok(Self { array })
91    }
92
93    /// Returns a Decimal vector with the same data as self, with the
94    /// specified precision and scale(should in Decimal128 range),
95    /// and return null if value is out of precision bound.
96    ///
97    /// For example:
98    /// value = 12345, precision = 3, the value will be casted to null.
99    pub fn with_precision_and_scale_to_null(self, precision: u8, scale: i8) -> Result<Self> {
100        self.null_if_overflow_precision(precision)
101            .with_precision_and_scale(precision, scale)
102    }
103
104    /// Return decimal value as string
105    pub fn value_as_string(&self, idx: usize) -> String {
106        self.array.value_as_string(idx)
107    }
108
109    /// Return decimal128 vector precision
110    pub fn precision(&self) -> u8 {
111        self.array.precision()
112    }
113
114    /// Return decimal128 vector scale
115    pub fn scale(&self) -> i8 {
116        self.array.scale()
117    }
118
119    /// Return decimal128 vector inner array
120    pub(crate) fn as_arrow(&self) -> &dyn Array {
121        &self.array
122    }
123
124    /// Validate decimal precision, if precision is invalid, return error.
125    pub fn validate_decimal_precision(&self, precision: u8) -> Result<()> {
126        self.array
127            .validate_decimal_precision(precision)
128            .context(ValueExceedsPrecisionSnafu { precision })
129    }
130
131    /// Values that exceed the precision bounds will be casted to Null.
132    fn null_if_overflow_precision(&self, precision: u8) -> Self {
133        Self {
134            array: self.array.null_if_overflow_precision(precision),
135        }
136    }
137
138    /// Get decimal128 Value from array by index.
139    fn get_decimal128_value_from_array(&self, index: usize) -> Option<Decimal128> {
140        if self.array.is_valid(index) {
141            // Safety: The index have been checked by `is_valid()`.
142            let value = unsafe { self.array.value_unchecked(index) };
143            // Safety: The precision and scale have been checked by Vector.
144            Some(Decimal128::new(value, self.precision(), self.scale()))
145        } else {
146            None
147        }
148    }
149}
150
151impl Vector for Decimal128Vector {
152    fn data_type(&self) -> ConcreteDataType {
153        if let ArrowDataType::Decimal128(p, s) = self.array.data_type() {
154            ConcreteDataType::decimal128_datatype(*p, *s)
155        } else {
156            ConcreteDataType::decimal128_default_datatype()
157        }
158    }
159
160    fn vector_type_name(&self) -> String {
161        "Decimal128Vector".to_string()
162    }
163
164    fn as_any(&self) -> &dyn Any {
165        self
166    }
167
168    fn len(&self) -> usize {
169        self.array.len()
170    }
171
172    fn to_arrow_array(&self) -> ArrayRef {
173        Arc::new(self.array.clone())
174    }
175
176    fn to_boxed_arrow_array(&self) -> Box<dyn Array> {
177        Box::new(self.array.clone())
178    }
179
180    fn validity(&self) -> Validity {
181        vectors::impl_validity_for_vector!(self.array)
182    }
183
184    fn memory_size(&self) -> usize {
185        self.array.get_buffer_memory_size()
186    }
187
188    fn null_count(&self) -> usize {
189        self.array.null_count()
190    }
191
192    fn is_null(&self, row: usize) -> bool {
193        self.array.is_null(row)
194    }
195
196    fn slice(&self, offset: usize, length: usize) -> VectorRef {
197        Arc::new(self.get_slice(offset, length))
198    }
199
200    fn get(&self, index: usize) -> Value {
201        if let Some(decimal) = self.get_decimal128_value_from_array(index) {
202            Value::Decimal128(decimal)
203        } else {
204            Value::Null
205        }
206    }
207
208    fn get_ref(&self, index: usize) -> ValueRef<'_> {
209        if let Some(decimal) = self.get_decimal128_value_from_array(index) {
210            ValueRef::Decimal128(decimal)
211        } else {
212            ValueRef::Null
213        }
214    }
215}
216
217impl From<Decimal128Array> for Decimal128Vector {
218    fn from(array: Decimal128Array) -> Self {
219        Self { array }
220    }
221}
222
223impl From<Vec<Option<i128>>> for Decimal128Vector {
224    fn from(vec: Vec<Option<i128>>) -> Self {
225        let array = Decimal128Array::from_iter(vec);
226        Self { array }
227    }
228}
229
230impl Serializable for Decimal128Vector {
231    fn serialize_to_json(&self) -> Result<Vec<serde_json::Value>> {
232        self.iter_data()
233            .map(|v| match v {
234                None => Ok(serde_json::Value::Null), // if decimal vector not present, map to NULL
235                Some(d) => serde_json::to_value(d),
236            })
237            .collect::<serde_json::Result<_>>()
238            .context(error::SerializeSnafu)
239    }
240}
241
242pub struct Decimal128Iter<'a> {
243    precision: u8,
244    scale: i8,
245    iter: ArrayIter<&'a Decimal128Array>,
246}
247
248impl Iterator for Decimal128Iter<'_> {
249    type Item = Option<Decimal128>;
250
251    fn next(&mut self) -> Option<Self::Item> {
252        self.iter
253            .next()
254            .map(|item| item.map(|v| Decimal128::new(v, self.precision, self.scale)))
255    }
256
257    fn size_hint(&self) -> (usize, Option<usize>) {
258        self.iter.size_hint()
259    }
260}
261
262impl ScalarVector for Decimal128Vector {
263    type OwnedItem = Decimal128;
264
265    type RefItem<'a> = Decimal128;
266
267    type Iter<'a> = Decimal128Iter<'a>;
268
269    type Builder = Decimal128VectorBuilder;
270
271    fn get_data(&self, idx: usize) -> Option<Self::RefItem<'_>> {
272        self.get_decimal128_value_from_array(idx)
273    }
274
275    fn iter_data(&self) -> Self::Iter<'_> {
276        Self::Iter {
277            precision: self.precision(),
278            scale: self.scale(),
279            iter: self.array.iter(),
280        }
281    }
282}
283
284pub struct Decimal128VectorBuilder {
285    precision: u8,
286    scale: i8,
287    mutable_array: Decimal128Builder,
288}
289
290impl MutableVector for Decimal128VectorBuilder {
291    fn data_type(&self) -> ConcreteDataType {
292        ConcreteDataType::decimal128_datatype(self.precision, self.scale)
293    }
294
295    fn len(&self) -> usize {
296        self.mutable_array.len()
297    }
298
299    fn as_any(&self) -> &dyn Any {
300        self
301    }
302
303    fn as_mut_any(&mut self) -> &mut dyn Any {
304        self
305    }
306
307    fn to_vector(&mut self) -> VectorRef {
308        Arc::new(self.finish())
309    }
310
311    fn to_vector_cloned(&self) -> VectorRef {
312        Arc::new(self.finish_cloned())
313    }
314
315    fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()> {
316        let decimal_val = value.try_into_decimal128()?.map(|v| v.val());
317        self.mutable_array.append_option(decimal_val);
318        Ok(())
319    }
320
321    fn push_null(&mut self) {
322        self.mutable_array.append_null();
323    }
324
325    fn extend_slice_of(&mut self, vector: &dyn Vector, offset: usize, length: usize) -> Result<()> {
326        let decimal_vector =
327            vector
328                .as_any()
329                .downcast_ref::<Decimal128Vector>()
330                .context(CastTypeSnafu {
331                    msg: format!(
332                        "Failed to cast vector from {} to Decimal128Vector",
333                        vector.vector_type_name(),
334                    ),
335                })?;
336        let slice = decimal_vector.get_slice(offset, length);
337        self.mutable_array
338            .extend(slice.iter_data().map(|v| v.map(|d| d.val())));
339        Ok(())
340    }
341}
342
343impl ScalarVectorBuilder for Decimal128VectorBuilder {
344    type VectorType = Decimal128Vector;
345
346    fn with_capacity(capacity: usize) -> Self {
347        Self {
348            precision: DECIMAL128_MAX_PRECISION,
349            scale: DECIMAL128_DEFAULT_SCALE,
350            mutable_array: Decimal128Builder::with_capacity(capacity),
351        }
352    }
353
354    fn push(&mut self, value: Option<<Self::VectorType as ScalarVector>::RefItem<'_>>) {
355        self.mutable_array.append_option(value.map(|v| v.val()));
356    }
357
358    fn finish(&mut self) -> Self::VectorType {
359        Decimal128Vector {
360            array: self.mutable_array.finish(),
361        }
362    }
363
364    fn finish_cloned(&self) -> Self::VectorType {
365        Decimal128Vector {
366            array: self.mutable_array.finish_cloned(),
367        }
368    }
369}
370
371impl Decimal128VectorBuilder {
372    /// Change the precision and scale of the Decimal128VectorBuilder.
373    pub fn with_precision_and_scale(self, precision: u8, scale: i8) -> Result<Self> {
374        let mutable_array = self
375            .mutable_array
376            .with_precision_and_scale(precision, scale)
377            .context(InvalidPrecisionOrScaleSnafu { precision, scale })?;
378        Ok(Self {
379            precision,
380            scale,
381            mutable_array,
382        })
383    }
384}
385
386vectors::impl_try_from_arrow_array_for_vector!(Decimal128Array, Decimal128Vector);
387
388#[cfg(test)]
389pub mod tests {
390    use arrow_array::Decimal128Array;
391    use common_decimal::Decimal128;
392
393    use super::*;
394    use crate::vectors::Int8Vector;
395    use crate::vectors::operations::VectorOp;
396
397    #[test]
398    fn test_from_arrow_decimal128_array() {
399        let decimal_array = Decimal128Array::from(vec![Some(123), Some(456)]);
400        let decimal_vector = Decimal128Vector::from(decimal_array);
401        let expect = Decimal128Vector::from_values(vec![123, 456]);
402        assert_eq!(decimal_vector, expect);
403
404        let decimal_array = Decimal128Array::from(vec![Some(123), Some(456)])
405            .with_precision_and_scale(10, 2)
406            .unwrap();
407        let decimal_vector = Decimal128Vector::from(decimal_array);
408        let expect = Decimal128Vector::from_values(vec![123, 456])
409            .with_precision_and_scale(10, 2)
410            .unwrap();
411        assert_eq!(decimal_vector, expect);
412
413        let decimal_array: ArrayRef = Arc::new(
414            Decimal128Array::from(vec![Some(123), Some(456)])
415                .with_precision_and_scale(3, 2)
416                .unwrap(),
417        );
418        let decimal_vector = Decimal128Vector::try_from_arrow_array(decimal_array).unwrap();
419        let expect = Decimal128Vector::from_values(vec![123, 456])
420            .with_precision_and_scale(3, 2)
421            .unwrap();
422        assert_eq!(decimal_vector, expect);
423    }
424
425    #[test]
426    fn test_from_slice() {
427        let decimal_vector = Decimal128Vector::from_slice([123, 456]);
428        let decimal_vector2 = Decimal128Vector::from_wrapper_slice([
429            Decimal128::new(123, 10, 2),
430            Decimal128::new(456, 10, 2),
431        ]);
432        let expect = Decimal128Vector::from_values(vec![123, 456]);
433
434        assert_eq!(decimal_vector, expect);
435        assert_eq!(decimal_vector2, expect);
436    }
437
438    #[test]
439    fn test_decimal128_vector_slice() {
440        let data = vec![100, 200, 300];
441        // create a decimal vector
442        let decimal_vector = Decimal128Vector::from_values(data.clone())
443            .with_precision_and_scale(10, 2)
444            .unwrap();
445        let decimal_vector2 = decimal_vector.slice(1, 2);
446        assert_eq!(decimal_vector2.len(), 2);
447        assert_eq!(
448            decimal_vector2.get(0),
449            Value::Decimal128(Decimal128::new(200, 10, 2))
450        );
451        assert_eq!(
452            decimal_vector2.get(1),
453            Value::Decimal128(Decimal128::new(300, 10, 2))
454        );
455    }
456
457    #[test]
458    fn test_decimal128_vector_basic() {
459        let data = vec![100, 200, 300];
460        // create a decimal vector
461        let decimal_vector = Decimal128Vector::from_values(data.clone())
462            .with_precision_and_scale(10, 2)
463            .unwrap();
464
465        // can use value_of_string(idx) get a decimal string
466        assert_eq!(decimal_vector.value_as_string(0), "1.00");
467
468        // iterator for-loop
469        for i in 0..data.len() {
470            assert_eq!(
471                decimal_vector.get_data(i),
472                Some(Decimal128::new((i + 1) as i128 * 100, 10, 2))
473            );
474            assert_eq!(
475                decimal_vector.get(i),
476                Value::Decimal128(Decimal128::new((i + 1) as i128 * 100, 10, 2))
477            );
478            assert_eq!(
479                decimal_vector.get_ref(i),
480                ValueRef::Decimal128(Decimal128::new((i + 1) as i128 * 100, 10, 2))
481            );
482        }
483
484        // directly convert vector with precision = 2 and scale = 1,
485        // then all of value will be null because of precision.
486        let decimal_vector = decimal_vector
487            .with_precision_and_scale_to_null(2, 1)
488            .unwrap();
489        assert_eq!(decimal_vector.len(), 3);
490        assert!(decimal_vector.is_null(0));
491        assert!(decimal_vector.is_null(1));
492        assert!(decimal_vector.is_null(2));
493    }
494
495    #[test]
496    fn test_decimal128_vector_builder() {
497        let mut decimal_builder = Decimal128VectorBuilder::with_capacity(3)
498            .with_precision_and_scale(10, 2)
499            .unwrap();
500        decimal_builder.push(Some(Decimal128::new(100, 10, 2)));
501        decimal_builder.push(Some(Decimal128::new(200, 10, 2)));
502        decimal_builder.push(Some(Decimal128::new(300, 10, 2)));
503        let decimal_vector = decimal_builder.finish();
504        assert_eq!(decimal_vector.len(), 3);
505        assert_eq!(decimal_vector.precision(), 10);
506        assert_eq!(decimal_vector.scale(), 2);
507        assert_eq!(
508            decimal_vector.get(0),
509            Value::Decimal128(Decimal128::new(100, 10, 2))
510        );
511        assert_eq!(
512            decimal_vector.get(1),
513            Value::Decimal128(Decimal128::new(200, 10, 2))
514        );
515        assert_eq!(
516            decimal_vector.get(2),
517            Value::Decimal128(Decimal128::new(300, 10, 2))
518        );
519
520        // push value error
521        let mut decimal_builder = Decimal128VectorBuilder::with_capacity(3);
522        decimal_builder.push(Some(Decimal128::new(123, 38, 10)));
523        decimal_builder.push(Some(Decimal128::new(1234, 38, 10)));
524        decimal_builder.push(Some(Decimal128::new(12345, 38, 10)));
525        let decimal_vector = decimal_builder.finish();
526        assert_eq!(decimal_vector.precision(), 38);
527        assert_eq!(decimal_vector.scale(), 10);
528        let result = decimal_vector
529            .with_precision_and_scale(3, 2)
530            .and_then(|x| x.validate_decimal_precision(3));
531        assert_eq!(
532            "Value exceeds the precision 3 bound",
533            result.unwrap_err().to_string()
534        );
535    }
536
537    #[test]
538    fn test_cast_to_decimal128() {
539        let vector = Int8Vector::from_values(vec![1, 2, 3, 4, 100]);
540        let casted_vector = vector.cast(&ConcreteDataType::decimal128_datatype(3, 1));
541        assert!(casted_vector.is_ok());
542        let vector = casted_vector.unwrap();
543        let array = vector.as_any().downcast_ref::<Decimal128Vector>().unwrap();
544        // because 100 is out of Decimal(3, 1) range, so it will be null
545        assert!(array.is_null(4));
546    }
547
548    #[test]
549    fn test_decimal28_vector_iter_data() {
550        let vector = Decimal128Vector::from_values(vec![1, 2, 3, 4])
551            .with_precision_and_scale(3, 1)
552            .unwrap();
553        let mut iter = vector.iter_data();
554        assert_eq!(iter.next(), Some(Some(Decimal128::new(1, 3, 1))));
555        assert_eq!(iter.next(), Some(Some(Decimal128::new(2, 3, 1))));
556        assert_eq!(iter.next(), Some(Some(Decimal128::new(3, 3, 1))));
557        assert_eq!(iter.next(), Some(Some(Decimal128::new(4, 3, 1))));
558        assert_eq!(iter.next(), None);
559
560        let values = vector
561            .iter_data()
562            .filter_map(|v| v.map(|x| x.val() * 2))
563            .collect::<Vec<_>>();
564        assert_eq!(values, vec![2, 4, 6, 8]);
565    }
566
567    #[test]
568    fn test_decimal128_vector_builder_finish_cloned() {
569        let mut builder = Decimal128VectorBuilder::with_capacity(1024);
570        builder.push(Some(Decimal128::new(1, 3, 1)));
571        builder.push(Some(Decimal128::new(1, 3, 1)));
572        builder.push(Some(Decimal128::new(1, 3, 1)));
573        builder.push(Some(Decimal128::new(1, 3, 1)));
574        let vector = builder.finish_cloned();
575        assert_eq!(vector.len(), 4);
576        assert_eq!(builder.len(), 4);
577    }
578}
datatypes/vectors/decimal.rs

datatypes/vectors/
decimal.rs
