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

Merged
myrrc merged 1 commit into from
Jul 7, 2026
Merged

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
149 changes: 133 additions & 16 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 @@ -18,10 +18,16 @@ use crate::aggregate_fn::combined::CombinedOptions;
use crate::aggregate_fn::combined::PairOptions;
use crate::aggregate_fn::fns::count::Count;
use crate::aggregate_fn::fns::sum::Sum;
use crate::aggregate_fn::fns::sum::sum_decimal_dtype;
use crate::builtins::ArrayBuiltins;
use crate::dtype::DType;
use crate::dtype::DecimalDType;
use crate::dtype::MAX_PRECISION;
use crate::dtype::MAX_SCALE;
use crate::dtype::Nullability;
use crate::dtype::PType;
use crate::dtype::i256;
use crate::scalar::DecimalValue;
use crate::scalar::Scalar;
use crate::scalar_fn::fns::operators::Operator;

Expand All @@ -45,10 +51,7 @@ pub fn mean(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
///
/// 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 and primitive numeric types are cast to f64. Decimals stay decimals.
#[derive(Clone, Debug)]
pub struct Mean;

Expand Down Expand Up @@ -88,18 +91,18 @@ 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),
};
if let DType::Decimal(..) = sum.dtype() {
vortex_bail!("grouped mean over decimals is not yet supported");
}
let target = DType::Primitive(PType::F64, Nullability::Nullable);
Comment thread
myrrc marked this conversation as resolved.
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");
if let DType::Decimal(decimal_dtype, _) = *left_scalar.dtype() {
return finalize_decimal_scalar(&left_scalar, &right_scalar, decimal_dtype);
}

let target = DType::Primitive(PType::F64, Nullability::Nullable);
Expand Down Expand Up @@ -140,13 +143,12 @@ 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
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,
}
}
Expand All @@ -156,13 +158,59 @@ fn mean_output_dtype(input_dtype: &DType) -> Option<DType> {
DType::Bool(_) | DType::Primitive(..) => {
Some(DType::Primitive(PType::F64, Nullability::Nullable))
}
DType::Decimal(..) => {
unimplemented!("mean for decimals is not yet implemented");
}
DType::Decimal(decimal_dtype, _) => Some(DType::Decimal(
mean_decimal_dtype(&sum_decimal_dtype(decimal_dtype)),
Nullability::Nullable,
)),
_ => None,
}
}

/// mean() output decimal type mimicking Spark/DataFusion/MySQL: decimal(p+4, s+4)
fn mean_decimal_dtype(sum: &DecimalDType) -> DecimalDType {
DecimalDType::new(
u8::min(MAX_PRECISION, sum.precision().saturating_sub(6)),
i8::min(MAX_SCALE, sum.scale() + 4),
)
}

fn finalize_decimal_scalar(
sum: &Scalar,
count: &Scalar,
sum_decimal: DecimalDType,
) -> VortexResult<Scalar> {
let target_decimal_dtype = mean_decimal_dtype(&sum_decimal);
let target_dtype = DType::Decimal(target_decimal_dtype, Nullability::Nullable);

// overflow
let Some(sum_value) = sum.as_decimal().decimal_value() else {
return Ok(Scalar::null(target_dtype));
};
// empty input
let count = count.as_primitive().typed_value::<u64>().unwrap_or(0);
if count == 0 {
return Ok(Scalar::null(target_dtype));
}

let Ok(sum) = DecimalValue::rescale_i256(
sum_value.as_i256(),
sum_decimal.scale(),
target_decimal_dtype.scale(),
) else {
return Ok(Scalar::null(target_dtype));
};
let mean = sum / i256::from_i128(i128::from(count));

let Ok(mean) = DecimalValue::try_from_i256(mean, target_decimal_dtype) else {
return Ok(Scalar::null(target_dtype));
};
Comment thread
robert3005 marked this conversation as resolved.
Ok(Scalar::decimal(
mean,
target_decimal_dtype,
Nullability::Nullable,
))
}

#[cfg(test)]
mod tests {
use vortex_buffer::buffer;
Expand All @@ -175,7 +223,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 +323,73 @@ 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::Decimal(DecimalDType::new(10, 6), Nullability::Nullable)
);
// mean(1.00, 2.00, 3.00) = 2.000000
assert_eq!(
result.as_decimal().decimal_value(),
Some(DecimalValue::I256(i256::from_i128(2_000_000)))
);
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)?;
// mean(1.50, 4.50) = 3.000000
assert_eq!(
result.as_decimal().decimal_value(),
Some(DecimalValue::I256(i256::from_i128(3_000_000)))
);
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)?;
// mean(1.00, 2.00, 3.00, 4.00, 5.00) = 3.000000
assert_eq!(
result.as_decimal().decimal_value(),
Some(DecimalValue::I256(i256::from_i128(3_000_000)))
);
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)?;
// mean(1.00, 0.00, 0.00) = 1/3 => 0.333333
assert_eq!(
result.as_decimal().decimal_value(),
Some(DecimalValue::I256(i256::from_i128(333_333)))
);
Ok(())
}

#[test]
fn mean_multi_batch() -> VortexResult<()> {
let mut ctx = array_session().create_execution_ctx();
Expand Down
19 changes: 11 additions & 8 deletions vortex-array/src/aggregate_fn/fns/sum/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,16 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
#[derive(Clone, Debug)]
pub struct Sum;

// Both Spark and DataFusion use this heuristic.
// - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala#L66
// - https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L188
pub(crate) fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType {
DecimalDType::new(
u8::min(MAX_PRECISION, input.precision() + 10),
input.scale(),
)
}

impl AggregateFnVTable for Sum {
type Options = NumericalAggregateOpts;
type Partial = SumPartial;
Expand Down Expand Up @@ -118,14 +128,7 @@ impl AggregateFnVTable for Sum {
}
},
DType::Decimal(decimal_dtype, _) => {
// Both Spark and DataFusion use this heuristic.
// - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala#L66
// - https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L188
let precision = u8::min(MAX_PRECISION, decimal_dtype.precision() + 10);
DType::Decimal(
DecimalDType::new(precision, decimal_dtype.scale()),
Nullable,
)
DType::Decimal(sum_decimal_dtype(decimal_dtype), Nullable)
}
// Unsupported types
_ => return None,
Expand Down
Loading
Loading