Skip to content

Commit e4f7d07

Browse files
tlm365tlm365
authored
Minor: Add unit tests for ceil/floor functions (#1728)
1 parent 1f3af8f commit e4f7d07

File tree

2 files changed

+326
-0
lines changed

2 files changed

+326
-0
lines changed

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

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

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

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

0 commit comments

Comments
 (0)