andygrove commented on code in PR #63:
URL: https://github.com/apache/datafusion-spark/pull/63#discussion_r4566763715
##########
crates/sqllogictest/test_files/array/array_position.slt:
##########
@@ -0,0 +1,197 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under
one or more contributor license agreements
+
+# Tests for Spark-compatible array_position.
+# Returns the 1-based position of the first match, 0 if not found, NULL if the
array or element is NULL.
+# Expected values verified with pyspark==4.2.0.
+
+query I
+SELECT array_position(array(1, 2, 3, 4), 3);
+----
+3
+
+# Not found returns 0 (DataFusion's built-in returns NULL)
+query I
+SELECT array_position(array(1, 2, 3, 4), 5);
+----
+0
+
+query I
+SELECT array_position(array('a', 'b', 'c'), 'b');
+----
+2
+
+# NULL elements in the array do not match and still count towards the position
+query I
+SELECT array_position(array(1, 2, NULL, 3), 3);
+----
+4
+
+query I
+SELECT array_position(array(1, 2, NULL, 3), 5);
+----
+0
+
+# First occurrence wins
+query I
+SELECT array_position(array(1, 2, 1, 3), 1);
+----
+1
+
+query I
+SELECT array_position(array(1, 2, 3), CAST(NULL AS INT));
+----
+NULL
+
+query I
+SELECT array_position(CAST(NULL AS ARRAY<INT>), 1);
+----
+NULL
+
+query I
+SELECT array_position(array(1, 2, 3), NULL);
Review Comment:
I checked this in PySpark 4.2.0, and `array_position(array(1, 2, 3), NULL)`
returns NULL there: the untyped NULL element is converted to the array's
element type before the type check. `NULL_TYPE` applies when the array argument
is an untyped NULL, and `array_position(NULL, 1)` fails with
`DATATYPE_MISMATCH.NULL_TYPE`. We error there too, so I kept this case with a
comment and added a `query error` test for `array_position(NULL, 1)` in 98787b5.
##########
crates/spark/src/function/array/array_position.rs:
##########
@@ -0,0 +1,219 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+use arrow::array::{
+ Array, ArrayRef, AsArray, GenericListArray, Int64Array, OffsetSizeTrait,
+ make_comparator,
+};
+use arrow::buffer::{NullBuffer, ScalarBuffer};
+use arrow::compute::SortOptions;
+use arrow::datatypes::{DataType, Float32Type, Float64Type};
+use datafusion_common::{Result, ScalarValue, exec_err, not_impl_err};
+use datafusion_expr::{
+ ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
+ Volatility,
+};
+use datafusion_macros::user_doc;
+use std::sync::Arc;
+
+/// Spark-compatible `array_position` expression.
+///
+/// <https://spark.apache.org/docs/latest/api/sql/index.html#array_position>
+#[user_doc(
+ doc_section(label = "Array Functions"),
+ description = "Returns the 1-based position of the first occurrence of the
element in the \
+ array, as a `BIGINT`. Unlike the DataFusion built-in, which returns
NULL when the element \
+ is not found, this returns 0. The result is NULL only when the array
or the element is \
+ NULL. NULL elements inside the array never match. For floating point
values, NaN is \
+ equal to NaN and -0.0 is equal to 0.0.",
+ syntax_example = "array_position(array, element)",
+ sql_example = r#"```sql
+> SELECT array_position(array(3, 2, 1), 1);
+3
+> SELECT array_position(array(3, 2, 1), 4);
+0
+```"#,
+ argument(name = "array", description = "Array expression. NULL returns
NULL."),
+ argument(
+ name = "element",
+ description = "Element to search for. NULL returns NULL."
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SparkArrayPosition {
+ signature: Signature,
+}
+
+impl Default for SparkArrayPosition {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SparkArrayPosition {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::array_and_element(Volatility::Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for SparkArrayPosition {
+ fn name(&self) -> &str {
+ "array_position"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
Result<ColumnarValue> {
+ let is_scalar = args
+ .args
+ .iter()
+ .all(|a| matches!(a, ColumnarValue::Scalar(_)));
+ let arrays = ColumnarValue::values_to_arrays(&args.args)?;
+ let [array, element] = arrays.as_slice() else {
+ return exec_err!("array_position expects exactly two arguments");
+ };
+
+ let result = match array.data_type() {
+ DataType::List(_) => position::<i32>(array, element)?,
+ DataType::LargeList(_) => position::<i64>(array, element)?,
+ other => return exec_err!("array_position does not support type
{other}"),
+ };
+
+ if is_scalar {
+ Ok(ColumnarValue::Scalar(ScalarValue::try_from_array(
+ &result, 0,
+ )?))
+ } else {
+ Ok(ColumnarValue::Array(result))
+ }
+ }
+
+ fn documentation(&self) -> Option<&Documentation> {
+ self.doc()
+ }
+}
+
+type Equality = Box<dyn Fn(usize, usize) -> bool + Send + Sync>;
+
+fn has_float(data_type: &DataType) -> bool {
+ match data_type {
+ DataType::Float32 | DataType::Float64 => true,
+ DataType::List(f) | DataType::LargeList(f) |
DataType::FixedSizeList(f, _) => {
+ has_float(f.data_type())
+ }
+ DataType::Struct(fields) => fields.iter().any(|f|
has_float(f.data_type())),
+ DataType::Map(f, _) => has_float(f.data_type()),
+ _ => false,
+ }
+}
+
+/// Wraps a non-null equality so that NULL equals NULL and NULL differs from
any value, which
+/// is how Spark compares values nested inside arrays and structs.
+fn null_aware(
+ left: &ArrayRef,
+ right: &ArrayRef,
+ eq: impl Fn(usize, usize) -> bool + Send + Sync + 'static,
+) -> Equality {
+ let (left, right) = (Arc::clone(left), Arc::clone(right));
+ Box::new(move |i, j| match (left.is_null(i), right.is_null(j)) {
+ (true, true) => true,
+ (false, false) => eq(i, j),
+ _ => false,
+ })
+}
+
+/// Builds `eq(left_index, right_index)` in Spark's equality, where NaN equals
NaN and -0.0
+/// equals 0.0 at any depth. Subtrees without a float use Arrow's comparator.
+fn equality(left: &ArrayRef, right: &ArrayRef) -> Result<Equality> {
Review Comment:
Done in 98787b5. `normalize_array` is now `pub(super)`. `position`
normalizes values and element once per batch and uses `make_comparator` for
every type, which removes the custom equality code and adds support for
fixed-size lists of floats.
##########
crates/sqllogictest/test_files/array/array_position.slt:
##########
@@ -0,0 +1,197 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under
one or more contributor license agreements
+
+# Tests for Spark-compatible array_position.
+# Returns the 1-based position of the first match, 0 if not found, NULL if the
array or element is NULL.
+# Expected values verified with pyspark==4.2.0.
+
+query I
+SELECT array_position(array(1, 2, 3, 4), 3);
+----
+3
+
+# Not found returns 0 (DataFusion's built-in returns NULL)
+query I
+SELECT array_position(array(1, 2, 3, 4), 5);
+----
+0
+
+query I
+SELECT array_position(array('a', 'b', 'c'), 'b');
+----
+2
+
+# NULL elements in the array do not match and still count towards the position
+query I
+SELECT array_position(array(1, 2, NULL, 3), 3);
+----
+4
+
+query I
+SELECT array_position(array(1, 2, NULL, 3), 5);
+----
+0
+
+# First occurrence wins
+query I
+SELECT array_position(array(1, 2, 1, 3), 1);
+----
+1
+
+query I
+SELECT array_position(array(1, 2, 3), CAST(NULL AS INT));
+----
+NULL
+
+query I
+SELECT array_position(CAST(NULL AS ARRAY<INT>), 1);
+----
+NULL
+
+query I
+SELECT array_position(array(1, 2, 3), NULL);
+----
+NULL
+
+query I
+SELECT array_position(array(), 1);
+----
+0
+
+# NaN equals NaN
+query I
+SELECT array_position(array(CAST('NaN' AS DOUBLE), CAST(1.0 AS DOUBLE)),
CAST('NaN' AS DOUBLE));
+----
+1
+
+query I
+SELECT array_position(array(CAST('NaN' AS FLOAT), CAST(1.0 AS FLOAT)),
CAST('NaN' AS FLOAT));
+----
+1
+
+query I
+SELECT array_position(array(true, false), false);
+----
+2
+
+query I
+SELECT array_position(array(CAST(1.5 AS DECIMAL(3,1)), CAST(2.5 AS
DECIMAL(3,1))), CAST(2.5 AS DECIMAL(3,1)));
+----
+2
+
+query I
+SELECT array_position(array(DATE '2020-01-01', DATE '2021-01-01'), DATE
'2021-01-01');
+----
+2
+
+query I
+SELECT array_position(array(X'01', X'02'), X'02');
+----
+2
+
+# Nested arrays
+query I
+SELECT array_position(array(array(1, 2), array(3)), array(3));
+----
+2
+
+# Element is coerced to the array's element type
+query I
+SELECT array_position(array(1, 2, 3), 2L);
+----
+2
+
+query I
+SELECT array_position(array(1, 2, 3), CAST(2.0 AS DOUBLE));
+----
+2
+
+# Signed zero: -0.0 equals 0.0
+query I
+SELECT array_position(array(CAST('-0.0' AS DOUBLE), CAST(1.0 AS DOUBLE)),
CAST(0.0 AS DOUBLE));
+----
+1
+
+query I
+SELECT array_position(array(CAST(0.0 AS DOUBLE), CAST(1.0 AS DOUBLE)),
CAST('-0.0' AS DOUBLE));
+----
+1
+
+query I
+SELECT array_position(array(CAST(1.0 AS DOUBLE), CAST('-0.0' AS DOUBLE)),
CAST(0.0 AS DOUBLE));
+----
+2
+
+query I
+SELECT array_position(array(CAST('-0.0' AS FLOAT), CAST(1.0 AS FLOAT)),
CAST(0.0 AS FLOAT));
+----
+1
+
+# Nested floats follow Spark's ordering: NaN equals NaN, -0.0 equals 0.0
+query I
+SELECT array_position(array(array(CAST('NaN' AS DOUBLE)), array(CAST(1.0 AS
DOUBLE))), array(CAST('NaN' AS DOUBLE)));
+----
+1
+
+query I
+SELECT array_position(array(array(CAST('-0.0' AS DOUBLE)), array(CAST(1.0 AS
DOUBLE))), array(CAST(0.0 AS DOUBLE)));
+----
+1
+
+query I
+SELECT array_position(array(array(CAST('-0.0' AS FLOAT)), array(CAST(1.0 AS
FLOAT))), array(CAST(0.0 AS FLOAT)));
+----
+1
+
+query I
+SELECT array_position(array(named_struct('a', CAST('NaN' AS DOUBLE)),
named_struct('a', CAST(1.0 AS DOUBLE))), named_struct('a', CAST('NaN' AS
DOUBLE)));
+----
+1
+
+query I
+SELECT array_position(array(named_struct('a', CAST('-0.0' AS DOUBLE)),
named_struct('a', CAST(1.0 AS DOUBLE))), named_struct('a', CAST(0.0 AS
DOUBLE)));
+----
+1
+
+# NULLs nested inside the compared values are equal to each other
+query I
+SELECT array_position(array(array(1, CAST(NULL AS INT)), array(1, 2)),
array(1, CAST(NULL AS INT)));
+----
+1
+
+query I
+SELECT array_position(array(array(CAST(NULL AS INT)), array(1)),
array(CAST(NULL AS INT)));
+----
+1
+
+query I
+SELECT array_position(array(named_struct('a', CAST(NULL AS INT)),
named_struct('a', CAST(1 AS INT))), named_struct('a', CAST(NULL AS INT)));
+----
+1
+
+query I
+SELECT array_position(array(array(CAST(1.0 AS DOUBLE), CAST(NULL AS DOUBLE))),
array(CAST(1.0 AS DOUBLE), CAST(NULL AS DOUBLE)));
+----
+1
+
+###
+### Column inputs
+###
+
+statement ok
+CREATE TABLE t AS VALUES
+ (make_array(1, 2, 3), 3),
+ (make_array(1, NULL, 3), 3),
+ (make_array(1, 2), 5),
+ (NULL, 1),
+ (make_array(1, 2), NULL);
+
+query I
+SELECT array_position(column1, column2) FROM t;
+----
+3
+3
+0
+NULL
+NULL
Review Comment:
Added in 98787b5: `LargeList(Int64)` and `LargeList(Float64)` cases, a
sliced batch, and column tables for floats (NaN, `-0.0`), strings, nested
floats and structs with signed zero. The PySpark runner passes on the whole
file. I wrote the nested-float and struct cases directly rather than porting
Comet's tests.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
