rluvaton commented on code in PR #6913:
URL: https://github.com/apache/arrow-rs/pull/6913#discussion_r1901661568
##########
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>>>,
Review Comment:
Thanks, removed
##########
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,
+ ];
Review Comment:
Thanks, updated
--
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]
