native/spark-expr/src/array_funcs/array_position.rs

// 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::{
    make_comparator, Array, ArrayRef, AsArray, BooleanArray, GenericListArray, Int64Array,
    OffsetSizeTrait,
};
use arrow::buffer::{NullBuffer, ScalarBuffer};
use arrow::compute::SortOptions;
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 std::cmp::Ordering;
use std::sync::Arc;

use crate::float_semantics::{compare_floats, has_float_leaf, normalize_nested_floats};

/// 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 Arrow's comparator 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);
        for i in start..end {
            if !values_typed.is_null(i) {
                let v = values_typed.value(i);
                if compare_floats(v, search_val).is_eq() {
                    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 Arrow's comparator.
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];
    let values_normalized =
        has_float_leaf(values.data_type()).then(|| normalize_nested_floats(values));
    let element_normalized =
        has_float_leaf(element.data_type()).then(|| normalize_nested_floats(element));
    let comparator = make_comparator(
        values_normalized.as_ref().unwrap_or(values).as_ref(),
        element_normalized.as_ref().unwrap_or(element).as_ref(),
        SortOptions::default(),
    )?;

    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();
        for i in start..end {
            if !values.is_null(i) && comparator(i, row_index) == Ordering::Equal {
                result[row_index] = (i - start + 1) as i64;
                break;
            }
        }
    }

    Ok(Arc::new(Int64Array::new(ScalarBuffer::from(result), nulls)))
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow::array::ListArray;
    use arrow::buffer::OffsetBuffer;
    use arrow::datatypes::{Field, Float64Type};

    #[test]
    fn test_nested_float_and_null_position() -> DataFusionResult<()> {
        let values = ListArray::from_iter_primitive::<Float64Type, _, _>([
            Some(vec![Some(1.0)]),
            Some(vec![Some(f64::NAN)]),
            Some(vec![Some(-0.0)]),
            Some(vec![Some(0.0)]),
            Some(vec![Some(1.0), None]),
        ]);
        let array = ListArray::new(
            Arc::new(Field::new("item", values.data_type().clone(), true)),
            OffsetBuffer::new(vec![0, 2, 4, 5].into()),
            Arc::new(values),
            None,
        );
        let element = ListArray::from_iter_primitive::<Float64Type, _, _>([
            Some(vec![Some(f64::NAN)]),
            Some(vec![Some(0.0)]),
            Some(vec![Some(1.0), None]),
        ]);

        let result = array_position_inner(&[Arc::new(array), Arc::new(element)])?;
        let result = result.as_any().downcast_ref::<Int64Array>().unwrap();
        assert_eq!(result, &Int64Array::from(vec![2, 1, 1]));
        Ok(())
    }

    // array_position over array<struct<...>> currently falls back to Spark
    // (ArraysBase.isTypeSupported rejects StructType, see #1307), so this only
    // exercises position_fallback directly and isn't reachable from a SQL query.
    #[test]
    fn test_struct_float_field_signed_zero_position() -> DataFusionResult<()> {
        use arrow::array::{Float64Builder, StructBuilder};

        let fields = vec![Arc::new(Field::new("a", DataType::Float64, true))];
        let mut values_builder =
            StructBuilder::new(fields.clone(), vec![Box::new(Float64Builder::new())]);
        for v in [-0.0, 1.0] {
            values_builder
                .field_builder::<Float64Builder>(0)
                .unwrap()
                .append_value(v);
            values_builder.append(true);
        }
        let values = Arc::new(values_builder.finish());
        let array = ListArray::new(
            Arc::new(Field::new("item", values.data_type().clone(), true)),
            OffsetBuffer::new(vec![0, 2].into()),
            values,
            None,
        );

        let mut element_builder = StructBuilder::new(fields, vec![Box::new(Float64Builder::new())]);
        element_builder
            .field_builder::<Float64Builder>(0)
            .unwrap()
            .append_value(0.0);
        element_builder.append(true);
        let element = element_builder.finish();

        let result = array_position_inner(&[Arc::new(array), Arc::new(element)])?;
        let result = result.as_any().downcast_ref::<Int64Array>().unwrap();
        // {-0.0} is the first element and now matches {0.0}, matching Spark.
        assert_eq!(result, &Int64Array::from(vec![1]));
        Ok(())
    }
}

#[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)
    }
}