Skip to content

Commit

Permalink
Add take_arrays util for getting entries from 2d arrays (#6475)
Browse files Browse the repository at this point in the history
* Add take_arrays util function

* Update comments

* Minor changes
  • Loading branch information
akurmustafa authored Sep 29, 2024
1 parent f41c258 commit f0e39cc
Showing 1 changed file with 59 additions and 0 deletions.
59 changes: 59 additions & 0 deletions arrow-select/src/take.rs
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,65 @@ pub fn take(
}
}

/// For each [ArrayRef] in the [`Vec<ArrayRef>`], take elements by index and create a new
/// [`Vec<ArrayRef>`] from those indices.
///
/// ```text
/// ┌────────┬────────┐
/// │ │ │ ┌────────┐ ┌────────┬────────┐
/// │ A │ 1 │ │ │ │ │ │
/// ├────────┼────────┤ │ 0 │ │ A │ 1 │
/// │ │ │ ├────────┤ ├────────┼────────┤
/// │ D │ 4 │ │ │ │ │ │
/// ├────────┼────────┤ │ 2 │ take_arrays(values,indices) │ B │ 2 │
/// │ │ │ ├────────┤ ├────────┼────────┤
/// │ B │ 2 │ │ │ ───────────────────────────► │ │ │
/// ├────────┼────────┤ │ 3 │ │ C │ 3 │
/// │ │ │ ├────────┤ ├────────┼────────┤
/// │ C │ 3 │ │ │ │ │ │
/// ├────────┼────────┤ │ 1 │ │ D │ 4 │
/// │ │ │ └────────┘ └────────┼────────┘
/// │ E │ 5 │
/// └────────┴────────┘
/// values arrays indices array result
/// ```
///
/// # Errors
/// This function errors whenever:
/// * An index cannot be casted to `usize` (typically 32 bit architectures)
/// * An index is out of bounds and `options` is set to check bounds.
///
/// # Safety
///
/// When `options` is not set to check bounds, taking indexes after `len` will panic.
///
/// # Examples
/// ```
/// # use std::sync::Arc;
/// # use arrow_array::{StringArray, UInt32Array, cast::AsArray};
/// # use arrow_select::take::{take, take_arrays};
/// let string_values = Arc::new(StringArray::from(vec!["zero", "one", "two"]));
/// let values = Arc::new(UInt32Array::from(vec![0, 1, 2]));
///
/// // Take items at index 2, and 1:
/// let indices = UInt32Array::from(vec![2, 1]);
/// let taken_arrays = take_arrays(&[string_values, values], &indices, None).unwrap();
/// let taken_string = taken_arrays[0].as_string::<i32>();
/// assert_eq!(*taken_string, StringArray::from(vec!["two", "one"]));
/// let taken_values = taken_arrays[1].as_primitive();
/// assert_eq!(*taken_values, UInt32Array::from(vec![2, 1]));
/// ```
pub fn take_arrays(
arrays: &[ArrayRef],
indices: &dyn Array,
options: Option<TakeOptions>,
) -> Result<Vec<ArrayRef>, ArrowError> {
arrays
.iter()
.map(|array| take(array.as_ref(), indices, options.clone()))
.collect()
}

/// Verifies that the non-null values of `indices` are all `< len`
fn check_bounds<T: ArrowPrimitiveType>(
len: usize,
Expand Down

0 comments on commit f0e39cc

Please sign in to comment.