Skip to content

mean() for decimals #8649

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
myrrc wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

mean() for decimals #8649

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 58 additions & 22 deletions vortex-array/src/aggregate_fn/fns/mean/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -42,13 +41,8 @@ pub fn mean(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
}

/// 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;

Expand Down Expand Up @@ -88,20 +82,13 @@ impl BinaryCombined for Mean {
}

fn finalize(&self, sum: ArrayRef, count: ArrayRef) -> VortexResult<ArrayRef> {
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);

@robert3005 robert3005 Jul 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should decimal mean be f64 as well? The code previously kept them decimals, I think they should stay decimals

@myrrc myrrc Jul 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hm, okay, that makes sense, I'll rewrite to keep mean as decimal

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<Scalar> {
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)?;
Expand Down Expand Up @@ -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<DType> {
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<DType> {
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,
}
}
Expand All @@ -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]
Expand Down Expand Up @@ -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_::<f64>(), 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_::<f64>(), 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_::<f64>(), 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_::<f64>(), Some(1.0f64 / 3.0));
Ok(())
}

#[test]
fn mean_multi_batch() -> VortexResult<()> {
let mut ctx = array_session().create_execution_ctx();
Expand Down
134 changes: 130 additions & 4 deletions vortex-array/src/arrays/decimal/compute/cast.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -77,16 +82,19 @@ impl CastKernel for Decimal {
dtype: &DType,
ctx: &mut ExecutionCtx,
) -> VortexResult<Option<ArrayRef>> {
// 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 {
Expand Down Expand Up @@ -143,6 +151,57 @@ impl CastKernel for Decimal {
}
}

fn cast_to_f64(
array: ArrayView<'_, Decimal>,
scale: i8,
nullability: Nullability,
ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
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::<f64>::zeroed(n).freeze(),
Mask::AllTrue(_) => match_each_decimal_value_type!(array.values_type(), |F| {
let values = array.buffer::<F>();
let values = values.as_slice();
let mut out = BufferMut::<f64>::with_capacity(n);
values.map_into(&mut out.spare_capacity_mut()[..n], |v: F| {
to_f64_lossy::<F>(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::<F>();
let values = values.as_slice();
let mut out = BufferMut::<f64>::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::<F>(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<F: NativeDecimalType>(value: F) -> f64 {
NumToPrimitive::to_f64(&value).unwrap_or(f64::NAN)
}

fn cast_decimal_values<F, T>(
array: ArrayView<'_, Decimal>,
from_decimal_dtype: DecimalDType,
Expand Down Expand Up @@ -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;
Expand All @@ -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]
Expand Down Expand Up @@ -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::<Canonical>(&mut ctx)
.unwrap()
.into_primitive();
assert_eq!(primitive.to_buffer::<f64>().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::<Canonical>(&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::<Canonical>(&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::<f64>().as_ref(), &[0.0, 0.0, 0.0]);
}
}
Loading