datafusion_spark/function/string/
elt.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::sync::Arc;
19
20use arrow::array::{
21    Array, ArrayRef, AsArray, PrimitiveArray, StringArray, StringBuilder,
22};
23use arrow::compute::{can_cast_types, cast};
24use arrow::datatypes::DataType::{Int64, Utf8};
25use arrow::datatypes::{DataType, Int64Type};
26use datafusion_common::cast::as_string_array;
27use datafusion_common::{DataFusionError, Result, exec_err, plan_datafusion_err};
28use datafusion_expr::{
29    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
30};
31use datafusion_functions::utils::make_scalar_function;
32
33#[derive(Debug, PartialEq, Eq, Hash)]
34pub struct SparkElt {
35    signature: Signature,
36}
37
38impl Default for SparkElt {
39    fn default() -> Self {
40        SparkElt::new()
41    }
42}
43
44impl SparkElt {
45    pub fn new() -> Self {
46        Self {
47            signature: Signature::user_defined(Volatility::Immutable),
48        }
49    }
50}
51
52impl ScalarUDFImpl for SparkElt {
53    fn name(&self) -> &str {
54        "elt"
55    }
56
57    fn signature(&self) -> &Signature {
58        &self.signature
59    }
60
61    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
62        Ok(Utf8)
63    }
64
65    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
66        let enable_ansi_mode = args.config_options.execution.enable_ansi_mode;
67        make_scalar_function(
68            move |arrays: &[ArrayRef]| elt(arrays, enable_ansi_mode),
69            vec![],
70        )(&args.args)
71    }
72
73    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
74        let length = arg_types.len();
75        if length < 2 {
76            return Err(plan_datafusion_err!(
77                "ELT function expects at least 2 arguments: index, value1"
78            ));
79        }
80
81        let idx_dt: &DataType = &arg_types[0];
82        if *idx_dt != Int64 && !can_cast_types(idx_dt, &Int64) {
83            return Err(DataFusionError::Plan(format!(
84                "ELT index must be Int64 (or castable to Int64), got {idx_dt:?}"
85            )));
86        }
87        let mut coerced = vec![Utf8; length];
88        coerced[0] = Int64;
89        Ok(coerced)
90    }
91}
92
93fn elt(args: &[ArrayRef], enable_ansi_mode: bool) -> Result<ArrayRef> {
94    let n_rows = args[0].len();
95
96    let idx: &PrimitiveArray<Int64Type> =
97        args[0].as_primitive_opt::<Int64Type>().ok_or_else(|| {
98            DataFusionError::Plan(format!(
99                "ELT function: first argument must be Int64 (got {:?})",
100                args[0].data_type()
101            ))
102        })?;
103
104    let num_values = args.len() - 1;
105    let mut cols: Vec<StringArray> = Vec::with_capacity(num_values);
106    for a in args.iter().skip(1) {
107        let casted = cast(a, &Utf8)?;
108        cols.push(as_string_array(&casted)?.clone());
109    }
110
111    let mut builder = StringBuilder::new();
112
113    for i in 0..n_rows {
114        if idx.is_null(i) {
115            builder.append_null();
116            continue;
117        }
118
119        let index = idx.value(i);
120
121        if index < 1 || (index as usize) > num_values {
122            if enable_ansi_mode {
123                return exec_err!(
124                    "The index {index} is out of bounds. The array has {num_values} elements."
125                );
126            }
127            builder.append_null();
128            continue;
129        }
130
131        let value_idx = (index as usize) - 1;
132        let col = &cols[value_idx];
133
134        if col.is_null(i) {
135            builder.append_null();
136        } else {
137            builder.append_value(col.value(i));
138        }
139    }
140
141    Ok(Arc::new(builder.finish()) as ArrayRef)
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147    use arrow::array::Int64Array;
148
149    fn run_elt_arrays(arrs: Vec<ArrayRef>) -> Result<StringArray> {
150        run_elt_arrays_with(arrs, false)
151    }
152
153    fn run_elt_arrays_with(arrs: Vec<ArrayRef>, ansi: bool) -> Result<StringArray> {
154        let arr = elt(&arrs, ansi)?;
155        Ok(as_string_array(&arr)?.clone())
156    }
157
158    #[test]
159    fn elt_utf8_basic() -> Result<()> {
160        let idx = Arc::new(Int64Array::from(vec![
161            Some(1),
162            Some(2),
163            Some(3),
164            Some(4),
165            Some(0),
166            None,
167        ]));
168        let v1 = Arc::new(StringArray::from(vec![
169            Some("a1"),
170            Some("a2"),
171            Some("a3"),
172            Some("a4"),
173            Some("a5"),
174            Some("a6"),
175        ]));
176        let v2 = Arc::new(StringArray::from(vec![
177            Some("b1"),
178            Some("b2"),
179            None,
180            Some("b4"),
181            Some("b5"),
182            Some("b6"),
183        ]));
184        let v3 = Arc::new(StringArray::from(vec![
185            Some("c1"),
186            Some("c2"),
187            Some("c3"),
188            None,
189            Some("c5"),
190            Some("c6"),
191        ]));
192
193        let out = run_elt_arrays(vec![idx, v1, v2, v3])?;
194        assert_eq!(out.len(), 6);
195        assert_eq!(out.value(0), "a1");
196        assert_eq!(out.value(1), "b2");
197        assert_eq!(out.value(2), "c3");
198        assert!(out.is_null(3));
199        assert!(out.is_null(4));
200        assert!(out.is_null(5));
201        Ok(())
202    }
203
204    #[test]
205    fn elt_int64_basic() -> Result<()> {
206        let idx = Arc::new(Int64Array::from(vec![Some(2), Some(1), Some(2)]));
207        let v1 = Arc::new(Int64Array::from(vec![Some(10), Some(20), Some(30)]));
208        let v2 = Arc::new(Int64Array::from(vec![Some(100), None, Some(300)]));
209
210        let out = run_elt_arrays(vec![idx, v1, v2])?;
211        assert_eq!(out.len(), 3);
212        assert_eq!(out.value(0), "100");
213        assert_eq!(out.value(1), "20");
214        assert_eq!(out.value(2), "300");
215        Ok(())
216    }
217
218    #[test]
219    fn elt_out_of_range_all_null() -> Result<()> {
220        let idx = Arc::new(Int64Array::from(vec![Some(5), Some(-1), Some(0)]));
221        let v1 = Arc::new(StringArray::from(vec![Some("x"), Some("y"), Some("z")]));
222        let v2 = Arc::new(StringArray::from(vec![Some("a"), Some("b"), Some("c")]));
223
224        let out = run_elt_arrays(vec![idx, v1, v2])?;
225        assert!(out.is_null(0));
226        assert!(out.is_null(1));
227        assert!(out.is_null(2));
228        Ok(())
229    }
230
231    #[test]
232    fn elt_utf8_returns_utf8() -> Result<()> {
233        let idx = Arc::new(Int64Array::from(vec![Some(1)]));
234        let v1 = Arc::new(StringArray::from(vec![Some("scala")]));
235        let v2 = Arc::new(StringArray::from(vec![Some("java")]));
236
237        let out = run_elt_arrays(vec![idx, v1, v2])?;
238        assert_eq!(out.data_type(), &Utf8);
239        Ok(())
240    }
241}
datafusion_spark/function/string/elt.rs

datafusion_spark/function/string/
elt.rs
