Skip to content

Commit 25c34f9

Browse files
authored
cleanup array_has (#12460)
* cleanup "array_has" * revert array_has_dispatch_for_scalar to use vec
1 parent 5c7721b commit 25c34f9

File tree

1 file changed

+25
-37
lines changed

1 file changed

+25
-37
lines changed

datafusion/functions-nested/src/array_has.rs

Lines changed: 25 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -96,44 +96,33 @@ impl ScalarUDFImpl for ArrayHas {
9696
}
9797

9898
fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
99-
// Always return null if the second argumet is null
100-
// i.e. array_has(array, null) -> null
101-
if let ColumnarValue::Scalar(s) = &args[1] {
102-
if s.is_null() {
103-
return Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)));
99+
match &args[1] {
100+
ColumnarValue::Array(array_needle) => {
101+
// the needle is already an array, convert the haystack to an array of the same length
102+
let haystack = args[0].to_owned().into_array(array_needle.len())?;
103+
let array = array_has_inner_for_array(&haystack, array_needle)?;
104+
Ok(ColumnarValue::Array(array))
104105
}
105-
}
106-
107-
// first, identify if any of the arguments is an Array. If yes, store its `len`,
108-
// as any scalar will need to be converted to an array of len `len`.
109-
let len = args
110-
.iter()
111-
.fold(Option::<usize>::None, |acc, arg| match arg {
112-
ColumnarValue::Scalar(_) => acc,
113-
ColumnarValue::Array(a) => Some(a.len()),
114-
});
115-
116-
let is_scalar = len.is_none();
117-
118-
let result = match args[1] {
119-
ColumnarValue::Array(_) => {
120-
let args = ColumnarValue::values_to_arrays(args)?;
121-
array_has_inner_for_array(&args[0], &args[1])
122-
}
123-
ColumnarValue::Scalar(_) => {
106+
ColumnarValue::Scalar(scalar_needle) => {
107+
// Always return null if the second argument is null
108+
// i.e. array_has(array, null) -> null
109+
if scalar_needle.is_null() {
110+
return Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)));
111+
}
112+
113+
// since the needle is a scalar, convert it to an array of size 1
124114
let haystack = args[0].to_owned().into_array(1)?;
125-
let needle = args[1].to_owned().into_array(1)?;
115+
let needle = scalar_needle.to_array_of_size(1)?;
126116
let needle = Scalar::new(needle);
127-
array_has_inner_for_scalar(&haystack, &needle)
117+
let array = array_has_inner_for_scalar(&haystack, &needle)?;
118+
if let ColumnarValue::Scalar(_) = &args[0] {
119+
// If both inputs are scalar, keeps output as scalar
120+
let scalar_value = ScalarValue::try_from_array(&array, 0)?;
121+
Ok(ColumnarValue::Scalar(scalar_value))
122+
} else {
123+
Ok(ColumnarValue::Array(array))
124+
}
128125
}
129-
};
130-
131-
if is_scalar {
132-
// If all inputs are scalar, keeps output as scalar
133-
let result = result.and_then(|arr| ScalarValue::try_from_array(&arr, 0));
134-
result.map(ColumnarValue::Scalar)
135-
} else {
136-
result.map(ColumnarValue::Array)
137126
}
138127
}
139128

@@ -218,10 +207,9 @@ fn array_has_dispatch_for_scalar<O: OffsetSizeTrait>(
218207
}
219208
let sliced_array = eq_array.slice(start, length);
220209
// For nested list, check number of nulls
221-
if sliced_array.null_count() == length {
222-
continue;
210+
if sliced_array.null_count() != length {
211+
final_contained[i] = Some(sliced_array.true_count() > 0);
223212
}
224-
final_contained[i] = Some(sliced_array.true_count() > 0);
225213
}
226214

227215
Ok(Arc::new(BooleanArray::from(final_contained)))

0 commit comments

Comments
 (0)