comphead commented on code in PR #20770:
URL: https://github.com/apache/datafusion/pull/20770#discussion_r2942203676
##########
datafusion/functions-nested/src/position.rs:
##########
@@ -436,36 +435,70 @@ impl ScalarUDFImpl for ArrayPositions {
}
}
-fn array_positions_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
- let [array, element] = take_function_args("array_positions", args)?;
+/// Attempts the scalar-needle fast path for `array_positions`.
+fn try_array_positions_scalar(args: &[ColumnarValue]) ->
Result<Option<ColumnarValue>> {
+ let [haystack_arg, needle_arg] = take_function_args("array_positions",
args)?;
+
+ let scalar_needle = match needle_arg {
+ ColumnarValue::Scalar(s) => s,
+ ColumnarValue::Array(_) => return Ok(None),
+ };
+
+ // `not_distinct` doesn't support nested types (List, Struct, etc.),
+ // so fall back to the per-row path for those.
+ if scalar_needle.data_type().is_nested() {
+ return Ok(None);
+ }
+
+ let (num_rows, all_inputs_scalar) = match haystack_arg {
+ ColumnarValue::Array(a) => (a.len(), false),
+ ColumnarValue::Scalar(_) => (1, true),
+ };
- match &array.data_type() {
+ let needle = scalar_needle.to_array_of_size(1)?;
+ let haystack = haystack_arg.to_array(num_rows)?;
+
+ let result = match haystack.data_type() {
List(_) => {
- let arr = as_list_array(&array)?;
- crate::utils::check_datatypes("array_positions", &[arr.values(),
element])?;
- general_positions::<i32>(arr, element)
+ let list = as_list_array(&haystack)?;
+ array_positions_scalar::<i32>(list, &needle)
}
LargeList(_) => {
- let arr = as_large_list_array(&array)?;
- crate::utils::check_datatypes("array_positions", &[arr.values(),
element])?;
- general_positions::<i64>(arr, element)
- }
- array_type => {
- exec_err!("array_positions does not support type '{array_type}'.")
+ let list = as_large_list_array(&haystack)?;
+ array_positions_scalar::<i64>(list, &needle)
}
+ t => exec_err!("array_positions does not support type '{t}'"),
+ }?;
+
+ if all_inputs_scalar {
+ Ok(Some(ColumnarValue::Scalar(ScalarValue::try_from_array(
+ &result, 0,
+ )?)))
+ } else {
+ Ok(Some(ColumnarValue::Array(result)))
}
}
-fn general_positions<OffsetSize: OffsetSizeTrait>(
- list_array: &GenericListArray<OffsetSize>,
- element_array: &ArrayRef,
+fn array_positions_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
+ let [haystack, needle] = take_function_args("array_positions", args)?;
+
+ match &haystack.data_type() {
+ List(_) => general_positions::<i32>(as_list_array(&haystack)?, needle),
+ LargeList(_) =>
general_positions::<i64>(as_large_list_array(&haystack)?, needle),
+ t => exec_err!("array_positions does not support type '{t}'"),
Review Comment:
```suggestion
dt => exec_err!("array_positions does not support type '{dt}'"),
```
--
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]
