Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

cleanup array_has #12460

Merged
merged 2 commits into from
Sep 16, 2024
Merged

cleanup array_has #12460

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d4ab9d7
cleanup "array_has"
samuelcolvin Sep 13, 2024
f68a3da
revert array_has_dispatch_for_scalar to use vec
samuelcolvin Sep 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 25 additions & 37 deletions datafusion/functions-nested/src/array_has.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,44 +95,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)
}
}

Expand Down Expand Up @@ -214,10 +203,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)))
Expand Down