Skip to content

Commit

Permalink
Minor: Move Sum aggregate function test to slt (#10382)
Browse files Browse the repository at this point in the history 
* move test to slt

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

* remove sum distinct

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

---------

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>
  • Loading branch information
@jayzhan211
jayzhan211 committed May 5, 2024
1 parent 1f23703 commit b412dba
Show file tree
Hide file tree
Showing 3 changed files with 121 additions and 192 deletions.
111 changes: 0 additions & 111 deletions datafusion/physical-expr/src/aggregate/sum.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -289,114 +289,3 @@ impl<T: ArrowNumericType> Accumulator for SlidingSumAccumulator<T> {
true
}
}

#[cfg(test)]
mod tests {
use super::*;
use crate::expressions::tests::assert_aggregate;
use arrow_array::*;
use datafusion_expr::AggregateFunction;

#[test]
fn sum_decimal() {
// test agg
let array: ArrayRef = Arc::new(
(1..6)
.map(Some)
.collect::<Decimal128Array>()
.with_precision_and_scale(10, 0)
.unwrap(),
);

assert_aggregate(
array,
AggregateFunction::Sum,
false,
ScalarValue::Decimal128(Some(15), 20, 0),
);
}

#[test]
fn sum_decimal_with_nulls() {
// test agg
let array: ArrayRef = Arc::new(
(1..6)
.map(|i| if i == 2 { None } else { Some(i) })
.collect::<Decimal128Array>()
.with_precision_and_scale(35, 0)
.unwrap(),
);

assert_aggregate(
array,
AggregateFunction::Sum,
false,
ScalarValue::Decimal128(Some(13), 38, 0),
);
}

#[test]
fn sum_decimal_all_nulls() {
// test with batch
let array: ArrayRef = Arc::new(
std::iter::repeat::<Option<i128>>(None)
.take(6)
.collect::<Decimal128Array>()
.with_precision_and_scale(10, 0)
.unwrap(),
);

// test agg
assert_aggregate(
array,
AggregateFunction::Sum,
false,
ScalarValue::Decimal128(None, 20, 0),
);
}

#[test]
fn sum_i32() {
let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]));
assert_aggregate(a, AggregateFunction::Sum, false, ScalarValue::from(15i64));
}

#[test]
fn sum_i32_with_nulls() {
let a: ArrayRef = Arc::new(Int32Array::from(vec![
Some(1),
None,
Some(3),
Some(4),
Some(5),
]));
assert_aggregate(a, AggregateFunction::Sum, false, ScalarValue::from(13i64));
}

#[test]
fn sum_i32_all_nulls() {
let a: ArrayRef = Arc::new(Int32Array::from(vec![None, None]));
assert_aggregate(a, AggregateFunction::Sum, false, ScalarValue::Int64(None));
}

#[test]
fn sum_u32() {
let a: ArrayRef =
Arc::new(UInt32Array::from(vec![1_u32, 2_u32, 3_u32, 4_u32, 5_u32]));
assert_aggregate(a, AggregateFunction::Sum, false, ScalarValue::from(15u64));
}

#[test]
fn sum_f32() {
let a: ArrayRef =
Arc::new(Float32Array::from(vec![1_f32, 2_f32, 3_f32, 4_f32, 5_f32]));
assert_aggregate(a, AggregateFunction::Sum, false, ScalarValue::from(15_f64));
}

#[test]
fn sum_f64() {
let a: ArrayRef =
Arc::new(Float64Array::from(vec![1_f64, 2_f64, 3_f64, 4_f64, 5_f64]));
assert_aggregate(a, AggregateFunction::Sum, false, ScalarValue::from(15_f64));
}
}
81 changes: 0 additions & 81 deletions datafusion/physical-expr/src/aggregate/sum_distinct.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -200,84 +200,3 @@ impl<T: ArrowPrimitiveType> Accumulator for DistinctSumAccumulator<T> {
+ self.values.capacity() * std::mem::size_of::<T::Native>()
}
}

#[cfg(test)]
mod tests {
use super::*;
use crate::expressions::tests::assert_aggregate;
use arrow::array::*;
use datafusion_expr::AggregateFunction;

fn run_update_batch(
return_type: DataType,
arrays: &[ArrayRef],
) -> Result<(Vec<ScalarValue>, ScalarValue)> {
let agg = DistinctSum::new(vec![], String::from("__col_name__"), return_type);

let mut accum = agg.create_accumulator()?;
accum.update_batch(arrays)?;

Ok((accum.state()?, accum.evaluate()?))
}

#[test]
fn sum_distinct_update_batch() -> Result<()> {
let array_int64: ArrayRef = Arc::new(Int64Array::from(vec![1, 1, 3]));
let arrays = vec![array_int64];
let (states, result) = run_update_batch(DataType::Int64, &arrays)?;

assert_eq!(states.len(), 1);
assert_eq!(result, ScalarValue::Int64(Some(4)));

Ok(())
}

#[test]
fn sum_distinct_i32_with_nulls() {
let array = Arc::new(Int32Array::from(vec![
Some(1),
Some(1),
None,
Some(2),
Some(2),
Some(3),
]));
assert_aggregate(array, AggregateFunction::Sum, true, 6_i64.into());
}

#[test]
fn sum_distinct_u32_with_nulls() {
let array: ArrayRef = Arc::new(UInt32Array::from(vec![
Some(1_u32),
Some(1_u32),
Some(3_u32),
Some(3_u32),
None,
]));
assert_aggregate(array, AggregateFunction::Sum, true, 4_u64.into());
}

#[test]
fn sum_distinct_f64() {
let array: ArrayRef =
Arc::new(Float64Array::from(vec![1_f64, 1_f64, 3_f64, 3_f64, 3_f64]));
assert_aggregate(array, AggregateFunction::Sum, true, 4_f64.into());
}

#[test]
fn sum_distinct_decimal_with_nulls() {
let array: ArrayRef = Arc::new(
(1..6)
.map(|i| if i == 2 { None } else { Some(i % 2) })
.collect::<Decimal128Array>()
.with_precision_and_scale(35, 0)
.unwrap(),
);
assert_aggregate(
array,
AggregateFunction::Sum,
true,
ScalarValue::Decimal128(Some(1), 38, 0),
);
}
}
121 changes: 121 additions & 0 deletions datafusion/sqllogictest/test_files/aggregate.slt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1625,7 +1625,53 @@ select mean(c1) from test
----
1.75

# aggregate sum distinct, coerced result from i32 to i64
statement ok
create table t (c int) as values (1), (2), (1), (3), (null), (null), (-3), (-3);

query IT
select sum(distinct c), arrow_typeof(sum(distinct c)) from t;
----
3 Int64

statement ok
drop table t;

# aggregate sum distinct, coerced result from u32 to u64
statement ok
create table t (c int unsigned) as values (1), (2), (1), (3), (null), (null), (3);

query IT
select sum(distinct c), arrow_typeof(sum(distinct c)) from t;
----
6 UInt64

statement ok
drop table t;

# aggregate sum distinct, coerced result from f32 to f64
statement ok
create table t (c float) as values (1.0), (2.2), (1.0), (3.3), (null), (null), (3.3), (-2.0);

query RT
select sum(distinct c), arrow_typeof(sum(distinct c)) from t;
----
4.5 Float64

statement ok
drop table t;

# aggregate sum distinct with decimal
statement ok
create table t (c decimal(35, 0)) as values (1), (2), (1), (3), (null), (null), (3), (-2);

query RT
select sum(distinct c), arrow_typeof(sum(distinct c)) from t;
----
4 Decimal128(38, 0)

statement ok
drop table t;

# query_sum_distinct - 2 different aggregate functions: avg and sum(distinct)
query RI
Expand Down Expand Up @@ -2316,6 +2362,81 @@ select sum(c1), arrow_typeof(sum(c1)) from d_table;
----
100 Decimal128(20, 3)

# aggregate sum with deciaml
statement ok
create table t (c decimal(35, 3)) as values (10), (null), (20);

query RT
select sum(c), arrow_typeof(sum(c)) from t;
----
30 Decimal128(38, 3)

statement ok
drop table t;

# aggregate sum with i32, sum coerced result to i64
statement ok
create table t (c int) as values (1), (-1), (10), (null), (-11);

query IT
select sum(c), arrow_typeof(sum(c)) from t;
----
-1 Int64

statement ok
drop table t;

# aggregate sum with all nulls
statement ok
create table t (c1 decimal(10, 0), c2 int) as values (null, null), (null, null), (null, null);

query RTIT
select
sum(c1), arrow_typeof(sum(c1)),
sum(c2), arrow_typeof(sum(c2))
from t;
----
NULL Decimal128(20, 0) NULL Int64

statement ok
drop table t;

# aggregate sum with u32, sum coerced result to u64
statement ok
create table t (c int unsigned) as values (1), (0), (10), (null), (4);

query IT
select sum(c), arrow_typeof(sum(c)) from t;
----
15 UInt64

statement ok
drop table t;

# aggregate sum with f32, sum coerced result to f64
statement ok
create table t (c float) as values (1.2), (0.2), (-1.2), (null), (-1.0);

query RT
select sum(c), arrow_typeof(sum(c)) from t;
----
-0.79999999702 Float64

statement ok
drop table t;

# aggregate sum with f64
statement ok
create table t (c double) as values (1.2), (0.2), (-1.2), (null), (-1.0);

query RT
select sum(c), arrow_typeof(sum(c)) from t;
----
-0.8 Float64

statement ok
drop table t;

query TRT
select c2, sum(c1), arrow_typeof(sum(c1)) from d_table GROUP BY c2 ORDER BY c2;
----
Expand Down

0 comments on commit b412dba

Please sign in to comment.