diff --git a/datafusion/spark/src/function/string/elt.rs b/datafusion/spark/src/function/string/elt.rs index e58faf0c40f93..b88477a7720f3 100644 --- a/datafusion/spark/src/function/string/elt.rs +++ b/datafusion/spark/src/function/string/elt.rs @@ -24,7 +24,7 @@ use arrow::compute::{can_cast_types, cast}; use arrow::datatypes::DataType::{Int64, Utf8}; use arrow::datatypes::{DataType, Int64Type}; use datafusion_common::cast::as_string_array; -use datafusion_common::{DataFusionError, Result, plan_datafusion_err}; +use datafusion_common::{DataFusionError, Result, exec_err, plan_datafusion_err}; use datafusion_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; @@ -63,7 +63,11 @@ impl ScalarUDFImpl for SparkElt { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(elt, vec![])(&args.args) + let enable_ansi_mode = args.config_options.execution.enable_ansi_mode; + make_scalar_function( + move |arrays: &[ArrayRef]| elt(arrays, enable_ansi_mode), + vec![], + )(&args.args) } fn coerce_types(&self, arg_types: &[DataType]) -> Result> { @@ -80,18 +84,13 @@ impl ScalarUDFImpl for SparkElt { "ELT index must be Int64 (or castable to Int64), got {idx_dt:?}" ))); } - let mut coerced = Vec::with_capacity(arg_types.len()); - coerced.push(Int64); - - for _ in 1..length { - coerced.push(Utf8); - } - + let mut coerced = vec![Utf8; length]; + coerced[0] = Int64; Ok(coerced) } } -fn elt(args: &[ArrayRef]) -> Result { +fn elt(args: &[ArrayRef], enable_ansi_mode: bool) -> Result { let n_rows = args[0].len(); let idx: &PrimitiveArray = @@ -103,11 +102,10 @@ fn elt(args: &[ArrayRef]) -> Result { })?; let num_values = args.len() - 1; - let mut cols: Vec> = Vec::with_capacity(num_values); + let mut cols: Vec = Vec::with_capacity(num_values); for a in args.iter().skip(1) { let casted = cast(a, &Utf8)?; - let sa = as_string_array(&casted)?; - cols.push(Arc::new(sa.clone())); + cols.push(as_string_array(&casted)?.clone()); } let mut builder = StringBuilder::new(); @@ -120,10 +118,12 @@ fn elt(args: &[ArrayRef]) -> Result { let index = idx.value(i); - // TODO: if spark.sql.ansi.enabled is true, - // throw ArrayIndexOutOfBoundsException for invalid indices; - // if false, return NULL instead (current behavior). if index < 1 || (index as usize) > num_values { + if enable_ansi_mode { + return exec_err!( + "The index {index} is out of bounds. The array has {num_values} elements." + ); + } builder.append_null(); continue; } @@ -146,13 +146,13 @@ mod tests { use super::*; use arrow::array::Int64Array; - fn run_elt_arrays(arrs: Vec) -> Result> { - let arr = elt(&arrs)?; - let string_array = arr - .as_any() - .downcast_ref::() - .ok_or_else(|| DataFusionError::Internal("expected Utf8".into()))?; - Ok(Arc::new(string_array.clone())) + fn run_elt_arrays(arrs: Vec) -> Result { + run_elt_arrays_with(arrs, false) + } + + fn run_elt_arrays_with(arrs: Vec, ansi: bool) -> Result { + let arr = elt(&arrs, ansi)?; + Ok(as_string_array(&arr)?.clone()) } #[test] diff --git a/datafusion/sqllogictest/test_files/spark/string/elt.slt b/datafusion/sqllogictest/test_files/spark/string/elt.slt index 12917d17e1e47..9f0348324aadc 100644 --- a/datafusion/sqllogictest/test_files/spark/string/elt.slt +++ b/datafusion/sqllogictest/test_files/spark/string/elt.slt @@ -59,3 +59,143 @@ query T SELECT elt(1, 10, null) ---- 10 + +######################################## +# ANSI mode = false (default): invalid indices return NULL +######################################## + +# Index 0 -> NULL (Spark returns NULL when ANSI is off) +query T +SELECT elt(0::int, 'a', 'b'); +---- +NULL + +# Negative index -> NULL +query T +SELECT elt(-1::int, 'a', 'b'); +---- +NULL + +# Index far beyond the input list -> NULL +query T +SELECT elt(100::int, 'a', 'b', 'c'); +---- +NULL + +# NULL index -> NULL regardless of mode +query T +SELECT elt(NULL::int, 'a', 'b'); +---- +NULL + +# NULL value at the selected index -> NULL +query T +SELECT elt(2::int, 'a', NULL); +---- +NULL + +# Three-argument list, pick middle element +query T +SELECT elt(2::int, 'scala', 'java', 'python'); +---- +java + +# Three-argument list, pick last element +query T +SELECT elt(3::int, 'scala', 'java', 'python'); +---- +python + +# Mixed types get cast to string (Spark returns string) +query T +SELECT elt(2::int, 1, 2, 3); +---- +2 + +# Vectorized: mix of valid, out-of-range, and NULL indices in ANSI-off mode +statement ok +CREATE TABLE elt_rows(idx INT, a STRING, b STRING, c STRING) AS VALUES + (1, 'a1', 'b1', 'c1'), + (2, 'a2', 'b2', 'c2'), + (3, 'a3', 'b3', 'c3'), + (0, 'a4', 'b4', 'c4'), + (-1, 'a5', 'b5', 'c5'), + (4, 'a6', 'b6', 'c6'), + (NULL, 'a7', 'b7', 'c7'); + +query T +SELECT elt(idx, a, b, c) FROM elt_rows ORDER BY a; +---- +a1 +b2 +c3 +NULL +NULL +NULL +NULL + +statement ok +DROP TABLE elt_rows; + +######################################## +# ANSI mode = true: invalid indices raise ArrayIndexOutOfBoundsException +######################################## + +statement ok +set datafusion.execution.enable_ansi_mode = true; + +# Valid indices still work +query T +SELECT elt(1::int, 'scala', 'java'); +---- +scala + +query T +SELECT elt(2::int, 'scala', 'java'); +---- +java + +# NULL index still returns NULL (matches Spark: no error when index itself is NULL) +query T +SELECT elt(NULL::int, 'a', 'b'); +---- +NULL + +# NULL value at valid index still returns NULL (only invalid indices error) +query T +SELECT elt(1::int, NULL, 'b'); +---- +NULL + +# Out-of-range positive index errors +statement error DataFusion error: Execution error: The index 3 is out of bounds\. The array has 2 elements\. +SELECT elt(3::int, 'scala', 'java'); + +# Zero index errors +statement error DataFusion error: Execution error: The index 0 is out of bounds\. The array has 2 elements\. +SELECT elt(0::int, 'scala', 'java'); + +# Negative index errors +statement error DataFusion error: Execution error: The index -1 is out of bounds\. The array has 2 elements\. +SELECT elt(-1::int, 'scala', 'java'); + +# Large positive index errors +statement error DataFusion error: Execution error: The index 100 is out of bounds\. The array has 3 elements\. +SELECT elt(100::int, 'a', 'b', 'c'); + +# Vectorized: a batch that contains any invalid index errors in ANSI mode +statement ok +CREATE TABLE elt_ansi(idx INT, a STRING, b STRING) AS VALUES + (1, 'a1', 'b1'), + (2, 'a2', 'b2'), + (3, 'a3', 'b3'); + +statement error DataFusion error: Execution error: The index 3 is out of bounds\. The array has 2 elements\. +SELECT elt(idx, a, b) FROM elt_ansi; + +statement ok +DROP TABLE elt_ansi; + +# Reset ANSI mode +statement ok +set datafusion.execution.enable_ansi_mode = false;