|
15 | 15 | // specific language governing permissions and limitations |
16 | 16 | // under the License. |
17 | 17 |
|
18 | | -use arrow::array::{Scalar, new_null_array}; |
| 18 | +use arrow::array::{ArrayRef, BooleanArray, Scalar, new_null_array}; |
19 | 19 | use arrow::compute::kernels::numeric::add; |
20 | 20 | use arrow::compute::kernels::{ |
| 21 | + boolean::{and, is_not_null, or}, |
21 | 22 | cmp::{eq, lt}, |
22 | | - numeric::rem, |
| 23 | + numeric::{neg, rem}, |
23 | 24 | zip::zip, |
24 | 25 | }; |
25 | 26 | use arrow::datatypes::DataType; |
| 27 | +use arrow::error::ArrowError; |
26 | 28 | use datafusion_common::{Result, ScalarValue, assert_eq_or_internal_err}; |
27 | 29 | use datafusion_expr::{ |
28 | 30 | ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, |
29 | 31 | }; |
30 | 32 |
|
| 33 | +/// Returns a one element array holding negative zero, for the floating point |
| 34 | +/// types only. |
| 35 | +/// |
| 36 | +/// Arrow's comparison kernels order floating point values totally, so `-0.0` |
| 37 | +/// compares as distinct from, and less than, `0.0`. Java, and therefore Spark, |
| 38 | +/// treats `-0.0` as equal to zero. The helper below uses this to restore the |
| 39 | +/// IEEE 754 answer. |
| 40 | +fn negative_zero(data_type: &DataType) -> Result<Option<ArrayRef>> { |
| 41 | + match data_type { |
| 42 | + DataType::Float16 | DataType::Float32 | DataType::Float64 => { |
| 43 | + let zero = ScalarValue::new_zero(data_type)?.to_array()?; |
| 44 | + Ok(Some(neg(zero.as_ref())?)) |
| 45 | + } |
| 46 | + _ => Ok(None), |
| 47 | + } |
| 48 | +} |
| 49 | + |
| 50 | +/// Rows of `values` that equal zero, counting `-0.0` as zero. |
| 51 | +fn is_zero(values: &ArrayRef) -> Result<BooleanArray> { |
| 52 | + let zero = ScalarValue::new_zero(values.data_type())?.to_array()?; |
| 53 | + let mask = eq(values, &Scalar::new(zero))?; |
| 54 | + match negative_zero(values.data_type())? { |
| 55 | + Some(negative_zero) => Ok(or(&mask, &eq(values, &Scalar::new(negative_zero))?)?), |
| 56 | + None => Ok(mask), |
| 57 | + } |
| 58 | +} |
| 59 | + |
31 | 60 | /// Computes `rem(left, right)` with divide-by-zero handling. |
32 | | -/// In ANSI mode, any zero divisor causes an error. |
33 | | -/// In legacy mode (ANSI off), zero divisors are replaced with NULL before |
34 | | -/// computing the remainder, so those positions return NULL while others |
35 | | -/// compute normally. |
| 61 | +/// In ANSI mode, a zero divisor of any numeric type causes an error, with |
| 62 | +/// `-0.0` counting as zero; a row whose dividend is NULL never raises, to |
| 63 | +/// match Spark's null-intolerant remainder. In legacy mode (ANSI off), zero |
| 64 | +/// divisors are replaced with NULL before computing the remainder, so those |
| 65 | +/// positions return NULL while others compute normally. |
36 | 66 | fn try_rem( |
37 | | - left: &arrow::array::ArrayRef, |
38 | | - right: &arrow::array::ArrayRef, |
| 67 | + left: &ArrayRef, |
| 68 | + right: &ArrayRef, |
39 | 69 | enable_ansi_mode: bool, |
40 | | -) -> Result<arrow::array::ArrayRef> { |
| 70 | +) -> Result<ArrayRef> { |
| 71 | + let divisor_is_zero = is_zero(right)?; |
| 72 | + // Null out zero divisors so that the remainder kernels never see one: |
| 73 | + // division by zero then returns NULL instead of erroring (integers) or |
| 74 | + // returning NaN (floats). ANSI mode reports the error itself below, so |
| 75 | + // this substitution is harmless on rows that must raise. |
| 76 | + let null = Scalar::new(new_null_array(right.data_type(), 1)); |
| 77 | + let safe_right = zip(&divisor_is_zero, &null, right)?; |
41 | 78 | if enable_ansi_mode { |
42 | | - Ok(rem(left, right)?) |
43 | | - } else { |
44 | | - // In legacy mode, null out zero divisors so that division by zero |
45 | | - // returns NULL instead of erroring (integers) or returning NaN (floats). |
46 | | - let zero = ScalarValue::new_zero(right.data_type())?.to_array()?; |
47 | | - let zero = Scalar::new(zero); |
48 | | - let null = Scalar::new(new_null_array(right.data_type(), 1)); |
49 | | - let is_zero = eq(right, &zero)?; |
50 | | - let safe_right = zip(&is_zero, &null, right)?; |
51 | | - Ok(rem(left, &safe_right)?) |
| 79 | + // Spark's remainder expressions are null intolerant, so a row whose |
| 80 | + // dividend is NULL evaluates to NULL and never raises, even when the |
| 81 | + // divisor on that row is zero. Mask the check by the validity of the |
| 82 | + // dividend to match. |
| 83 | + let raises = and(&divisor_is_zero, &is_not_null(left.as_ref())?)?; |
| 84 | + if raises.iter().flatten().any(|raises| raises) { |
| 85 | + return Err(ArrowError::DivideByZero.into()); |
| 86 | + } |
52 | 87 | } |
| 88 | + Ok(rem(left, &safe_right)?) |
53 | 89 | } |
54 | 90 |
|
55 | 91 | /// Spark-compatible `mod` function |
@@ -418,6 +454,101 @@ mod test { |
418 | 454 | assert!(result.is_err()); |
419 | 455 | } |
420 | 456 |
|
| 457 | + #[test] |
| 458 | + fn test_mod_zero_division_ansi_float() { |
| 459 | + // In ANSI mode a zero divisor of any numeric type must raise, |
| 460 | + // including floating point, where Arrow's `rem` follows IEEE 754 |
| 461 | + // and quietly returns NaN (#23894) |
| 462 | + let left = Float64Array::from(vec![Some(10.5), Some(7.2)]); |
| 463 | + let right = Float64Array::from(vec![Some(0.0), Some(2.0)]); |
| 464 | + |
| 465 | + let left_value = ColumnarValue::Array(Arc::new(left)); |
| 466 | + let right_value = ColumnarValue::Array(Arc::new(right)); |
| 467 | + |
| 468 | + let result = spark_mod(&[left_value, right_value], true); |
| 469 | + assert!(result.is_err()); |
| 470 | + } |
| 471 | + |
| 472 | + #[test] |
| 473 | + fn test_mod_negative_zero_divisor_legacy() { |
| 474 | + // `-0.0` counts as a zero divisor, so it returns NULL in legacy |
| 475 | + // mode rather than NaN (#23894) |
| 476 | + let left = Float64Array::from(vec![Some(10.5), Some(7.5)]); |
| 477 | + let right = Float64Array::from(vec![Some(-0.0), Some(2.0)]); |
| 478 | + |
| 479 | + let left_value = ColumnarValue::Array(Arc::new(left)); |
| 480 | + let right_value = ColumnarValue::Array(Arc::new(right)); |
| 481 | + |
| 482 | + let result = spark_mod(&[left_value, right_value], false).unwrap(); |
| 483 | + |
| 484 | + if let ColumnarValue::Array(result_array) = result { |
| 485 | + let result_float64 = result_array |
| 486 | + .as_any() |
| 487 | + .downcast_ref::<Float64Array>() |
| 488 | + .unwrap(); |
| 489 | + assert!(result_float64.is_null(0)); // 10.5 % -0.0 = NULL |
| 490 | + assert_eq!(result_float64.value(1), 1.5); // 7.5 % 2.0 = 1.5 |
| 491 | + } else { |
| 492 | + panic!("Expected array result"); |
| 493 | + } |
| 494 | + } |
| 495 | + |
| 496 | + #[test] |
| 497 | + fn test_mod_negative_zero_divisor_ansi() { |
| 498 | + // `-0.0` counts as a zero divisor, so it raises in ANSI mode (#23894) |
| 499 | + let left = Float64Array::from(vec![Some(10.5)]); |
| 500 | + let right = Float64Array::from(vec![Some(-0.0)]); |
| 501 | + |
| 502 | + let left_value = ColumnarValue::Array(Arc::new(left)); |
| 503 | + let right_value = ColumnarValue::Array(Arc::new(right)); |
| 504 | + |
| 505 | + let result = spark_mod(&[left_value, right_value], true); |
| 506 | + assert!(result.is_err()); |
| 507 | + } |
| 508 | + |
| 509 | + #[test] |
| 510 | + fn test_mod_zero_division_ansi_null_dividend() { |
| 511 | + // Spark's remainder expressions are null intolerant: a NULL dividend |
| 512 | + // short-circuits to NULL before the divisor is validated, so a zero |
| 513 | + // divisor on such a row must not raise, even in ANSI mode (#23894) |
| 514 | + let left = Int32Array::from(vec![None, Some(10)]); |
| 515 | + let right = Int32Array::from(vec![Some(0), Some(3)]); |
| 516 | + |
| 517 | + let left_value = ColumnarValue::Array(Arc::new(left)); |
| 518 | + let right_value = ColumnarValue::Array(Arc::new(right)); |
| 519 | + |
| 520 | + let result = spark_mod(&[left_value, right_value], true).unwrap(); |
| 521 | + |
| 522 | + if let ColumnarValue::Array(result_array) = result { |
| 523 | + let result_int32 = |
| 524 | + result_array.as_any().downcast_ref::<Int32Array>().unwrap(); |
| 525 | + assert!(result_int32.is_null(0)); // NULL % 0 = NULL (no error) |
| 526 | + assert_eq!(result_int32.value(1), 1); // 10 % 3 = 1 |
| 527 | + } else { |
| 528 | + panic!("Expected array result"); |
| 529 | + } |
| 530 | + |
| 531 | + // Same for floating point |
| 532 | + let left = Float64Array::from(vec![None, Some(10.5)]); |
| 533 | + let right = Float64Array::from(vec![Some(0.0), Some(2.0)]); |
| 534 | + |
| 535 | + let left_value = ColumnarValue::Array(Arc::new(left)); |
| 536 | + let right_value = ColumnarValue::Array(Arc::new(right)); |
| 537 | + |
| 538 | + let result = spark_mod(&[left_value, right_value], true).unwrap(); |
| 539 | + |
| 540 | + if let ColumnarValue::Array(result_array) = result { |
| 541 | + let result_float64 = result_array |
| 542 | + .as_any() |
| 543 | + .downcast_ref::<Float64Array>() |
| 544 | + .unwrap(); |
| 545 | + assert!(result_float64.is_null(0)); // NULL % 0.0 = NULL (no error) |
| 546 | + assert!((result_float64.value(1) - 0.5).abs() < f64::EPSILON); // 10.5 % 2.0 = 0.5 |
| 547 | + } else { |
| 548 | + panic!("Expected array result"); |
| 549 | + } |
| 550 | + } |
| 551 | + |
421 | 552 | // PMOD tests |
422 | 553 | #[test] |
423 | 554 | fn test_pmod_int32() { |
|
0 commit comments