diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index b421d69a966..179409a2136 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_session::registry::CachedId; use crate::ArrayRef; @@ -42,13 +41,8 @@ pub fn mean(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { } /// Compute the arithmetic mean of an array. -/// /// Implemented as `Sum / Count` via [`BinaryCombined`]. -/// -/// Coercion / return type: -/// - Booleans and primitive numeric types are coerced to `f64` and the result -/// is a nullable `f64`. -/// - Decimals are kept as decimals but not implemented currently +/// Booleans, primitive numeric types and decimals are coerced to f64. #[derive(Clone, Debug)] pub struct Mean; @@ -88,20 +82,13 @@ impl BinaryCombined for Mean { } fn finalize(&self, sum: ArrayRef, count: ArrayRef) -> VortexResult { - let target = match sum.dtype() { - DType::Decimal(..) => sum.dtype().with_nullability(Nullability::Nullable), - _ => DType::Primitive(PType::F64, Nullability::Nullable), - }; + let target = DType::Primitive(PType::F64, Nullability::Nullable); let sum_cast = sum.cast(target.clone())?; let count_cast = count.cast(target)?; sum_cast.binary(count_cast, Operator::Div) } fn finalize_scalar(&self, left_scalar: Scalar, right_scalar: Scalar) -> VortexResult { - if let DType::Decimal(..) = left_scalar.dtype() { - vortex_bail!("mean::finalize_scalar not yet implemented for decimal inputs"); - } - let target = DType::Primitive(PType::F64, Nullability::Nullable); let sum_cast = left_scalar.cast(&target)?; let count_cast = right_scalar.cast(&target)?; @@ -140,25 +127,21 @@ impl BinaryCombined for Mean { /// - Bool stays as bool — `Sum` has a native bool path and bool → f64 isn't /// currently a direct cast in vortex. /// - Primitive numerics → `f64` so the sum and finalize work without overflow. +/// - Decimals stay as decimals. finalize() casts to f64 fn coerced_input_dtype(input_dtype: &DType) -> Option { match input_dtype { DType::Bool(_) => Some(input_dtype.clone()), DType::Primitive(_, n) => Some(DType::Primitive(PType::F64, *n)), - DType::Decimal(..) => { - unimplemented!("mean is not implemented for decimals yet") - } + DType::Decimal(..) => Some(input_dtype.clone()), _ => None, } } fn mean_output_dtype(input_dtype: &DType) -> Option { match input_dtype { - DType::Bool(_) | DType::Primitive(..) => { + DType::Bool(_) | DType::Primitive(..) | DType::Decimal(..) => { Some(DType::Primitive(PType::F64, Nullability::Nullable)) } - DType::Decimal(..) => { - unimplemented!("mean for decimals is not yet implemented"); - } _ => None, } } @@ -175,7 +158,9 @@ mod tests { use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; + use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; + use crate::dtype::DecimalDType; use crate::validity::Validity; #[test] @@ -273,6 +258,57 @@ mod tests { Ok(()) } + #[test] + fn mean_decimal() -> VortexResult<()> { + let dtype = DecimalDType::new(6, 2); + let array = + DecimalArray::new(buffer![100i32, 200, 300], dtype, Validity::NonNullable).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let result = mean(&array, &mut ctx)?; + assert_eq!( + result.dtype(), + &DType::Primitive(PType::F64, Nullability::Nullable) + ); + assert_eq!(result.as_primitive().as_::(), Some(2.0)); + Ok(()) + } + + #[test] + fn mean_decimal_null() -> VortexResult<()> { + let dtype = DecimalDType::new(6, 2); + let validity = Validity::from_iter([true, false, true]); + let array = DecimalArray::new(buffer![150i32, 0, 450], dtype, validity).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let result = mean(&array, &mut ctx)?; + assert_eq!(result.as_primitive().as_::(), Some(3.0)); + Ok(()) + } + + #[test] + fn mean_decimal_chunked() -> VortexResult<()> { + let dtype = DecimalDType::new(6, 2); + let validity = Validity::NonNullable; + let chunk1 = DecimalArray::new(buffer![100i32, 200], dtype, validity.clone()).into_array(); + let chunk2 = DecimalArray::new(buffer![300i32, 400, 500], dtype, validity).into_array(); + let dtype = chunk1.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![chunk1, chunk2], dtype)?; + let mut ctx = array_session().create_execution_ctx(); + let result = mean(&chunked.into_array(), &mut ctx)?; + assert_eq!(result.as_primitive().as_::(), Some(3.0)); + Ok(()) + } + + #[test] + fn mean_decimal_33() -> VortexResult<()> { + let dtype = DecimalDType::new(6, 2); + let buf = buffer![100i32, 0, 0]; + let array = DecimalArray::new(buf, dtype, Validity::NonNullable).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let result = mean(&array, &mut ctx)?; + assert_eq!(result.as_primitive().as_::(), Some(1.0f64 / 3.0)); + Ok(()) + } + #[test] fn mean_multi_batch() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); diff --git a/vortex-array/src/arrays/decimal/compute/cast.rs b/vortex-array/src/arrays/decimal/compute/cast.rs index 80fafde5059..b6f061d9e8a 100644 --- a/vortex-array/src/arrays/decimal/compute/cast.rs +++ b/vortex-array/src/arrays/decimal/compute/cast.rs @@ -1,7 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use num_traits::AsPrimitive; use num_traits::CheckedMul; +use num_traits::ToPrimitive as NumToPrimitive; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_compute::lane_kernels::IndexedSourceExt; @@ -19,11 +21,14 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Decimal; use crate::arrays::DecimalArray; +use crate::arrays::PrimitiveArray; use crate::dtype::BigCast; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::DecimalType; use crate::dtype::NativeDecimalType; +use crate::dtype::Nullability; +use crate::dtype::PType; use crate::dtype::i256; use crate::match_each_decimal_value_type; use crate::scalar::DecimalValue; @@ -77,16 +82,19 @@ impl CastKernel for Decimal { dtype: &DType, ctx: &mut ExecutionCtx, ) -> VortexResult> { - // Early return if not casting to decimal - let DType::Decimal(to_decimal_dtype, to_nullability) = dtype else { - return Ok(None); - }; let DType::Decimal(from_decimal_dtype, _) = array.dtype() else { vortex_panic!( "DecimalArray must have decimal dtype, got {:?}", array.dtype() ); }; + if let DType::Primitive(PType::F64, nullability) = dtype { + let scale = from_decimal_dtype.scale(); + return cast_to_f64(array, scale, *nullability, ctx).map(Some); + } + let DType::Decimal(to_decimal_dtype, to_nullability) = dtype else { + return Ok(None); + }; // If the dtype is exactly the same, return self if array.dtype() == dtype { @@ -143,6 +151,57 @@ impl CastKernel for Decimal { } } +fn cast_to_f64( + array: ArrayView<'_, Decimal>, + scale: i8, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let source_validity = array.validity()?; + let n = array.len(); + let mask = source_validity.execute_mask(n, ctx)?; + let validity = source_validity.cast_nullability(nullability, n, ctx)?; + + let scale: i32 = scale.as_(); + let inv_factor: f64 = 10f64.powi(-scale); + + let buffer = match &mask { + Mask::AllFalse(_) => BufferMut::::zeroed(n).freeze(), + Mask::AllTrue(_) => match_each_decimal_value_type!(array.values_type(), |F| { + let values = array.buffer::(); + let values = values.as_slice(); + let mut out = BufferMut::::with_capacity(n); + values.map_into(&mut out.spare_capacity_mut()[..n], |v: F| { + to_f64_lossy::(v) * inv_factor + }); + // SAFETY: map_into wrote every lane before returning. + unsafe { out.set_len(n) }; + out.freeze() + }), + Mask::Values(mask_values) => match_each_decimal_value_type!(array.values_type(), |F| { + let values = array.buffer::(); + let values = values.as_slice(); + let mut out = BufferMut::::with_capacity(n); + let write_result = values.try_map_masked_into( + mask_values.bit_buffer(), + &mut out.spare_capacity_mut()[..n], + |v: F| Some(to_f64_lossy::(v) * inv_factor), + ); + debug_assert!(write_result.is_ok()); + // SAFETY: try_map_masked_into wrote every lane before returning Ok. + unsafe { out.set_len(n) }; + out.freeze() + }), + }; + + Ok(PrimitiveArray::new(buffer, validity).into_array()) +} + +#[inline] +fn to_f64_lossy(value: F) -> f64 { + NumToPrimitive::to_f64(&value).unwrap_or(f64::NAN) +} + fn cast_decimal_values( array: ArrayView<'_, Decimal>, from_decimal_dtype: DecimalDType, @@ -386,6 +445,7 @@ mod tests { use vortex_buffer::buffer; use super::upcast_decimal_values; + use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; @@ -398,6 +458,8 @@ mod tests { use crate::dtype::DecimalDType; use crate::dtype::DecimalType; use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::scalar::Scalar; use crate::validity::Validity; #[test] @@ -769,4 +831,68 @@ mod tests { .contains("Cannot downcast decimal values") ); } + + #[test] + fn cast_decimal_f64() { + let dtype = DecimalDType::new(6, 2); + let array = DecimalArray::new(buffer![125i32, 250, -375], dtype, Validity::NonNullable); + let target = DType::Primitive(PType::F64, Nullability::NonNullable); + let casted = array.into_array().cast(target.clone()).unwrap(); + assert_eq!(casted.dtype(), &target); + + let mut ctx = array_session().create_execution_ctx(); + let primitive = casted + .execute::(&mut ctx) + .unwrap() + .into_primitive(); + assert_eq!(primitive.to_buffer::().as_ref(), &[1.25, 2.5, -3.75]); + } + + #[test] + fn cast_decimal_f64_null() { + let values = [Some(100i32), None, Some(-200)]; + let array = DecimalArray::from_option_iter(values, DecimalDType::new(6, 2)); + let target = DType::Primitive(PType::F64, Nullability::Nullable); + let casted = array.into_array().cast(target.clone()).unwrap(); + assert_eq!(casted.dtype(), &target); + + let mut ctx = array_session().create_execution_ctx(); + let primitive = casted + .execute::(&mut ctx) + .unwrap() + .into_primitive(); + + assert!(primitive.is_valid(0, &mut ctx).unwrap()); + assert!(!primitive.is_valid(1, &mut ctx).unwrap()); + assert!(primitive.is_valid(2, &mut ctx).unwrap()); + + assert_eq!( + primitive.execute_scalar(0, &mut ctx).unwrap(), + Scalar::from(1f64) + ); + assert_eq!( + primitive.execute_scalar(2, &mut ctx).unwrap(), + Scalar::from(-2f64) + ); + } + + #[test] + fn cast_decimal_f64_all_null() { + let dtype = DecimalDType::new(6, 2); + let buf = buffer![i32::MAX, i32::MIN, 12345]; + let array = DecimalArray::new(buf, dtype, Validity::AllInvalid); + let target = DType::Primitive(PType::F64, Nullability::Nullable); + let casted = array.into_array().cast(target.clone()).unwrap(); + assert_eq!(casted.dtype(), &target); + + let primitive = casted + .execute::(&mut array_session().create_execution_ctx()) + .unwrap() + .into_primitive(); + let mut ctx = array_session().create_execution_ctx(); + for i in 0..2 { + assert!(!primitive.is_valid(i, &mut ctx).unwrap()); + } + assert_eq!(primitive.to_buffer::().as_ref(), &[0.0, 0.0, 0.0]); + } }