feat: optimize array_position with typed array comparison and address review feedback

andygrove · andygrove · commit 6cf8f6a597d4 · 2026-03-16T12:25:59.000-06:00
- Use typed array downcasting instead of ScalarValue for element
  comparison, improving performance from 0.4X to 0.7-0.8X of Spark
- Add getSupportLevel override marking as Incompatible (NaN equality)
- Add NaN edge case tests for float/double arrays
- Add CometArrayExpressionBenchmark microbenchmark
- Make spark_array_position function private
- Update docs to mark array_position as supported
diff --git a/docs/spark_expressions_support.md b/docs/spark_expressions_support.md
@@ -93,7 +93,7 @@
 - [x] array_join
 - [x] array_max
 - [ ] array_min
-- [ ] array_position
+- [x] array_position
 - [x] array_remove
 - [x] array_repeat
 - [x] array_union
diff --git a/native/spark-expr/src/array_funcs/array_position.rs b/native/spark-expr/src/array_funcs/array_position.rs
@@ -15,8 +15,13 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::array::{Array, ArrayRef, GenericListArray, Int64Array, OffsetSizeTrait};
-use arrow::datatypes::DataType;
+use arrow::array::{
+    Array, ArrayRef, AsArray, BooleanArray, GenericListArray, Int64Array, OffsetSizeTrait,
+};
+use arrow::datatypes::{
+    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, Volatility,
@@ -26,7 +31,7 @@ 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).
-pub fn spark_array_position(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
+fn spark_array_position(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
     if args.len() != 2 {
         return exec_err!("array_position function takes exactly two arguments");
     }
@@ -63,6 +68,105 @@ fn array_position_inner(args: &[ArrayRef]) -> Result<ArrayRef, DataFusionError>
     }
 }
 
+/// Find the 1-based position of `search_val` in a typed primitive array.
+/// Returns 0 if not found.
+macro_rules! find_position_primitive {
+    ($list_items:expr, $element:expr, $row_index:expr, $arrow_type:ty) => {{
+        let items = $list_items.as_primitive::<$arrow_type>();
+        let search = $element.as_primitive::<$arrow_type>();
+        let search_val = search.value($row_index);
+        let mut pos: i64 = 0;
+        for i in 0..items.len() {
+            if !items.is_null(i) && items.value(i) == search_val {
+                pos = (i + 1) as i64;
+                break;
+            }
+        }
+        pos
+    }};
+}
+
+fn find_position_in_row(
+    list_items: &ArrayRef,
+    element: &ArrayRef,
+    row_index: usize,
+) -> Result<i64, DataFusionError> {
+    let pos = match list_items.data_type() {
+        DataType::Boolean => {
+            let items = list_items.as_any().downcast_ref::<BooleanArray>().unwrap();
+            let search = element.as_any().downcast_ref::<BooleanArray>().unwrap();
+            let search_val = search.value(row_index);
+            let mut pos: i64 = 0;
+            for i in 0..items.len() {
+                if !items.is_null(i) && items.value(i) == search_val {
+                    pos = (i + 1) as i64;
+                    break;
+                }
+            }
+            pos
+        }
+        DataType::Int8 => find_position_primitive!(list_items, element, row_index, Int8Type),
+        DataType::Int16 => find_position_primitive!(list_items, element, row_index, Int16Type),
+        DataType::Int32 => find_position_primitive!(list_items, element, row_index, Int32Type),
+        DataType::Int64 => find_position_primitive!(list_items, element, row_index, Int64Type),
+        DataType::Float32 => {
+            find_position_primitive!(list_items, element, row_index, Float32Type)
+        }
+        DataType::Float64 => {
+            find_position_primitive!(list_items, element, row_index, Float64Type)
+        }
+        DataType::Decimal128(_, _) => {
+            find_position_primitive!(list_items, element, row_index, Decimal128Type)
+        }
+        DataType::Date32 => {
+            find_position_primitive!(list_items, element, row_index, Date32Type)
+        }
+        DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, _) => {
+            find_position_primitive!(list_items, element, row_index, TimestampMicrosecondType)
+        }
+        DataType::Utf8 => {
+            let items = list_items.as_string::<i32>();
+            let search = element.as_string::<i32>();
+            let search_val = search.value(row_index);
+            let mut pos: i64 = 0;
+            for i in 0..items.len() {
+                if !items.is_null(i) && items.value(i) == search_val {
+                    pos = (i + 1) as i64;
+                    break;
+                }
+            }
+            pos
+        }
+        DataType::LargeUtf8 => {
+            let items = list_items.as_string::<i64>();
+            let search = element.as_string::<i64>();
+            let search_val = search.value(row_index);
+            let mut pos: i64 = 0;
+            for i in 0..items.len() {
+                if !items.is_null(i) && items.value(i) == search_val {
+                    pos = (i + 1) as i64;
+                    break;
+                }
+            }
+            pos
+        }
+        // Fallback to ScalarValue for complex types (nested arrays, etc.)
+        _ => {
+            let element_scalar = ScalarValue::try_from_array(element, row_index)?;
+            let mut pos: i64 = 0;
+            for i in 0..list_items.len() {
+                let item_scalar = ScalarValue::try_from_array(list_items, i)?;
+                if !item_scalar.is_null() && element_scalar == item_scalar {
+                    pos = (i + 1) as i64;
+                    break;
+                }
+            }
+            pos
+        }
+    };
+    Ok(pos)
+}
+
 fn generic_array_position<O: OffsetSizeTrait>(
     array: &ArrayRef,
     element: &ArrayRef,
@@ -75,30 +179,11 @@ fn generic_array_position<O: OffsetSizeTrait>(
     let mut data = Vec::with_capacity(list_array.len());
 
     for row_index in 0..list_array.len() {
-        if list_array.is_null(row_index) {
-            // Null array returns null position (same as Spark)
-            data.push(None);
-        } else if element.is_null(row_index) {
-            // Searching for null element returns null in Spark
+        if list_array.is_null(row_index) || element.is_null(row_index) {
             data.push(None);
         } else {
             let list_array_row = list_array.value(row_index);
-
-            // Get the search element as a scalar
-            let element_scalar = ScalarValue::try_from_array(element, row_index)?;
-
-            // Compare element to each item in the list
-            let mut position: i64 = 0;
-            for i in 0..list_array_row.len() {
-                let list_item_scalar = ScalarValue::try_from_array(&list_array_row, i)?;
-
-                // null != anything in Spark array_position
-                if !list_item_scalar.is_null() && element_scalar == list_item_scalar {
-                    position = (i + 1) as i64; // 1-indexed
-                    break;
-                }
-            }
-
+            let position = find_position_in_row(&list_array_row, element, row_index)?;
             data.push(Some(position));
         }
     }
diff --git a/native/spark-expr/src/array_funcs/mod.rs b/native/spark-expr/src/array_funcs/mod.rs
@@ -22,7 +22,7 @@ mod list_extract;
 mod size;
 
 pub use array_insert::ArrayInsert;
-pub use array_position::{spark_array_position, SparkArrayPositionFunc};
+pub use array_position::SparkArrayPositionFunc;
 pub use get_array_struct_fields::GetArrayStructFields;
 pub use list_extract::ListExtract;
 pub use size::{spark_size, SparkSizeFunc};
diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala
@@ -664,6 +664,8 @@ object CometSize extends CometExpressionSerde[Size] {
 
 object CometArrayPosition extends CometExpressionSerde[ArrayPosition] with ArraysBase {
 
+  override def getSupportLevel(expr: ArrayPosition): SupportLevel = Incompatible(None)
+
   override def convert(
       expr: ArrayPosition,
       inputs: Seq[Attribute],
diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_position.sql b/spark/src/test/resources/sql-tests/expressions/array/array_position.sql
@@ -15,6 +15,7 @@
 -- specific language governing permissions and limitations
 -- under the License.
 
+-- Config: spark.comet.expression.ArrayPosition.allowIncompatible=true
 -- ConfigMatrix: parquet.enable.dictionary=false,true
 
 statement
@@ -163,6 +164,33 @@ INSERT INTO test_ap_double VALUES
 query
 SELECT array_position(arr, val) FROM test_ap_double
 
+-- NaN handling for float arrays (Spark treats NaN == NaN)
+query spark_answer_only
+SELECT array_position(array(cast('NaN' as float), cast(1.0 as float)), cast('NaN' as float))
+
+query spark_answer_only
+SELECT array_position(array(cast(1.0 as float), cast('NaN' as float), cast(2.0 as float)), cast('NaN' as float))
+
+-- NaN handling for double arrays (Spark treats NaN == NaN)
+query spark_answer_only
+SELECT array_position(array(cast('NaN' as double), 1.0), cast('NaN' as double))
+
+query spark_answer_only
+SELECT array_position(array(1.0, cast('NaN' as double), 2.0), cast('NaN' as double))
+
+-- NaN handling with column data
+statement
+CREATE TABLE test_ap_nan(arr array<float>, val float) USING parquet
+
+statement
+INSERT INTO test_ap_nan VALUES
+  (array(cast('NaN' as float), cast(1.0 as float)), cast('NaN' as float)),
+  (array(cast(1.0 as float), cast('NaN' as float), cast(2.0 as float)), cast('NaN' as float)),
+  (array(cast(1.0 as float), cast(2.0 as float)), cast('NaN' as float))
+
+query ignore(NaN equality: IEEE 754 says NaN != NaN but Spark treats NaN == NaN)
+SELECT array_position(arr, val) FROM test_ap_nan
+
 -- decimal arrays
 statement
 CREATE TABLE test_ap_decimal(arr array<decimal(10,2)>, val decimal(10,2)) USING parquet
diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometArrayExpressionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometArrayExpressionBenchmark.scala
@@ -0,0 +1,106 @@
+/*
+ * 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.
+ */
+
+package org.apache.spark.sql.benchmark
+
+// spotless:off
+/**
+ * Benchmark to measure performance of Comet array expressions. To run this benchmark:
+ * {{{
+ *   SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometArrayExpressionBenchmark
+ * }}}
+ * Results will be written to "spark/benchmarks/CometArrayExpressionBenchmark-**results.txt".
+ */
+// spotless:on
+object CometArrayExpressionBenchmark extends CometBenchmarkBase {
+
+  def arrayPositionBenchmark(values: Int): Unit = {
+    withTempPath { dir =>
+      withTempTable("parquetV1Table") {
+        // Create a table with int arrays of size 10 and a search value
+        prepareTable(
+          dir,
+          spark.sql(s"""SELECT
+               |  array(
+               |    cast(value % 100 as int),
+               |    cast((value + 1) % 100 as int),
+               |    cast((value + 2) % 100 as int),
+               |    cast((value + 3) % 100 as int),
+               |    cast((value + 4) % 100 as int),
+               |    cast((value + 5) % 100 as int),
+               |    cast((value + 6) % 100 as int),
+               |    cast((value + 7) % 100 as int),
+               |    cast((value + 8) % 100 as int),
+               |    cast((value + 9) % 100 as int)
+               |  ) as int_arr,
+               |  cast((value + 5) % 100 as int) as search_val
+               |FROM $tbl""".stripMargin))
+
+        val extraConfigs =
+          Map("spark.comet.expression.ArrayPosition.allowIncompatible" -> "true")
+
+        runExpressionBenchmark(
+          "array_position - int array",
+          values,
+          "SELECT array_position(int_arr, search_val) FROM parquetV1Table",
+          extraConfigs)
+      }
+    }
+
+    withTempPath { dir =>
+      withTempTable("parquetV1Table") {
+        // Create a table with string arrays of size 10 and a search value
+        prepareTable(
+          dir,
+          spark.sql(s"""SELECT
+               |  array(
+               |    cast(value % 100 as string),
+               |    cast((value + 1) % 100 as string),
+               |    cast((value + 2) % 100 as string),
+               |    cast((value + 3) % 100 as string),
+               |    cast((value + 4) % 100 as string),
+               |    cast((value + 5) % 100 as string),
+               |    cast((value + 6) % 100 as string),
+               |    cast((value + 7) % 100 as string),
+               |    cast((value + 8) % 100 as string),
+               |    cast((value + 9) % 100 as string)
+               |  ) as str_arr,
+               |  cast((value + 5) % 100 as string) as search_val
+               |FROM $tbl""".stripMargin))
+
+        val extraConfigs =
+          Map("spark.comet.expression.ArrayPosition.allowIncompatible" -> "true")
+
+        runExpressionBenchmark(
+          "array_position - string array",
+          values,
+          "SELECT array_position(str_arr, search_val) FROM parquetV1Table",
+          extraConfigs)
+      }
+    }
+  }
+
+  override def runCometBenchmark(mainArgs: Array[String]): Unit = {
+    val values = 1024 * 1024
+
+    runBenchmarkWithTable("ArrayPosition", values) { v =>
+      arrayPositionBenchmark(v)
+    }
+  }
+}
