feat: support ansi for elt

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<ColumnarValue> {
-        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<Vec<DataType>> {
@@ -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<ArrayRef, DataFusionError> {
+fn elt(args: &[ArrayRef], enable_ansi_mode: bool) -> Result<ArrayRef> {
     let n_rows = args[0].len();
 
     let idx: &PrimitiveArray<Int64Type> =
@@ -103,11 +102,10 @@ fn elt(args: &[ArrayRef]) -> Result<ArrayRef, DataFusionError> {
         })?;
 
     let num_values = args.len() - 1;
-    let mut cols: Vec<Arc<StringArray>> = Vec::with_capacity(num_values);
+    let mut cols: Vec<StringArray> = 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<ArrayRef, DataFusionError> {
 
         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<ArrayRef>) -> Result<Arc<StringArray>> {
-        let arr = elt(&arrs)?;
-        let string_array = arr
-            .as_any()
-            .downcast_ref::<StringArray>()
-            .ok_or_else(|| DataFusionError::Internal("expected Utf8".into()))?;
-        Ok(Arc::new(string_array.clone()))
+    fn run_elt_arrays(arrs: Vec<ArrayRef>) -> Result<StringArray> {
+        run_elt_arrays_with(arrs, false)
+    }
+
+    fn run_elt_arrays_with(arrs: Vec<ArrayRef>, ansi: bool) -> Result<StringArray> {
+        let arr = elt(&arrs, ansi)?;
+        Ok(as_string_array(&arr)?.clone())
     }
 
     #[test]

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;