Function spark_round

datafusion/spark/src/function/math/round.rs:421–654

(args: &[ColumnarValue], enable_ansi_mode: bool)

// ---------------------------------------------------------------------------

fn spark_round(args: &[ColumnarValue], enable_ansi_mode: bool) -> Result<ColumnarValue> {
    if args.is_empty() || args.len() > 2 {
        return exec_err!("round requires 1 or 2 arguments, got {}", args.len());
    }

    let scale = match get_scale(args)? {
        Some(s) => s,
        None => {
            // NULL scale → return NULL with the same data type as the first argument
            return Ok(ColumnarValue::Scalar(ScalarValue::try_from(
                args[0].data_type(),
            )?));
        }
    };

    match &args[0] {
        ColumnarValue::Array(array) => match array.data_type() {
            DataType::Null => Ok(args[0].clone()),

            // Integer types
            DataType::Int8 => {
                impl_integer_array_round!(array, Int8Type, scale, enable_ansi_mode)
            }
            DataType::Int16 => {
                impl_integer_array_round!(array, Int16Type, scale, enable_ansi_mode)
            }
            DataType::Int32 => {
                impl_integer_array_round!(array, Int32Type, scale, enable_ansi_mode)
            }
            DataType::Int64 => {
                impl_integer_array_round!(array, Int64Type, scale, enable_ansi_mode)
            }

            // Unsigned integer types
            DataType::UInt8 => {
                impl_integer_array_round!(array, UInt8Type, scale, enable_ansi_mode)
            }
            DataType::UInt16 => {
                impl_integer_array_round!(array, UInt16Type, scale, enable_ansi_mode)
            }
            DataType::UInt32 => {
                impl_integer_array_round!(array, UInt32Type, scale, enable_ansi_mode)
            }
            DataType::UInt64 => {
                let array = array.as_primitive::<UInt64Type>();
                let result: PrimitiveArray<UInt64Type> = array.try_unary(|x| {
                    let v_i64 = i64::try_from(x).map_err(|_| {
                        (exec_err!(
                            "round: UInt64 value {x} exceeds i64::MAX and cannot be rounded"
                        ) as Result<(), _>)
                            .unwrap_err()
                    })?;
                    round_integer(v_i64, scale, enable_ansi_mode)
                        .map(|v| v as u64)
                })?;
                Ok(ColumnarValue::Array(Arc::new(result)))
            }