This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 25c34f9ecc cleanup `array_has` (#12460)
25c34f9ecc is described below
commit 25c34f9ecc225aa6e38a1534ab362a239135b703
Author: Samuel Colvin <[email protected]>
AuthorDate: Mon Sep 16 19:32:54 2024 +0100
cleanup `array_has` (#12460)
* cleanup "array_has"
* revert array_has_dispatch_for_scalar to use vec
---
datafusion/functions-nested/src/array_has.rs | 62 +++++++++++-----------------
1 file changed, 25 insertions(+), 37 deletions(-)
diff --git a/datafusion/functions-nested/src/array_has.rs
b/datafusion/functions-nested/src/array_has.rs
index dec964df21..8f8d123bf5 100644
--- a/datafusion/functions-nested/src/array_has.rs
+++ b/datafusion/functions-nested/src/array_has.rs
@@ -96,44 +96,33 @@ impl ScalarUDFImpl for ArrayHas {
}
fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
- // Always return null if the second argumet is null
- // i.e. array_has(array, null) -> null
- if let ColumnarValue::Scalar(s) = &args[1] {
- if s.is_null() {
- return Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)));
+ match &args[1] {
+ ColumnarValue::Array(array_needle) => {
+ // the needle is already an array, convert the haystack to an
array of the same length
+ let haystack =
args[0].to_owned().into_array(array_needle.len())?;
+ let array = array_has_inner_for_array(&haystack,
array_needle)?;
+ Ok(ColumnarValue::Array(array))
}
- }
-
- // first, identify if any of the arguments is an Array. If yes, store
its `len`,
- // as any scalar will need to be converted to an array of len `len`.
- 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 result = match args[1] {
- ColumnarValue::Array(_) => {
- let args = ColumnarValue::values_to_arrays(args)?;
- array_has_inner_for_array(&args[0], &args[1])
- }
- ColumnarValue::Scalar(_) => {
+ ColumnarValue::Scalar(scalar_needle) => {
+ // Always return null if the second argument is null
+ // i.e. array_has(array, null) -> null
+ if scalar_needle.is_null() {
+ return
Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)));
+ }
+
+ // since the needle is a scalar, convert it to an array of
size 1
let haystack = args[0].to_owned().into_array(1)?;
- let needle = args[1].to_owned().into_array(1)?;
+ let needle = scalar_needle.to_array_of_size(1)?;
let needle = Scalar::new(needle);
- array_has_inner_for_scalar(&haystack, &needle)
+ let array = array_has_inner_for_scalar(&haystack, &needle)?;
+ if let ColumnarValue::Scalar(_) = &args[0] {
+ // If both inputs are scalar, keeps output as scalar
+ let scalar_value = ScalarValue::try_from_array(&array, 0)?;
+ Ok(ColumnarValue::Scalar(scalar_value))
+ } else {
+ Ok(ColumnarValue::Array(array))
+ }
}
- };
-
- if is_scalar {
- // If all inputs are scalar, keeps output as scalar
- let result = result.and_then(|arr|
ScalarValue::try_from_array(&arr, 0));
- result.map(ColumnarValue::Scalar)
- } else {
- result.map(ColumnarValue::Array)
}
}
@@ -218,10 +207,9 @@ fn array_has_dispatch_for_scalar<O: OffsetSizeTrait>(
}
let sliced_array = eq_array.slice(start, length);
// For nested list, check number of nulls
- if sliced_array.null_count() == length {
- continue;
+ if sliced_array.null_count() != length {
+ final_contained[i] = Some(sliced_array.true_count() > 0);
}
- final_contained[i] = Some(sliced_array.true_count() > 0);
}
Ok(Arc::new(BooleanArray::from(final_contained)))
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
