diff --git a/datafusion/spark/src/function/array/shuffle.rs b/datafusion/spark/src/function/array/shuffle.rs index eaeff6538..8051825ac 100644 --- a/datafusion/spark/src/function/array/shuffle.rs +++ b/datafusion/spark/src/function/array/shuffle.rs @@ -105,11 +105,8 @@ impl ScalarUDFImpl for SparkShuffle { &self, args: datafusion_expr::ScalarFunctionArgs, ) -> Result { - if args.args.is_empty() { - return exec_err!("shuffle expects at least 1 argument"); - } - if args.args.len() > 2 { - return exec_err!("shuffle expects at most 2 arguments"); + if args.args.is_empty() || args.args.len() > 2 { + return exec_err!("shuffle expects 1 or 2 argument(s)"); } // Extract seed from second argument if present @@ -131,10 +128,10 @@ fn extract_seed(seed_arg: &ColumnarValue) -> Result> { ColumnarValue::Scalar(scalar) => { let seed = match scalar { ScalarValue::Int64(Some(v)) => Some(*v as u64), - ScalarValue::Null => None, + ScalarValue::Null | ScalarValue::Int64(None) => None, _ => { return exec_err!( - "shuffle seed must be Int64 type, got '{}'", + "shuffle seed must be Int64 type but got '{}'", scalar.data_type() ); } @@ -164,7 +161,10 @@ fn array_shuffle_with_seed(arg: &[ArrayRef], seed: Option) -> Result Ok(Arc::clone(input_array)), - array_type => exec_err!("shuffle does not support type '{array_type}'."), + array_type => exec_err!( + "shuffle does not support type '{array_type}'; \ + expected types: List, LargeList, FixedSizeList or Null." + ), } } diff --git a/datafusion/sqllogictest/test_files/spark/array/shuffle.slt b/datafusion/sqllogictest/test_files/spark/array/shuffle.slt index 5e4865332..01d319b61 100644 --- a/datafusion/sqllogictest/test_files/spark/array/shuffle.slt +++ b/datafusion/sqllogictest/test_files/spark/array/shuffle.slt @@ -107,6 +107,16 @@ SELECT shuffle([1, 2, 3, 4], CAST('2' AS INT)); ---- [1, 4, 2, 3] +query ? +SELECT shuffle(['ab'], NULL); +---- +[ab] + +query ? +SELECT shuffle(shuffle([3, 3], NULL), NULL); +---- +[3, 3] + # Clean up statement ok DROP TABLE test_shuffle_list_types;