Skip to content

Commit

Permalink
remove unnecessary wraps in sortk
Browse files Browse the repository at this point in the history
  • Loading branch information
Jiayu Liu committed Jun 11, 2021
1 parent 0c00776 commit 4b42ec2
Showing 1 changed file with 51 additions and 45 deletions.
96 changes: 51 additions & 45 deletions arrow/src/compute/kernels/sort.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,7 @@ pub fn sort_to_indices(

let (v, n) = partition_validity(values);

match values.data_type() {
Ok(match values.data_type() {
DataType::Boolean => sort_boolean(values, v, n, &options, limit),
DataType::Int8 => {
sort_primitive::<Int8Type, _>(values, v, n, cmp, &options, limit)
Expand Down Expand Up @@ -278,10 +278,12 @@ pub fn sort_to_indices(
DataType::Float64 => {
sort_list::<i32, Float64Type>(values, v, n, &options, limit)
}
t => Err(ArrowError::ComputeError(format!(
"Sort not supported for list type {:?}",
t
))),
t => {
return Err(ArrowError::ComputeError(format!(
"Sort not supported for list type {:?}",
t
)))
}
},
DataType::LargeList(field) => match field.data_type() {
DataType::Int8 => sort_list::<i64, Int8Type>(values, v, n, &options, limit),
Expand All @@ -304,10 +306,12 @@ pub fn sort_to_indices(
DataType::Float64 => {
sort_list::<i64, Float64Type>(values, v, n, &options, limit)
}
t => Err(ArrowError::ComputeError(format!(
"Sort not supported for list type {:?}",
t
))),
t => {
return Err(ArrowError::ComputeError(format!(
"Sort not supported for list type {:?}",
t
)))
}
},
DataType::FixedSizeList(field, _) => match field.data_type() {
DataType::Int8 => sort_list::<i32, Int8Type>(values, v, n, &options, limit),
Expand All @@ -330,10 +334,12 @@ pub fn sort_to_indices(
DataType::Float64 => {
sort_list::<i32, Float64Type>(values, v, n, &options, limit)
}
t => Err(ArrowError::ComputeError(format!(
"Sort not supported for list type {:?}",
t
))),
t => {
return Err(ArrowError::ComputeError(format!(
"Sort not supported for list type {:?}",
t
)))
}
},
DataType::Dictionary(key_type, value_type)
if *value_type.as_ref() == DataType::Utf8 =>
Expand Down Expand Up @@ -363,17 +369,21 @@ pub fn sort_to_indices(
DataType::UInt64 => {
sort_string_dictionary::<UInt64Type>(values, v, n, &options, limit)
}
t => Err(ArrowError::ComputeError(format!(
"Sort not supported for dictionary key type {:?}",
t
))),
t => {
return Err(ArrowError::ComputeError(format!(
"Sort not supported for dictionary key type {:?}",
t
)))
}
}
}
t => Err(ArrowError::ComputeError(format!(
"Sort not supported for data type {:?}",
t
))),
}
t => {
return Err(ArrowError::ComputeError(format!(
"Sort not supported for data type {:?}",
t
)))
}
})
}

/// Options that define how sort kernels should behave
Expand All @@ -396,14 +406,13 @@ impl Default for SortOptions {
}

/// Sort primitive values
#[allow(clippy::unnecessary_wraps)]
fn sort_boolean(
values: &ArrayRef,
value_indices: Vec<u32>,
null_indices: Vec<u32>,
options: &SortOptions,
limit: Option<usize>,
) -> Result<UInt32Array> {
) -> UInt32Array {
let values = values
.as_any()
.downcast_ref::<BooleanArray>()
Expand Down Expand Up @@ -469,19 +478,18 @@ fn sort_boolean(
vec![],
);

Ok(UInt32Array::from(result_data))
UInt32Array::from(result_data)
}

/// Sort primitive values
#[allow(clippy::unnecessary_wraps)]
fn sort_primitive<T, F>(
values: &ArrayRef,
value_indices: Vec<u32>,
null_indices: Vec<u32>,
cmp: F,
options: &SortOptions,
limit: Option<usize>,
) -> Result<UInt32Array>
) -> UInt32Array
where
T: ArrowPrimitiveType,
T::Native: std::cmp::PartialOrd,
Expand Down Expand Up @@ -549,7 +557,7 @@ where
vec![],
);

Ok(UInt32Array::from(result_data))
UInt32Array::from(result_data)
}

// insert valid and nan values in the correct order depending on the descending flag
Expand All @@ -574,7 +582,7 @@ fn sort_string<Offset: StringOffsetSizeTrait>(
null_indices: Vec<u32>,
options: &SortOptions,
limit: Option<usize>,
) -> Result<UInt32Array> {
) -> UInt32Array {
let values = values
.as_any()
.downcast_ref::<GenericStringArray<Offset>>()
Expand All @@ -597,7 +605,7 @@ fn sort_string_dictionary<T: ArrowDictionaryKeyType>(
null_indices: Vec<u32>,
options: &SortOptions,
limit: Option<usize>,
) -> Result<UInt32Array> {
) -> UInt32Array {
let values: &DictionaryArray<T> = as_dictionary_array::<T>(values);

let keys: &PrimitiveArray<T> = &values.keys_array();
Expand All @@ -620,15 +628,14 @@ fn sort_string_dictionary<T: ArrowDictionaryKeyType>(

/// shared implementation between dictionary encoded and plain string arrays
#[inline]
#[allow(clippy::unnecessary_wraps)]
fn sort_string_helper<'a, A: Array, F>(
values: &'a A,
value_indices: Vec<u32>,
null_indices: Vec<u32>,
options: &SortOptions,
limit: Option<usize>,
value_fn: F,
) -> Result<UInt32Array>
) -> UInt32Array
where
F: Fn(&'a A, u32) -> &str,
{
Expand Down Expand Up @@ -661,23 +668,22 @@ where
if options.nulls_first {
nulls.append(&mut valid_indices);
nulls.truncate(len);
return Ok(UInt32Array::from(nulls));
UInt32Array::from(nulls)
} else {
// no need to sort nulls as they are in the correct order already
valid_indices.append(&mut nulls);
valid_indices.truncate(len);
UInt32Array::from(valid_indices)
}

// no need to sort nulls as they are in the correct order already
valid_indices.append(&mut nulls);
valid_indices.truncate(len);
Ok(UInt32Array::from(valid_indices))
}

#[allow(clippy::unnecessary_wraps)]
fn sort_list<S, T>(
values: &ArrayRef,
value_indices: Vec<u32>,
mut null_indices: Vec<u32>,
options: &SortOptions,
limit: Option<usize>,
) -> Result<UInt32Array>
) -> UInt32Array
where
S: OffsetSizeTrait,
T: ArrowPrimitiveType,
Expand Down Expand Up @@ -727,12 +733,12 @@ where
if options.nulls_first {
null_indices.append(&mut valid_indices);
null_indices.truncate(len);
return Ok(UInt32Array::from(null_indices));
UInt32Array::from(null_indices)
} else {
valid_indices.append(&mut null_indices);
valid_indices.truncate(len);
UInt32Array::from(valid_indices)
}

valid_indices.append(&mut null_indices);
valid_indices.truncate(len);
Ok(UInt32Array::from(valid_indices))
}

/// Compare two `Array`s based on the ordering defined in [ord](crate::array::ord).
Expand Down

0 comments on commit 4b42ec2

Please sign in to comment.