Skip to content

Commit 9e68728

Browse files
taeppertaepper
committed
add short-hand for range begins/ends in NullLikePartition
1 parent a91e3a5 commit 9e68728

4 files changed

Lines changed: 44 additions & 51 deletions

File tree

‎cpp/src/arrow/compute/kernels/vector_array_sort.cc‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -80,8 +80,8 @@ struct PartitionNthToIndices {
8080
const auto p = PartitionNullsAndNans<ArrayType, NonStablePartitioner>(
8181
out_span, arr, 0, options.null_placement);
8282
auto nth_begin = out_span.data() + pivot;
83-
auto non_null_begin = p.non_null_like_range.data();
84-
auto non_null_end = p.non_null_like_range.data() + p.non_null_like_range.size();
83+
auto non_null_begin = p.non_null_like_begin();
84+
auto non_null_end = p.non_null_like_end();
8585
if (nth_begin >= non_null_begin && nth_begin < non_null_end) {
8686
std::nth_element(non_null_begin, nth_begin, non_null_end,
8787
[&arr](uint64_t left, uint64_t right) {
@@ -156,17 +156,15 @@ class ArrayCompareSorter {
156156
indices, values, offset, options.null_placement);
157157
if (options.order == SortOrder::Ascending) {
158158
std::stable_sort(
159-
p.non_null_like_range.data(),
160-
p.non_null_like_range.data() + p.non_null_like_range.size(),
159+
p.non_null_like_begin(), p.non_null_like_end(),
161160
[&values, &offset](uint64_t left, uint64_t right) {
162161
const auto lhs = GetView::LogicalValue(values.GetView(left - offset));
163162
const auto rhs = GetView::LogicalValue(values.GetView(right - offset));
164163
return lhs < rhs;
165164
});
166165
} else {
167166
std::stable_sort(
168-
p.non_null_like_range.data(),
169-
p.non_null_like_range.data() + p.non_null_like_range.size(),
167+
p.non_null_like_begin(), p.non_null_like_end(),
170168
[&values, &offset](uint64_t left, uint64_t right) {
171169
const auto lhs = GetView::LogicalValue(values.GetView(left - offset));
172170
const auto rhs = GetView::LogicalValue(values.GetView(right - offset));

‎cpp/src/arrow/compute/kernels/vector_select_k.cc‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -226,9 +226,9 @@ class ArraySelector : public TypeVisitor {
226226

227227
HeapSortNonNullsToOutput<InType, sort_order>(p.non_null_like_range, arr,
228228
output.non_null_like_range);
229-
std::copy(p.nan_range.begin(), p.nan_range.begin() + output.nan_range.size(),
229+
std::copy(p.nan_begin(), p.nan_begin() + output.nan_range.size(),
230230
output.nan_range.begin());
231-
std::copy(p.null_range.begin(), p.null_range.begin() + output.null_range.size(),
231+
std::copy(p.null_begin(), p.null_begin() + output.null_range.size(),
232232
output.null_range.begin());
233233

234234
*output_ = Datum(take_indices);
@@ -492,12 +492,12 @@ class RecordBatchSelector {
492492
}
493493
if (output.nan_range.size() > 0) {
494494
// We have the last sort_key, can just copy over the null values
495-
std::copy(p.nan_range.begin(), p.nan_range.begin() + output.nan_range.size(),
495+
std::copy(p.nan_begin(), p.nan_begin() + output.nan_range.size(),
496496
output.nan_range.begin());
497497
}
498498
if (output.null_range.size() > 0) {
499499
// We have the last sort_key, can just copy over the null values
500-
std::copy(p.null_range.begin(), p.null_range.begin() + output.null_range.size(),
500+
std::copy(p.null_begin(), p.null_begin() + output.null_range.size(),
501501
output.null_range.begin());
502502
}
503503
} else {

‎cpp/src/arrow/compute/kernels/vector_sort.cc‎

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -297,17 +297,15 @@ class ConcreteRecordBatchColumnSorter : public RecordBatchColumnSorter {
297297
// a counting sort compatible with indirect indexing.
298298
if (order_ == SortOrder::Ascending) {
299299
std::stable_sort(
300-
partitions.non_null_like_range.data(),
301-
partitions.non_null_like_range.data() + partitions.non_null_like_range.size(),
300+
partitions.non_null_like_begin(), partitions.non_null_like_end(),
302301
[&](uint64_t left, uint64_t right) {
303302
const auto lhs = GetView::LogicalValue(array_.GetView(left - offset));
304303
const auto rhs = GetView::LogicalValue(array_.GetView(right - offset));
305304
return lhs < rhs;
306305
});
307306
} else {
308307
std::stable_sort(
309-
partitions.non_null_like_range.data(),
310-
partitions.non_null_like_range.data() + partitions.non_null_like_range.size(),
308+
partitions.non_null_like_begin(), partitions.non_null_like_end(),
311309
[&](uint64_t left, uint64_t right) {
312310
// We don't use 'left > right' here to reduce required operator.
313311
// If we use 'right < left' here, '<' is only required.
@@ -526,8 +524,7 @@ class MultipleKeyRecordBatchSorter : public TypeVisitor {
526524
const auto p = PartitionNullsInternal<Type>(first_sort_key);
527525

528526
// Sort first-key non-nulls
529-
std::stable_sort(p.non_null_like_range.data(),
530-
p.non_null_like_range.data() + p.non_null_like_range.size(),
527+
std::stable_sort(p.non_null_like_begin(), p.non_null_like_end(),
531528
[&](uint64_t left, uint64_t right) {
532529
// Both values are never null nor NaN
533530
// (otherwise they've been partitioned away above).
@@ -574,7 +571,7 @@ class MultipleKeyRecordBatchSorter : public TypeVisitor {
574571
// Sort all NaNs by the second and following sort keys.
575572
// TODO: could we instead run an independent sort from the second key on
576573
// this slice?
577-
std::stable_sort(p.nan_range.data(), p.nan_range.data() + p.nan_range.size(),
574+
std::stable_sort(p.nan_begin(), p.nan_end(),
578575
[&comparator](uint64_t left, uint64_t right) {
579576
return comparator.Compare(left, right, 1);
580577
});
@@ -583,7 +580,7 @@ class MultipleKeyRecordBatchSorter : public TypeVisitor {
583580
// Sort all nulls by the second and following sort keys.
584581
// TODO: could we instead run an independent sort from the second key on
585582
// this slice?
586-
std::stable_sort(p.null_range.data(), p.null_range.data() + p.null_range.size(),
583+
std::stable_sort(p.null_begin(), p.null_end(),
587584
[&comparator](uint64_t left, uint64_t right) {
588585
return comparator.Compare(left, right, 1);
589586
});

‎cpp/src/arrow/compute/kernels/vector_sort_internal.h‎

Lines changed: 31 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -114,13 +114,23 @@ struct GenericNullLikePartition {
114114
std::span<IndexType> nan_range;
115115
std::span<IndexType> null_range;
116116

117+
IndexType* non_null_like_begin() const { return non_null_like_range.data(); }
118+
IndexType* non_null_like_end() const {
119+
return non_null_like_range.data() + non_null_like_range.size();
120+
}
121+
IndexType* nan_begin() const { return nan_range.data(); }
122+
IndexType* nan_end() const { return nan_range.data() + nan_range.size(); }
123+
IndexType* null_begin() const { return null_range.data(); }
124+
IndexType* null_end() const { return null_range.data() + null_range.size(); }
125+
117126
IndexType* overall_begin() const {
118-
return std::min(non_null_like_range.data(), null_range.data());
127+
// nans are always in the middle
128+
return std::min(non_null_like_begin(), null_begin());
119129
}
120130

121131
IndexType* overall_end() const {
122-
return std::max(non_null_like_range.data() + non_null_like_range.size(),
123-
null_range.data() + null_range.size());
132+
// nans are always in the middle
133+
return std::max(non_null_like_end(), null_end());
124134
}
125135

126136
// Note that "_begin" is not actually the begin of the stored ranges, but can be much
@@ -310,27 +320,21 @@ struct ChunkedMergeImpl {
310320
const ChunkedNullLikePartition& left, const ChunkedNullLikePartition& right) const {
311321
// Input layout:
312322
// [left nul .. left nan .. left non-nul .. right nul .. right nan .. right non-nul]
313-
ARROW_DCHECK_EQ(left.null_range.data() + left.null_range.size(),
314-
left.nan_range.data());
315-
ARROW_DCHECK_EQ(left.nan_range.data() + left.nan_range.size(),
316-
left.non_null_like_range.data());
317-
ARROW_DCHECK_EQ(left.non_null_like_range.data() + left.non_null_like_range.size(),
318-
right.null_range.data());
319-
ARROW_DCHECK_EQ(right.null_range.data() + right.null_range.size(),
320-
right.nan_range.data());
321-
ARROW_DCHECK_EQ(right.nan_range.data() + right.nan_range.size(),
322-
right.non_null_like_range.data());
323+
ARROW_DCHECK_EQ(left.null_end(), left.nan_begin());
324+
ARROW_DCHECK_EQ(left.nan_end(), left.non_null_like_begin());
325+
ARROW_DCHECK_EQ(left.non_null_like_end(), right.null_begin());
326+
ARROW_DCHECK_EQ(right.null_end(), right.nan_begin());
327+
ARROW_DCHECK_EQ(right.nan_end(), right.non_null_like_begin());
323328

324329
// Mutate the input, stably in two steps, to obtain the following layouts:
325330
// [left nul .. left nan .. right nul .. right nan .. left non-nul .. right non-nus]
326-
std::rotate(left.non_null_like_range.data(), right.null_range.data(),
327-
right.nan_range.data() + right.nan_range.size());
331+
std::rotate(left.non_null_like_begin(), right.null_begin(), right.nan_end());
328332

329333
// only use sizes of ranges that are at a different position now
330334
// [left nul .. right nul .. left nan .. right nan .. left non-nulls .. right
331335
// non-nulls] this is a no-op if no nan values are present
332-
std::rotate(left.nan_range.data(), left.nan_range.data() + left.nan_range.size(),
333-
left.nan_range.data() + left.nan_range.size() + right.null_range.size());
336+
std::rotate(left.nan_begin(), left.nan_end(),
337+
left.nan_end() + right.null_range.size());
334338

335339
std::span<CompressedChunkLocation> full_span{left.overall_begin(),
336340
right.overall_end()};
@@ -356,28 +360,22 @@ struct ChunkedMergeImpl {
356360
const ChunkedNullLikePartition& right) const {
357361
// Input layout:
358362
// [left non-nul .. left nan .. left nul .. right non-nul .. right nan .. right nulls]
359-
ARROW_DCHECK_EQ(left.non_null_like_range.data() + left.non_null_like_range.size(),
360-
left.nan_range.data());
361-
ARROW_DCHECK_EQ(left.nan_range.data() + left.nan_range.size(),
362-
left.null_range.data());
363-
ARROW_DCHECK_EQ(left.null_range.data() + left.null_range.size(),
364-
right.non_null_like_range.data());
365-
ARROW_DCHECK_EQ(right.non_null_like_range.data() + right.non_null_like_range.size(),
366-
right.nan_range.data());
367-
ARROW_DCHECK_EQ(right.nan_range.data() + right.nan_range.size(),
368-
right.null_range.data());
363+
ARROW_DCHECK_EQ(left.non_null_like_end(), left.nan_begin());
364+
ARROW_DCHECK_EQ(left.nan_end(), left.null_begin());
365+
ARROW_DCHECK_EQ(left.null_end(), right.non_null_like_begin());
366+
ARROW_DCHECK_EQ(right.non_null_like_end(), right.nan_begin());
367+
ARROW_DCHECK_EQ(right.nan_end(), right.null_begin());
369368

370369
// Mutate the input, stably in two steps, to obtain the following layouts:
371370
// [left non-nul .. right non-nul .. left nan .. left nul .. right nan .. right nul]
372-
std::rotate(left.nan_range.data(), right.non_null_like_range.data(),
373-
right.non_null_like_range.data() + right.non_null_like_range.size());
371+
std::rotate(left.nan_begin(), right.non_null_like_begin(), right.non_null_like_end());
374372

375373
// only use sizes of ranges that are at a different position now
376-
// [left non-nul .. right non-nul .. left nan .. left nul .. right nan .. right nul]
374+
// [left non-nul .. right non-nul .. left nan .. right nan .. left null .. right nul]
377375
// this is a no-op if no nan values are present
378-
auto new_left_null_range_begin =
379-
left.non_null_like_range.data() + left.non_null_like_range.size() +
380-
right.non_null_like_range.size() + left.nan_range.size();
376+
auto new_left_null_range_begin = left.non_null_like_end() +
377+
right.non_null_like_range.size() +
378+
left.nan_range.size();
381379
std::rotate(
382380
new_left_null_range_begin, new_left_null_range_begin + left.null_range.size(),
383381
new_left_null_range_begin + left.null_range.size() + right.nan_range.size());

0 commit comments

Comments
 (0)