Skip to content

Commit e1e1d09

Browse files
committed
Bring over apache/arrow-rs/9683, integrate into sorts, add heuristic to choose between sort implementations.
1 parent 8d91fb0 commit e1e1d09

5 files changed

Lines changed: 390 additions & 9 deletions

File tree

‎datafusion/physical-plan/src/aggregates/row_hash.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1162,6 +1162,7 @@ impl GroupedHashAggregateStream {
11621162
emit,
11631163
self.spill_state.spill_expr.clone(),
11641164
self.batch_size,
1165+
false,
11651166
);
11661167
let spillfile = self
11671168
.spill_state

‎datafusion/physical-plan/src/sorts/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ mod cursor;
2222
mod merge;
2323
mod multi_level_merge;
2424
pub mod partial_sort;
25+
mod radix;
2526
pub mod sort;
2627
pub mod sort_preserving_merge;
2728
mod stream;
Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
// TODO: replace with arrow_row::radix::radix_sort_to_indices once
19+
// available in arrow-rs (see https://github.com/apache/arrow-rs/pull/9683)
20+
21+
//! MSD radix sort on row-encoded keys.
22+
23+
use arrow::array::UInt32Array;
24+
use arrow::row::{RowConverter, Rows, SortField};
25+
use arrow_ord::sort::SortColumn;
26+
use std::sync::Arc;
27+
28+
/// 256-bucket histogram + scatter costs more than comparison sort at small n.
29+
const FALLBACK_THRESHOLD: usize = 64;
30+
31+
/// 8 bytes covers the discriminating prefix of most key layouts; deeper
32+
/// recursion hits diminishing returns as buckets become sparse.
33+
const MAX_DEPTH: usize = 8;
34+
35+
/// Sort row indices using MSD radix sort on row-encoded keys.
36+
///
37+
/// Returns a `UInt32Array` of row indices in sorted order.
38+
pub(crate) fn radix_sort_to_indices(
39+
sort_columns: &[SortColumn],
40+
) -> arrow::error::Result<UInt32Array> {
41+
let sort_fields: Vec<SortField> = sort_columns
42+
.iter()
43+
.map(|col| {
44+
SortField::new_with_options(
45+
col.values.data_type().clone(),
46+
col.options.unwrap_or_default(),
47+
)
48+
})
49+
.collect();
50+
51+
let arrays: Vec<_> = sort_columns
52+
.iter()
53+
.map(|col| Arc::clone(&col.values))
54+
.collect();
55+
56+
let converter = RowConverter::new(sort_fields)?;
57+
let rows = converter.convert_columns(&arrays)?;
58+
59+
let n = rows.num_rows();
60+
let mut indices: Vec<u32> = (0..n as u32).collect();
61+
let mut temp = vec![0u32; n];
62+
msd_radix_sort(&mut indices, &mut temp, &rows, 0);
63+
Ok(UInt32Array::from(indices))
64+
}
65+
66+
fn msd_radix_sort(indices: &mut [u32], temp: &mut [u32], rows: &Rows, byte_pos: usize) {
67+
let n = indices.len();
68+
69+
if n <= FALLBACK_THRESHOLD || byte_pos >= MAX_DEPTH {
70+
indices.sort_unstable_by(|&a, &b| {
71+
let ra = unsafe { rows.row_unchecked(a as usize) };
72+
let rb = unsafe { rows.row_unchecked(b as usize) };
73+
ra.cmp(&rb)
74+
});
75+
return;
76+
}
77+
78+
let mut counts = [0u32; 256];
79+
for &idx in &*indices {
80+
let row = unsafe { rows.row_unchecked(idx as usize) };
81+
let byte = row.data().get(byte_pos).copied().unwrap_or(0);
82+
counts[byte as usize] += 1;
83+
}
84+
85+
// Skip scatter when all rows share the same byte
86+
if counts.iter().filter(|&&c| c > 0).count() == 1 {
87+
msd_radix_sort(indices, temp, rows, byte_pos + 1);
88+
return;
89+
}
90+
91+
let mut offsets = [0u32; 257];
92+
for i in 0..256 {
93+
offsets[i + 1] = offsets[i] + counts[i];
94+
}
95+
96+
let temp = &mut temp[..n];
97+
let mut write_pos = offsets;
98+
for &idx in &*indices {
99+
let row = unsafe { rows.row_unchecked(idx as usize) };
100+
let byte = row.data().get(byte_pos).copied().unwrap_or(0) as usize;
101+
temp[write_pos[byte] as usize] = idx;
102+
write_pos[byte] += 1;
103+
}
104+
indices.copy_from_slice(temp);
105+
106+
for bucket in 0..256 {
107+
let start = offsets[bucket] as usize;
108+
let end = offsets[bucket + 1] as usize;
109+
if end - start > 1 {
110+
msd_radix_sort(
111+
&mut indices[start..end],
112+
&mut temp[start..end],
113+
rows,
114+
byte_pos + 1,
115+
);
116+
}
117+
}
118+
}

0 commit comments

Comments
 (0)