Function spark_array_repeat

datafusion/spark/src/function/array/repeat.rs:92–121

This is a Spark-specific wrapper around DataFusion's array_repeat that returns NULL if the count argument is NULL (Spark behavior), whereas DataFusion's array_repeat ignores NULLs.

(args: ScalarFunctionArgs)

/// This is a Spark-specific wrapper around DataFusion's array_repeat that returns NULL
/// if the count argument is NULL (Spark behavior), whereas DataFusion's array_repeat ignores NULLs.
fn spark_array_repeat(args: ScalarFunctionArgs) -> Result<ColumnarValue> {
    let ScalarFunctionArgs {
        args: arg_values,
        arg_fields,
        number_rows,
        return_field,
        config_options,
    } = args;
    let return_type = return_field.data_type().clone();

    // A NULL element should be repeated into the array, not cause a NULL result.
    let null_mask = compute_null_mask(&arg_values[1..]);

    // If count is null then return NULL immediately
    if matches!(null_mask, NullMaskResolution::ReturnNull) {
        return Ok(ColumnarValue::Scalar(ScalarValue::try_from(return_type)?));
    }

    let array_repeat_func = ArrayRepeat::new();
    let func_args = ScalarFunctionArgs {
        args: arg_values,
        arg_fields,
        number_rows,
        return_field,
        config_options,
    };
    let result = array_repeat_func.invoke_with_args(func_args)?;

    apply_null_mask(result, null_mask, &return_type)
}