Skip to content

Commit ce60579

Browse files
author
Tai Le Manh
committed
Minor: Add unit tests for ceil/floor functions
up
1 parent 25e39ab commit ce60579

File tree

2 files changed

+308
-0
lines changed

2 files changed

+308
-0
lines changed

‎native/spark-expr/src/math_funcs/ceil.rs

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,3 +81,157 @@ fn decimal_ceil_f(scale: &i8) -> impl Fn(i128) -> i128 {
8181
let div = 10_i128.pow_wrapping(*scale as u32);
8282
move |x: i128| div_ceil(x, div)
8383
}
84+
85+
#[cfg(test)]
86+
mod test {
87+
use crate::spark_ceil;
88+
use arrow::array::{Decimal128Array, Float32Array, Float64Array, Int64Array};
89+
use arrow::datatypes::DataType;
90+
use datafusion::common::cast::as_int64_array;
91+
use datafusion::common::{Result, ScalarValue};
92+
use datafusion::physical_plan::ColumnarValue;
93+
use std::sync::Arc;
94+
95+
#[test]
96+
fn test_ceil_f32_array() -> Result<()> {
97+
let input = Float32Array::from(vec![
98+
Some(125.2345),
99+
Some(15.0001),
100+
Some(0.1),
101+
Some(-0.9),
102+
Some(-1.1),
103+
Some(123.0),
104+
None,
105+
]);
106+
let args = vec![ColumnarValue::Array(Arc::new(input))];
107+
let ColumnarValue::Array(result) = spark_ceil(&args, &DataType::Float32)? else {
108+
unreachable!()
109+
};
110+
let actual = as_int64_array(&result)?;
111+
let expected = Int64Array::from(vec![
112+
Some(126),
113+
Some(16),
114+
Some(1),
115+
Some(0),
116+
Some(-1),
117+
Some(123),
118+
None,
119+
]);
120+
assert_eq!(actual, &expected);
121+
Ok(())
122+
}
123+
124+
#[test]
125+
fn test_ceil_f64_array() -> Result<()> {
126+
let input = Float64Array::from(vec![
127+
Some(125.2345),
128+
Some(15.0001),
129+
Some(0.1),
130+
Some(-0.9),
131+
Some(-1.1),
132+
Some(123.0),
133+
None,
134+
]);
135+
let args = vec![ColumnarValue::Array(Arc::new(input))];
136+
let ColumnarValue::Array(result) = spark_ceil(&args, &DataType::Float64)? else {
137+
unreachable!()
138+
};
139+
let actual = as_int64_array(&result)?;
140+
let expected = Int64Array::from(vec![
141+
Some(126),
142+
Some(16),
143+
Some(1),
144+
Some(0),
145+
Some(-1),
146+
Some(123),
147+
None,
148+
]);
149+
assert_eq!(actual, &expected);
150+
Ok(())
151+
}
152+
153+
#[test]
154+
fn test_ceil_i64_array() -> Result<()> {
155+
let input = Int64Array::from(vec![Some(-1), Some(0), Some(1), None]);
156+
let args = vec![ColumnarValue::Array(Arc::new(input))];
157+
let ColumnarValue::Array(result) = spark_ceil(&args, &DataType::Int64)? else {
158+
unreachable!()
159+
};
160+
let actual = as_int64_array(&result)?;
161+
let expected = Int64Array::from(vec![Some(-1), Some(0), Some(1), None]);
162+
assert_eq!(actual, &expected);
163+
Ok(())
164+
}
165+
166+
#[test]
167+
fn test_ceil_decimal128_array() -> Result<()> {
168+
let array = Decimal128Array::from(vec![
169+
Some(12345), // 123.45
170+
Some(12500), // 125.00
171+
Some(-12999), // -129.99
172+
None,
173+
])
174+
.with_precision_and_scale(5, 2)?;
175+
let args = vec![ColumnarValue::Array(Arc::new(array))];
176+
let ColumnarValue::Array(result) = spark_ceil(&args, &DataType::Decimal128(5, 2))? else {
177+
unreachable!()
178+
};
179+
let expected = Decimal128Array::from(vec![Some(12400), Some(12500), Some(-12900), None])
180+
.with_precision_and_scale(5, 2)?;
181+
let actual = result.as_any().downcast_ref::<Decimal128Array>().unwrap();
182+
assert_eq!(actual, &expected);
183+
Ok(())
184+
}
185+
186+
#[test]
187+
fn test_ceil_f32_scalar() -> Result<()> {
188+
let args = vec![ColumnarValue::Scalar(ScalarValue::Float32(Some(125.2345)))];
189+
let ColumnarValue::Scalar(ScalarValue::Int64(Some(result))) =
190+
spark_ceil(&args, &DataType::Float32)?
191+
else {
192+
unreachable!()
193+
};
194+
assert_eq!(result, 126);
195+
Ok(())
196+
}
197+
198+
#[test]
199+
fn test_ceil_f64_scalar() -> Result<()> {
200+
let args = vec![ColumnarValue::Scalar(ScalarValue::Float64(Some(-1.1)))];
201+
let ColumnarValue::Scalar(ScalarValue::Int64(Some(result))) =
202+
spark_ceil(&args, &DataType::Float64)?
203+
else {
204+
unreachable!()
205+
};
206+
assert_eq!(result, -1);
207+
Ok(())
208+
}
209+
210+
#[test]
211+
fn test_ceil_i64_scalar() -> Result<()> {
212+
let args = vec![ColumnarValue::Scalar(ScalarValue::Int64(Some(48)))];
213+
let ColumnarValue::Scalar(ScalarValue::Int64(Some(result))) =
214+
spark_ceil(&args, &DataType::Int64)?
215+
else {
216+
unreachable!()
217+
};
218+
assert_eq!(result, 48);
219+
Ok(())
220+
}
221+
222+
#[test]
223+
fn test_ceil_decimal128_scalar() -> Result<()> {
224+
let args = vec![ColumnarValue::Scalar(ScalarValue::Decimal128(
225+
Some(567),
226+
3,
227+
1,
228+
))]; // 56.7
229+
let ColumnarValue::Scalar(ScalarValue::Decimal128(Some(result), 3, 1)) =
230+
spark_ceil(&args, &DataType::Decimal128(3, 1))?
231+
else {
232+
unreachable!()
233+
};
234+
assert_eq!(result, 570); // 57.0
235+
Ok(())
236+
}
237+
}

‎native/spark-expr/src/math_funcs/floor.rs

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,3 +81,157 @@ fn decimal_floor_f(scale: &i8) -> impl Fn(i128) -> i128 {
8181
let div = 10_i128.pow_wrapping(*scale as u32);
8282
move |x: i128| div_floor(x, div)
8383
}
84+
85+
#[cfg(test)]
86+
mod test {
87+
use crate::spark_floor;
88+
use arrow::array::{Decimal128Array, Float32Array, Float64Array, Int64Array};
89+
use arrow::datatypes::DataType;
90+
use datafusion::common::cast::as_int64_array;
91+
use datafusion::common::{Result, ScalarValue};
92+
use datafusion::physical_plan::ColumnarValue;
93+
use std::sync::Arc;
94+
95+
#[test]
96+
fn test_floor_f32_array() -> Result<()> {
97+
let input = Float32Array::from(vec![
98+
Some(125.9345),
99+
Some(15.9999),
100+
Some(0.9),
101+
Some(-0.1),
102+
Some(-1.999),
103+
Some(123.0),
104+
None,
105+
]);
106+
let args = vec![ColumnarValue::Array(Arc::new(input))];
107+
let ColumnarValue::Array(result) = spark_floor(&args, &DataType::Float32)? else {
108+
unreachable!()
109+
};
110+
let actual = as_int64_array(&result)?;
111+
let expected = Int64Array::from(vec![
112+
Some(125),
113+
Some(15),
114+
Some(0),
115+
Some(-1),
116+
Some(-2),
117+
Some(123),
118+
None,
119+
]);
120+
assert_eq!(actual, &expected);
121+
Ok(())
122+
}
123+
124+
#[test]
125+
fn test_floor_f64_array() -> Result<()> {
126+
let input = Float64Array::from(vec![
127+
Some(125.9345),
128+
Some(15.9999),
129+
Some(0.9),
130+
Some(-0.1),
131+
Some(-1.999),
132+
Some(123.0),
133+
None,
134+
]);
135+
let args = vec![ColumnarValue::Array(Arc::new(input))];
136+
let ColumnarValue::Array(result) = spark_floor(&args, &DataType::Float64)? else {
137+
unreachable!()
138+
};
139+
let actual = as_int64_array(&result)?;
140+
let expected = Int64Array::from(vec![
141+
Some(125),
142+
Some(15),
143+
Some(0),
144+
Some(-1),
145+
Some(-2),
146+
Some(123),
147+
None,
148+
]);
149+
assert_eq!(actual, &expected);
150+
Ok(())
151+
}
152+
153+
#[test]
154+
fn test_floor_i64_array() -> Result<()> {
155+
let input = Int64Array::from(vec![Some(-1), Some(0), Some(1), None]);
156+
let args = vec![ColumnarValue::Array(Arc::new(input))];
157+
let ColumnarValue::Array(result) = spark_floor(&args, &DataType::Int64)? else {
158+
unreachable!()
159+
};
160+
let actual = as_int64_array(&result)?;
161+
let expected = Int64Array::from(vec![Some(-1), Some(0), Some(1), None]);
162+
assert_eq!(actual, &expected);
163+
Ok(())
164+
}
165+
166+
#[test]
167+
fn test_floor_decimal128_array() -> Result<()> {
168+
let array = Decimal128Array::from(vec![
169+
Some(12345), // 123.45
170+
Some(12500), // 125.00
171+
Some(-12999), // -129.99
172+
None,
173+
])
174+
.with_precision_and_scale(5, 2)?;
175+
let args = vec![ColumnarValue::Array(Arc::new(array))];
176+
let ColumnarValue::Array(result) = spark_floor(&args, &DataType::Decimal128(5, 2))? else {
177+
unreachable!()
178+
};
179+
let expected = Decimal128Array::from(vec![Some(12300), Some(12500), Some(-13000), None])
180+
.with_precision_and_scale(5, 2)?;
181+
let actual = result.as_any().downcast_ref::<Decimal128Array>().unwrap();
182+
assert_eq!(actual, &expected);
183+
Ok(())
184+
}
185+
186+
#[test]
187+
fn test_floor_f32_scalar() -> Result<()> {
188+
let args = vec![ColumnarValue::Scalar(ScalarValue::Float32(Some(125.9345)))];
189+
let ColumnarValue::Scalar(ScalarValue::Int64(Some(result))) =
190+
spark_floor(&args, &DataType::Float32)?
191+
else {
192+
unreachable!()
193+
};
194+
assert_eq!(result, 125);
195+
Ok(())
196+
}
197+
198+
#[test]
199+
fn test_floor_f64_scalar() -> Result<()> {
200+
let args = vec![ColumnarValue::Scalar(ScalarValue::Float64(Some(-1.999)))];
201+
let ColumnarValue::Scalar(ScalarValue::Int64(Some(result))) =
202+
spark_floor(&args, &DataType::Float64)?
203+
else {
204+
unreachable!()
205+
};
206+
assert_eq!(result, -2);
207+
Ok(())
208+
}
209+
210+
#[test]
211+
fn test_floor_i64_scalar() -> Result<()> {
212+
let args = vec![ColumnarValue::Scalar(ScalarValue::Int64(Some(48)))];
213+
let ColumnarValue::Scalar(ScalarValue::Int64(Some(result))) =
214+
spark_floor(&args, &DataType::Int64)?
215+
else {
216+
unreachable!()
217+
};
218+
assert_eq!(result, 48);
219+
Ok(())
220+
}
221+
222+
#[test]
223+
fn test_floor_decimal128_scalar() -> Result<()> {
224+
let args = vec![ColumnarValue::Scalar(ScalarValue::Decimal128(
225+
Some(567),
226+
3,
227+
1,
228+
))]; // 56.7
229+
let ColumnarValue::Scalar(ScalarValue::Decimal128(Some(result), 3, 1)) =
230+
spark_floor(&args, &DataType::Decimal128(3, 1))?
231+
else {
232+
unreachable!()
233+
};
234+
assert_eq!(result, 560); // 56.0
235+
Ok(())
236+
}
237+
}

0 commit comments

Comments
 (0)