Decimal Vector

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-array/src/builders/decimal.rs

@@ -2,16 +2,17 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::any::Any;
+
 use vortex_buffer::BufferMut;
 use vortex_dtype::{BigCast, DType, DecimalDType, NativeDecimalType, Nullability};
-use vortex_error::{vortex_ensure, vortex_panic, VortexExpect, VortexResult};
+use vortex_error::{VortexExpect, VortexResult, vortex_ensure, vortex_panic};
 use vortex_mask::Mask;
 use vortex_scalar::{
-    i256, match_each_decimal_value, match_each_decimal_value_type, DecimalValue, Scalar,
+    DecimalValue, Scalar, i256, match_each_decimal_value, match_each_decimal_value_type,
 };
 
 use crate::arrays::DecimalArray;
-use crate::builders::{ArrayBuilder, LazyBitBufferBuilder, DEFAULT_BUILDER_CAPACITY};
+use crate::builders::{ArrayBuilder, DEFAULT_BUILDER_CAPACITY, LazyBitBufferBuilder};
 use crate::canonical::Canonical;
 use crate::{Array, ArrayRef, IntoArray, ToCanonical};
 

vortex-array/src/stats/mod.rs

@@ -8,12 +8,12 @@ use std::hash::Hash;
 
 use arrow_buffer::bit_iterator::BitIterator;
 use arrow_buffer::{BooleanBufferBuilder, MutableBuffer};
-use enum_iterator::{all, last, Sequence};
+use enum_iterator::{Sequence, all, last};
 use log::debug;
 use num_enum::{IntoPrimitive, TryFromPrimitive};
 pub use stats_set::*;
 use vortex_dtype::Nullability::{NonNullable, Nullable};
-use vortex_dtype::{i256, DType, DecimalDType, NativeDecimalType, PType};
+use vortex_dtype::{DType, DecimalDType, NativeDecimalType, PType, i256};
 
 mod array;
 mod bound;

vortex-compute/src/mask/mod.rs

@@ -5,11 +5,11 @@
 
 use std::ops::BitAnd;
 
-use vortex_dtype::NativePType;
+use vortex_dtype::{NativeDecimalType, NativePType};
 use vortex_mask::Mask;
 use vortex_vector::{
-    BoolVector, NullVector, PVector, PrimitiveVector, StructVector, VarBinType, VarBinVector,
-    Vector, match_each_pvector, match_each_vector,
+    BoolVector, DVector, DecimalVector, NullVector, PVector, PrimitiveVector, StructVector,
+    VarBinType, VarBinVector, Vector, match_each_dvector, match_each_pvector, match_each_vector,
 };
 
 /// Trait for masking the validity of an array or vector.
@@ -42,6 +42,20 @@ impl MaskValidity for BoolVector {
     }
 }
 
+impl MaskValidity for DecimalVector {
+    fn mask_validity(self, mask: &Mask) -> Self {
+        match_each_dvector!(self, |v| { MaskValidity::mask_validity(v, mask).into() })
+    }
+}
+
+impl<D: NativeDecimalType> MaskValidity for DVector<D> {
+    fn mask_validity(self, mask: &Mask) -> Self {
+        let (ps, elements, validity) = self.into_parts();
+        // SAFETY: we are preserving the original elements buffer and only modifying the validity.
+        unsafe { Self::new_unchecked(ps, elements, validity.bitand(mask)) }
+    }
+}
+
 impl MaskValidity for PrimitiveVector {
     fn mask_validity(self, mask: &Mask) -> Self {
         match_each_pvector!(self, |v| { MaskValidity::mask_validity(v, mask).into() })

vortex-datafusion/src/convert/scalars.rs

@@ -6,11 +6,11 @@ use std::sync::Arc;
 use datafusion_common::ScalarValue;
 use vortex::buffer::ByteBuffer;
 use vortex::dtype::datetime::arrow::make_temporal_ext_dtype;
-use vortex::dtype::datetime::{is_temporal_ext_type, TemporalMetadata, TimeUnit};
+use vortex::dtype::datetime::{TemporalMetadata, TimeUnit, is_temporal_ext_type};
 use vortex::dtype::half::f16;
 use vortex::dtype::{DType, DecimalDType, NativeDecimalType, Nullability, PType};
-use vortex::error::{vortex_bail, VortexResult};
-use vortex::scalar::{i256, DecimalValue, Scalar};
+use vortex::error::{VortexResult, vortex_bail};
+use vortex::scalar::{DecimalValue, Scalar, i256};
 
 use crate::convert::{FromDataFusion, TryToDataFusion};
 
@@ -244,12 +244,12 @@ impl FromDataFusion<ScalarValue> for Scalar {
 
 #[cfg(test)]
 mod tests {
-    use datafusion_common::arrow::datatypes::i256 as arrow_i256;
     use datafusion_common::ScalarValue;
+    use datafusion_common::arrow::datatypes::i256 as arrow_i256;
     use rstest::rstest;
     use vortex::buffer::ByteBuffer;
     use vortex::dtype::{DType, DecimalDType, Nullability, PType};
-    use vortex::scalar::{i256, DecimalValue, Scalar};
+    use vortex::scalar::{DecimalValue, Scalar, i256};
 
     use super::*;
 

vortex-dtype/src/arbitrary/mod.rs

@@ -7,8 +7,8 @@ use arbitrary::{Arbitrary, Result, Unstructured};
 use vortex_error::VortexExpect;
 
 use crate::{
-    i256, DType, DecimalDType, FieldName, FieldNames, NativeDecimalType, Nullability,
-    PType, StructFields,
+    DType, DecimalDType, FieldName, FieldNames, NativeDecimalType, Nullability, PType,
+    StructFields, i256,
 };
 
 mod decimal;

vortex-dtype/src/decimal/mod.rs

@@ -1,17 +1,18 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+mod max_precision;
 mod precision;
 mod types;
-mod max_precision;
 
-use crate::{i256, DType};
-use num_traits::ToPrimitive;
 use std::fmt::{Display, Formatter};
-use vortex_error::{vortex_bail, vortex_panic, VortexError, VortexExpect, VortexResult};
 
+use num_traits::ToPrimitive;
 pub use precision::*;
 pub use types::*;
+use vortex_error::{VortexError, VortexExpect, VortexResult, vortex_bail, vortex_panic};
+
+use crate::{DType, i256};
 
 const MAX_PRECISION: u8 = <i256 as NativeDecimalType>::MAX_PRECISION;
 const MAX_SCALE: i8 = <i256 as NativeDecimalType>::MAX_SCALE;

vortex-dtype/src/decimal/precision.rs

@@ -1,9 +1,13 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use crate::{DecimalDType, NativeDecimalType};
+use std::any::type_name;
+use std::fmt::Display;
 use std::marker::PhantomData;
-use vortex_error::{vortex_bail, VortexExpect, VortexResult};
+
+use vortex_error::{VortexExpect, VortexResult, vortex_bail};
+
+use crate::{DecimalDType, NativeDecimalType};
 
 /// A struct representing the precision and scale of a decimal type, to be represented
 /// by the native type `D`.
@@ -14,6 +18,18 @@ pub struct PrecisionScale<D> {
     phantom: PhantomData<D>,
 }
 
+impl<D: NativeDecimalType> Display for PrecisionScale<D> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "decimal({}, p={}, s={})",
+            type_name::<D>(),
+            self.precision,
+            self.scale
+        )
+    }
+}
+
 impl<D: NativeDecimalType> PrecisionScale<D> {
     /// Create a new [`PrecisionScale`] with the given precision and scale.
     ///

vortex-dtype/src/decimal/types.rs

@@ -1,15 +1,16 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::fmt::{Debug, Display};
+use std::panic::RefUnwindSafe;
+
+use num_traits::{ConstOne, ConstZero};
+use paste::paste;
+
 use crate::decimal::max_precision::{
     MAX_DECIMAL256_FOR_EACH_PRECISION, MIN_DECIMAL256_FOR_EACH_PRECISION,
 };
-use crate::{i256, BigCast};
-use num_traits::ConstOne;
-use num_traits::ConstZero;
-use paste::paste;
-use std::fmt::{Debug, Display};
-use std::panic::RefUnwindSafe;
+use crate::{BigCast, i256};
 
 /// Type of the decimal values.
 ///

vortex-scalar/src/arbitrary.rs

@@ -12,7 +12,7 @@ use std::sync::Arc;
 use arbitrary::{Result, Unstructured};
 use vortex_buffer::{BufferString, ByteBuffer};
 use vortex_dtype::half::f16;
-use vortex_dtype::{i256, DType, DecimalDType, NativeDecimalType, PType};
+use vortex_dtype::{DType, DecimalDType, NativeDecimalType, PType, i256};
 
 use crate::{DecimalValue, InnerScalarValue, PValue, Scalar, ScalarValue};
 

vortex-scalar/src/scalar.rs

@@ -6,8 +6,8 @@ use std::hash::Hash;
 use std::sync::Arc;
 
 use vortex_buffer::Buffer;
-use vortex_dtype::{i256, DType, NativeDType, NativeDecimalType, Nullability};
-use vortex_error::{vortex_bail, vortex_err, VortexError, VortexExpect, VortexResult};
+use vortex_dtype::{DType, NativeDType, NativeDecimalType, Nullability, i256};
+use vortex_error::{VortexError, VortexExpect, VortexResult, vortex_bail, vortex_err};
 
 use super::*;