You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "psvri (via GitHub)" <gi...@apache.org> on 2023/07/01 18:53:53 UTC
[GitHub] [arrow-rs] psvri opened a new pull request, #4473: Improve primitive type sorts by 13-67%
psvri opened a new pull request, #4473:
URL: https://github.com/apache/arrow-rs/pull/4473
# Which issue does this PR close?
Closes #.
# Rationale for this change
The current sort implementation for primitive types first sorts by indices and then performs a take operation. The kernel can be improved by directly sorting.
The results for i64 on my laptop are as follows
```
sort 2^10 time: [9.1287 µs 9.1711 µs 9.2279 µs]
change: [-27.180% -25.980% -24.697%] (p = 0.00 < 0.05)
Performance has improved.
Found 17 outliers among 100 measurements (17.00%)
1 (1.00%) low severe
2 (2.00%) low mild
3 (3.00%) high mild
11 (11.00%) high severe
sort 2^12 time: [70.191 µs 70.437 µs 70.741 µs]
change: [-15.190% -13.941% -12.614%] (p = 0.00 < 0.05)
Performance has improved.
Found 12 outliers among 100 measurements (12.00%)
6 (6.00%) high mild
6 (6.00%) high severe
sort nulls 2^10 time: [4.9898 µs 5.0080 µs 5.0319 µs]
change: [-68.212% -67.754% -67.288%] (p = 0.00 < 0.05)
Performance has improved.
Found 17 outliers among 100 measurements (17.00%)
2 (2.00%) low mild
5 (5.00%) high mild
10 (10.00%) high severe
sort nulls 2^12 time: [34.333 µs 34.641 µs 34.983 µs]
change: [-58.256% -57.089% -56.063%] (p = 0.00 < 0.05)
Performance has improved.
Found 16 outliers among 100 measurements (16.00%)
3 (3.00%) high mild
13 (13.00%) high severe
```
# What changes are included in this PR?
I reworked the sort kernel so that primitive types are sorted directly without using sort_by_indices . I have also included a new primitive benchmark `sort_kernel_primitives.rs` .
# Are there any user-facing changes?
No
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1251265451
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,137 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn sort_native_type<T>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ <T as arrow_array::ArrowPrimitiveType>::Native: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let primitive_values = values.as_primitive::<T>();
+
+ let result_capacity = values.len()
+ * std::mem::size_of::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+
+ let mutable_slice =
+ mutable_buffer.typed_data_mut::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+
+ let array_data = values.to_data();
+ let input_values = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+
+ let mut validity_buffer = BooleanBufferBuilder::new(values.len());
+ let values_slice;
+
+ if sort_options.nulls_first {
+ values_slice = &mut mutable_slice[nulls_count..];
+ validity_buffer.append_n(nulls_count, false);
+ validity_buffer.append_n(valid_count, true);
+ } else {
+ values_slice = &mut mutable_slice[..valid_count];
+ validity_buffer.append_n(valid_count, true);
+ validity_buffer.append_n(nulls_count, false);
+ }
+
+ for (write_index, index) in nulls.valid_indices().enumerate() {
+ values_slice[write_index] = primitive_values.value(index);
+ }
+
+ values_slice.sort_unstable_by(|a, b| a.compare(*b));
+ if sort_options.descending {
+ values_slice.reverse();
+ }
+
+ null_bit_buffer = Some(validity_buffer.finish().into());
+ } else {
+ mutable_slice.copy_from_slice(&input_values[..values.len()]);
+ mutable_slice.sort_unstable_by(|a, b| a.compare(*b));
Review Comment:
Not for floats, we use total ordering not the default partial ordering
--
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
[GitHub] [arrow-rs] Dandandan commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1251266182
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,137 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn sort_native_type<T>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ <T as arrow_array::ArrowPrimitiveType>::Native: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let primitive_values = values.as_primitive::<T>();
+
+ let result_capacity = values.len()
+ * std::mem::size_of::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+
+ let mutable_slice =
+ mutable_buffer.typed_data_mut::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+
+ let array_data = values.to_data();
+ let input_values = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+
+ let mut validity_buffer = BooleanBufferBuilder::new(values.len());
+ let values_slice;
+
+ if sort_options.nulls_first {
+ values_slice = &mut mutable_slice[nulls_count..];
+ validity_buffer.append_n(nulls_count, false);
+ validity_buffer.append_n(valid_count, true);
+ } else {
+ values_slice = &mut mutable_slice[..valid_count];
+ validity_buffer.append_n(valid_count, true);
+ validity_buffer.append_n(nulls_count, false);
+ }
+
+ for (write_index, index) in nulls.valid_indices().enumerate() {
+ values_slice[write_index] = primitive_values.value(index);
+ }
+
+ values_slice.sort_unstable_by(|a, b| a.compare(*b));
+ if sort_options.descending {
+ values_slice.reverse();
+ }
+
+ null_bit_buffer = Some(validity_buffer.finish().into());
+ } else {
+ mutable_slice.copy_from_slice(&input_values[..values.len()]);
+ mutable_slice.sort_unstable_by(|a, b| a.compare(*b));
Review Comment:
Ah forgot about that :)
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1250585716
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +60,211 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type, i8>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type, i16>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type, i32>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type, i64>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type, u8>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type, u16>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type, u32>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type, u64>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type, f32>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type, f64>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type, i32>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type, i64>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType, i32>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType, i32>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType, i64>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType, i32>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType, i128>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType, i64>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn compress_store<U>(input: *const U, mut output: *mut U, mask: u8) -> isize
+where
+ U: ArrowNativeType,
+{
+ let mut offset = 0;
+ if mask != 0 {
+ for i in 0..8 {
+ if (mask & (1 << i)) != 0 {
+ // This is safe since a valid bit i.e bit set to 1 indicates a valid value
+ unsafe {
+ *output = *input.offset(i);
+ offset += 1;
+ output = output.offset(1);
+ }
+ }
+ }
+ }
+ offset
+}
+
+fn create_null_buffer(
+ valid_count: usize,
+ nulls_count: usize,
+ length: usize,
+ sort_options: SortOptions,
+) -> Option<Buffer> {
+ let null_capacity = (length / 8) + (length % 8 != 0) as usize;
+ let mut mutable_null_buffer = MutableBuffer::new(null_capacity * 8);
+ mutable_null_buffer.resize(null_capacity, 0);
+
+ let mutable_null_buffer_slice = mutable_null_buffer.as_slice_mut();
+
+ if valid_count > 0 {
+ let mut count = valid_count;
+ let mut index = 0;
+ if sort_options.nulls_first {
+ let remaining_nulls = nulls_count % 8;
+ index = nulls_count / 8;
+
+ if remaining_nulls != 0 {
+ let valid_values_count = min(8 - remaining_nulls, valid_count);
+ mutable_null_buffer_slice[index] =
+ ((1 << valid_values_count) - 1) << remaining_nulls;
+ count -= valid_values_count;
+ index += 1;
+ }
+ }
+ while count >= 8 {
+ mutable_null_buffer_slice[index] = u8::MAX;
+ index += 1;
+ count -= 8;
+ }
+ if count != 0 {
+ mutable_null_buffer_slice[index] = (1 << count) - 1;
+ }
+ }
+
+ Some(mutable_null_buffer.into())
+}
+
+fn sort_native_type<T, U>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ U: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let values = values.as_primitive::<T>();
+
+ let result_capacity = values.len() * std::mem::size_of::<U>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+ let mutable_slice: &mut [U] = mutable_buffer.typed_data_mut();
+
+ let array_data = values.to_data();
+ let input_values: &[U] = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+ let null_buffer = nulls.buffer().as_slice();
+
+ let mut mutable_slice_ptr = mutable_slice.as_mut_ptr();
+ let mut input_values_ptr = input_values.as_ptr();
+
+ if sort_options.nulls_first {
+ // This is safe since the offset in in bounds
+ unsafe {
+ mutable_slice_ptr = mutable_slice_ptr.add(values.null_count());
+ }
+ }
+
+ // This is safe since we are in bounds
+ let values_slice =
+ unsafe { slice::from_raw_parts_mut(mutable_slice_ptr, valid_count) };
+
+ for mask in null_buffer {
Review Comment:
I believe this is incorrect as it fails to correctly account for any offset
--
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
[GitHub] [arrow-rs] psvri commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "psvri (via GitHub)" <gi...@apache.org>.
psvri commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1251193105
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,137 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn sort_native_type<T>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ <T as arrow_array::ArrowPrimitiveType>::Native: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let primitive_values = values.as_primitive::<T>();
+
+ let result_capacity = values.len()
+ * std::mem::size_of::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+
+ let mutable_slice =
+ mutable_buffer.typed_data_mut::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+
+ let array_data = values.to_data();
+ let input_values = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+
+ let mut validity_buffer = BooleanBufferBuilder::new(values.len());
+ let values_slice;
+
+ if sort_options.nulls_first {
+ values_slice = &mut mutable_slice[nulls_count..];
+ validity_buffer.append_n(nulls_count, false);
+ validity_buffer.append_n(valid_count, true);
+ } else {
+ values_slice = &mut mutable_slice[..valid_count];
+ validity_buffer.append_n(valid_count, true);
+ validity_buffer.append_n(nulls_count, false);
+ }
+
+ for (write_index, index) in nulls.valid_indices().enumerate() {
+ values_slice[write_index] = primitive_values.value(index);
+ }
+
+ values_slice.sort_unstable_by(|a, b| a.compare(*b));
+ if sort_options.descending {
+ values_slice.reverse();
+ }
+
+ null_bit_buffer = Some(validity_buffer.finish().into());
+ } else {
+ mutable_slice.copy_from_slice(&input_values[..values.len()]);
+ mutable_slice.sort_unstable_by(|a, b| a.compare(*b));
+ if sort_options.descending {
+ mutable_slice.reverse();
+ }
Review Comment:
Tested it just now on my laptop. The difference is only b/n 2-3% .
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1251442526
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,137 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
Review Comment:
If you changed sort_native_type to take PrimitiveArray instead of dyn Array you could use the downcast_primitive_array macro here
--
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
[GitHub] [arrow-rs] psvri commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "psvri (via GitHub)" <gi...@apache.org>.
psvri commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1251109415
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +60,211 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type, i8>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type, i16>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type, i32>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type, i64>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type, u8>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type, u16>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type, u32>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type, u64>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type, f32>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type, f64>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type, i32>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type, i64>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType, i32>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType, i32>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType, i64>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType, i32>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType, i128>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType, i64>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn compress_store<U>(input: *const U, mut output: *mut U, mask: u8) -> isize
+where
+ U: ArrowNativeType,
+{
+ let mut offset = 0;
+ if mask != 0 {
+ for i in 0..8 {
+ if (mask & (1 << i)) != 0 {
+ // This is safe since a valid bit i.e bit set to 1 indicates a valid value
+ unsafe {
+ *output = *input.offset(i);
+ offset += 1;
+ output = output.offset(1);
+ }
+ }
+ }
+ }
+ offset
+}
+
+fn create_null_buffer(
+ valid_count: usize,
+ nulls_count: usize,
+ length: usize,
+ sort_options: SortOptions,
+) -> Option<Buffer> {
+ let null_capacity = (length / 8) + (length % 8 != 0) as usize;
+ let mut mutable_null_buffer = MutableBuffer::new(null_capacity * 8);
+ mutable_null_buffer.resize(null_capacity, 0);
+
+ let mutable_null_buffer_slice = mutable_null_buffer.as_slice_mut();
+
+ if valid_count > 0 {
+ let mut count = valid_count;
+ let mut index = 0;
+ if sort_options.nulls_first {
+ let remaining_nulls = nulls_count % 8;
+ index = nulls_count / 8;
+
+ if remaining_nulls != 0 {
+ let valid_values_count = min(8 - remaining_nulls, valid_count);
+ mutable_null_buffer_slice[index] =
+ ((1 << valid_values_count) - 1) << remaining_nulls;
+ count -= valid_values_count;
+ index += 1;
+ }
+ }
+ while count >= 8 {
+ mutable_null_buffer_slice[index] = u8::MAX;
+ index += 1;
+ count -= 8;
+ }
+ if count != 0 {
+ mutable_null_buffer_slice[index] = (1 << count) - 1;
+ }
+ }
+
+ Some(mutable_null_buffer.into())
+}
+
+fn sort_native_type<T, U>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ U: ArrowNativeTypeOp,
Review Comment:
Done
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1251253583
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,137 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn sort_native_type<T>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ <T as arrow_array::ArrowPrimitiveType>::Native: ArrowNativeTypeOp,
Review Comment:
This shouldn't be necessary given the constraints on ArrowPrimitiveType::Native
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,137 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn sort_native_type<T>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ <T as arrow_array::ArrowPrimitiveType>::Native: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let primitive_values = values.as_primitive::<T>();
+
+ let result_capacity = values.len()
+ * std::mem::size_of::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
Review Comment:
Have you considered just using Vec here?
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,137 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn sort_native_type<T>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ <T as arrow_array::ArrowPrimitiveType>::Native: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let primitive_values = values.as_primitive::<T>();
+
+ let result_capacity = values.len()
+ * std::mem::size_of::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+
+ let mutable_slice =
+ mutable_buffer.typed_data_mut::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+
+ let array_data = values.to_data();
+ let input_values = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
Review Comment:
```suggestion
if let Some(nulls) = array.nulls().filter(|n| n.null_count() > 0) {
```
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,137 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn sort_native_type<T>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ <T as arrow_array::ArrowPrimitiveType>::Native: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let primitive_values = values.as_primitive::<T>();
+
+ let result_capacity = values.len()
+ * std::mem::size_of::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+
+ let mutable_slice =
+ mutable_buffer.typed_data_mut::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+
+ let array_data = values.to_data();
+ let input_values = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
Review Comment:
It might be nicer to use an expression style here, rather than using mut
e.g.
let nulls = match array.nulls().filter(|n| n.null_count() > 0) {
Some(nulls) => ...,
None => ...
}
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,137 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn sort_native_type<T>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ <T as arrow_array::ArrowPrimitiveType>::Native: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let primitive_values = values.as_primitive::<T>();
+
+ let result_capacity = values.len()
+ * std::mem::size_of::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+
+ let mutable_slice =
+ mutable_buffer.typed_data_mut::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+
+ let array_data = values.to_data();
+ let input_values = array_data.buffer(0);
Review Comment:
```suggestion
let array = array.as_primitive::<T>();
let input_values = array.values().as_ref();
```
This not only avoids marshaling to ArrayData, but also the code is technically exploiting an implementation detail that PrimitiveArray returns ArrayData with a zero offset
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,137 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn sort_native_type<T>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ <T as arrow_array::ArrowPrimitiveType>::Native: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let primitive_values = values.as_primitive::<T>();
+
+ let result_capacity = values.len()
+ * std::mem::size_of::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+
+ let mutable_slice =
+ mutable_buffer.typed_data_mut::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+
+ let array_data = values.to_data();
+ let input_values = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+
+ let mut validity_buffer = BooleanBufferBuilder::new(values.len());
+ let values_slice;
Review Comment:
I personally prefer the expression style, e.g.
```
let values_slice = match sort_options.nulls_first {
true => ...,
false => ...
}
```
It makes it easier to see what is going on and where the value is being set
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,137 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn sort_native_type<T>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ <T as arrow_array::ArrowPrimitiveType>::Native: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let primitive_values = values.as_primitive::<T>();
+
+ let result_capacity = values.len()
+ * std::mem::size_of::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+
+ let mutable_slice =
+ mutable_buffer.typed_data_mut::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+
+ let array_data = values.to_data();
+ let input_values = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+
+ let mut validity_buffer = BooleanBufferBuilder::new(values.len());
+ let values_slice;
+
+ if sort_options.nulls_first {
+ values_slice = &mut mutable_slice[nulls_count..];
+ validity_buffer.append_n(nulls_count, false);
+ validity_buffer.append_n(valid_count, true);
+ } else {
+ values_slice = &mut mutable_slice[..valid_count];
+ validity_buffer.append_n(valid_count, true);
+ validity_buffer.append_n(nulls_count, false);
+ }
+
+ for (write_index, index) in nulls.valid_indices().enumerate() {
+ values_slice[write_index] = primitive_values.value(index);
+ }
+
+ values_slice.sort_unstable_by(|a, b| a.compare(*b));
+ if sort_options.descending {
+ values_slice.reverse();
+ }
+
+ null_bit_buffer = Some(validity_buffer.finish().into());
+ } else {
+ mutable_slice.copy_from_slice(&input_values[..values.len()]);
+ mutable_slice.sort_unstable_by(|a, b| a.compare(*b));
+ if sort_options.descending {
+ mutable_slice.reverse();
+ }
Review Comment:
I think given the marginal speed difference it makes sense to save on codegen by using reverse :+1:
--
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
[GitHub] [arrow-rs] psvri commented on pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "psvri (via GitHub)" <gi...@apache.org>.
psvri commented on PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#issuecomment-1620347007
Thanks for the above comments and the latest commit addresses them. It simplified the code a lot.
--
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
[GitHub] [arrow-rs] psvri commented on a diff in pull request #4473: Improve primitive type sorts by 13-67%
Posted by "psvri (via GitHub)" <gi...@apache.org>.
psvri commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1248931317
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +60,211 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type, i8>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type, i16>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type, i32>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type, i64>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type, u8>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type, u16>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type, u32>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type, u64>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type, f32>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type, f64>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type, i32>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type, i64>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType, i32>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType, i32>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType, i64>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType, i32>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType, i128>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType, i64>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn compress_store<U>(input: *const U, mut output: *mut U, mask: u8) -> isize
Review Comment:
This approach is preferred as we can simd optimise this fairly easily as shown [here](https://github.com/psvri/rust-simd-sort/blob/main/simd_sort/src/platform/x86/avx512/bit64.rs#L82)
--
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
[GitHub] [arrow-rs] psvri commented on pull request #4473: Improve primitive type sorts by 13-67%
Posted by "psvri (via GitHub)" <gi...@apache.org>.
psvri commented on PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#issuecomment-1616048975
I havent created an issue for this. Let me know if its required.
--
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
[GitHub] [arrow-rs] psvri commented on a diff in pull request #4473: Improve primitive type sorts by 13-67%
Posted by "psvri (via GitHub)" <gi...@apache.org>.
psvri commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1248931317
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +60,211 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type, i8>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type, i16>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type, i32>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type, i64>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type, u8>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type, u16>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type, u32>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type, u64>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type, f32>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type, f64>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type, i32>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type, i64>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType, i32>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType, i32>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType, i64>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType, i32>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType, i128>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType, i64>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn compress_store<U>(input: *const U, mut output: *mut U, mask: u8) -> isize
Review Comment:
This approach is preferred as we can simd optimise this fairly easily as shown [here](https://github.com/psvri/rust-simd-sort/blob/main/simd_sort/src/platform/x86/avx512/bit64.rs#L82) for some architectures
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1250585716
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +60,211 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type, i8>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type, i16>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type, i32>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type, i64>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type, u8>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type, u16>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type, u32>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type, u64>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type, f32>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type, f64>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type, i32>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type, i64>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType, i32>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType, i32>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType, i64>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType, i32>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType, i128>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType, i64>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn compress_store<U>(input: *const U, mut output: *mut U, mask: u8) -> isize
+where
+ U: ArrowNativeType,
+{
+ let mut offset = 0;
+ if mask != 0 {
+ for i in 0..8 {
+ if (mask & (1 << i)) != 0 {
+ // This is safe since a valid bit i.e bit set to 1 indicates a valid value
+ unsafe {
+ *output = *input.offset(i);
+ offset += 1;
+ output = output.offset(1);
+ }
+ }
+ }
+ }
+ offset
+}
+
+fn create_null_buffer(
+ valid_count: usize,
+ nulls_count: usize,
+ length: usize,
+ sort_options: SortOptions,
+) -> Option<Buffer> {
+ let null_capacity = (length / 8) + (length % 8 != 0) as usize;
+ let mut mutable_null_buffer = MutableBuffer::new(null_capacity * 8);
+ mutable_null_buffer.resize(null_capacity, 0);
+
+ let mutable_null_buffer_slice = mutable_null_buffer.as_slice_mut();
+
+ if valid_count > 0 {
+ let mut count = valid_count;
+ let mut index = 0;
+ if sort_options.nulls_first {
+ let remaining_nulls = nulls_count % 8;
+ index = nulls_count / 8;
+
+ if remaining_nulls != 0 {
+ let valid_values_count = min(8 - remaining_nulls, valid_count);
+ mutable_null_buffer_slice[index] =
+ ((1 << valid_values_count) - 1) << remaining_nulls;
+ count -= valid_values_count;
+ index += 1;
+ }
+ }
+ while count >= 8 {
+ mutable_null_buffer_slice[index] = u8::MAX;
+ index += 1;
+ count -= 8;
+ }
+ if count != 0 {
+ mutable_null_buffer_slice[index] = (1 << count) - 1;
+ }
+ }
+
+ Some(mutable_null_buffer.into())
+}
+
+fn sort_native_type<T, U>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ U: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let values = values.as_primitive::<T>();
+
+ let result_capacity = values.len() * std::mem::size_of::<U>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+ let mutable_slice: &mut [U] = mutable_buffer.typed_data_mut();
+
+ let array_data = values.to_data();
+ let input_values: &[U] = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+ let null_buffer = nulls.buffer().as_slice();
+
+ let mut mutable_slice_ptr = mutable_slice.as_mut_ptr();
+ let mut input_values_ptr = input_values.as_ptr();
+
+ if sort_options.nulls_first {
+ // This is safe since the offset in in bounds
+ unsafe {
+ mutable_slice_ptr = mutable_slice_ptr.add(values.null_count());
+ }
+ }
+
+ // This is safe since we are in bounds
+ let values_slice =
+ unsafe { slice::from_raw_parts_mut(mutable_slice_ptr, valid_count) };
+
+ for mask in null_buffer {
Review Comment:
This is incorrect as it fails to correctly account for any offset present in the null buffer, you likely need to use bit_chunks or iter_valid
--
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
[GitHub] [arrow-rs] psvri commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "psvri (via GitHub)" <gi...@apache.org>.
psvri commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1251109791
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +60,211 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type, i8>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type, i16>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type, i32>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type, i64>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type, u8>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type, u16>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type, u32>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type, u64>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type, f32>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type, f64>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type, i32>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type, i64>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType, i32>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType, i32>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType, i64>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType, i32>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType, i128>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType, i64>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn compress_store<U>(input: *const U, mut output: *mut U, mask: u8) -> isize
+where
+ U: ArrowNativeType,
+{
+ let mut offset = 0;
+ if mask != 0 {
+ for i in 0..8 {
+ if (mask & (1 << i)) != 0 {
+ // This is safe since a valid bit i.e bit set to 1 indicates a valid value
+ unsafe {
+ *output = *input.offset(i);
+ offset += 1;
+ output = output.offset(1);
+ }
+ }
+ }
+ }
+ offset
+}
+
+fn create_null_buffer(
+ valid_count: usize,
+ nulls_count: usize,
+ length: usize,
+ sort_options: SortOptions,
+) -> Option<Buffer> {
+ let null_capacity = (length / 8) + (length % 8 != 0) as usize;
+ let mut mutable_null_buffer = MutableBuffer::new(null_capacity * 8);
+ mutable_null_buffer.resize(null_capacity, 0);
+
+ let mutable_null_buffer_slice = mutable_null_buffer.as_slice_mut();
+
+ if valid_count > 0 {
+ let mut count = valid_count;
+ let mut index = 0;
+ if sort_options.nulls_first {
+ let remaining_nulls = nulls_count % 8;
+ index = nulls_count / 8;
+
+ if remaining_nulls != 0 {
+ let valid_values_count = min(8 - remaining_nulls, valid_count);
+ mutable_null_buffer_slice[index] =
+ ((1 << valid_values_count) - 1) << remaining_nulls;
+ count -= valid_values_count;
+ index += 1;
+ }
+ }
+ while count >= 8 {
+ mutable_null_buffer_slice[index] = u8::MAX;
+ index += 1;
+ count -= 8;
+ }
+ if count != 0 {
+ mutable_null_buffer_slice[index] = (1 << count) - 1;
+ }
+ }
+
+ Some(mutable_null_buffer.into())
+}
+
+fn sort_native_type<T, U>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ U: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let values = values.as_primitive::<T>();
+
+ let result_capacity = values.len() * std::mem::size_of::<U>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+ let mutable_slice: &mut [U] = mutable_buffer.typed_data_mut();
+
+ let array_data = values.to_data();
+ let input_values: &[U] = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+ let null_buffer = nulls.buffer().as_slice();
+
+ let mut mutable_slice_ptr = mutable_slice.as_mut_ptr();
+ let mut input_values_ptr = input_values.as_ptr();
+
+ if sort_options.nulls_first {
+ // This is safe since the offset in in bounds
+ unsafe {
+ mutable_slice_ptr = mutable_slice_ptr.add(values.null_count());
+ }
+ }
+
+ // This is safe since we are in bounds
+ let values_slice =
+ unsafe { slice::from_raw_parts_mut(mutable_slice_ptr, valid_count) };
+
+ for mask in null_buffer {
Review Comment:
I had forgotten about the presence of offset in buffers . Should be fixed in latest commit
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1252128677
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,74 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ downcast_primitive_array!(
+ values => return sort_native_type(values, options),
+ DataType::RunEndEncoded(_, _) => return sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ return take(values, &indices, None)
Review Comment:
```suggestion
values => sort_native_type(values, options),
DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
_ => {
let indices = sort_to_indices(values, options, None)?;
take(values, &indices, None)
```
--
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
[GitHub] [arrow-rs] tustvold merged pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold merged PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1250580989
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +60,211 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type, i8>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type, i16>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type, i32>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type, i64>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type, u8>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type, u16>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type, u32>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type, u64>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type, f32>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type, f64>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type, i32>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type, i64>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType, i32>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType, i32>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType, i64>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType, i32>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType, i128>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType, i64>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn compress_store<U>(input: *const U, mut output: *mut U, mask: u8) -> isize
+where
+ U: ArrowNativeType,
+{
+ let mut offset = 0;
+ if mask != 0 {
+ for i in 0..8 {
+ if (mask & (1 << i)) != 0 {
+ // This is safe since a valid bit i.e bit set to 1 indicates a valid value
+ unsafe {
+ *output = *input.offset(i);
+ offset += 1;
+ output = output.offset(1);
+ }
+ }
+ }
+ }
+ offset
+}
+
+fn create_null_buffer(
+ valid_count: usize,
+ nulls_count: usize,
+ length: usize,
+ sort_options: SortOptions,
+) -> Option<Buffer> {
+ let null_capacity = (length / 8) + (length % 8 != 0) as usize;
+ let mut mutable_null_buffer = MutableBuffer::new(null_capacity * 8);
+ mutable_null_buffer.resize(null_capacity, 0);
+
+ let mutable_null_buffer_slice = mutable_null_buffer.as_slice_mut();
+
+ if valid_count > 0 {
+ let mut count = valid_count;
+ let mut index = 0;
+ if sort_options.nulls_first {
+ let remaining_nulls = nulls_count % 8;
+ index = nulls_count / 8;
+
+ if remaining_nulls != 0 {
+ let valid_values_count = min(8 - remaining_nulls, valid_count);
+ mutable_null_buffer_slice[index] =
+ ((1 << valid_values_count) - 1) << remaining_nulls;
+ count -= valid_values_count;
+ index += 1;
+ }
+ }
+ while count >= 8 {
+ mutable_null_buffer_slice[index] = u8::MAX;
+ index += 1;
+ count -= 8;
+ }
+ if count != 0 {
+ mutable_null_buffer_slice[index] = (1 << count) - 1;
+ }
+ }
+
+ Some(mutable_null_buffer.into())
+}
+
+fn sort_native_type<T, U>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ U: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let values = values.as_primitive::<T>();
+
+ let result_capacity = values.len() * std::mem::size_of::<U>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+ let mutable_slice: &mut [U] = mutable_buffer.typed_data_mut();
+
+ let array_data = values.to_data();
+ let input_values: &[U] = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+ let null_buffer = nulls.buffer().as_slice();
+
+ let mut mutable_slice_ptr = mutable_slice.as_mut_ptr();
+ let mut input_values_ptr = input_values.as_ptr();
+
+ if sort_options.nulls_first {
+ // This is safe since the offset in in bounds
+ unsafe {
+ mutable_slice_ptr = mutable_slice_ptr.add(values.null_count());
+ }
+ }
+
+ // This is safe since we are in bounds
+ let values_slice =
+ unsafe { slice::from_raw_parts_mut(mutable_slice_ptr, valid_count) };
+
+ for mask in null_buffer {
+ let written_count =
+ compress_store::<U>(input_values_ptr, mutable_slice_ptr, *mask);
+ // This is safe as the offset increments are within bounds
+ unsafe {
+ input_values_ptr = input_values_ptr.offset(8);
+ mutable_slice_ptr = mutable_slice_ptr.offset(written_count);
+ }
+ }
+
+ values_slice.sort_unstable_by(|a, b| a.compare(*b));
+ if sort_options.descending {
+ values_slice.reverse();
+ }
+
+ null_bit_buffer =
+ create_null_buffer(valid_count, nulls_count, values.len(), sort_options);
+ } else {
+ mutable_slice.copy_from_slice(input_values);
+ mutable_slice.sort_unstable_by(|a, b| a.compare(*b));
+ if sort_options.descending {
+ mutable_slice.reverse();
+ }
+ }
+ // This is safe since data types match
Review Comment:
Perhaps we could use `PrimitiveArray::new().with_data_type()` instead
--
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
[GitHub] [arrow-rs] psvri commented on pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "psvri (via GitHub)" <gi...@apache.org>.
psvri commented on PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#issuecomment-1618879088
I have reworked the code to have no unsafe and reused
> Is there some way we might reduce the amount of unsafe code here, given this is a rare special case (where you don't need the indices to sort other columns) I'm keen to keep the maintenance overheads down.
I have removed all unsafe code in the latest commit.
--
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
[GitHub] [arrow-rs] Dandandan commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1251252335
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,137 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn sort_native_type<T>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ <T as arrow_array::ArrowPrimitiveType>::Native: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let primitive_values = values.as_primitive::<T>();
+
+ let result_capacity = values.len()
+ * std::mem::size_of::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+
+ let mutable_slice =
+ mutable_buffer.typed_data_mut::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+
+ let array_data = values.to_data();
+ let input_values = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+
+ let mut validity_buffer = BooleanBufferBuilder::new(values.len());
+ let values_slice;
+
+ if sort_options.nulls_first {
+ values_slice = &mut mutable_slice[nulls_count..];
+ validity_buffer.append_n(nulls_count, false);
+ validity_buffer.append_n(valid_count, true);
+ } else {
+ values_slice = &mut mutable_slice[..valid_count];
+ validity_buffer.append_n(valid_count, true);
+ validity_buffer.append_n(nulls_count, false);
+ }
+
+ for (write_index, index) in nulls.valid_indices().enumerate() {
+ values_slice[write_index] = primitive_values.value(index);
+ }
+
+ values_slice.sort_unstable_by(|a, b| a.compare(*b));
+ if sort_options.descending {
+ values_slice.reverse();
+ }
+
+ null_bit_buffer = Some(validity_buffer.finish().into());
+ } else {
+ mutable_slice.copy_from_slice(&input_values[..values.len()]);
+ mutable_slice.sort_unstable_by(|a, b| a.compare(*b));
+ if sort_options.descending {
+ mutable_slice.reverse();
+ }
Review Comment:
Thanks for checking. I would argue it's also a bit more simple :)
--
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
[GitHub] [arrow-rs] psvri commented on pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "psvri (via GitHub)" <gi...@apache.org>.
psvri commented on PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#issuecomment-1619379917
I will make these changes today.
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1250585716
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +60,211 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type, i8>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type, i16>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type, i32>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type, i64>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type, u8>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type, u16>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type, u32>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type, u64>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type, f32>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type, f64>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type, i32>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type, i64>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType, i32>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType, i32>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType, i64>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType, i32>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType, i128>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType, i64>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn compress_store<U>(input: *const U, mut output: *mut U, mask: u8) -> isize
+where
+ U: ArrowNativeType,
+{
+ let mut offset = 0;
+ if mask != 0 {
+ for i in 0..8 {
+ if (mask & (1 << i)) != 0 {
+ // This is safe since a valid bit i.e bit set to 1 indicates a valid value
+ unsafe {
+ *output = *input.offset(i);
+ offset += 1;
+ output = output.offset(1);
+ }
+ }
+ }
+ }
+ offset
+}
+
+fn create_null_buffer(
+ valid_count: usize,
+ nulls_count: usize,
+ length: usize,
+ sort_options: SortOptions,
+) -> Option<Buffer> {
+ let null_capacity = (length / 8) + (length % 8 != 0) as usize;
+ let mut mutable_null_buffer = MutableBuffer::new(null_capacity * 8);
+ mutable_null_buffer.resize(null_capacity, 0);
+
+ let mutable_null_buffer_slice = mutable_null_buffer.as_slice_mut();
+
+ if valid_count > 0 {
+ let mut count = valid_count;
+ let mut index = 0;
+ if sort_options.nulls_first {
+ let remaining_nulls = nulls_count % 8;
+ index = nulls_count / 8;
+
+ if remaining_nulls != 0 {
+ let valid_values_count = min(8 - remaining_nulls, valid_count);
+ mutable_null_buffer_slice[index] =
+ ((1 << valid_values_count) - 1) << remaining_nulls;
+ count -= valid_values_count;
+ index += 1;
+ }
+ }
+ while count >= 8 {
+ mutable_null_buffer_slice[index] = u8::MAX;
+ index += 1;
+ count -= 8;
+ }
+ if count != 0 {
+ mutable_null_buffer_slice[index] = (1 << count) - 1;
+ }
+ }
+
+ Some(mutable_null_buffer.into())
+}
+
+fn sort_native_type<T, U>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ U: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let values = values.as_primitive::<T>();
+
+ let result_capacity = values.len() * std::mem::size_of::<U>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+ let mutable_slice: &mut [U] = mutable_buffer.typed_data_mut();
+
+ let array_data = values.to_data();
+ let input_values: &[U] = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+ let null_buffer = nulls.buffer().as_slice();
+
+ let mut mutable_slice_ptr = mutable_slice.as_mut_ptr();
+ let mut input_values_ptr = input_values.as_ptr();
+
+ if sort_options.nulls_first {
+ // This is safe since the offset in in bounds
+ unsafe {
+ mutable_slice_ptr = mutable_slice_ptr.add(values.null_count());
+ }
+ }
+
+ // This is safe since we are in bounds
+ let values_slice =
+ unsafe { slice::from_raw_parts_mut(mutable_slice_ptr, valid_count) };
+
+ for mask in null_buffer {
Review Comment:
This is incorrect as it fails to correctly account for any offset present in the null buffer, you likely need to use bit_chunks or [valid_indices](https://docs.rs/arrow-buffer/latest/arrow_buffer/buffer/struct.NullBuffer.html#method.valid_indices)
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1250585716
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +60,211 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type, i8>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type, i16>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type, i32>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type, i64>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type, u8>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type, u16>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type, u32>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type, u64>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type, f32>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type, f64>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type, i32>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type, i64>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType, i32>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType, i32>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType, i64>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType, i32>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType, i128>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType, i64>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn compress_store<U>(input: *const U, mut output: *mut U, mask: u8) -> isize
+where
+ U: ArrowNativeType,
+{
+ let mut offset = 0;
+ if mask != 0 {
+ for i in 0..8 {
+ if (mask & (1 << i)) != 0 {
+ // This is safe since a valid bit i.e bit set to 1 indicates a valid value
+ unsafe {
+ *output = *input.offset(i);
+ offset += 1;
+ output = output.offset(1);
+ }
+ }
+ }
+ }
+ offset
+}
+
+fn create_null_buffer(
+ valid_count: usize,
+ nulls_count: usize,
+ length: usize,
+ sort_options: SortOptions,
+) -> Option<Buffer> {
+ let null_capacity = (length / 8) + (length % 8 != 0) as usize;
+ let mut mutable_null_buffer = MutableBuffer::new(null_capacity * 8);
+ mutable_null_buffer.resize(null_capacity, 0);
+
+ let mutable_null_buffer_slice = mutable_null_buffer.as_slice_mut();
+
+ if valid_count > 0 {
+ let mut count = valid_count;
+ let mut index = 0;
+ if sort_options.nulls_first {
+ let remaining_nulls = nulls_count % 8;
+ index = nulls_count / 8;
+
+ if remaining_nulls != 0 {
+ let valid_values_count = min(8 - remaining_nulls, valid_count);
+ mutable_null_buffer_slice[index] =
+ ((1 << valid_values_count) - 1) << remaining_nulls;
+ count -= valid_values_count;
+ index += 1;
+ }
+ }
+ while count >= 8 {
+ mutable_null_buffer_slice[index] = u8::MAX;
+ index += 1;
+ count -= 8;
+ }
+ if count != 0 {
+ mutable_null_buffer_slice[index] = (1 << count) - 1;
+ }
+ }
+
+ Some(mutable_null_buffer.into())
+}
+
+fn sort_native_type<T, U>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ U: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let values = values.as_primitive::<T>();
+
+ let result_capacity = values.len() * std::mem::size_of::<U>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+ let mutable_slice: &mut [U] = mutable_buffer.typed_data_mut();
+
+ let array_data = values.to_data();
+ let input_values: &[U] = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+ let null_buffer = nulls.buffer().as_slice();
+
+ let mut mutable_slice_ptr = mutable_slice.as_mut_ptr();
+ let mut input_values_ptr = input_values.as_ptr();
+
+ if sort_options.nulls_first {
+ // This is safe since the offset in in bounds
+ unsafe {
+ mutable_slice_ptr = mutable_slice_ptr.add(values.null_count());
+ }
+ }
+
+ // This is safe since we are in bounds
+ let values_slice =
+ unsafe { slice::from_raw_parts_mut(mutable_slice_ptr, valid_count) };
+
+ for mask in null_buffer {
Review Comment:
This is incorrect as it fails to correctly account for any offset present in the null buffer
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1250589882
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +60,211 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type, i8>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type, i16>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type, i32>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type, i64>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type, u8>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type, u16>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type, u32>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type, u64>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type, f32>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type, f64>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type, i32>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type, i64>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType, i32>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType, i32>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType, i64>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType, i32>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType, i128>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType, i64>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn compress_store<U>(input: *const U, mut output: *mut U, mask: u8) -> isize
+where
+ U: ArrowNativeType,
+{
+ let mut offset = 0;
+ if mask != 0 {
+ for i in 0..8 {
+ if (mask & (1 << i)) != 0 {
+ // This is safe since a valid bit i.e bit set to 1 indicates a valid value
+ unsafe {
+ *output = *input.offset(i);
+ offset += 1;
+ output = output.offset(1);
+ }
+ }
+ }
+ }
+ offset
+}
+
+fn create_null_buffer(
+ valid_count: usize,
+ nulls_count: usize,
+ length: usize,
+ sort_options: SortOptions,
+) -> Option<Buffer> {
+ let null_capacity = (length / 8) + (length % 8 != 0) as usize;
+ let mut mutable_null_buffer = MutableBuffer::new(null_capacity * 8);
+ mutable_null_buffer.resize(null_capacity, 0);
+
+ let mutable_null_buffer_slice = mutable_null_buffer.as_slice_mut();
+
+ if valid_count > 0 {
+ let mut count = valid_count;
+ let mut index = 0;
+ if sort_options.nulls_first {
+ let remaining_nulls = nulls_count % 8;
+ index = nulls_count / 8;
+
+ if remaining_nulls != 0 {
+ let valid_values_count = min(8 - remaining_nulls, valid_count);
+ mutable_null_buffer_slice[index] =
+ ((1 << valid_values_count) - 1) << remaining_nulls;
+ count -= valid_values_count;
+ index += 1;
+ }
+ }
+ while count >= 8 {
+ mutable_null_buffer_slice[index] = u8::MAX;
+ index += 1;
+ count -= 8;
+ }
+ if count != 0 {
+ mutable_null_buffer_slice[index] = (1 << count) - 1;
+ }
+ }
+
+ Some(mutable_null_buffer.into())
+}
+
+fn sort_native_type<T, U>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ U: ArrowNativeTypeOp,
Review Comment:
Could we just use `T::Native` instead of a separate type parameter?
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1250583243
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +60,211 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type, i8>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type, i16>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type, i32>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type, i64>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type, u8>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type, u16>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type, u32>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type, u64>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type, f32>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type, f64>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type, i32>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type, i64>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType, i32>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType, i32>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType, i64>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType, i32>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType, i128>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType, i64>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn compress_store<U>(input: *const U, mut output: *mut U, mask: u8) -> isize
+where
+ U: ArrowNativeType,
+{
+ let mut offset = 0;
+ if mask != 0 {
+ for i in 0..8 {
+ if (mask & (1 << i)) != 0 {
+ // This is safe since a valid bit i.e bit set to 1 indicates a valid value
+ unsafe {
+ *output = *input.offset(i);
+ offset += 1;
+ output = output.offset(1);
+ }
+ }
+ }
+ }
+ offset
+}
+
+fn create_null_buffer(
+ valid_count: usize,
+ nulls_count: usize,
+ length: usize,
+ sort_options: SortOptions,
+) -> Option<Buffer> {
+ let null_capacity = (length / 8) + (length % 8 != 0) as usize;
Review Comment:
This appears to duplicate logic found in BooleanBufferBuilder? In particular https://docs.rs/arrow-buffer/latest/src/arrow_buffer/builder/boolean.rs.html#147-170
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1250583243
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +60,211 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type, i8>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type, i16>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type, i32>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type, i64>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type, u8>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type, u16>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type, u32>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type, u64>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type, f32>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type, f64>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type, i32>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type, i64>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType, i32>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType, i32>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType, i64>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType, i32>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType, i128>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType, i64>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn compress_store<U>(input: *const U, mut output: *mut U, mask: u8) -> isize
+where
+ U: ArrowNativeType,
+{
+ let mut offset = 0;
+ if mask != 0 {
+ for i in 0..8 {
+ if (mask & (1 << i)) != 0 {
+ // This is safe since a valid bit i.e bit set to 1 indicates a valid value
+ unsafe {
+ *output = *input.offset(i);
+ offset += 1;
+ output = output.offset(1);
+ }
+ }
+ }
+ }
+ offset
+}
+
+fn create_null_buffer(
+ valid_count: usize,
+ nulls_count: usize,
+ length: usize,
+ sort_options: SortOptions,
+) -> Option<Buffer> {
+ let null_capacity = (length / 8) + (length % 8 != 0) as usize;
Review Comment:
This appears to duplicate logic found in BooleanBufferBuilder?
--
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
[GitHub] [arrow-rs] tustvold commented on pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#issuecomment-1617740976
Is there some way we might reduce the amount of unsafe code here, given this is a rare special case (where you don't need the indices to sort other columns) I'm keen to keep the maintenance overheads down.
--
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
[GitHub] [arrow-rs] psvri commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "psvri (via GitHub)" <gi...@apache.org>.
psvri commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1251111511
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +60,211 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type, i8>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type, i16>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type, i32>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type, i64>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type, u8>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type, u16>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type, u32>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type, u64>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type, f32>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type, f64>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type, i32>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type, i64>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType, i32>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType, i32>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType, i64>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType, i32>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType, i128>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType, i64>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn compress_store<U>(input: *const U, mut output: *mut U, mask: u8) -> isize
+where
+ U: ArrowNativeType,
+{
+ let mut offset = 0;
+ if mask != 0 {
+ for i in 0..8 {
+ if (mask & (1 << i)) != 0 {
+ // This is safe since a valid bit i.e bit set to 1 indicates a valid value
+ unsafe {
+ *output = *input.offset(i);
+ offset += 1;
+ output = output.offset(1);
+ }
+ }
+ }
+ }
+ offset
+}
+
+fn create_null_buffer(
+ valid_count: usize,
+ nulls_count: usize,
+ length: usize,
+ sort_options: SortOptions,
+) -> Option<Buffer> {
+ let null_capacity = (length / 8) + (length % 8 != 0) as usize;
Review Comment:
Thanks for this pointer. I have used it in the latest commit.
--
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
[GitHub] [arrow-rs] Dandandan commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1251177488
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,137 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn sort_native_type<T>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ <T as arrow_array::ArrowPrimitiveType>::Native: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let primitive_values = values.as_primitive::<T>();
+
+ let result_capacity = values.len()
+ * std::mem::size_of::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+
+ let mutable_slice =
+ mutable_buffer.typed_data_mut::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+
+ let array_data = values.to_data();
+ let input_values = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+
+ let mut validity_buffer = BooleanBufferBuilder::new(values.len());
+ let values_slice;
+
+ if sort_options.nulls_first {
+ values_slice = &mut mutable_slice[nulls_count..];
+ validity_buffer.append_n(nulls_count, false);
+ validity_buffer.append_n(valid_count, true);
+ } else {
+ values_slice = &mut mutable_slice[..valid_count];
+ validity_buffer.append_n(valid_count, true);
+ validity_buffer.append_n(nulls_count, false);
+ }
+
+ for (write_index, index) in nulls.valid_indices().enumerate() {
+ values_slice[write_index] = primitive_values.value(index);
+ }
+
+ values_slice.sort_unstable_by(|a, b| a.compare(*b));
+ if sort_options.descending {
+ values_slice.reverse();
+ }
+
+ null_bit_buffer = Some(validity_buffer.finish().into());
+ } else {
+ mutable_slice.copy_from_slice(&input_values[..values.len()]);
+ mutable_slice.sort_unstable_by(|a, b| a.compare(*b));
+ if sort_options.descending {
+ mutable_slice.reverse();
+ }
Review Comment:
This looks like it should be faster for the descending case?
```suggestion
if sort_options.descending {
mutable_slice.sort_unstable_by(|a, b| b.compare(*a));
} else {
mutable_slice.sort_unstable_by(|a, b| a.compare(*b));
}
```
--
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
[GitHub] [arrow-rs] psvri commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "psvri (via GitHub)" <gi...@apache.org>.
psvri commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1251110673
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +60,211 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type, i8>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type, i16>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type, i32>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type, i64>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type, u8>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type, u16>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type, u32>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type, u64>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type, f32>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type, f64>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type, i32>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type, i64>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType, i32>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType, i32>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType, i64>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType, i32>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType, i128>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType, i64>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn compress_store<U>(input: *const U, mut output: *mut U, mask: u8) -> isize
+where
+ U: ArrowNativeType,
+{
+ let mut offset = 0;
+ if mask != 0 {
+ for i in 0..8 {
+ if (mask & (1 << i)) != 0 {
+ // This is safe since a valid bit i.e bit set to 1 indicates a valid value
+ unsafe {
+ *output = *input.offset(i);
+ offset += 1;
+ output = output.offset(1);
+ }
+ }
+ }
+ }
+ offset
+}
+
+fn create_null_buffer(
+ valid_count: usize,
+ nulls_count: usize,
+ length: usize,
+ sort_options: SortOptions,
+) -> Option<Buffer> {
+ let null_capacity = (length / 8) + (length % 8 != 0) as usize;
+ let mut mutable_null_buffer = MutableBuffer::new(null_capacity * 8);
+ mutable_null_buffer.resize(null_capacity, 0);
+
+ let mutable_null_buffer_slice = mutable_null_buffer.as_slice_mut();
+
+ if valid_count > 0 {
+ let mut count = valid_count;
+ let mut index = 0;
+ if sort_options.nulls_first {
+ let remaining_nulls = nulls_count % 8;
+ index = nulls_count / 8;
+
+ if remaining_nulls != 0 {
+ let valid_values_count = min(8 - remaining_nulls, valid_count);
+ mutable_null_buffer_slice[index] =
+ ((1 << valid_values_count) - 1) << remaining_nulls;
+ count -= valid_values_count;
+ index += 1;
+ }
+ }
+ while count >= 8 {
+ mutable_null_buffer_slice[index] = u8::MAX;
+ index += 1;
+ count -= 8;
+ }
+ if count != 0 {
+ mutable_null_buffer_slice[index] = (1 << count) - 1;
+ }
+ }
+
+ Some(mutable_null_buffer.into())
+}
+
+fn sort_native_type<T, U>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ U: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let values = values.as_primitive::<T>();
+
+ let result_capacity = values.len() * std::mem::size_of::<U>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+ let mutable_slice: &mut [U] = mutable_buffer.typed_data_mut();
+
+ let array_data = values.to_data();
+ let input_values: &[U] = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+ let null_buffer = nulls.buffer().as_slice();
+
+ let mut mutable_slice_ptr = mutable_slice.as_mut_ptr();
+ let mut input_values_ptr = input_values.as_ptr();
+
+ if sort_options.nulls_first {
+ // This is safe since the offset in in bounds
+ unsafe {
+ mutable_slice_ptr = mutable_slice_ptr.add(values.null_count());
+ }
+ }
+
+ // This is safe since we are in bounds
+ let values_slice =
+ unsafe { slice::from_raw_parts_mut(mutable_slice_ptr, valid_count) };
+
+ for mask in null_buffer {
+ let written_count =
+ compress_store::<U>(input_values_ptr, mutable_slice_ptr, *mask);
+ // This is safe as the offset increments are within bounds
+ unsafe {
+ input_values_ptr = input_values_ptr.offset(8);
+ mutable_slice_ptr = mutable_slice_ptr.offset(written_count);
+ }
+ }
+
+ values_slice.sort_unstable_by(|a, b| a.compare(*b));
+ if sort_options.descending {
+ values_slice.reverse();
+ }
+
+ null_bit_buffer =
+ create_null_buffer(valid_count, nulls_count, values.len(), sort_options);
+ } else {
+ mutable_slice.copy_from_slice(input_values);
+ mutable_slice.sort_unstable_by(|a, b| a.compare(*b));
+ if sort_options.descending {
+ mutable_slice.reverse();
+ }
+ }
+ // This is safe since data types match
Review Comment:
Done
--
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
[GitHub] [arrow-rs] Dandandan commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1251263331
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,137 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn sort_native_type<T>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ <T as arrow_array::ArrowPrimitiveType>::Native: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let primitive_values = values.as_primitive::<T>();
+
+ let result_capacity = values.len()
+ * std::mem::size_of::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+
+ let mutable_slice =
+ mutable_buffer.typed_data_mut::<<T as arrow_array::ArrowPrimitiveType>::Native>();
+
+ let array_data = values.to_data();
+ let input_values = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+
+ let mut validity_buffer = BooleanBufferBuilder::new(values.len());
+ let values_slice;
+
+ if sort_options.nulls_first {
+ values_slice = &mut mutable_slice[nulls_count..];
+ validity_buffer.append_n(nulls_count, false);
+ validity_buffer.append_n(valid_count, true);
+ } else {
+ values_slice = &mut mutable_slice[..valid_count];
+ validity_buffer.append_n(valid_count, true);
+ validity_buffer.append_n(nulls_count, false);
+ }
+
+ for (write_index, index) in nulls.valid_indices().enumerate() {
+ values_slice[write_index] = primitive_values.value(index);
+ }
+
+ values_slice.sort_unstable_by(|a, b| a.compare(*b));
+ if sort_options.descending {
+ values_slice.reverse();
+ }
+
+ null_bit_buffer = Some(validity_buffer.finish().into());
+ } else {
+ mutable_slice.copy_from_slice(&input_values[..values.len()]);
+ mutable_slice.sort_unstable_by(|a, b| a.compare(*b));
Review Comment:
```suggestion
mutable_slice.sort_unstable();
```
Should be the same?
--
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
[GitHub] [arrow-rs] psvri commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "psvri (via GitHub)" <gi...@apache.org>.
psvri commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1252138507
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +58,74 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ downcast_primitive_array!(
+ values => return sort_native_type(values, options),
+ DataType::RunEndEncoded(_, _) => return sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ return take(values, &indices, None)
Review Comment:
Done
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4473: Improve in-place primitive sorts by 13-67%
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4473:
URL: https://github.com/apache/arrow-rs/pull/4473#discussion_r1250585716
##########
arrow-ord/src/sort.rs:
##########
@@ -57,11 +60,211 @@ pub fn sort(
values: &dyn Array,
options: Option<SortOptions>,
) -> Result<ArrayRef, ArrowError> {
- if let DataType::RunEndEncoded(_, _) = values.data_type() {
- return sort_run(values, options, None);
+ match values.data_type() {
+ DataType::Int8 => sort_native_type::<Int8Type, i8>(values, options),
+ DataType::Int16 => sort_native_type::<Int16Type, i16>(values, options),
+ DataType::Int32 => sort_native_type::<Int32Type, i32>(values, options),
+ DataType::Int64 => sort_native_type::<Int64Type, i64>(values, options),
+ DataType::UInt8 => sort_native_type::<UInt8Type, u8>(values, options),
+ DataType::UInt16 => sort_native_type::<UInt16Type, u16>(values, options),
+ DataType::UInt32 => sort_native_type::<UInt32Type, u32>(values, options),
+ DataType::UInt64 => sort_native_type::<UInt64Type, u64>(values, options),
+ DataType::Float32 => sort_native_type::<Float32Type, f32>(values, options),
+ DataType::Float64 => sort_native_type::<Float64Type, f64>(values, options),
+ DataType::Date32 => sort_native_type::<Date32Type, i32>(values, options),
+ DataType::Date64 => sort_native_type::<Date64Type, i64>(values, options),
+ DataType::Time32(TimeUnit::Second) => {
+ sort_native_type::<Time32SecondType, i32>(values, options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ sort_native_type::<Time32MillisecondType, i32>(values, options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ sort_native_type::<Time64MicrosecondType, i64>(values, options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ sort_native_type::<Time64NanosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Second, _) => {
+ sort_native_type::<TimestampSecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ sort_native_type::<TimestampMillisecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ sort_native_type::<TimestampMicrosecondType, i64>(values, options)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ sort_native_type::<TimestampNanosecondType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::YearMonth) => {
+ sort_native_type::<IntervalYearMonthType, i32>(values, options)
+ }
+ DataType::Interval(IntervalUnit::DayTime) => {
+ sort_native_type::<IntervalDayTimeType, i64>(values, options)
+ }
+ DataType::Interval(IntervalUnit::MonthDayNano) => {
+ sort_native_type::<IntervalMonthDayNanoType, i128>(values, options)
+ }
+ DataType::Duration(TimeUnit::Second) => {
+ sort_native_type::<DurationSecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Millisecond) => {
+ sort_native_type::<DurationMillisecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Microsecond) => {
+ sort_native_type::<DurationMicrosecondType, i64>(values, options)
+ }
+ DataType::Duration(TimeUnit::Nanosecond) => {
+ sort_native_type::<DurationNanosecondType, i64>(values, options)
+ }
+ DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
+ _ => {
+ let indices = sort_to_indices(values, options, None)?;
+ take(values, &indices, None)
+ }
}
- let indices = sort_to_indices(values, options, None)?;
- take(values, &indices, None)
+}
+
+fn compress_store<U>(input: *const U, mut output: *mut U, mask: u8) -> isize
+where
+ U: ArrowNativeType,
+{
+ let mut offset = 0;
+ if mask != 0 {
+ for i in 0..8 {
+ if (mask & (1 << i)) != 0 {
+ // This is safe since a valid bit i.e bit set to 1 indicates a valid value
+ unsafe {
+ *output = *input.offset(i);
+ offset += 1;
+ output = output.offset(1);
+ }
+ }
+ }
+ }
+ offset
+}
+
+fn create_null_buffer(
+ valid_count: usize,
+ nulls_count: usize,
+ length: usize,
+ sort_options: SortOptions,
+) -> Option<Buffer> {
+ let null_capacity = (length / 8) + (length % 8 != 0) as usize;
+ let mut mutable_null_buffer = MutableBuffer::new(null_capacity * 8);
+ mutable_null_buffer.resize(null_capacity, 0);
+
+ let mutable_null_buffer_slice = mutable_null_buffer.as_slice_mut();
+
+ if valid_count > 0 {
+ let mut count = valid_count;
+ let mut index = 0;
+ if sort_options.nulls_first {
+ let remaining_nulls = nulls_count % 8;
+ index = nulls_count / 8;
+
+ if remaining_nulls != 0 {
+ let valid_values_count = min(8 - remaining_nulls, valid_count);
+ mutable_null_buffer_slice[index] =
+ ((1 << valid_values_count) - 1) << remaining_nulls;
+ count -= valid_values_count;
+ index += 1;
+ }
+ }
+ while count >= 8 {
+ mutable_null_buffer_slice[index] = u8::MAX;
+ index += 1;
+ count -= 8;
+ }
+ if count != 0 {
+ mutable_null_buffer_slice[index] = (1 << count) - 1;
+ }
+ }
+
+ Some(mutable_null_buffer.into())
+}
+
+fn sort_native_type<T, U>(
+ values: &dyn Array,
+ options: Option<SortOptions>,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowPrimitiveType,
+ U: ArrowNativeTypeOp,
+{
+ let sort_options = options.unwrap_or_default();
+ let values = values.as_primitive::<T>();
+
+ let result_capacity = values.len() * std::mem::size_of::<U>();
+ let mut mutable_buffer = MutableBuffer::new(result_capacity);
+ mutable_buffer.resize(result_capacity, 0);
+ let mutable_slice: &mut [U] = mutable_buffer.typed_data_mut();
+
+ let array_data = values.to_data();
+ let input_values: &[U] = array_data.buffer(0);
+
+ let mut null_bit_buffer = None;
+
+ let nulls_count = values.null_count();
+ let valid_count = values.len() - nulls_count;
+
+ if values.null_count() > 0 {
+ let nulls = array_data.nulls().unwrap();
+ let null_buffer = nulls.buffer().as_slice();
+
+ let mut mutable_slice_ptr = mutable_slice.as_mut_ptr();
+ let mut input_values_ptr = input_values.as_ptr();
+
+ if sort_options.nulls_first {
+ // This is safe since the offset in in bounds
+ unsafe {
+ mutable_slice_ptr = mutable_slice_ptr.add(values.null_count());
+ }
+ }
+
+ // This is safe since we are in bounds
+ let values_slice =
+ unsafe { slice::from_raw_parts_mut(mutable_slice_ptr, valid_count) };
+
+ for mask in null_buffer {
Review Comment:
This is incorrect as it fails to correctly account for any offset present in the null buffer, you likely need to use bit_chunks or similar
--
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