jayzhan211 commented on code in PR #8170:
URL: https://github.com/apache/arrow-datafusion/pull/8170#discussion_r1393513521
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -953,114 +1025,68 @@ fn general_list_repeat(
)?))
}
-macro_rules! position {
- ($ARRAY:expr, $ELEMENT:expr, $INDEX:expr, $ARRAY_TYPE:ident) => {{
- let element = downcast_arg!($ELEMENT, $ARRAY_TYPE);
- $ARRAY
- .iter()
- .zip(element.iter())
- .zip($INDEX.iter())
- .map(|((arr, el), i)| {
- let index = match i {
- Some(i) => {
- if i <= 0 {
- 0
- } else {
- i - 1
- }
- }
- None => return exec_err!("initial position must not be
null"),
- };
-
- match arr {
- Some(arr) => {
- let child_array = downcast_arg!(arr, $ARRAY_TYPE);
-
- match child_array
- .iter()
- .skip(index as usize)
- .position(|x| x == el)
- {
- Some(value) => Ok(Some(value as u64 + index as u64
+ 1u64)),
- None => Ok(None),
- }
- }
- None => Ok(None),
- }
- })
- .collect::<Result<UInt64Array>>()?
- }};
-}
-
/// Array_position SQL function
pub fn array_position(args: &[ArrayRef]) -> Result<ArrayRef> {
- let arr = as_list_array(&args[0])?;
- let element = &args[1];
+ let list_array = as_list_array(&args[0])?;
+ let element_array = &args[1];
+
+ check_datatypes("array_position", &[list_array.values(), element_array])?;
- let index = if args.len() == 3 {
- as_int64_array(&args[2])?.clone()
+ let arr_from = if args.len() == 3 {
+ as_int64_array(&args[2])?
+ .values()
+ .to_vec()
+ .iter()
+ .map(|&x| x - 1)
+ .collect::<Vec<_>>()
} else {
- Int64Array::from_value(0, arr.len())
+ vec![0; list_array.len()]
};
- check_datatypes("array_position", &[arr.values(), element])?;
- macro_rules! array_function {
- ($ARRAY_TYPE:ident) => {
- position!(arr, element, index, $ARRAY_TYPE)
- };
+ // if `start_from` index is out of bounds, return error
+ for (arr, &from) in list_array.iter().zip(arr_from.iter()) {
+ if let Some(arr) = arr {
+ if from < 0 || from as usize >= arr.len() {
+ return internal_err!("start_from index out of bounds");
+ }
+ } else {
+ // We will get null if we got null in the array, so we don't need
to check
+ }
}
- let res = call_array_function!(arr.value_type(), true);
- Ok(Arc::new(res))
+ general_position::<i32>(list_array, element_array, arr_from)
}
-macro_rules! positions {
- ($ARRAY:expr, $ELEMENT:expr, $ARRAY_TYPE:ident) => {{
- let element = downcast_arg!($ELEMENT, $ARRAY_TYPE);
- let mut offsets: Vec<i32> = vec![0];
- let mut values =
- downcast_arg!(new_empty_array(&DataType::UInt64),
UInt64Array).clone();
- for comp in $ARRAY
- .iter()
- .zip(element.iter())
- .map(|(arr, el)| match arr {
- Some(arr) => {
- let child_array = downcast_arg!(arr, $ARRAY_TYPE);
- let res = child_array
- .iter()
- .enumerate()
- .filter(|(_, x)| *x == el)
- .flat_map(|(i, _)| Some((i + 1) as u64))
- .collect::<UInt64Array>();
+fn general_position<OffsetSize: OffsetSizeTrait>(
Review Comment:
no, but we will need it if we want to extend it for large list, I'm just
lazy to remove it for now :)
--
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]
