You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "alamb (via GitHub)" <gi...@apache.org> on 2023/12/08 17:04:28 UTC
Re: [PR] Remove `define_array_slice` and reuse `array_slice` for `array_pop_front/back` [arrow-datafusion]
alamb commented on code in PR #8401:
URL: https://github.com/apache/arrow-datafusion/pull/8401#discussion_r1420786795
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -370,135 +369,57 @@ 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)
- }
-}
-
-fn list_slice<T: Array + 'static>(
- array: &dyn Array,
- i: i64,
- j: i64,
- return_element: bool,
-) -> ArrayRef {
- let array = array.as_any().downcast_ref::<T>().unwrap();
-
- let array_type = array.data_type().clone();
+pub fn array_element(args: &[ArrayRef]) -> Result<ArrayRef> {
Review Comment:
Does this need to be `pub`? If it does, could you please add a documentation comment explaining what the arguments are (namely that the first is a list and the second is an array of indexes?
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -581,45 +502,118 @@ pub fn array_except(args: &[ArrayRef]) -> Result<ArrayRef> {
pub fn array_slice(args: &[ArrayRef]) -> Result<ArrayRef> {
let list_array = as_list_array(&args[0])?;
- let key = as_int64_array(&args[1])?;
- let extra_key = as_int64_array(&args[2])?;
- define_array_slice(list_array, key, extra_key, false)
-}
+ let from_array = as_int64_array(&args[1])?;
+ let to_array = as_int64_array(&args[2])?;
-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()];
- // Attention: `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 {
- // Attention: 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))
+ let values = list_array.values();
+ let original_data = values.to_data();
+ let capacity = Capacities::Array(original_data.len());
+
+ // use_nulls: false, we don't need nulls but empty array for array_slice, so we don't need explicit nulls but adjust offset to indicate nulls.
+ let mut mutable =
+ MutableArrayData::with_capacities(vec![&original_data], false, capacity);
+
+ // We have the slice syntax compatible with DuckDB v0.8.1.
+ // The rule `adjusted_from_index` and `adjusted_to_index` follows the rule of array_slice in duckdb.
+
+ fn adjusted_from_index(index: i64, len: usize) -> Option<i64> {
+ // 0 ~ len - 1
+ let adjusted_zero_index = if index < 0 {
+ index + len as i64
+ } else {
+ // array_slice(arr, 1, to) is the same as array_slice(arr, 0, to)
+ std::cmp::max(index - 1, 0)
+ };
+
+ if 0 <= adjusted_zero_index && adjusted_zero_index < len as i64 {
+ Some(adjusted_zero_index)
+ } else {
+ // Out of bounds
+ None
+ }
+ }
+
+ fn adjusted_to_index(index: i64, len: usize) -> Option<i64> {
+ // 0 ~ len - 1
+ let adjusted_zero_index = if index < 0 {
+ // array_slice in duckdb with negative to_index is python-like, so index itself is exclusive
+ index + len as i64 - 1
+ } else {
+ // array_slice(arr, from, len + 1) is the same as array_slice(arr, from, len)
+ std::cmp::min(index - 1, len as i64 - 1)
+ };
+
+ if 0 <= adjusted_zero_index && adjusted_zero_index < len as i64 {
+ Some(adjusted_zero_index)
+ } else {
+ // Out of bounds
+ None
+ }
+ }
+
+ let mut offsets = vec![0];
+
+ for (row_index, offset_window) in list_array.offsets().windows(2).enumerate() {
+ let start = offset_window[0] as usize;
+ let end = offset_window[1] as usize;
+ let len = end - start;
+
+ // len 0 indicate array is null, return empty array in this row.
+ if len == 0 {
+ offsets.push(offsets[row_index]);
+ continue;
+ }
+
+ // If index is null, we consider it as the minimum / maximum index of the array.
+ let from_index = if from_array.is_null(row_index) {
+ Some(0)
+ } else {
+ adjusted_from_index(from_array.value(row_index), len)
+ };
+
+ let to_index = if to_array.is_null(row_index) {
+ Some(len as i64 - 1)
+ } else {
+ adjusted_to_index(to_array.value(row_index), len)
+ };
+
+ if let (Some(from), Some(to)) = (from_index, to_index) {
+ if from <= to {
+ assert!(start + to as usize <= end);
+ mutable.extend(0, start + from as usize, start + to as usize + 1);
+ offsets.push(offsets[row_index] + (to - from + 1) as i32);
+ } else {
+ // invalid range, return empty array
+ offsets.push(offsets[row_index]);
+ }
+ } else {
+ // invalid range, return empty array
+ offsets.push(offsets[row_index]);
+ }
}
+
+ let data = mutable.freeze();
+
+ Ok(Arc::new(ListArray::try_new(
+ Arc::new(Field::new("item", list_array.value_type(), true)),
+ OffsetBuffer::new(offsets.into()),
+ arrow_array::make_array(data),
+ None,
+ )?))
}
+/// 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(
Review Comment:
this is very clever 👍
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -581,45 +502,118 @@ pub fn array_except(args: &[ArrayRef]) -> Result<ArrayRef> {
pub fn array_slice(args: &[ArrayRef]) -> Result<ArrayRef> {
Review Comment:
Also here I think it would help to add a comment explaining what the arguments are to this function and what the expected output is
Key details might be if the index values are 1 based or 0 based, for example
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -581,45 +502,118 @@ pub fn array_except(args: &[ArrayRef]) -> Result<ArrayRef> {
pub fn array_slice(args: &[ArrayRef]) -> Result<ArrayRef> {
let list_array = as_list_array(&args[0])?;
- let key = as_int64_array(&args[1])?;
- let extra_key = as_int64_array(&args[2])?;
- define_array_slice(list_array, key, extra_key, false)
-}
+ let from_array = as_int64_array(&args[1])?;
+ let to_array = as_int64_array(&args[2])?;
-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()];
- // Attention: `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 {
- // Attention: 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))
+ let values = list_array.values();
+ let original_data = values.to_data();
+ let capacity = Capacities::Array(original_data.len());
+
+ // use_nulls: false, we don't need nulls but empty array for array_slice, so we don't need explicit nulls but adjust offset to indicate nulls.
+ let mut mutable =
+ MutableArrayData::with_capacities(vec![&original_data], false, capacity);
+
+ // We have the slice syntax compatible with DuckDB v0.8.1.
+ // The rule `adjusted_from_index` and `adjusted_to_index` follows the rule of array_slice in duckdb.
+
+ fn adjusted_from_index(index: i64, len: usize) -> Option<i64> {
+ // 0 ~ len - 1
+ let adjusted_zero_index = if index < 0 {
+ index + len as i64
+ } else {
+ // array_slice(arr, 1, to) is the same as array_slice(arr, 0, to)
+ std::cmp::max(index - 1, 0)
+ };
+
+ if 0 <= adjusted_zero_index && adjusted_zero_index < len as i64 {
+ Some(adjusted_zero_index)
+ } else {
+ // Out of bounds
+ None
+ }
+ }
+
+ fn adjusted_to_index(index: i64, len: usize) -> Option<i64> {
+ // 0 ~ len - 1
+ let adjusted_zero_index = if index < 0 {
+ // array_slice in duckdb with negative to_index is python-like, so index itself is exclusive
+ index + len as i64 - 1
+ } else {
+ // array_slice(arr, from, len + 1) is the same as array_slice(arr, from, len)
+ std::cmp::min(index - 1, len as i64 - 1)
+ };
+
+ if 0 <= adjusted_zero_index && adjusted_zero_index < len as i64 {
+ Some(adjusted_zero_index)
+ } else {
+ // Out of bounds
+ None
+ }
+ }
+
+ let mut offsets = vec![0];
+
+ for (row_index, offset_window) in list_array.offsets().windows(2).enumerate() {
+ let start = offset_window[0] as usize;
+ let end = offset_window[1] as usize;
+ let len = end - start;
+
+ // len 0 indicate array is null, return empty array in this row.
+ if len == 0 {
+ offsets.push(offsets[row_index]);
+ continue;
+ }
+
+ // If index is null, we consider it as the minimum / maximum index of the array.
+ let from_index = if from_array.is_null(row_index) {
+ Some(0)
+ } else {
+ adjusted_from_index(from_array.value(row_index), len)
+ };
+
+ let to_index = if to_array.is_null(row_index) {
+ Some(len as i64 - 1)
+ } else {
+ adjusted_to_index(to_array.value(row_index), len)
+ };
+
+ if let (Some(from), Some(to)) = (from_index, to_index) {
+ if from <= to {
+ assert!(start + to as usize <= end);
+ mutable.extend(0, start + from as usize, start + to as usize + 1);
+ offsets.push(offsets[row_index] + (to - from + 1) as i32);
+ } else {
+ // invalid range, return empty array
+ offsets.push(offsets[row_index]);
+ }
+ } else {
+ // invalid range, return empty array
+ offsets.push(offsets[row_index]);
+ }
}
+
+ let data = mutable.freeze();
+
+ Ok(Arc::new(ListArray::try_new(
+ Arc::new(Field::new("item", list_array.value_type(), true)),
+ OffsetBuffer::new(offsets.into()),
+ arrow_array::make_array(data),
+ None,
+ )?))
}
+/// 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(
Review Comment:
this is very clever 👍
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org