Skip to content

Commit b511dcf

Browse files
myrrcmyrrc
committed
mean() for decimals
Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>
1 parent c7e8ae9 commit b511dcf

3 files changed

Lines changed: 274 additions & 28 deletions

File tree

‎vortex-array/src/aggregate_fn/fns/mean/mod.rs‎

Lines changed: 133 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,16 @@ use crate::aggregate_fn::combined::CombinedOptions;
1818
use crate::aggregate_fn::combined::PairOptions;
1919
use crate::aggregate_fn::fns::count::Count;
2020
use crate::aggregate_fn::fns::sum::Sum;
21+
use crate::aggregate_fn::fns::sum::sum_decimal_dtype;
2122
use crate::builtins::ArrayBuiltins;
2223
use crate::dtype::DType;
24+
use crate::dtype::DecimalDType;
25+
use crate::dtype::MAX_PRECISION;
26+
use crate::dtype::MAX_SCALE;
2327
use crate::dtype::Nullability;
2428
use crate::dtype::PType;
29+
use crate::dtype::i256;
30+
use crate::scalar::DecimalValue;
2531
use crate::scalar::Scalar;
2632
use crate::scalar_fn::fns::operators::Operator;
2733

@@ -45,10 +51,7 @@ pub fn mean(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
4551
///
4652
/// Implemented as `Sum / Count` via [`BinaryCombined`].
4753
///
48-
/// Coercion / return type:
49-
/// - Booleans and primitive numeric types are coerced to `f64` and the result
50-
/// is a nullable `f64`.
51-
/// - Decimals are kept as decimals but not implemented currently
54+
/// Booleans and primitive numeric types are cast to f64. Decimals stay decimals.
5255
#[derive(Clone, Debug)]
5356
pub struct Mean;
5457

@@ -88,18 +91,18 @@ impl BinaryCombined for Mean {
8891
}
8992

9093
fn finalize(&self, sum: ArrayRef, count: ArrayRef) -> VortexResult<ArrayRef> {
91-
let target = match sum.dtype() {
92-
DType::Decimal(..) => sum.dtype().with_nullability(Nullability::Nullable),
93-
_ => DType::Primitive(PType::F64, Nullability::Nullable),
94-
};
94+
if let DType::Decimal(..) = sum.dtype() {
95+
vortex_bail!("grouped mean over decimals is not yet supported");
96+
}
97+
let target = DType::Primitive(PType::F64, Nullability::Nullable);
9598
let sum_cast = sum.cast(target.clone())?;
9699
let count_cast = count.cast(target)?;
97100
sum_cast.binary(count_cast, Operator::Div)
98101
}
99102

100103
fn finalize_scalar(&self, left_scalar: Scalar, right_scalar: Scalar) -> VortexResult<Scalar> {
101-
if let DType::Decimal(..) = left_scalar.dtype() {
102-
vortex_bail!("mean::finalize_scalar not yet implemented for decimal inputs");
104+
if let DType::Decimal(decimal_dtype, _) = *left_scalar.dtype() {
105+
return finalize_decimal_scalar(&left_scalar, &right_scalar, decimal_dtype);
103106
}
104107

105108
let target = DType::Primitive(PType::F64, Nullability::Nullable);
@@ -140,13 +143,12 @@ impl BinaryCombined for Mean {
140143
/// - Bool stays as bool — `Sum` has a native bool path and bool → f64 isn't
141144
/// currently a direct cast in vortex.
142145
/// - Primitive numerics → `f64` so the sum and finalize work without overflow.
146+
/// - Decimals stay as decimals
143147
fn coerced_input_dtype(input_dtype: &DType) -> Option<DType> {
144148
match input_dtype {
145149
DType::Bool(_) => Some(input_dtype.clone()),
146150
DType::Primitive(_, n) => Some(DType::Primitive(PType::F64, *n)),
147-
DType::Decimal(..) => {
148-
unimplemented!("mean is not implemented for decimals yet")
149-
}
151+
DType::Decimal(..) => Some(input_dtype.clone()),
150152
_ => None,
151153
}
152154
}
@@ -156,13 +158,59 @@ fn mean_output_dtype(input_dtype: &DType) -> Option<DType> {
156158
DType::Bool(_) | DType::Primitive(..) => {
157159
Some(DType::Primitive(PType::F64, Nullability::Nullable))
158160
}
159-
DType::Decimal(..) => {
160-
unimplemented!("mean for decimals is not yet implemented");
161-
}
161+
DType::Decimal(decimal_dtype, _) => Some(DType::Decimal(
162+
mean_decimal_dtype(&sum_decimal_dtype(decimal_dtype)),
163+
Nullability::Nullable,
164+
)),
162165
_ => None,
163166
}
164167
}
165168

169+
/// mean() output decimal type mimicking Spark/DataFusion/MySQL: decimal(p+4, s+4)
170+
fn mean_decimal_dtype(sum: &DecimalDType) -> DecimalDType {
171+
DecimalDType::new(
172+
u8::min(MAX_PRECISION, sum.precision().saturating_sub(6)),
173+
i8::min(MAX_SCALE, sum.scale() + 4),
174+
)
175+
}
176+
177+
fn finalize_decimal_scalar(
178+
sum: &Scalar,
179+
count: &Scalar,
180+
sum_decimal: DecimalDType,
181+
) -> VortexResult<Scalar> {
182+
let target_decimal_dtype = mean_decimal_dtype(&sum_decimal);
183+
let target_dtype = DType::Decimal(target_decimal_dtype, Nullability::Nullable);
184+
185+
// overflow
186+
let Some(sum_value) = sum.as_decimal().decimal_value() else {
187+
return Ok(Scalar::null(target_dtype));
188+
};
189+
// empty input
190+
let count = count.as_primitive().typed_value::<u64>().unwrap_or(0);
191+
if count == 0 {
192+
return Ok(Scalar::null(target_dtype));
193+
}
194+
195+
let Ok(sum) = DecimalValue::rescale_i256(
196+
sum_value.as_i256(),
197+
sum_decimal.scale(),
198+
target_decimal_dtype.scale(),
199+
) else {
200+
return Ok(Scalar::null(target_dtype));
201+
};
202+
let mean = sum / i256::from_i128(i128::from(count));
203+
204+
let Ok(mean) = DecimalValue::try_from_i256(mean, target_decimal_dtype) else {
205+
return Ok(Scalar::null(target_dtype));
206+
};
207+
Ok(Scalar::decimal(
208+
mean,
209+
target_decimal_dtype,
210+
Nullability::Nullable,
211+
))
212+
}
213+
166214
#[cfg(test)]
167215
mod tests {
168216
use vortex_buffer::buffer;
@@ -175,7 +223,9 @@ mod tests {
175223
use crate::arrays::BoolArray;
176224
use crate::arrays::ChunkedArray;
177225
use crate::arrays::ConstantArray;
226+
use crate::arrays::DecimalArray;
178227
use crate::arrays::PrimitiveArray;
228+
use crate::dtype::DecimalDType;
179229
use crate::validity::Validity;
180230

181231
#[test]
@@ -273,6 +323,73 @@ mod tests {
273323
Ok(())
274324
}
275325

326+
#[test]
327+
fn mean_decimal() -> VortexResult<()> {
328+
let dtype = DecimalDType::new(6, 2);
329+
let array =
330+
DecimalArray::new(buffer![100i32, 200, 300], dtype, Validity::NonNullable).into_array();
331+
let mut ctx = array_session().create_execution_ctx();
332+
let result = mean(&array, &mut ctx)?;
333+
assert_eq!(
334+
result.dtype(),
335+
&DType::Decimal(DecimalDType::new(10, 6), Nullability::Nullable)
336+
);
337+
// mean(1.00, 2.00, 3.00) = 2.000000
338+
assert_eq!(
339+
result.as_decimal().decimal_value(),
340+
Some(DecimalValue::I256(i256::from_i128(2_000_000)))
341+
);
342+
Ok(())
343+
}
344+
345+
#[test]
346+
fn mean_decimal_null() -> VortexResult<()> {
347+
let dtype = DecimalDType::new(6, 2);
348+
let validity = Validity::from_iter([true, false, true]);
349+
let array = DecimalArray::new(buffer![150i32, 0, 450], dtype, validity).into_array();
350+
let mut ctx = array_session().create_execution_ctx();
351+
let result = mean(&array, &mut ctx)?;
352+
// mean(1.50, 4.50) = 3.000000
353+
assert_eq!(
354+
result.as_decimal().decimal_value(),
355+
Some(DecimalValue::I256(i256::from_i128(3_000_000)))
356+
);
357+
Ok(())
358+
}
359+
360+
#[test]
361+
fn mean_decimal_chunked() -> VortexResult<()> {
362+
let dtype = DecimalDType::new(6, 2);
363+
let validity = Validity::NonNullable;
364+
let chunk1 = DecimalArray::new(buffer![100i32, 200], dtype, validity.clone()).into_array();
365+
let chunk2 = DecimalArray::new(buffer![300i32, 400, 500], dtype, validity).into_array();
366+
let dtype = chunk1.dtype().clone();
367+
let chunked = ChunkedArray::try_new(vec![chunk1, chunk2], dtype)?;
368+
let mut ctx = array_session().create_execution_ctx();
369+
let result = mean(&chunked.into_array(), &mut ctx)?;
370+
// mean(1.00, 2.00, 3.00, 4.00, 5.00) = 3.000000
371+
assert_eq!(
372+
result.as_decimal().decimal_value(),
373+
Some(DecimalValue::I256(i256::from_i128(3_000_000)))
374+
);
375+
Ok(())
376+
}
377+
378+
#[test]
379+
fn mean_decimal_33() -> VortexResult<()> {
380+
let dtype = DecimalDType::new(6, 2);
381+
let buf = buffer![100i32, 0, 0];
382+
let array = DecimalArray::new(buf, dtype, Validity::NonNullable).into_array();
383+
let mut ctx = array_session().create_execution_ctx();
384+
let result = mean(&array, &mut ctx)?;
385+
// mean(1.00, 0.00, 0.00) = 1/3 => 0.333333
386+
assert_eq!(
387+
result.as_decimal().decimal_value(),
388+
Some(DecimalValue::I256(i256::from_i128(333_333)))
389+
);
390+
Ok(())
391+
}
392+
276393
#[test]
277394
fn mean_multi_batch() -> VortexResult<()> {
278395
let mut ctx = array_session().create_execution_ctx();

‎vortex-array/src/aggregate_fn/fns/sum/mod.rs‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,16 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
7777
#[derive(Clone, Debug)]
7878
pub struct Sum;
7979

80+
// Both Spark and DataFusion use this heuristic.
81+
// - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala#L66
82+
// - https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L188
83+
pub(crate) fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType {
84+
DecimalDType::new(
85+
u8::min(MAX_PRECISION, input.precision() + 10),
86+
input.scale(),
87+
)
88+
}
89+
8090
impl AggregateFnVTable for Sum {
8191
type Options = NumericalAggregateOpts;
8292
type Partial = SumPartial;
@@ -118,14 +128,7 @@ impl AggregateFnVTable for Sum {
118128
}
119129
},
120130
DType::Decimal(decimal_dtype, _) => {
121-
// Both Spark and DataFusion use this heuristic.
122-
// - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala#L66
123-
// - https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L188
124-
let precision = u8::min(MAX_PRECISION, decimal_dtype.precision() + 10);
125-
DType::Decimal(
126-
DecimalDType::new(precision, decimal_dtype.scale()),
127-
Nullable,
128-
)
131+
DType::Decimal(sum_decimal_dtype(decimal_dtype), Nullable)
129132
}
130133
// Unsupported types
131134
_ => return None,

0 commit comments

Comments
 (0)