Jefffrey commented on code in PR #20461:
URL: https://github.com/apache/datafusion/pull/20461#discussion_r2887682226


##########
datafusion/spark/src/function/math/modulus.rs:
##########
@@ -15,37 +15,78 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use arrow::array::new_null_array;
 use arrow::compute::kernels::numeric::add;
-use arrow::compute::kernels::{cmp::lt, numeric::rem, zip::zip};
+use arrow::compute::kernels::{
+    cmp::{eq, lt},
+    numeric::rem,
+    zip::zip,
+};
 use arrow::datatypes::DataType;
 use datafusion_common::{Result, ScalarValue, assert_eq_or_internal_err};
 use datafusion_expr::{
     ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
 };
 use std::any::Any;
 
+/// Attempts `rem(left, right)` with per-element divide-by-zero handling.
+/// In ANSI mode, any zero divisor causes an error.
+/// In legacy mode (ANSI off), positions where the divisor is zero return NULL
+/// while other positions compute normally.
+fn try_rem(
+    left: &arrow::array::ArrayRef,
+    right: &arrow::array::ArrayRef,
+    enable_ansi_mode: bool,
+) -> Result<arrow::array::ArrayRef> {
+    match rem(left, right) {
+        Ok(result) => Ok(result),
+        Err(arrow::error::ArrowError::DivideByZero) if !enable_ansi_mode => {
+            // Integer rem fails when ANY divisor element is zero.
+            // Handle per-element: replace zeros with 1, compute rem, then 
null out zero positions.
+            let zero = ScalarValue::new_zero(right.data_type())?
+                .to_array_of_size(right.len())?;
+            let is_zero = eq(right, &zero)?;
+            let one =
+                
ScalarValue::new_one(right.data_type())?.to_array_of_size(right.len())?;
+            let safe_right = zip(&is_zero, &one, right)?;
+            let result = rem(left, &safe_right)?;
+            let null_array = new_null_array(result.data_type(), result.len());
+            Ok(zip(&is_zero, &null_array, &result)?)
+        }

Review Comment:
   ```suggestion
           Err(arrow::error::ArrowError::DivideByZero) if !enable_ansi_mode => {
               // Integer rem fails when ANY divisor element is zero.
               // Handle per-element: null out zero divisors
               let zero = ScalarValue::new_zero(right.data_type())?.to_array()?;
               let zero = Scalar::new(zero);
               let null = Scalar::new(new_null_array(right.data_type(), 1));
               let is_zero = eq(right, &zero)?;
               let safe_right = zip(&is_zero, &null, right)?;
               Ok(rem(left, &safe_right)?)
           }
   ```
   
   We don't to zip in the ones, we can simply null out the zeros as `rem` 
doesn't perform the division on null elements anyway.
   
   - Unfortunately we don't have a way to replace the null buffer given just an 
`ArrayRef` so it seems this zip approach really is the only way



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to