Skip to content

Commit

Permalink
replace pop front and pop back
Browse files Browse the repository at this point in the history 
Signed-off-by: jayzhan211 <[email protected]>
  • Loading branch information
@jayzhan211
jayzhan211 committed Dec 3, 2023
1 parent cc850d6 commit 8c7e323
Showing 1 changed file with 20 additions and 169 deletions.
189 changes: 20 additions & 169 deletions datafusion/physical-expr/src/array_expressions.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -360,135 +360,6 @@ pub fn make_array(arrays: &[ArrayRef]) -> Result<ArrayRef> {
}
}

fn return_empty(return_null: bool, data_type: DataType) -> Arc<dyn Array> {
if return_null {
new_null_array(&data_type, 1)
} else {
new_empty_array(&data_type)
}
}

macro_rules! list_slice {
($ARRAY:expr, $I:expr, $J:expr, $RETURN_ELEMENT:expr, $ARRAY_TYPE:ident) => {{
let array = $ARRAY.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap();
if $I == 0 && $J == 0 || $ARRAY.is_empty() {
return return_empty($RETURN_ELEMENT, $ARRAY.data_type().clone());
}

let i = if $I < 0 {
if $I.abs() as usize > array.len() {
return return_empty(true, $ARRAY.data_type().clone());
}

(array.len() as i64 + $I + 1) as usize
} else {
if $I == 0 {
1
} else {
$I as usize
}
};
let j = if $J < 0 {
if $J.abs() as usize > array.len() {
return return_empty(true, $ARRAY.data_type().clone());
}

if $RETURN_ELEMENT {
(array.len() as i64 + $J + 1) as usize
} else {
(array.len() as i64 + $J) as usize
}
} else {
if $J == 0 {
1
} else {
if $J as usize > array.len() {
array.len()
} else {
$J as usize
}
}
};

if i > j || i as usize > $ARRAY.len() {
return_empty($RETURN_ELEMENT, $ARRAY.data_type().clone())
} else {
Arc::new(array.slice((i - 1), (j + 1 - i)))
}
}};
}

macro_rules! slice {
($ARRAY:expr, $KEY:expr, $EXTRA_KEY:expr, $RETURN_ELEMENT:expr, $ARRAY_TYPE:ident) => {{
let sliced_array: Vec<Arc<dyn Array>> = $ARRAY
.iter()
.zip($KEY.iter())
.zip($EXTRA_KEY.iter())
.map(|((arr, i), j)| match (arr, i, j) {
(Some(arr), Some(i), Some(j)) => {
list_slice!(arr, i, j, $RETURN_ELEMENT, $ARRAY_TYPE)
}
(Some(arr), None, Some(j)) => {
list_slice!(arr, 1i64, j, $RETURN_ELEMENT, $ARRAY_TYPE)
}
(Some(arr), Some(i), None) => {
list_slice!(arr, i, arr.len() as i64, $RETURN_ELEMENT, $ARRAY_TYPE)
}
(Some(arr), None, None) if !$RETURN_ELEMENT => arr,
_ => return_empty($RETURN_ELEMENT, $ARRAY.value_type().clone()),
})
.collect();

// concat requires input of at least one array
if sliced_array.is_empty() {
Ok(return_empty($RETURN_ELEMENT, $ARRAY.value_type()))
} else {
let vec = sliced_array
.iter()
.map(|a| a.as_ref())
.collect::<Vec<&dyn Array>>();
let mut i: i32 = 0;
let mut offsets = vec![i];
offsets.extend(
vec.iter()
.map(|a| {
i += a.len() as i32;
i
})
.collect::<Vec<_>>(),
);
let values = compute::concat(vec.as_slice()).unwrap();

if $RETURN_ELEMENT {
Ok(values)
} else {
let field =
Arc::new(Field::new("item", $ARRAY.value_type().clone(), true));
Ok(Arc::new(ListArray::try_new(
field,
OffsetBuffer::new(offsets.into()),
values,
None,
)?))
}
}
}};
}

fn define_array_slice(
list_array: &ListArray,
key: &Int64Array,
extra_key: &Int64Array,
return_element: bool,
) -> Result<ArrayRef> {
macro_rules! array_function {
($ARRAY_TYPE:ident) => {
slice!(list_array, key, extra_key, return_element, $ARRAY_TYPE)
};
}
call_array_function!(list_array.value_type(), true)
}

pub fn array_element(args: &[ArrayRef]) -> Result<ArrayRef> {
let list_array = as_list_array(&args[0])?;
let indexes = as_int64_array(&args[1])?;
Expand Down Expand Up @@ -722,40 +593,18 @@ pub fn array_slice(args: &[ArrayRef]) -> Result<ArrayRef> {
)?))
}

fn general_array_pop(
list_array: &GenericListArray<i32>,
from_back: bool,
) -> Result<(Vec<i64>, Vec<i64>)> {
if from_back {
let key = vec![0; list_array.len()];
// Atttetion: `arr.len() - 1` in extra key defines the last element position (position = index + 1, not inclusive) we want in the new array.
let extra_key: Vec<_> = list_array
.iter()
.map(|x| x.map_or(0, |arr| arr.len() as i64 - 1))
.collect();
Ok((key, extra_key))
} else {
// Atttetion: 2 in the `key`` defines the first element position (position = index + 1) we want in the new array.
// We only handle two cases of the first element index: if the old array has any elements, starts from 2 (index + 1), or starts from initial.
let key: Vec<_> = list_array.iter().map(|x| x.map_or(0, |_| 2)).collect();
let extra_key: Vec<_> = list_array
.iter()
.map(|x| x.map_or(0, |arr| arr.len() as i64))
.collect();
Ok((key, extra_key))
}
}

/// array_pop_back SQL function
pub fn array_pop_back(args: &[ArrayRef]) -> Result<ArrayRef> {
let list_array = as_list_array(&args[0])?;
let (key, extra_key) = general_array_pop(list_array, true)?;

define_array_slice(
list_array,
&Int64Array::from(key),
&Int64Array::from(extra_key),
false,
)
let from_array = Int64Array::from(vec![1; list_array.len()]);
let to_array = Int64Array::from(
list_array
.iter()
.map(|arr| arr.map_or(0, |arr| arr.len() as i64 - 1))
.collect::<Vec<i64>>(),
);
let args = vec![args[0].clone(), Arc::new(from_array), Arc::new(to_array)];
array_slice(args.as_slice())
}

/// Appends or prepends elements to a ListArray.
Expand Down Expand Up @@ -879,16 +728,18 @@ pub fn gen_range(args: &[ArrayRef]) -> Result<ArrayRef> {
Ok(arr)
}

/// array_pop_front SQL function
pub fn array_pop_front(args: &[ArrayRef]) -> Result<ArrayRef> {
let list_array = as_list_array(&args[0])?;
let (key, extra_key) = general_array_pop(list_array, false)?;

define_array_slice(
list_array,
&Int64Array::from(key),
&Int64Array::from(extra_key),
false,
)
let from_array = Int64Array::from(vec![2; list_array.len()]);
let to_array = Int64Array::from(
list_array
.iter()
.map(|arr| arr.map_or(0, |arr| arr.len() as i64))
.collect::<Vec<i64>>(),
);
let args = vec![args[0].clone(), Arc::new(from_array), Arc::new(to_array)];
array_slice(args.as_slice())
}

/// Array_append SQL function
Expand Down

0 comments on commit 8c7e323

Please sign in to comment.