Jefffrey commented on code in PR #6913:
URL: https://github.com/apache/arrow-rs/pull/6913#discussion_r1902277707
##########
arrow-ord/src/sort.rs:
##########
@@ -801,19 +801,186 @@ mod tests {
use half::f16;
use rand::rngs::StdRng;
use rand::{Rng, RngCore, SeedableRng};
+ use std::ops::Deref;
+
+ mod build_arrays_helper {
+ use crate::sort::tests::get_same_lists_length;
+ use arrow_array::{
+ ArrayRef, ArrowPrimitiveType, Decimal128Array, Decimal256Array,
FixedSizeBinaryArray,
+ FixedSizeListArray, GenericBinaryArray, GenericListArray,
LargeStringArray,
+ OffsetSizeTrait, PrimitiveArray, StringArray, StringViewArray,
+ };
+ use arrow_buffer::i256;
+ use std::sync::Arc;
+
+ pub fn decimal128_array(data: Vec<Option<i128>>) -> Vec<ArrayRef> {
+ vec![Arc::new(
+ data.into_iter()
+ .collect::<Decimal128Array>()
+ .with_precision_and_scale(23, 6)
+ .unwrap(),
+ )]
+ }
- fn create_decimal128_array(data: Vec<Option<i128>>) -> Decimal128Array {
- data.into_iter()
- .collect::<Decimal128Array>()
- .with_precision_and_scale(23, 6)
- .unwrap()
+ pub fn decimal256_array(data: Vec<Option<i256>>) -> Vec<ArrayRef> {
+ vec![Arc::new(
+ data.into_iter()
+ .collect::<Decimal256Array>()
+ .with_precision_and_scale(53, 6)
+ .unwrap(),
+ )]
+ }
+
+ pub fn primitive_arrays<T>(data: Vec<Option<T::Native>>) ->
Vec<ArrayRef>
+ where
+ T: ArrowPrimitiveType,
+ PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
+ {
+ vec![Arc::new(PrimitiveArray::<T>::from(data)) as ArrayRef]
+ }
+
+ pub fn string_arrays(data: Vec<Option<&str>>) -> Vec<ArrayRef> {
+ vec![
+ Arc::new(StringArray::from(data.clone())) as ArrayRef,
+ Arc::new(LargeStringArray::from(data.clone())) as ArrayRef,
+ Arc::new(StringViewArray::from(data)) as ArrayRef,
+ ]
+ }
+
+ pub fn binary_arrays(data: Vec<Option<Vec<u8>>>) -> Vec<ArrayRef> {
+ // Generic size binary array
+ fn generic_binary_array<S: OffsetSizeTrait>(
+ data: &[Option<Vec<u8>>],
+ ) -> Arc<GenericBinaryArray<S>> {
+ Arc::new(GenericBinaryArray::<S>::from_opt_vec(
+ data.iter()
+ .map(|binary| binary.as_ref().map(Vec::as_slice))
+ .collect(),
+ ))
+ }
+
+ let mut arrays = vec![
+ generic_binary_array::<i32>(&data) as ArrayRef,
+ generic_binary_array::<i64>(&data) as ArrayRef,
+ ];
+
+ if let Some(first_length) = get_same_lists_length(&data) {
+ arrays.push(Arc::new(
+ FixedSizeBinaryArray::try_from_sparse_iter_with_size(
+ data.iter().cloned(),
+ first_length as i32,
+ )
+ .unwrap(),
+ ));
+ }
+
+ arrays
+ }
+
+ pub fn primitive_generic_list_array<OffsetSize, T>(
+ data: &[Option<Vec<Option<T::Native>>>],
+ ) -> ArrayRef
+ where
+ OffsetSize: OffsetSizeTrait,
+ T: ArrowPrimitiveType,
+ PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
+ {
+ Arc::new(GenericListArray::<OffsetSize>::from_iter_primitive::<T,
_, _>(data.to_vec()))
+ }
+
+ pub fn primitive_fixed_list_array<T>(
+ data: &[Option<Vec<Option<T::Native>>>],
+ length: i32,
+ ) -> ArrayRef
+ where
+ T: ArrowPrimitiveType,
+ PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
+ {
+ Arc::new(FixedSizeListArray::from_iter_primitive::<T, _, _>(
+ data.to_vec(),
+ length,
+ ))
+ }
+
+ pub fn primitive_list_arrays<T>(data:
Vec<Option<Vec<Option<T::Native>>>>) -> Vec<ArrayRef>
+ where
+ T: ArrowPrimitiveType,
+ PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
+ {
+ let mut arrays = vec![
+ primitive_generic_list_array::<i32, T>(&data),
+ primitive_generic_list_array::<i64, T>(&data),
+ ];
+
+ if let Some(first_length) = get_same_lists_length(&data) {
+ arrays.push(primitive_fixed_list_array::<T>(&data,
first_length as i32));
+ }
+
+ arrays
+ }
}
- fn create_decimal256_array(data: Vec<Option<i256>>) -> Decimal256Array {
- data.into_iter()
- .collect::<Decimal256Array>()
- .with_precision_and_scale(53, 6)
- .unwrap()
+ /// Return Some(length) if all lists have the same length, None otherwise
+ fn get_same_lists_length<T>(data: &[Option<Vec<T>>]) -> Option<usize> {
+ let first_length = data
+ .iter()
+ .find_map(|x| x.as_ref().map(|x| x.len()))
+ .unwrap_or(0);
+ let first_non_match_length = data
+ .iter()
+ .map(|x| x.as_ref().map(|x| x.len()).unwrap_or(first_length))
+ .any(|x| x != first_length);
+
+ if first_non_match_length {
+ None
+ } else {
+ Some(first_length)
+ }
+ }
+
+ /// Test sorting arrays
+ ///
+ /// the `into_lists_fn` function is used to convert the data into the list
variants wanted to be tested, (e.g. ListArray, LargeListArray,
FixedSizeListArray or any other list variant)
+ /// The ordering of the returned vector must be consistent as it will be
called with the data and the expected data
+ fn test_sort_arrays<
+ ListItemType: Clone,
+ IntoListFn: Fn(Vec<Option<ListItemType>>) -> Vec<ArrayRef>,
+ >(
+ into_lists_fn: IntoListFn,
+ data: Vec<Option<ListItemType>>,
+ options: Option<SortOptions>,
+ limit: Option<usize>,
+ expected_data: Vec<Option<ListItemType>>,
+ ) {
Review Comment:
I think this was missed from the changes you applied
--
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]
