native/spark-expr/src/math_funcs/checked_arithmetic.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use arrow::array::{Array, ArrowNativeTypeOp, BooleanBufferBuilder, PrimitiveArray};
use arrow::array::{ArrayRef, AsArray};

use crate::{divide_by_zero_error, EvalMode, SparkError};
use array::{Datum, Scalar};
use arrow::array;
use arrow::buffer::NullBuffer;
use arrow::compute::kernels::{arity, numeric};
use arrow::datatypes::{
    ArrowPrimitiveType, DataType, Float16Type, Float32Type, Float64Type, Int16Type, Int32Type,
    Int64Type, Int8Type,
};
use arrow::error::ArrowError;
use datafusion::common::DataFusionError;
use datafusion::common::ScalarValue;
use datafusion::physical_plan::ColumnarValue;
use std::sync::Arc;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum MathOp {
    Add,
    Sub,
    Mul,
    Div,
}

impl MathOp {
    fn function_name(self) -> &'static str {
        match self {
            Self::Add => "try_add",
            Self::Sub => "try_subtract",
            Self::Mul => "try_multiply",
            Self::Div => "try_divide",
        }
    }
}

/// Arrow's error does not carry the failing row. Only after its fast kernel reports an
/// overflow, recover the first non-null overflowing pair for Spark's Byte/Short error.
/// Keep scalar broadcasting and array offsets intact; never inspect values beneath NULLs.
fn binary_overflow_error<T: ArrowPrimitiveType>(
    left: &dyn Datum,
    right: &dyn Datum,
    op: MathOp,
) -> Option<SparkError>
where
    T::Native: std::fmt::Display,
{
    let (left, left_scalar) = left.get();
    let (right, right_scalar) = right.get();
    let left = left.as_primitive::<T>();
    let right = right.as_primitive::<T>();
    let len = if left_scalar { right.len() } else { left.len() };
    for row in 0..len {
        let l = if left_scalar { 0 } else { row };
        let r = if right_scalar { 0 } else { row };
        if left.is_null(l) || right.is_null(r) {
            continue;
        }
        let l = left.value(l);
        let r = right.value(r);
        let (result, symbol) = match op {
            MathOp::Add => (l.add_checked(r), "+"),
            MathOp::Sub => (l.sub_checked(r), "-"),
            MathOp::Mul => (l.mul_checked(r), "*"),
            MathOp::Div => return None,
        };
        if result.is_err() {
            return Some(SparkError::BinaryArithmeticOverflow {
                value1: l.to_string(),
                symbol: symbol.to_string(),
                value2: r.to_string(),
                function_name: op.function_name().to_string(),
            });
        }
    }
    None
}

fn try_arithmetic_kernel<T>(
    left: &PrimitiveArray<T>,
    right: &PrimitiveArray<T>,
    op: MathOp,
) -> Result<ArrayRef, DataFusionError>
where
    T: ArrowPrimitiveType,
{
    match op {
        MathOp::Add => checked_binary(left, right, |l, r| l.add_checked(r)),
        MathOp::Sub => checked_binary(left, right, |l, r| l.sub_checked(r)),
        MathOp::Mul => checked_binary(left, right, |l, r| l.mul_checked(r)),
        MathOp::Div => checked_binary(left, right, |l, r| l.div_checked(r)),
    }
}

fn integer_type_name(data_type: &DataType) -> Result<&'static str, DataFusionError> {
    match data_type {
        DataType::Int64 => Ok("long"),
        DataType::Int32 => Ok("integer"),
        DataType::Int16 => Ok("short"),
        DataType::Int8 => Ok("byte"),
        _ => Err(DataFusionError::Internal(format!(
            "Unsupported integer data type: {:?}",
            data_type
        ))),
    }
}

fn float_type_name(data_type: &DataType) -> Result<&'static str, DataFusionError> {
    match data_type {
        DataType::Float64 => Ok("double"),
        DataType::Float32 | DataType::Float16 => Ok("float"),
        _ => Err(DataFusionError::Internal(format!(
            "Unsupported float data type: {:?}",
            data_type
        ))),
    }
}

fn ansi_arithmetic_kernel(
    left: &ColumnarValue,
    right: &ColumnarValue,
    op: MathOp,
    data_type: &DataType,
) -> Result<ColumnarValue, DataFusionError> {
    let from_type = integer_type_name(data_type)?;
    let run_kernel = |l: &dyn Datum, r: &dyn Datum| {
        let result = match op {
            MathOp::Add => numeric::add(l, r),
            MathOp::Sub => numeric::sub(l, r),
            MathOp::Mul => numeric::mul(l, r),
            MathOp::Div => numeric::div(l, r),
        };
        result.map_err(|error| match error {
            ArrowError::DivideByZero => divide_by_zero_error().into(),
            ArrowError::ArithmeticOverflow(_) => {
                let binary_error = if op != MathOp::Div {
                    match data_type {
                        DataType::Int8 => binary_overflow_error::<Int8Type>(l, r, op),
                        DataType::Int16 => binary_overflow_error::<Int16Type>(l, r, op),
                        _ => None,
                    }
                } else {
                    None
                };
                DataFusionError::from(binary_error.unwrap_or_else(|| {
                    SparkError::ArithmeticOverflow {
                        from_type: from_type.to_string(),
                        function_name: op.function_name().to_string(),
                    }
                }))
            }
            _ => error.into(),
        })
    };

    let is_scalars = matches!(
        (left, right),
        (ColumnarValue::Scalar(_), ColumnarValue::Scalar(_))
    );

    let result_array = match (left, right) {
        (ColumnarValue::Array(l), ColumnarValue::Array(r)) => run_kernel(&l, &r),
        (ColumnarValue::Scalar(l), ColumnarValue::Array(r)) => {
            let l_arr = l.to_array()?;
            let l_scalar = Scalar::new(l_arr);
            run_kernel(&l_scalar, &r)
        }
        (ColumnarValue::Array(l), ColumnarValue::Scalar(r)) => {
            let r_arr = r.to_array()?;
            let r_scalar = Scalar::new(r_arr);
            run_kernel(&l, &r_scalar)
        }
        (ColumnarValue::Scalar(l), ColumnarValue::Scalar(r)) => {
            let l_arr = l.to_array()?;
            let r_arr = r.to_array()?;
            let l_scalar = Scalar::new(l_arr);
            let r_scalar = Scalar::new(r_arr);
            run_kernel(&l_scalar, &r_scalar)
        }
    };

    let array = result_array?;

    if is_scalars {
        let scalar_val = ScalarValue::try_from_array(array.as_ref(), 0)?;
        Ok(ColumnarValue::Scalar(scalar_val))
    } else {
        Ok(ColumnarValue::Array(array))
    }
}

fn ansi_float_div<T>(
    left: &PrimitiveArray<T>,
    right: &PrimitiveArray<T>,
) -> Result<ArrayRef, DataFusionError>
where
    T: ArrowPrimitiveType,
{
    let from_type = float_type_name(left.data_type())?;
    arity::try_binary::<_, _, _, T>(left, right, |l, r| l.div_checked(r))
        .map(|array| Arc::new(array) as ArrayRef)
        .map_err(|e| match e {
            ArrowError::DivideByZero => divide_by_zero_error().into(),
            _ => DataFusionError::from(SparkError::ArithmeticOverflow {
                from_type: String::from(from_type),
                function_name: String::new(),
            }),
        })
}

fn checked_binary<T, F>(
    left: &PrimitiveArray<T>,
    right: &PrimitiveArray<T>,
    op: F,
) -> Result<ArrayRef, DataFusionError>
where
    T: ArrowPrimitiveType,
    F: Fn(T::Native, T::Native) -> Result<T::Native, ArrowError>,
{
    let len = left.len();
    let lhs = &left.values()[..len];
    let rhs = &right.values()[..len];
    let nulls = NullBuffer::union(left.nulls(), right.nulls());

    let mut values = vec![T::Native::default(); len];
    let mut overflowed: Vec<usize> = Vec::new();

    for (i, (out, (&l, &r))) in values.iter_mut().zip(lhs.iter().zip(rhs)).enumerate() {
        match op(l, r) {
            Ok(v) => *out = v,
            Err(_) => {
                overflowed.push(i);
            }
        }
    }

    let nulls = if overflowed.is_empty() {
        nulls
    } else {
        let mut validity = BooleanBufferBuilder::new(len);
        match &nulls {
            Some(n) => validity.append_buffer(n.inner()),
            None => validity.append_n(len, true),
        }
        for i in overflowed {
            validity.set_bit(i, false);
        }
        Some(NullBuffer::new(validity.finish()))
    };

    if let Some(n) = &nulls {
        if n.null_count() > 0 {
            for (out, valid) in values.iter_mut().zip(n.iter()) {
                if !valid {
                    *out = T::Native::default();
                }
            }
        }
    }

    Ok(Arc::new(PrimitiveArray::<T>::new(values.into(), nulls)) as ArrayRef)
}

pub fn checked_add(
    args: &[ColumnarValue],
    data_type: &DataType,
    eval_mode: EvalMode,
) -> Result<ColumnarValue, DataFusionError> {
    checked_arithmetic_internal(args, data_type, MathOp::Add, eval_mode)
}

pub fn checked_sub(
    args: &[ColumnarValue],
    data_type: &DataType,
    eval_mode: EvalMode,
) -> Result<ColumnarValue, DataFusionError> {
    checked_arithmetic_internal(args, data_type, MathOp::Sub, eval_mode)
}

pub fn checked_mul(
    args: &[ColumnarValue],
    data_type: &DataType,
    eval_mode: EvalMode,
) -> Result<ColumnarValue, DataFusionError> {
    checked_arithmetic_internal(args, data_type, MathOp::Mul, eval_mode)
}

pub fn checked_div(
    args: &[ColumnarValue],
    data_type: &DataType,
    eval_mode: EvalMode,
) -> Result<ColumnarValue, DataFusionError> {
    checked_arithmetic_internal(args, data_type, MathOp::Div, eval_mode)
}

#[inline]
fn is_integer_type(data_type: &DataType) -> bool {
    matches!(
        data_type,
        DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64
    )
}

fn checked_arithmetic_internal(
    args: &[ColumnarValue],
    data_type: &DataType,
    op: MathOp,
    eval_mode: EvalMode,
) -> Result<ColumnarValue, DataFusionError> {
    let left = &args[0];
    let right = &args[1];

    let is_ansi_mode = match eval_mode {
        EvalMode::Try => false,
        EvalMode::Ansi => true,
        _ => {
            return Err(DataFusionError::Internal(format!(
                "Unsupported mode : {:?}",
                eval_mode
            )))
        }
    };

    // Early return for integer types in ANSI mode using the fast Datum/Scalar path
    if is_ansi_mode && is_integer_type(data_type) {
        return ansi_arithmetic_kernel(left, right, op, data_type);
    }

    // Materialize operands for Try-mode and float division
    let (left_arr, right_arr): (ArrayRef, ArrayRef) = match (left, right) {
        (ColumnarValue::Array(l), ColumnarValue::Array(r)) => (Arc::clone(l), Arc::clone(r)),
        (ColumnarValue::Scalar(l), ColumnarValue::Array(r)) => {
            (l.to_array_of_size(r.len())?, Arc::clone(r))
        }
        (ColumnarValue::Array(l), ColumnarValue::Scalar(r)) => {
            (Arc::clone(l), r.to_array_of_size(l.len())?)
        }
        (ColumnarValue::Scalar(l), ColumnarValue::Scalar(r)) => (l.to_array()?, r.to_array()?),
    };

    let result_array = match data_type {
        DataType::Int8 => try_arithmetic_kernel(
            left_arr.as_primitive::<Int8Type>(),
            right_arr.as_primitive::<Int8Type>(),
            op,
        ),
        DataType::Int16 => try_arithmetic_kernel(
            left_arr.as_primitive::<Int16Type>(),
            right_arr.as_primitive::<Int16Type>(),
            op,
        ),
        DataType::Int32 => try_arithmetic_kernel(
            left_arr.as_primitive::<Int32Type>(),
            right_arr.as_primitive::<Int32Type>(),
            op,
        ),
        DataType::Int64 => try_arithmetic_kernel(
            left_arr.as_primitive::<Int64Type>(),
            right_arr.as_primitive::<Int64Type>(),
            op,
        ),
        DataType::Float16 if op == MathOp::Div => {
            if is_ansi_mode {
                ansi_float_div(
                    left_arr.as_primitive::<Float16Type>(),
                    right_arr.as_primitive::<Float16Type>(),
                )
            } else {
                try_arithmetic_kernel(
                    left_arr.as_primitive::<Float16Type>(),
                    right_arr.as_primitive::<Float16Type>(),
                    op,
                )
            }
        }
        DataType::Float32 if op == MathOp::Div => {
            if is_ansi_mode {
                ansi_float_div(
                    left_arr.as_primitive::<Float32Type>(),
                    right_arr.as_primitive::<Float32Type>(),
                )
            } else {
                try_arithmetic_kernel(
                    left_arr.as_primitive::<Float32Type>(),
                    right_arr.as_primitive::<Float32Type>(),
                    op,
                )
            }
        }
        DataType::Float64 if op == MathOp::Div => {
            if is_ansi_mode {
                ansi_float_div(
                    left_arr.as_primitive::<Float64Type>(),
                    right_arr.as_primitive::<Float64Type>(),
                )
            } else {
                try_arithmetic_kernel(
                    left_arr.as_primitive::<Float64Type>(),
                    right_arr.as_primitive::<Float64Type>(),
                    op,
                )
            }
        }
        _ => Err(DataFusionError::Internal(format!(
            "Unsupported data type: {:?}",
            data_type
        ))),
    };

    Ok(ColumnarValue::Array(result_array?))
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow::array::{Float64Array, Int32Array};

    fn int32_args(left: Vec<Option<i32>>, right: Vec<Option<i32>>) -> Vec<ColumnarValue> {
        vec![
            ColumnarValue::Array(Arc::new(Int32Array::from(left))),
            ColumnarValue::Array(Arc::new(Int32Array::from(right))),
        ]
    }

    fn as_int32(value: ColumnarValue) -> Int32Array {
        let ColumnarValue::Array(array) = value else {
            unreachable!()
        };
        array.as_primitive::<Int32Type>().clone()
    }

    #[test]
    fn test_checked_add_propagates_nulls() {
        let args = int32_args(
            vec![Some(1), None, Some(3), None],
            vec![Some(10), Some(20), None, None],
        );
        let result = as_int32(checked_add(&args, &DataType::Int32, EvalMode::Ansi).unwrap());
        assert_eq!(result, Int32Array::from(vec![Some(11), None, None, None]));
    }

    #[test]
    fn test_checked_add_overflow_is_null_in_try_mode() {
        let args = int32_args(vec![Some(i32::MAX), Some(1)], vec![Some(1), Some(1)]);
        let result = as_int32(checked_add(&args, &DataType::Int32, EvalMode::Try).unwrap());
        assert_eq!(result, Int32Array::from(vec![None, Some(2)]));
    }

    #[test]
    fn test_checked_add_overflow_errors_in_ansi_mode() {
        let args = int32_args(vec![Some(i32::MAX)], vec![Some(1)]);
        let result = checked_add(&args, &DataType::Int32, EvalMode::Ansi);
        assert!(result.is_err());
        assert!(result
            .err()
            .unwrap()
            .message()
            .contains("[ARITHMETIC_OVERFLOW] integer overflow"));
    }

    #[test]
    fn test_checked_sub_and_mul_overflow() {
        let args = int32_args(vec![Some(i32::MIN), Some(5)], vec![Some(1), Some(3)]);
        let result = as_int32(checked_sub(&args, &DataType::Int32, EvalMode::Try).unwrap());
        assert_eq!(result, Int32Array::from(vec![None, Some(2)]));

        let args = int32_args(vec![Some(i32::MAX), Some(5)], vec![Some(2), Some(3)]);
        let result = as_int32(checked_mul(&args, &DataType::Int32, EvalMode::Try).unwrap());
        assert_eq!(result, Int32Array::from(vec![None, Some(15)]));
    }

    #[test]
    fn test_checked_div_by_zero() {
        let args = vec![
            ColumnarValue::Array(Arc::new(Float64Array::from(vec![Some(1.0), Some(8.0)]))),
            ColumnarValue::Array(Arc::new(Float64Array::from(vec![Some(0.0), Some(2.0)]))),
        ];
        let ColumnarValue::Array(array) =
            checked_div(&args, &DataType::Float64, EvalMode::Try).unwrap()
        else {
            unreachable!()
        };
        assert_eq!(
            array.as_primitive::<Float64Type>(),
            &Float64Array::from(vec![None, Some(4.0)])
        );

        assert!(checked_div(&args, &DataType::Float64, EvalMode::Ansi).is_err());
    }

    /// A null slot may hold any value in its underlying buffer (e.g. after a filter or a slice).
    /// Such a row must not overflow-error in ANSI mode, and its result slot must read back as a
    /// null holding the default value.
    #[test]
    fn test_null_row_with_garbage_value_does_not_error_in_ansi_mode() {
        let nulls = NullBuffer::from(vec![false, true]);
        let left = Int32Array::new(vec![i32::MAX, 1].into(), Some(nulls));
        let right = Int32Array::from(vec![Some(1), Some(1)]);
        let args = vec![
            ColumnarValue::Array(Arc::new(left)),
            ColumnarValue::Array(Arc::new(right)),
        ];

        let result = as_int32(checked_add(&args, &DataType::Int32, EvalMode::Ansi).unwrap());
        assert_eq!(result, Int32Array::from(vec![None, Some(2)]));
        assert_eq!(result.values()[0], 0);
    }

    #[test]
    fn test_ansi_integer_div_by_zero() {
        let args = int32_args(vec![Some(10)], vec![Some(0)]);
        let result = checked_div(&args, &DataType::Int32, EvalMode::Ansi);
        assert!(result.is_err());
    }

    #[test]
    fn test_ansi_scalar_operand() {
        let args = vec![
            ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(10), Some(20)]))),
            ColumnarValue::Scalar(ScalarValue::Int32(Some(5))),
        ];
        let result = as_int32(checked_add(&args, &DataType::Int32, EvalMode::Ansi).unwrap());
        assert_eq!(result, Int32Array::from(vec![Some(15), Some(25)]));
    }

    #[test]
    fn test_ansi_int64_overflow() {
        let args = vec![
            ColumnarValue::Array(Arc::new(array::Int64Array::from(vec![Some(i64::MAX)]))),
            ColumnarValue::Array(Arc::new(array::Int64Array::from(vec![Some(1)]))),
        ];
        let result = checked_add(&args, &DataType::Int64, EvalMode::Ansi);
        assert!(result.is_err());
        assert!(result
            .err()
            .unwrap()
            .message()
            .contains("[ARITHMETIC_OVERFLOW] long overflow"));
    }

    #[test]
    fn test_ansi_scalar_operands() {
        let args_right = vec![
            ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(10), Some(20)]))),
            ColumnarValue::Scalar(ScalarValue::Int32(Some(5))),
        ];
        let res_right =
            as_int32(checked_add(&args_right, &DataType::Int32, EvalMode::Ansi).unwrap());
        assert_eq!(res_right, Int32Array::from(vec![Some(15), Some(25)]));

        let args_left = vec![
            ColumnarValue::Scalar(ScalarValue::Int32(Some(5))),
            ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(10), Some(20)]))),
        ];
        let res_left = as_int32(checked_add(&args_left, &DataType::Int32, EvalMode::Ansi).unwrap());
        assert_eq!(res_left, Int32Array::from(vec![Some(15), Some(25)]));

        let args_scalar = vec![
            ColumnarValue::Scalar(ScalarValue::Int32(Some(10))),
            ColumnarValue::Scalar(ScalarValue::Int32(Some(20))),
        ];
        let res_scalar = checked_add(&args_scalar, &DataType::Int32, EvalMode::Ansi).unwrap();
        match res_scalar {
            ColumnarValue::Scalar(ScalarValue::Int32(v)) => assert_eq!(v, Some(30)),
            _ => panic!("Expected scalar result"),
        }
    }

    fn overflow_json(error: DataFusionError) -> serde_json::Value {
        let DataFusionError::External(error) = error else {
            panic!("Expected structured Spark error, got {error:?}")
        };
        let error = error.downcast_ref::<SparkError>().unwrap();
        serde_json::from_str(&error.to_json()).unwrap()
    }

    #[test]
    fn test_ansi_integral_overflow_payloads() {
        // Both overflow signs and all scalar/array shapes. In particular, subtraction must
        // preserve operand order when the scalar is on the left.
        for (data_type, min, max) in [
            (DataType::Int8, i8::MIN as i64, i8::MAX as i64),
            (DataType::Int16, i16::MIN as i64, i16::MAX as i64),
            (DataType::Int32, i32::MIN as i64, i32::MAX as i64),
            (DataType::Int64, i64::MIN, i64::MAX),
        ] {
            let value = |v| match data_type {
                DataType::Int8 => ScalarValue::Int8(Some(v as i8)),
                DataType::Int16 => ScalarValue::Int16(Some(v as i16)),
                DataType::Int32 => ScalarValue::Int32(Some(v as i32)),
                DataType::Int64 => ScalarValue::Int64(Some(v)),
                _ => unreachable!(),
            };
            for (op, symbol, function, pairs) in [
                (MathOp::Add, "+", "try_add", [(max, 1), (min, -1)]),
                (MathOp::Sub, "-", "try_subtract", [(min, 1), (max, -1)]),
                (MathOp::Mul, "*", "try_multiply", [(max, 2), (min, -1)]),
            ] {
                for (l, r) in pairs {
                    for (left_scalar, right_scalar) in
                        [(false, false), (true, false), (false, true), (true, true)]
                    {
                        let operand = |v: ScalarValue, scalar| {
                            if scalar {
                                ColumnarValue::Scalar(v)
                            } else {
                                ColumnarValue::Array(v.to_array().unwrap())
                            }
                        };
                        let args = [
                            operand(value(l), left_scalar),
                            operand(value(r), right_scalar),
                        ];
                        let error = overflow_json(
                            checked_arithmetic_internal(&args, &data_type, op, EvalMode::Ansi)
                                .unwrap_err(),
                        );
                        let expected = if matches!(data_type, DataType::Int8 | DataType::Int16) {
                            serde_json::json!({
                                "value1": l.to_string(), "symbol": symbol,
                                "value2": r.to_string(), "functionName": function
                            })
                        } else {
                            serde_json::json!({
                                "fromType": integer_type_name(&data_type).unwrap(),
                                "functionName": function
                            })
                        };
                        let class = if matches!(data_type, DataType::Int8 | DataType::Int16) {
                            "BINARY_ARITHMETIC_OVERFLOW"
                        } else {
                            "ARITHMETIC_OVERFLOW"
                        };
                        assert_eq!(error["errorClass"], class, "{data_type:?} {op:?}");
                        assert_eq!(error["params"], expected, "{data_type:?} {op:?}");
                        // TRY must still return a typed NULL for the identical inputs.
                        let result =
                            checked_arithmetic_internal(&args, &data_type, op, EvalMode::Try)
                                .unwrap();
                        assert_eq!(result.into_array(1).unwrap().null_count(), 1);
                    }
                }
            }
        }
    }

    #[test]
    fn test_ansi_wide_integer_overflow_suggestion() {
        // A separate regression so the baseline also demonstrates the missing suggestion.
        let args = int32_args(vec![Some(i32::MAX)], vec![Some(1)]);
        let error =
            overflow_json(checked_add(&args, &DataType::Int32, EvalMode::Ansi).unwrap_err());
        assert_eq!(error["params"]["functionName"], "try_add");
    }

    #[test]
    fn test_ansi_small_integer_overflow_skips_nulls_and_respects_slices() {
        macro_rules! check {
            ($ty:ty, $native:ty) => {{
                // Ignore both the row before the slice and overflowing garbage beneath NULLs.
                let left = PrimitiveArray::<$ty>::new(
                    vec![<$native>::MAX, <$native>::MAX, 1, <$native>::MAX].into(),
                    Some(NullBuffer::from(vec![true, false, true, true])),
                )
                .slice(1, 3);
                let right = PrimitiveArray::<$ty>::from_iter_values([1, 1, 2]);
                let args = [
                    ColumnarValue::Array(Arc::new(left.clone())),
                    ColumnarValue::Array(Arc::new(right)),
                ];
                let error = overflow_json(
                    checked_add(&args, left.data_type(), EvalMode::Ansi).unwrap_err(),
                );
                assert_eq!(error["params"]["value1"], <$native>::MAX.to_string());
                assert_eq!(error["params"]["value2"], "2");

                let args = [
                    ColumnarValue::Array(Arc::new(left.slice(0, 2))),
                    ColumnarValue::Array(Arc::new(PrimitiveArray::<$ty>::from_iter_values([1, 1]))),
                ];
                let result = checked_add(&args, left.data_type(), EvalMode::Ansi)
                    .unwrap()
                    .into_array(2)
                    .unwrap();
                assert_eq!(
                    result.as_primitive::<$ty>(),
                    &PrimitiveArray::<$ty>::from_iter([None, Some(2)])
                );
            }};
        }
        check!(Int8Type, i8);
        check!(Int16Type, i16);
    }

    #[test]
    fn test_ansi_small_integer_null_and_empty_scalar_broadcasts() {
        for data_type in [DataType::Int8, DataType::Int16] {
            let null = ScalarValue::try_from(&data_type).unwrap();
            let max = if data_type == DataType::Int8 {
                ScalarValue::Int8(Some(i8::MAX))
            } else {
                ScalarValue::Int16(Some(i16::MAX))
            };
            for op in [MathOp::Add, MathOp::Sub, MathOp::Mul] {
                for args in [
                    [
                        ColumnarValue::Scalar(null.clone()),
                        ColumnarValue::Array(max.to_array_of_size(3).unwrap()),
                    ],
                    [
                        ColumnarValue::Array(max.to_array_of_size(3).unwrap()),
                        ColumnarValue::Scalar(null.clone()),
                    ],
                    [
                        ColumnarValue::Scalar(max.clone()),
                        ColumnarValue::Array(max.to_array_of_size(0).unwrap()),
                    ],
                    [
                        ColumnarValue::Array(max.to_array_of_size(0).unwrap()),
                        ColumnarValue::Scalar(max.clone()),
                    ],
                    [
                        ColumnarValue::Scalar(null.clone()),
                        ColumnarValue::Scalar(max.clone()),
                    ],
                ] {
                    let result =
                        checked_arithmetic_internal(&args, &data_type, op, EvalMode::Ansi).unwrap();
                    match result {
                        ColumnarValue::Array(array) => assert_eq!(array.len(), array.null_count()),
                        ColumnarValue::Scalar(value) => assert!(value.is_null()),
                    }
                }
            }
        }
    }
}