Skip to content

Commit

Permalink
move median test
Browse files Browse the repository at this point in the history 
Signed-off-by: jayzhan211 <[email protected]>
  • Loading branch information
@jayzhan211
jayzhan211 committed May 22, 2024
1 parent 65e281a commit b1df1b7
Show file tree
Hide file tree
Showing 2 changed files with 156 additions and 384 deletions.
384 changes: 0 additions & 384 deletions datafusion/physical-expr/src/aggregate/median.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -295,387 +295,3 @@ fn calculate_median<T: ArrowNumericType>(
Some(*median)
}
}

#[cfg(test)]
mod tests {
use super::*;
use crate::expressions::col;
use crate::expressions::tests::aggregate;
use crate::generic_test_distinct_op;
use arrow::{array::*, datatypes::*};

#[test]
fn median_decimal() -> Result<()> {
// test median
let array: ArrayRef = Arc::new(
(1..7)
.map(Some)
.collect::<Decimal128Array>()
.with_precision_and_scale(10, 4)?,
);

generic_test_distinct_op!(
array,
DataType::Decimal128(10, 4),
Median,
false,
ScalarValue::Decimal128(Some(3), 10, 4)
)
}

#[test]
fn median_decimal_with_nulls() -> Result<()> {
let array: ArrayRef = Arc::new(
(1..6)
.map(|i| if i == 2 { None } else { Some(i) })
.collect::<Decimal128Array>()
.with_precision_and_scale(10, 4)?,
);
generic_test_distinct_op!(
array,
DataType::Decimal128(10, 4),
Median,
false,
ScalarValue::Decimal128(Some(3), 10, 4)
)
}

#[test]
fn median_decimal_all_nulls() -> Result<()> {
// test median
let array: ArrayRef = Arc::new(
std::iter::repeat::<Option<i128>>(None)
.take(6)
.collect::<Decimal128Array>()
.with_precision_and_scale(10, 4)?,
);
generic_test_distinct_op!(
array,
DataType::Decimal128(10, 4),
Median,
false,
ScalarValue::Decimal128(None, 10, 4)
)
}

#[test]
fn median_i32_odd() -> Result<()> {
let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]));
generic_test_distinct_op!(
a,
DataType::Int32,
Median,
false,
ScalarValue::from(3_i32)
)
}

#[test]
fn median_i32_even() -> Result<()> {
let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6]));
generic_test_distinct_op!(
a,
DataType::Int32,
Median,
false,
ScalarValue::from(3_i32)
)
}

#[test]
fn median_i32_with_nulls() -> Result<()> {
let a: ArrayRef = Arc::new(Int32Array::from(vec![
Some(1),
None,
Some(3),
Some(4),
Some(5),
]));
generic_test_distinct_op!(
a,
DataType::Int32,
Median,
false,
ScalarValue::from(3i32)
)
}

#[test]
fn median_i32_all_nulls() -> Result<()> {
let a: ArrayRef = Arc::new(Int32Array::from(vec![None, None]));
generic_test_distinct_op!(
a,
DataType::Int32,
Median,
false,
ScalarValue::Int32(None)
)
}

#[test]
fn median_u32_odd() -> Result<()> {
let a: ArrayRef =
Arc::new(UInt32Array::from(vec![1_u32, 2_u32, 3_u32, 4_u32, 5_u32]));
generic_test_distinct_op!(
a,
DataType::UInt32,
Median,
false,
ScalarValue::from(3u32)
)
}

#[test]
fn median_u32_even() -> Result<()> {
let a: ArrayRef = Arc::new(UInt32Array::from(vec![
1_u32, 2_u32, 3_u32, 4_u32, 5_u32, 6_u32,
]));
generic_test_distinct_op!(
a,
DataType::UInt32,
Median,
false,
ScalarValue::from(3u32)
)
}

#[test]
fn median_f32_odd() -> Result<()> {
let a: ArrayRef =
Arc::new(Float32Array::from(vec![1_f32, 2_f32, 3_f32, 4_f32, 5_f32]));
generic_test_distinct_op!(
a,
DataType::Float32,
Median,
false,
ScalarValue::from(3_f32)
)
}

#[test]
fn median_f32_even() -> Result<()> {
let a: ArrayRef = Arc::new(Float32Array::from(vec![
1_f32, 2_f32, 3_f32, 4_f32, 5_f32, 6_f32,
]));
generic_test_distinct_op!(
a,
DataType::Float32,
Median,
false,
ScalarValue::from(3.5_f32)
)
}

#[test]
fn median_f64_odd() -> Result<()> {
let a: ArrayRef =
Arc::new(Float64Array::from(vec![1_f64, 2_f64, 3_f64, 4_f64, 5_f64]));
generic_test_distinct_op!(
a,
DataType::Float64,
Median,
false,
ScalarValue::from(3_f64)
)
}

#[test]
fn median_f64_even() -> Result<()> {
let a: ArrayRef = Arc::new(Float64Array::from(vec![
1_f64, 2_f64, 3_f64, 4_f64, 5_f64, 6_f64,
]));
generic_test_distinct_op!(
a,
DataType::Float64,
Median,
false,
ScalarValue::from(3.5_f64)
)
}

#[test]
fn distinct_median_decimal() -> Result<()> {
let array: ArrayRef = Arc::new(
vec![1, 1, 1, 1, 2, 3, 1, 1, 3]
.into_iter()
.map(Some)
.collect::<Decimal128Array>()
.with_precision_and_scale(10, 4)?,
);

generic_test_distinct_op!(
array,
DataType::Decimal128(10, 4),
Median,
true,
ScalarValue::Decimal128(Some(2), 10, 4)
)
}

#[test]
fn distinct_median_decimal_with_nulls() -> Result<()> {
let array: ArrayRef = Arc::new(
vec![Some(3), Some(1), None, Some(3), Some(2), Some(3), Some(3)]
.into_iter()
.collect::<Decimal128Array>()
.with_precision_and_scale(10, 4)?,
);
generic_test_distinct_op!(
array,
DataType::Decimal128(10, 4),
Median,
true,
ScalarValue::Decimal128(Some(2), 10, 4)
)
}

#[test]
fn distinct_median_decimal_all_nulls() -> Result<()> {
let array: ArrayRef = Arc::new(
std::iter::repeat::<Option<i128>>(None)
.take(6)
.collect::<Decimal128Array>()
.with_precision_and_scale(10, 4)?,
);
generic_test_distinct_op!(
array,
DataType::Decimal128(10, 4),
Median,
true,
ScalarValue::Decimal128(None, 10, 4)
)
}

#[test]
fn distinct_median_i32_odd() -> Result<()> {
let a: ArrayRef = Arc::new(Int32Array::from(vec![2, 1, 1, 2, 1, 3]));
generic_test_distinct_op!(
a,
DataType::Int32,
Median,
true,
ScalarValue::from(2_i32)
)
}

#[test]
fn distinct_median_i32_even() -> Result<()> {
let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 1, 3, 1, 1]));
generic_test_distinct_op!(
a,
DataType::Int32,
Median,
true,
ScalarValue::from(2_i32)
)
}

#[test]
fn distinct_median_i32_with_nulls() -> Result<()> {
let a: ArrayRef = Arc::new(Int32Array::from(vec![
Some(1),
None,
Some(1),
Some(1),
Some(3),
]));
generic_test_distinct_op!(
a,
DataType::Int32,
Median,
true,
ScalarValue::from(2i32)
)
}

#[test]
fn distinct_median_i32_all_nulls() -> Result<()> {
let a: ArrayRef = Arc::new(Int32Array::from(vec![None, None]));
generic_test_distinct_op!(
a,
DataType::Int32,
Median,
true,
ScalarValue::Int32(None)
)
}

#[test]
fn distinct_median_u32_odd() -> Result<()> {
let a: ArrayRef =
Arc::new(UInt32Array::from(vec![1_u32, 1_u32, 2_u32, 1_u32, 3_u32]));
generic_test_distinct_op!(
a,
DataType::UInt32,
Median,
true,
ScalarValue::from(2u32)
)
}

#[test]
fn distinct_median_u32_even() -> Result<()> {
let a: ArrayRef = Arc::new(UInt32Array::from(vec![
1_u32, 1_u32, 1_u32, 1_u32, 3_u32, 3_u32,
]));
generic_test_distinct_op!(
a,
DataType::UInt32,
Median,
true,
ScalarValue::from(2u32)
)
}

#[test]
fn distinct_median_f32_odd() -> Result<()> {
let a: ArrayRef =
Arc::new(Float32Array::from(vec![3_f32, 2_f32, 1_f32, 1_f32, 1_f32]));
generic_test_distinct_op!(
a,
DataType::Float32,
Median,
true,
ScalarValue::from(2_f32)
)
}

#[test]
fn distinct_median_f32_even() -> Result<()> {
let a: ArrayRef =
Arc::new(Float32Array::from(vec![1_f32, 1_f32, 1_f32, 1_f32, 2_f32]));
generic_test_distinct_op!(
a,
DataType::Float32,
Median,
true,
ScalarValue::from(1.5_f32)
)
}

#[test]
fn distinct_median_f64_odd() -> Result<()> {
let a: ArrayRef =
Arc::new(Float64Array::from(vec![1_f64, 1_f64, 1_f64, 2_f64, 3_f64]));
generic_test_distinct_op!(
a,
DataType::Float64,
Median,
true,
ScalarValue::from(2_f64)
)
}

#[test]
fn distinct_median_f64_even() -> Result<()> {
let a: ArrayRef =
Arc::new(Float64Array::from(vec![1_f64, 1_f64, 1_f64, 1_f64, 2_f64]));
generic_test_distinct_op!(
a,
DataType::Float64,
Median,
true,
ScalarValue::from(1.5_f64)
)
}
}
Loading

0 comments on commit b1df1b7

Please sign in to comment.