| // Licensed to the Apache Software Foundation (ASF) under one |
| // or more contributor license agreements. See the NOTICE file |
| // distributed with this work for additional information |
| // regarding copyright ownership. The ASF licenses this file |
| // to you under the Apache License, Version 2.0 (the |
| // "License"); you may not use this file except in compliance |
| // with the License. You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, |
| // software distributed under the License is distributed on an |
| // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| // KIND, either express or implied. See the License for the |
| // specific language governing permissions and limitations |
| // under the License. |
| |
| use arrow::array::{ |
| Array, ArrayRef, AsArray, BooleanArray, GenericListArray, Int64Array, OffsetSizeTrait, |
| }; |
| use arrow::buffer::{NullBuffer, ScalarBuffer}; |
| use arrow::datatypes::{ |
| ArrowPrimitiveType, DataType, Date32Type, Decimal128Type, Float32Type, Float64Type, Int16Type, |
| Int32Type, Int64Type, Int8Type, TimestampMicrosecondType, |
| }; |
| use datafusion::common::{exec_err, DataFusionError, Result as DataFusionResult, ScalarValue}; |
| use datafusion::logical_expr::{ |
| ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, |
| }; |
| use num::Float; |
| use std::sync::Arc; |
| |
| /// Spark array_position() function that returns the 1-based position of an element in an array. |
| /// Returns 0 if the element is not found (Spark behavior differs from DataFusion which returns null). |
| fn spark_array_position(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> { |
| if args.len() != 2 { |
| return exec_err!("array_position function takes exactly two arguments"); |
| } |
| |
| let len = args |
| .iter() |
| .fold(Option::<usize>::None, |acc, arg| match arg { |
| ColumnarValue::Scalar(_) => acc, |
| ColumnarValue::Array(a) => Some(a.len()), |
| }); |
| |
| let is_scalar = len.is_none(); |
| let arrays = ColumnarValue::values_to_arrays(args)?; |
| |
| let result = array_position_inner(&arrays)?; |
| |
| if is_scalar { |
| let scalar = ScalarValue::try_from_array(&result, 0)?; |
| Ok(ColumnarValue::Scalar(scalar)) |
| } else { |
| Ok(ColumnarValue::Array(result)) |
| } |
| } |
| |
| fn array_position_inner(args: &[ArrayRef]) -> Result<ArrayRef, DataFusionError> { |
| let array = &args[0]; |
| let element = &args[1]; |
| |
| match array.data_type() { |
| DataType::List(_) => generic_array_position::<i32>(array, element), |
| DataType::LargeList(_) => generic_array_position::<i64>(array, element), |
| other => exec_err!("array_position does not support type '{other:?}'"), |
| } |
| } |
| |
| /// Searches for an element in a list array using the flat values buffer and offsets directly, |
| /// avoiding per-row subarray allocation. Dispatches to typed fast paths by element data type. |
| fn generic_array_position<O: OffsetSizeTrait>( |
| array: &ArrayRef, |
| element: &ArrayRef, |
| ) -> Result<ArrayRef, DataFusionError> { |
| let list_array = array |
| .as_any() |
| .downcast_ref::<GenericListArray<O>>() |
| .ok_or_else(|| DataFusionError::Internal("expected list array".into()))?; |
| |
| let values = list_array.values(); |
| let offsets = list_array.offsets(); |
| let elem_type = values.data_type().clone(); |
| |
| match &elem_type { |
| DataType::Boolean => position_boolean::<O>(list_array, offsets, values, element), |
| DataType::Int8 => position_primitive::<O, Int8Type>(list_array, offsets, values, element), |
| DataType::Int16 => position_primitive::<O, Int16Type>(list_array, offsets, values, element), |
| DataType::Int32 => position_primitive::<O, Int32Type>(list_array, offsets, values, element), |
| DataType::Int64 => position_primitive::<O, Int64Type>(list_array, offsets, values, element), |
| DataType::Float32 => position_float::<O, Float32Type>(list_array, offsets, values, element), |
| DataType::Float64 => position_float::<O, Float64Type>(list_array, offsets, values, element), |
| DataType::Decimal128(_, _) => { |
| position_primitive::<O, Decimal128Type>(list_array, offsets, values, element) |
| } |
| DataType::Date32 => { |
| position_primitive::<O, Date32Type>(list_array, offsets, values, element) |
| } |
| DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, _) => { |
| position_primitive::<O, TimestampMicrosecondType>(list_array, offsets, values, element) |
| } |
| DataType::Utf8 => position_string::<O, i32>(list_array, offsets, values, element), |
| DataType::LargeUtf8 => position_string::<O, i64>(list_array, offsets, values, element), |
| // Fallback to ScalarValue for complex types (nested arrays, etc.) |
| _ => position_fallback::<O>(list_array, offsets, values, element), |
| } |
| } |
| |
| /// Compute the combined null buffer from list array and element nulls. |
| fn combined_nulls( |
| list_array_nulls: Option<&NullBuffer>, |
| element_nulls: Option<&NullBuffer>, |
| ) -> Option<NullBuffer> { |
| match (list_array_nulls, element_nulls) { |
| (Some(a), Some(b)) => NullBuffer::union(Some(a), Some(b)), |
| (Some(a), None) => Some(a.clone()), |
| (None, Some(b)) => Some(b.clone()), |
| (None, None) => None, |
| } |
| } |
| |
| /// Fast path for primitive types: downcast once, iterate using offsets into the flat buffer. |
| fn position_primitive<O: OffsetSizeTrait, T: ArrowPrimitiveType>( |
| list_array: &GenericListArray<O>, |
| offsets: &arrow::buffer::OffsetBuffer<O>, |
| values: &ArrayRef, |
| element: &ArrayRef, |
| ) -> Result<ArrayRef, DataFusionError> |
| where |
| T::Native: PartialEq, |
| { |
| let values_typed = values.as_primitive::<T>(); |
| let element_typed = element.as_primitive::<T>(); |
| let num_rows = list_array.len(); |
| let nulls = combined_nulls(list_array.nulls(), element.nulls()); |
| let mut result = vec![0i64; num_rows]; |
| |
| for (row_index, w) in offsets.windows(2).enumerate() { |
| if nulls.as_ref().is_some_and(|n| n.is_null(row_index)) { |
| continue; |
| } |
| let start = w[0].as_usize(); |
| let end = w[1].as_usize(); |
| let search_val = element_typed.value(row_index); |
| for i in start..end { |
| if !values_typed.is_null(i) && values_typed.value(i) == search_val { |
| result[row_index] = (i - start + 1) as i64; |
| break; |
| } |
| } |
| } |
| |
| Ok(Arc::new(Int64Array::new(ScalarBuffer::from(result), nulls))) |
| } |
| |
| /// Float path: same as primitive but treats NaN == NaN (Spark's ordering.equiv() semantics). |
| fn position_float<O: OffsetSizeTrait, T: ArrowPrimitiveType>( |
| list_array: &GenericListArray<O>, |
| offsets: &arrow::buffer::OffsetBuffer<O>, |
| values: &ArrayRef, |
| element: &ArrayRef, |
| ) -> Result<ArrayRef, DataFusionError> |
| where |
| T::Native: PartialEq + num::Float, |
| { |
| let values_typed = values.as_primitive::<T>(); |
| let element_typed = element.as_primitive::<T>(); |
| let num_rows = list_array.len(); |
| let nulls = combined_nulls(list_array.nulls(), element.nulls()); |
| let mut result = vec![0i64; num_rows]; |
| |
| for (row_index, w) in offsets.windows(2).enumerate() { |
| if nulls.as_ref().is_some_and(|n| n.is_null(row_index)) { |
| continue; |
| } |
| let start = w[0].as_usize(); |
| let end = w[1].as_usize(); |
| let search_val = element_typed.value(row_index); |
| let search_is_nan = search_val.is_nan(); |
| for i in start..end { |
| if !values_typed.is_null(i) { |
| let v = values_typed.value(i); |
| if (search_is_nan && v.is_nan()) || v == search_val { |
| result[row_index] = (i - start + 1) as i64; |
| break; |
| } |
| } |
| } |
| } |
| |
| Ok(Arc::new(Int64Array::new(ScalarBuffer::from(result), nulls))) |
| } |
| |
| /// Boolean path. |
| fn position_boolean<O: OffsetSizeTrait>( |
| list_array: &GenericListArray<O>, |
| offsets: &arrow::buffer::OffsetBuffer<O>, |
| values: &ArrayRef, |
| element: &ArrayRef, |
| ) -> Result<ArrayRef, DataFusionError> { |
| let values_typed = values |
| .as_any() |
| .downcast_ref::<BooleanArray>() |
| .ok_or_else(|| DataFusionError::Internal("expected boolean array".into()))?; |
| let element_typed = element |
| .as_any() |
| .downcast_ref::<BooleanArray>() |
| .ok_or_else(|| DataFusionError::Internal("expected boolean array".into()))?; |
| let num_rows = list_array.len(); |
| let nulls = combined_nulls(list_array.nulls(), element.nulls()); |
| let mut result = vec![0i64; num_rows]; |
| |
| for (row_index, w) in offsets.windows(2).enumerate() { |
| if nulls.as_ref().is_some_and(|n| n.is_null(row_index)) { |
| continue; |
| } |
| let start = w[0].as_usize(); |
| let end = w[1].as_usize(); |
| let search_val = element_typed.value(row_index); |
| for i in start..end { |
| if !values_typed.is_null(i) && values_typed.value(i) == search_val { |
| result[row_index] = (i - start + 1) as i64; |
| break; |
| } |
| } |
| } |
| |
| Ok(Arc::new(Int64Array::new(ScalarBuffer::from(result), nulls))) |
| } |
| |
| /// String path: downcast once, iterate using offsets into the flat string buffer. |
| fn position_string<O: OffsetSizeTrait, S: OffsetSizeTrait>( |
| list_array: &GenericListArray<O>, |
| offsets: &arrow::buffer::OffsetBuffer<O>, |
| values: &ArrayRef, |
| element: &ArrayRef, |
| ) -> Result<ArrayRef, DataFusionError> { |
| let values_typed = values.as_string::<S>(); |
| let element_typed = element.as_string::<S>(); |
| let num_rows = list_array.len(); |
| let nulls = combined_nulls(list_array.nulls(), element.nulls()); |
| let mut result = vec![0i64; num_rows]; |
| |
| for (row_index, w) in offsets.windows(2).enumerate() { |
| if nulls.as_ref().is_some_and(|n| n.is_null(row_index)) { |
| continue; |
| } |
| let start = w[0].as_usize(); |
| let end = w[1].as_usize(); |
| let search_val = element_typed.value(row_index); |
| for i in start..end { |
| if !values_typed.is_null(i) && values_typed.value(i) == search_val { |
| result[row_index] = (i - start + 1) as i64; |
| break; |
| } |
| } |
| } |
| |
| Ok(Arc::new(Int64Array::new(ScalarBuffer::from(result), nulls))) |
| } |
| |
| /// Fallback for complex types (nested arrays, structs, etc.) using ScalarValue comparison. |
| fn position_fallback<O: OffsetSizeTrait>( |
| list_array: &GenericListArray<O>, |
| offsets: &arrow::buffer::OffsetBuffer<O>, |
| values: &ArrayRef, |
| element: &ArrayRef, |
| ) -> Result<ArrayRef, DataFusionError> { |
| let num_rows = list_array.len(); |
| let nulls = combined_nulls(list_array.nulls(), element.nulls()); |
| let mut result = vec![0i64; num_rows]; |
| |
| for (row_index, w) in offsets.windows(2).enumerate() { |
| if nulls.as_ref().is_some_and(|n| n.is_null(row_index)) { |
| continue; |
| } |
| let start = w[0].as_usize(); |
| let end = w[1].as_usize(); |
| let search_scalar = ScalarValue::try_from_array(element, row_index)?; |
| for i in start..end { |
| if !values.is_null(i) { |
| let item_scalar = ScalarValue::try_from_array(values, i)?; |
| if search_scalar == item_scalar { |
| result[row_index] = (i - start + 1) as i64; |
| break; |
| } |
| } |
| } |
| } |
| |
| Ok(Arc::new(Int64Array::new(ScalarBuffer::from(result), nulls))) |
| } |
| |
| #[derive(Debug, Hash, Eq, PartialEq)] |
| pub struct SparkArrayPositionFunc { |
| signature: Signature, |
| } |
| |
| impl Default for SparkArrayPositionFunc { |
| fn default() -> Self { |
| Self::new() |
| } |
| } |
| |
| impl SparkArrayPositionFunc { |
| pub fn new() -> Self { |
| Self { |
| signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable), |
| } |
| } |
| } |
| |
| impl ScalarUDFImpl for SparkArrayPositionFunc { |
| fn name(&self) -> &str { |
| "spark_array_position" |
| } |
| |
| fn signature(&self) -> &Signature { |
| &self.signature |
| } |
| |
| fn return_type(&self, _arg_types: &[DataType]) -> DataFusionResult<DataType> { |
| Ok(DataType::Int64) |
| } |
| |
| fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> { |
| spark_array_position(&args.args) |
| } |
| } |
