comphead commented on code in PR #63:
URL: https://github.com/apache/datafusion-spark/pull/63#discussion_r4566026448
##########
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:
Spark rejects this. `ArrayPosition.checkInputDataTypes` returns
`DATATYPE_MISMATCH.NULL_TYPE` when either argument is `NullType`, in 3.5, 4.0
and 4.2. The typed `CAST(NULL AS INT)` case above is the one that returns NULL.
The header says the values were verified with pyspark 4.2.0, so could this case
go, or be marked as DataFusion-only behavior?
##########
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:
`array_contains` on main (#83) already has Spark float equality:
`normalize_floats` and `normalize_array` canonicalize NaN and `-0.0` at any
depth of list, large list, fixed-size list and struct. With those made
`pub(super)`, this function could normalize `values` and `element` once per
batch and use `make_comparator` for every type. That removes `has_float`,
`null_aware`, `equality`, `list_equality` and `Equality`, and also covers
`FixedSizeList` of floats, which this rejects as not supported. It costs one
copy of the float buffers per batch, but leaves the crate with one definition
of Spark float equality.
##########
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:
Nothing covers `LargeList` (the `position::<i64>` path), floats or strings
in column input, or a sliced list. Comet's unit tests for a nested float column
and struct signed zero were not ported either. Could we add `arrow_cast(...,
'LargeList(Int64)')` cases and a table with NaN, `-0.0` and string rows here?
--
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]
