This is an automated email from the ASF dual-hosted git repository.
pitrou pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new bc9d1349e0 GH-50247: [C++] Reuse abstraction for null partitions in
sorting functions (#50248)
bc9d1349e0 is described below
commit bc9d1349e091dead2044ed200d3e3e867578b354
Author: Alexander Taepper <[email protected]>
AuthorDate: Wed Jul 22 13:14:22 2026 +0200
GH-50247: [C++] Reuse abstraction for null partitions in sorting functions
(#50248)
### Rationale for this change
@ pitrou mentioned this as a follow-up in #46926
### What changes are included in this PR?
Refactoring sorting methods to reuse the helper methods avoid maintaining
two abstractions for null partitions. The new abstraction was very seamless to
implement in most cases, but a few spots required some care
In particular, these functions were severly simlpified by the new
abstraction:
- `MarkDuplicates`: duplicate nulls and nans were detected by checking
every single row for `Null` one additional time, after we already had (and
discarded) the nullness information
- `GenericMergeImpl`: merging of `null`-ranges involved repartitioning
`null` and `nan` values in every merge invocation. Now, we track this
distinction and do not need any merge function for `null` and `nan` blocks
(unless we merge by multiple sort-keys, where we will merge according to the
remaining sort keys)
### Are these changes tested?
Yes, the compute test suite passes as before
### Are there any user-facing changes?
No.
* GitHub Issue: #50247
Authored-by: Alexander Taepper <[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>
---
cpp/src/arrow/compute/kernels/chunked_internal.cc | 25 +-
cpp/src/arrow/compute/kernels/chunked_internal.h | 25 +-
cpp/src/arrow/compute/kernels/vector_array_sort.cc | 159 +++---
cpp/src/arrow/compute/kernels/vector_rank.cc | 112 ++---
cpp/src/arrow/compute/kernels/vector_select_k.cc | 83 +--
cpp/src/arrow/compute/kernels/vector_sort.cc | 501 +++++++++----------
.../arrow/compute/kernels/vector_sort_internal.h | 554 ++++++++++-----------
cpp/src/arrow/compute/kernels/vector_sort_test.cc | 27 +
8 files changed, 668 insertions(+), 818 deletions(-)
diff --git a/cpp/src/arrow/compute/kernels/chunked_internal.cc
b/cpp/src/arrow/compute/kernels/chunked_internal.cc
index 39495ea0e6..3a853cfbd4 100644
--- a/cpp/src/arrow/compute/kernels/chunked_internal.cc
+++ b/cpp/src/arrow/compute/kernels/chunked_internal.cc
@@ -50,8 +50,7 @@ std::vector<int64_t> ChunkedIndexMapper::GetChunkLengths(
return chunk_lengths;
}
-Result<std::pair<CompressedChunkLocation*, CompressedChunkLocation*>>
-ChunkedIndexMapper::LogicalToPhysical() {
+Result<std::span<CompressedChunkLocation>>
ChunkedIndexMapper::LogicalToPhysical() {
// Check that indices would fall in bounds for CompressedChunkLocation
if (ARROW_PREDICT_FALSE(chunk_lengths_.size() >
CompressedChunkLocation::kMaxChunkIndex + 1)) {
@@ -67,13 +66,13 @@ ChunkedIndexMapper::LogicalToPhysical() {
}
}
- const int64_t num_indices = static_cast<int64_t>(indices_end_ -
indices_begin_);
+ const int64_t num_indices = static_cast<int64_t>(indices_.size());
DCHECK_EQ(num_indices, std::accumulate(chunk_lengths_.begin(),
chunk_lengths_.end(),
static_cast<int64_t>(0)));
CompressedChunkLocation* physical_begin =
- reinterpret_cast<CompressedChunkLocation*>(indices_begin_);
- DCHECK_EQ(physical_begin + num_indices,
- reinterpret_cast<CompressedChunkLocation*>(indices_end_));
+ reinterpret_cast<CompressedChunkLocation*>(indices_.data());
+ DCHECK_EQ(physical_begin + num_indices,
reinterpret_cast<CompressedChunkLocation*>(
+ indices_.data() +
indices_.size()));
int64_t chunk_offset = 0;
for (int64_t chunk_index = 0; chunk_index <
static_cast<int64_t>(chunk_lengths_.size());
@@ -82,17 +81,17 @@ ChunkedIndexMapper::LogicalToPhysical() {
for (int64_t i = 0; i < chunk_length; ++i) {
// Logical indices are expected to be chunk-partitioned, which avoids
costly
// chunked index resolution.
- DCHECK_GE(indices_begin_[chunk_offset + i],
static_cast<uint64_t>(chunk_offset));
- DCHECK_LT(indices_begin_[chunk_offset + i],
+ DCHECK_GE(indices_[chunk_offset + i],
static_cast<uint64_t>(chunk_offset));
+ DCHECK_LT(indices_[chunk_offset + i],
static_cast<uint64_t>(chunk_offset + chunk_length));
physical_begin[chunk_offset + i] = CompressedChunkLocation{
static_cast<uint64_t>(chunk_index),
- indices_begin_[chunk_offset + i] -
static_cast<uint64_t>(chunk_offset)};
+ indices_[chunk_offset + i] - static_cast<uint64_t>(chunk_offset)};
}
chunk_offset += chunk_length;
}
- return std::pair{physical_begin, physical_begin + num_indices};
+ return std::span<CompressedChunkLocation>{physical_begin, physical_begin +
num_indices};
}
Status ChunkedIndexMapper::PhysicalToLogical() {
@@ -105,15 +104,15 @@ Status ChunkedIndexMapper::PhysicalToLogical() {
}
}
- const int64_t num_indices = static_cast<int64_t>(indices_end_ -
indices_begin_);
+ const int64_t num_indices = static_cast<int64_t>(indices_.size());
CompressedChunkLocation* physical_begin =
- reinterpret_cast<CompressedChunkLocation*>(indices_begin_);
+ reinterpret_cast<CompressedChunkLocation*>(indices_.data());
for (int64_t i = 0; i < num_indices; ++i) {
const auto loc = physical_begin[i];
DCHECK_LT(loc.chunk_index(), chunk_offsets.size());
DCHECK_LT(loc.index_in_chunk(),
static_cast<uint64_t>(chunk_lengths_[loc.chunk_index()]));
- indices_begin_[i] =
+ indices_[i] =
chunk_offsets[loc.chunk_index()] +
static_cast<int64_t>(loc.index_in_chunk());
}
diff --git a/cpp/src/arrow/compute/kernels/chunked_internal.h
b/cpp/src/arrow/compute/kernels/chunked_internal.h
index 2dcfa0047e..23b500a120 100644
--- a/cpp/src/arrow/compute/kernels/chunked_internal.h
+++ b/cpp/src/arrow/compute/kernels/chunked_internal.h
@@ -127,27 +127,19 @@ ARROW_EXPORT std::vector<const Array*>
GetArrayPointers(const ArrayVector& array
// and vice-versa.
class ARROW_EXPORT ChunkedIndexMapper {
public:
- ChunkedIndexMapper(const std::vector<const Array*>& chunks, uint64_t*
indices_begin,
- uint64_t* indices_end)
- : ChunkedIndexMapper(std::span(chunks), indices_begin, indices_end) {}
- ChunkedIndexMapper(std::span<const Array* const> chunks, uint64_t*
indices_begin,
- uint64_t* indices_end)
- : chunk_lengths_(GetChunkLengths(chunks)),
- indices_begin_(indices_begin),
- indices_end_(indices_end) {}
- ChunkedIndexMapper(const RecordBatchVector& chunks, uint64_t* indices_begin,
- uint64_t* indices_end)
- : chunk_lengths_(GetChunkLengths(chunks)),
- indices_begin_(indices_begin),
- indices_end_(indices_end) {}
+ ChunkedIndexMapper(const std::vector<const Array*>& chunks,
std::span<uint64_t> indices)
+ : ChunkedIndexMapper(std::span(chunks), indices) {}
+ ChunkedIndexMapper(std::span<const Array* const> chunks, std::span<uint64_t>
indices)
+ : chunk_lengths_(GetChunkLengths(chunks)), indices_(indices) {}
+ ChunkedIndexMapper(const RecordBatchVector& chunks, std::span<uint64_t>
indices)
+ : chunk_lengths_(GetChunkLengths(chunks)), indices_(indices) {}
// Turn the original uint64_t logical indices into physical. This reuses the
// same memory area, so the logical indices cannot be used anymore until
// PhysicalToLogical() is called.
//
// This assumes that the logical indices are originally chunk-partitioned.
- Result<std::pair<CompressedChunkLocation*, CompressedChunkLocation*>>
- LogicalToPhysical();
+ Result<std::span<CompressedChunkLocation>> LogicalToPhysical();
// Turn the physical indices back into logical, making the uint64_t indices
// usable again.
@@ -158,8 +150,7 @@ class ARROW_EXPORT ChunkedIndexMapper {
static std::vector<int64_t> GetChunkLengths(const RecordBatchVector& chunks);
std::vector<int64_t> chunk_lengths_;
- uint64_t* indices_begin_;
- uint64_t* indices_end_;
+ std::span<uint64_t> indices_;
};
} // namespace arrow::compute::internal
diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index 6e7068f6ec..ad9444037c 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -71,17 +71,19 @@ struct PartitionNthToIndices {
return Status::IndexError("NthToIndices index out of bound");
}
ArrayData* out_arr = out->array_data().get();
- uint64_t* out_begin = out_arr->GetMutableValues<uint64_t>(1);
- uint64_t* out_end = out_begin + arr.length();
- std::iota(out_begin, out_end, 0);
+ std::span<uint64_t> out_span{out_arr->GetMutableValues<uint64_t>(1),
+ static_cast<size_t>(arr.length())};
+ std::iota(out_span.begin(), out_span.end(), 0);
if (pivot == arr.length()) {
return Status::OK();
}
- const auto p = PartitionNulls<ArrayType, NonStablePartitioner>(
- out_begin, out_end, arr, 0, options.null_placement);
- auto nth_begin = out_begin + pivot;
- if (nth_begin >= p.non_nulls_begin && nth_begin < p.non_nulls_end) {
- std::nth_element(p.non_nulls_begin, nth_begin, p.non_nulls_end,
+ const auto p = PartitionNullsAndNans<ArrayType, NonStablePartitioner>(
+ out_span, arr, 0, options.null_placement);
+ auto nth_begin = out_span.data() + pivot;
+ auto non_null_begin = p.non_null_like_begin();
+ auto non_null_end = p.non_null_like_end();
+ if (nth_begin >= non_null_begin && nth_begin < non_null_end) {
+ std::nth_element(non_null_begin, nth_begin, non_null_end,
[&arr](uint64_t left, uint64_t right) {
const auto lval =
GetView::LogicalValue(arr.GetView(left));
const auto rval =
GetView::LogicalValue(arr.GetView(right));
@@ -145,16 +147,16 @@ class ArrayCompareSorter {
using GetView = GetViewType<ArrowType>;
public:
- Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t*
indices_end,
- const Array& array, int64_t offset,
- const ArraySortOptions& options,
ExecContext*) {
+ Result<NullLikePartition> operator()(std::span<uint64_t> indices, const
Array& array,
+ int64_t offset, const ArraySortOptions&
options,
+ ExecContext*) {
const auto& values = checked_cast<const ArrayType&>(array);
- const auto p = PartitionNulls<ArrayType, StablePartitioner>(
- indices_begin, indices_end, values, offset, options.null_placement);
+ const auto p = PartitionNullsAndNans<ArrayType, StablePartitioner>(
+ indices, values, offset, options.null_placement);
if (options.order == SortOrder::Ascending) {
std::stable_sort(
- p.non_nulls_begin, p.non_nulls_end,
+ p.non_null_like_begin(), p.non_null_like_end(),
[&values, &offset](uint64_t left, uint64_t right) {
const auto lhs = GetView::LogicalValue(values.GetView(left -
offset));
const auto rhs = GetView::LogicalValue(values.GetView(right -
offset));
@@ -162,7 +164,7 @@ class ArrayCompareSorter {
});
} else {
std::stable_sort(
- p.non_nulls_begin, p.non_nulls_end,
+ p.non_null_like_begin(), p.non_null_like_end(),
[&values, &offset](uint64_t left, uint64_t right) {
const auto lhs = GetView::LogicalValue(values.GetView(left -
offset));
const auto rhs = GetView::LogicalValue(values.GetView(right -
offset));
@@ -178,10 +180,9 @@ class ArrayCompareSorter {
template <>
class ArrayCompareSorter<DictionaryType> {
public:
- Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t*
indices_end,
- const Array& array, int64_t offset,
- const ArraySortOptions& options,
- ExecContext* ctx) {
+ Result<NullLikePartition> operator()(std::span<uint64_t> indices, const
Array& array,
+ int64_t offset, const ArraySortOptions&
options,
+ ExecContext* ctx) {
const auto& dict_array = checked_cast<const DictionaryArray&>(array);
const auto& dict_values = dict_array.dictionary();
const auto& dict_indices = dict_array.indices();
@@ -220,7 +221,7 @@ class ArrayCompareSorter<DictionaryType> {
DCHECK_EQ(decoded_ranks->length(), dict_array.length());
ARROW_ASSIGN_OR_RAISE(auto rank_sorter,
GetArraySorter(*decoded_ranks->type()));
- return rank_sorter(indices_begin, indices_end, *decoded_ranks, offset,
options, ctx);
+ return rank_sorter(indices, *decoded_ranks, offset, options, ctx);
}
private:
@@ -264,12 +265,11 @@ class ArrayCompareSorter<DictionaryType> {
template <>
class ArrayCompareSorter<StructType> {
public:
- Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t*
indices_end,
- const Array& array, int64_t offset,
- const ArraySortOptions& options,
- ExecContext* ctx) {
+ Result<NullLikePartition> operator()(std::span<uint64_t> indices, const
Array& array,
+ int64_t offset, const ArraySortOptions&
options,
+ ExecContext* ctx) {
const auto& struct_array = checked_cast<const StructArray&>(array);
- return SortStructArray(ctx, indices_begin, indices_end, struct_array,
options.order,
+ return SortStructArray(ctx, indices, struct_array, options.order,
options.null_placement);
}
};
@@ -290,17 +290,16 @@ class ArrayCountSorter {
value_range_ = static_cast<uint32_t>(max - min) + 1;
}
- Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t*
indices_end,
- const Array& array, int64_t offset,
- const ArraySortOptions& options,
- ExecContext*) const {
+ Result<NullLikePartition> operator()(std::span<uint64_t> indices, const
Array& array,
+ int64_t offset, const ArraySortOptions&
options,
+ ExecContext*) const {
const auto& values = checked_cast<const ArrayType&>(array);
// 32bit counter performs much better than 64bit one
if (values.length() < (1LL << 32)) {
- return SortInternal<uint32_t>(indices_begin, indices_end, values,
offset, options);
+ return SortInternal<uint32_t>(indices, values, offset, options);
} else {
- return SortInternal<uint64_t>(indices_begin, indices_end, values,
offset, options);
+ return SortInternal<uint64_t>(indices, values, offset, options);
}
}
@@ -309,14 +308,13 @@ class ArrayCountSorter {
uint32_t value_range_{0};
template <typename CounterType>
- NullPartitionResult SortInternal(uint64_t* indices_begin, uint64_t*
indices_end,
- const ArrayType& values, int64_t offset,
- const ArraySortOptions& options) const {
+ NullLikePartition SortInternal(std::span<uint64_t> indices, const ArrayType&
values,
+ int64_t offset, const ArraySortOptions&
options) const {
const uint32_t value_range = value_range_;
// first and last slot reserved for prefix sum (depending on sort order)
std::vector<CounterType> counts(2 + value_range);
- NullPartitionResult p;
+ NullLikePartition p;
if (options.order == SortOrder::Ascending) {
// counts will be increasing, starting with 0 and ending with (length -
null_count)
@@ -325,13 +323,10 @@ class ArrayCountSorter {
counts[i] += counts[i - 1];
}
- if (options.null_placement == NullPlacement::AtStart) {
- p = NullPartitionResult::NullsAtStart(indices_begin, indices_end,
- indices_end -
counts[value_range]);
- } else {
- p = NullPartitionResult::NullsAtEnd(indices_begin, indices_end,
- indices_begin +
counts[value_range]);
- }
+ p = NullLikePartition::FromCounts(indices, counts[value_range], 0,
+ indices.size() - counts[value_range],
+ options.null_placement);
+
EmitIndices(p, values, offset, &counts[0]);
} else {
// counts will be decreasing, starting with (length - null_count) and
ending with 0
@@ -340,13 +335,8 @@ class ArrayCountSorter {
counts[i - 1] += counts[i];
}
- if (options.null_placement == NullPlacement::AtStart) {
- p = NullPartitionResult::NullsAtStart(indices_begin, indices_end,
- indices_end - counts[0]);
- } else {
- p = NullPartitionResult::NullsAtEnd(indices_begin, indices_end,
- indices_begin + counts[0]);
- }
+ p = NullLikePartition::FromCounts(indices, counts[0], 0, indices.size()
- counts[0],
+ options.null_placement);
EmitIndices(p, values, offset, &counts[1]);
}
return p;
@@ -359,14 +349,15 @@ class ArrayCountSorter {
}
template <typename CounterType>
- void EmitIndices(const NullPartitionResult& p, const ArrayType& values,
int64_t offset,
+ void EmitIndices(const NullLikePartition& p, const ArrayType& values,
int64_t offset,
CounterType* counts) const {
+ DCHECK(p.nan_range.empty());
int64_t index = offset;
CounterType count_nulls = 0;
VisitRawValuesInline<c_type>(
*values.data(),
- [&](c_type v) { p.non_nulls_begin[counts[v - min_]++] = index++; },
- [&]() { p.nulls_begin[count_nulls++] = index++; });
+ [&](c_type v) { p.non_null_like_range[counts[v - min_]++] = index++; },
+ [&]() { p.null_range[count_nulls++] = index++; });
}
};
@@ -375,25 +366,21 @@ class ArrayCountSorter<BooleanType> {
public:
ArrayCountSorter() = default;
- Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t*
indices_end,
- const Array& array, int64_t offset,
- const ArraySortOptions& options,
ExecContext*) {
+ Result<NullLikePartition> operator()(std::span<uint64_t> indices, const
Array& array,
+ int64_t offset, const ArraySortOptions&
options,
+ ExecContext*) {
const auto& values = checked_cast<const BooleanArray&>(array);
std::array<int64_t, 3> counts{0, 0, 0}; // false, true, null
const int64_t nulls = values.null_count();
+ const int64_t non_nulls = values.length() - nulls;
+
const int64_t ones = values.true_count();
- const int64_t zeros = values.length() - ones - nulls;
+ const int64_t zeros = non_nulls - ones;
- NullPartitionResult p;
- if (options.null_placement == NullPlacement::AtStart) {
- p = NullPartitionResult::NullsAtStart(indices_begin, indices_end,
- indices_begin + nulls);
- } else {
- p = NullPartitionResult::NullsAtEnd(indices_begin, indices_end,
- indices_end - nulls);
- }
+ NullLikePartition p = NullLikePartition::FromCounts(indices, non_nulls, 0,
nulls,
+
options.null_placement);
if (options.order == SortOrder::Ascending) {
// ones start after zeros
@@ -405,8 +392,8 @@ class ArrayCountSorter<BooleanType> {
int64_t index = offset;
VisitRawValuesInline(
- *values.data(), [&](bool v) { p.non_nulls_begin[counts[v]++] =
index++; },
- [&]() { p.nulls_begin[counts[2]++] = index++; });
+ *values.data(), [&](bool v) { p.non_null_like_range[counts[v]++] =
index++; },
+ [&]() { p.null_range[counts[2]++] = index++; });
return p;
}
};
@@ -420,10 +407,9 @@ class ArrayCountOrCompareSorter {
using c_type = typename ArrowType::c_type;
public:
- Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t*
indices_end,
- const Array& array, int64_t offset,
- const ArraySortOptions& options,
- ExecContext* ctx) {
+ Result<NullLikePartition> operator()(std::span<uint64_t> indices, const
Array& array,
+ int64_t offset, const ArraySortOptions&
options,
+ ExecContext* ctx) {
const auto& values = checked_cast<const ArrayType&>(array);
if (values.length() >= countsort_min_len_ && values.length() >
values.null_count()) {
@@ -435,11 +421,11 @@ class ArrayCountOrCompareSorter {
if (static_cast<uint64_t>(max) - static_cast<uint64_t>(min) <=
countsort_max_range_) {
count_sorter_.SetMinMax(min, max);
- return count_sorter_(indices_begin, indices_end, values, offset,
options, ctx);
+ return count_sorter_(indices, values, offset, options, ctx);
}
}
- return compare_sorter_(indices_begin, indices_end, values, offset,
options, ctx);
+ return compare_sorter_(indices, values, offset, options, ctx);
}
private:
@@ -461,11 +447,11 @@ class ArrayCountOrCompareSorter {
class ArrayNullSorter {
public:
- Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t*
indices_end,
- const Array& values, int64_t offset,
- const ArraySortOptions& options,
ExecContext*) {
- return NullPartitionResult::NullsOnly(indices_begin, indices_end,
- options.null_placement);
+ Result<NullLikePartition> operator()(std::span<uint64_t> indices, const
Array& values,
+ int64_t offset, const ArraySortOptions&
options,
+ ExecContext*) {
+ return NullLikePartition::FromCounts(indices, 0, 0, indices.size(),
+ options.null_placement);
}
};
@@ -542,14 +528,14 @@ struct ArraySortIndices {
static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult*
out) {
const auto& options = ArraySortIndicesState::Get(ctx);
ArrayData* out_arr = out->array_data().get();
- uint64_t* out_begin = out_arr->GetMutableValues<uint64_t>(1);
- uint64_t* out_end = out_begin + out_arr->length;
- std::iota(out_begin, out_end, 0);
+ std::span<uint64_t> out_span{out_arr->GetMutableValues<uint64_t>(1),
+ static_cast<size_t>(out_arr->length)};
+ std::iota(out_span.begin(), out_span.end(), 0);
ArrayType arr(batch[0].array.ToArrayData());
ARROW_ASSIGN_OR_RAISE(auto sorter,
GetArraySorter(*GetPhysicalType(arr.type())));
- return sorter(out_begin, out_end, arr, 0, options,
ctx->exec_context()).status();
+ return sorter(out_span, arr, 0, options, ctx->exec_context()).status();
}
};
@@ -557,12 +543,11 @@ Status ArraySortIndicesChunked(KernelContext* ctx, const
ExecBatch& batch, Datum
const auto& options = ArraySortIndicesState::Get(ctx);
ArrayData* out_arr = out->mutable_array();
DCHECK_EQ(out_arr->length, batch.length);
- uint64_t* out_begin = out_arr->GetMutableValues<uint64_t>(1);
- uint64_t* out_end = out_begin + out_arr->length;
- std::iota(out_begin, out_end, 0);
- return SortChunkedArray(ctx->exec_context(), out_begin, out_end,
- *batch[0].chunked_array(), options.order,
- options.null_placement)
+ std::span<uint64_t> out_span{out_arr->GetMutableValues<uint64_t>(1),
+ static_cast<size_t>(out_arr->length)};
+ std::iota(out_span.begin(), out_span.end(), 0);
+ return SortChunkedArray(ctx->exec_context(), out_span,
*batch[0].chunked_array(),
+ options.order, options.null_placement)
.status();
}
diff --git a/cpp/src/arrow/compute/kernels/vector_rank.cc
b/cpp/src/arrow/compute/kernels/vector_rank.cc
index d32dd359c7..26655d8076 100644
--- a/cpp/src/arrow/compute/kernels/vector_rank.cc
+++ b/cpp/src/arrow/compute/kernels/vector_rank.cc
@@ -37,16 +37,15 @@ namespace {
// is the same as the value at the previous sort index.
constexpr uint64_t kDuplicateMask = 1ULL << 63;
-template <typename ValueSelector, typename IsNullSelector>
-void MarkDuplicates(const NullPartitionResult& sorted, ValueSelector&&
value_selector,
- IsNullSelector&& is_null_selector) {
+template <typename ValueSelector>
+void MarkDuplicates(const NullLikePartition& sorted, ValueSelector&&
value_selector) {
using T = decltype(value_selector(int64_t{}));
// Process non-nulls
- if (sorted.non_nulls_end != sorted.non_nulls_begin) {
- auto it = sorted.non_nulls_begin;
+ if (!sorted.non_null_like_range.empty()) {
+ auto it = sorted.non_null_like_range.begin();
T prev_value = value_selector(*it);
- while (++it < sorted.non_nulls_end) {
+ while (++it < sorted.non_null_like_range.end()) {
T curr_value = value_selector(*it);
if (curr_value == prev_value) {
*it |= kDuplicateMask;
@@ -55,23 +54,24 @@ void MarkDuplicates(const NullPartitionResult& sorted,
ValueSelector&& value_sel
}
}
+ // Process nans
+ if (!sorted.nan_range.empty()) {
+ for (auto& index : sorted.nan_range.subspan(1)) {
+ index |= kDuplicateMask;
+ }
+ }
+
// Process nulls
- if (sorted.nulls_end != sorted.nulls_begin) {
- auto it = sorted.nulls_begin;
- bool prev_is_null = is_null_selector(*it);
- while (++it < sorted.nulls_end) {
- bool curr_is_null = is_null_selector(*it);
- if (curr_is_null == prev_is_null) {
- *it |= kDuplicateMask;
- }
- prev_is_null = curr_is_null;
+ if (!sorted.null_range.empty()) {
+ for (auto& index : sorted.null_range.subspan(1)) {
+ index |= kDuplicateMask;
}
}
}
template <typename ArrowType>
-Result<NullPartitionResult> DoSortAndMarkDuplicate(
- ExecContext* ctx, uint64_t* indices_begin, uint64_t* indices_end, const
Array& input,
+Result<NullLikePartition> DoSortAndMarkDuplicate(
+ ExecContext* ctx, std::span<uint64_t> indices, const Array& input,
const std::shared_ptr<DataType>& physical_type, const SortOrder order,
const NullPlacement null_placement, bool needs_duplicates) {
using GetView = GetViewType<ArrowType>;
@@ -80,51 +80,38 @@ Result<NullPartitionResult> DoSortAndMarkDuplicate(
ARROW_ASSIGN_OR_RAISE(auto array_sorter, GetArraySorter(*physical_type));
ArrayType array(input.data());
- ARROW_ASSIGN_OR_RAISE(auto sorted,
- array_sorter(indices_begin, indices_end, array, 0,
- ArraySortOptions(order, null_placement),
ctx));
+ ARROW_ASSIGN_OR_RAISE(
+ auto sorted,
+ array_sorter(indices, array, 0, ArraySortOptions(order, null_placement),
ctx));
if (needs_duplicates) {
auto value_selector = [&array](int64_t index) {
return GetView::LogicalValue(array.GetView(index));
};
- if constexpr (has_null_like_values<ArrowType>()) {
- auto is_null_selector = [&array](int64_t index) { return
array.IsNull(index); };
- MarkDuplicates(sorted, value_selector, is_null_selector);
- } else {
- MarkDuplicates(sorted, value_selector, [](int64_t) { return true; });
- }
+ MarkDuplicates(sorted, value_selector);
}
return sorted;
}
template <typename ArrowType>
-Result<NullPartitionResult> DoSortAndMarkDuplicate(
- ExecContext* ctx, uint64_t* indices_begin, uint64_t* indices_end,
- const ChunkedArray& input, const std::shared_ptr<DataType>& physical_type,
- const SortOrder order, const NullPlacement null_placement, bool
needs_duplicates) {
+Result<NullLikePartition> DoSortAndMarkDuplicate(
+ ExecContext* ctx, std::span<uint64_t> indices, const ChunkedArray& input,
+ const std::shared_ptr<DataType>& physical_type, const SortOrder order,
+ const NullPlacement null_placement, bool needs_duplicates) {
auto physical_chunks = GetPhysicalChunks(input, physical_type);
if (physical_chunks.empty()) {
- return NullPartitionResult{};
+ return NullLikePartition::FromCounts(indices, 0, 0, 0, null_placement);
}
- ARROW_ASSIGN_OR_RAISE(auto sorted,
- SortChunkedArray(ctx, indices_begin, indices_end,
physical_type,
- physical_chunks, order,
null_placement));
+ ARROW_ASSIGN_OR_RAISE(
+ auto sorted, SortChunkedArray(ctx, indices, physical_type,
physical_chunks, order,
+ null_placement));
if (needs_duplicates) {
const auto arrays = GetArrayPointers(physical_chunks);
auto value_selector = [resolver =
ChunkedArrayResolver(std::span(arrays))](int64_t index) {
return resolver.Resolve(index).Value<ArrowType>();
};
- if constexpr (has_null_like_values<ArrowType>()) {
- auto is_null_selector =
- [resolver = ChunkedArrayResolver(std::span(arrays))](int64_t index) {
- return resolver.Resolve(index).IsNull();
- };
- MarkDuplicates(sorted, value_selector, is_null_selector);
- } else {
- MarkDuplicates(sorted, value_selector, [](int64_t) { return true; });
- }
+ MarkDuplicates(sorted, value_selector);
}
return sorted;
}
@@ -132,31 +119,29 @@ Result<NullPartitionResult> DoSortAndMarkDuplicate(
template <typename InputType>
class SortAndMarkDuplicate : public TypeVisitor {
public:
- SortAndMarkDuplicate(ExecContext* ctx, uint64_t* indices_begin, uint64_t*
indices_end,
+ SortAndMarkDuplicate(ExecContext* ctx, std::span<uint64_t> indices,
const InputType& input, const SortOrder order,
const NullPlacement null_placement, const bool
needs_duplicate)
: TypeVisitor(),
ctx_(ctx),
- indices_begin_(indices_begin),
- indices_end_(indices_end),
+ indices_(indices),
input_(input),
order_(order),
null_placement_(null_placement),
needs_duplicates_(needs_duplicate),
physical_type_(GetPhysicalType(input.type())) {}
- Result<NullPartitionResult> Run() {
+ Result<NullLikePartition> Run() {
RETURN_NOT_OK(physical_type_->Accept(this));
return sorted_;
}
-#define VISIT(TYPE)
\
- Status Visit(const TYPE& type) {
\
- ARROW_ASSIGN_OR_RAISE(
\
- sorted_, DoSortAndMarkDuplicate<TYPE>(ctx_, indices_begin_,
indices_end_, \
- input_, physical_type_, order_,
\
- null_placement_,
needs_duplicates_)); \
- return Status::OK();
\
+#define VISIT(TYPE)
\
+ Status Visit(const TYPE& type) {
\
+ ARROW_ASSIGN_OR_RAISE(sorted_, DoSortAndMarkDuplicate<TYPE>(
\
+ ctx_, indices_, input_, physical_type_,
order_, \
+ null_placement_, needs_duplicates_));
\
+ return Status::OK();
\
}
VISIT_SORTABLE_PHYSICAL_TYPES(VISIT)
@@ -165,20 +150,19 @@ class SortAndMarkDuplicate : public TypeVisitor {
private:
ExecContext* ctx_;
- uint64_t* indices_begin_;
- uint64_t* indices_end_;
+ std::span<uint64_t> indices_;
const InputType& input_;
const SortOrder order_;
const NullPlacement null_placement_;
const bool needs_duplicates_;
const std::shared_ptr<DataType> physical_type_;
- NullPartitionResult sorted_{};
+ NullLikePartition sorted_{};
};
// A CRTP-based helper class for "rank_normal" and "rank_quantile"
template <typename Derived>
struct BaseQuantileRanker {
- Result<Datum> CreateRankings(ExecContext* ctx, const NullPartitionResult&
sorted) {
+ Result<Datum> CreateRankings(ExecContext* ctx, const NullLikePartition&
sorted) {
const int64_t length = sorted.overall_end() - sorted.overall_begin();
ARROW_ASSIGN_OR_RAISE(auto rankings,
MakeMutableFloat64Array(length, ctx->memory_pool()));
@@ -228,7 +212,7 @@ struct NormalRanker : public
BaseQuantileRanker<NormalRanker> {
struct OrdinalRanker {
explicit OrdinalRanker(RankOptions::Tiebreaker tiebreaker) :
tiebreaker_(tiebreaker) {}
- Result<Datum> CreateRankings(ExecContext* ctx, const NullPartitionResult&
sorted) {
+ Result<Datum> CreateRankings(ExecContext* ctx, const NullLikePartition&
sorted) {
const int64_t length = sorted.overall_end() - sorted.overall_begin();
ARROW_ASSIGN_OR_RAISE(auto rankings,
MakeMutableUInt64Array(length, ctx->memory_pool()));
@@ -379,13 +363,13 @@ class RankMetaFunctionBase : public MetaFunction {
int64_t length = input.length();
ARROW_ASSIGN_OR_RAISE(auto indices,
MakeMutableUInt64Array(length, ctx->memory_pool()));
- auto* indices_begin = indices->GetMutableValues<uint64_t>(1);
- auto* indices_end = indices_begin + length;
- std::iota(indices_begin, indices_end, 0);
+ std::span<uint64_t> indices_span{indices->GetMutableValues<uint64_t>(1),
+ static_cast<size_t>(length)};
+ std::iota(indices_span.begin(), indices_span.end(), 0);
auto needs_duplicates = Derived::NeedsDuplicates(options);
ARROW_ASSIGN_OR_RAISE(
- auto sorted, SortAndMarkDuplicate(ctx, indices_begin, indices_end,
input, order,
- null_placement, needs_duplicates)
+ auto sorted, SortAndMarkDuplicate(ctx, indices_span, input, order,
null_placement,
+ needs_duplicates)
.Run());
auto ranker = Derived::GetRanker(options);
diff --git a/cpp/src/arrow/compute/kernels/vector_select_k.cc
b/cpp/src/arrow/compute/kernels/vector_select_k.cc
index ea072d179e..7d94fa2ba3 100644
--- a/cpp/src/arrow/compute/kernels/vector_select_k.cc
+++ b/cpp/src/arrow/compute/kernels/vector_select_k.cc
@@ -76,13 +76,10 @@ class SelectKComparator<SortOrder::Descending> {
}
};
-struct OutputRangesByNullLikeness {
- std::span<uint64_t> non_null_like_range;
- std::span<uint64_t> nan_range;
- std::span<uint64_t> null_range;
-};
-
-OutputRangesByNullLikeness CalculateOutputRangesByNullLikeness(
+// Clip the group counts to the k output slots and lay out the corresponding
+// ranges in `output_indices`. Because k was clipped to the input length, the
+// clipped counts always sum to k == output_indices.size().
+NullLikePartition CalculateOutputRangesByNullLikeness(
int64_t non_null_like_count, int64_t nan_count, int64_t null_count,
NullPlacement null_placement, std::span<uint64_t> output_indices) {
auto k = static_cast<int64_t>(output_indices.size());
@@ -93,21 +90,13 @@ OutputRangesByNullLikeness
CalculateOutputRangesByNullLikeness(
non_null_like_to_take = std::min(k, non_null_like_count);
nan_to_take = std::min(k - non_null_like_to_take, nan_count);
null_to_take = std::min(k - non_null_like_to_take - nan_to_take,
null_count);
- return OutputRangesByNullLikeness{
- .non_null_like_range = output_indices.subspan(0,
non_null_like_to_take),
- .nan_range = output_indices.subspan(non_null_like_to_take,
nan_to_take),
- .null_range =
- output_indices.subspan(non_null_like_to_take + nan_to_take,
null_to_take)};
} else {
null_to_take = std::min(k, null_count);
nan_to_take = std::min(k - null_to_take, nan_count);
non_null_like_to_take = std::min(k - null_to_take - nan_to_take,
non_null_like_count);
- return OutputRangesByNullLikeness{
- .non_null_like_range =
- output_indices.subspan(null_to_take + nan_to_take,
non_null_like_to_take),
- .nan_range = output_indices.subspan(null_to_take, nan_to_take),
- .null_range = output_indices.subspan(0, null_to_take)};
}
+ return NullLikePartition::FromCounts(output_indices, non_null_like_to_take,
nan_to_take,
+ null_to_take, null_placement);
}
template <typename Comparator>
@@ -165,29 +154,6 @@ void HeapSortNonNullsToOutput(std::span<uint64_t>
non_null_input_range,
}
}
-struct PartitionResultByNullLikeness {
- std::span<uint64_t> non_null_like_range;
- std::span<uint64_t> null_range;
- std::span<uint64_t> nan_range;
-};
-
-template <typename ArrayType, typename Partitioner>
-PartitionResultByNullLikeness PartitionNullsAndNans(uint64_t* indices_begin,
- uint64_t* indices_end,
- const ArrayType& values,
- int64_t offset,
- NullPlacement
null_placement) {
- // Partition nulls at start (resp. end), and null-like values just before
(resp. after)
- NullPartitionResult p = PartitionNullsOnly<Partitioner>(indices_begin,
indices_end,
- values, offset,
null_placement);
- NullPartitionResult q = PartitionNullLikes<ArrayType, Partitioner>(
- p.non_nulls_begin, p.non_nulls_end, values, offset, null_placement);
- return PartitionResultByNullLikeness{
- .non_null_like_range = {q.non_nulls_begin, q.non_nulls_end},
- .null_range = {p.nulls_begin, p.nulls_end},
- .nan_range = {q.nulls_begin, q.nulls_end}};
-}
-
class ArraySelector : public TypeVisitor {
public:
ArraySelector(ExecContext* ctx, const Array& array, const SelectKOptions&
options,
@@ -227,16 +193,14 @@ class ArraySelector : public TypeVisitor {
std::vector<uint64_t> indices(arr.length());
- uint64_t* indices_begin = indices.data();
- uint64_t* indices_end = indices_begin + indices.size();
- std::iota(indices_begin, indices_end, 0);
+ std::iota(indices.begin(), indices.end(), 0);
ARROW_ASSIGN_OR_RAISE(auto take_indices,
MakeMutableUInt64Array(k_, ctx_->memory_pool()));
auto* output_begin = take_indices->template GetMutableValues<uint64_t>(1);
const auto p = PartitionNullsAndNans<ArrayType, NonStablePartitioner>(
- indices_begin, indices_end, arr, 0, null_placement_);
+ indices, arr, 0, null_placement_);
// From k, calculate
// l = non_null_like elements to take from PartitionResult
@@ -251,10 +215,9 @@ class ArraySelector : public TypeVisitor {
HeapSortNonNullsToOutput<InType, sort_order>(p.non_null_like_range, arr,
output.non_null_like_range);
- std::copy(p.nan_range.begin(), p.nan_range.begin() +
output.nan_range.size(),
- output.nan_range.begin());
- std::copy(p.null_range.begin(), p.null_range.begin() +
output.null_range.size(),
- output.null_range.begin());
+ std::copy(p.nan_begin(), p.nan_begin() + output.nan_range.size(),
output.nan_begin());
+ std::copy(p.null_begin(), p.null_begin() + output.null_range.size(),
+ output.null_begin());
*output_ = Datum(take_indices);
return Status::OK();
@@ -323,7 +286,7 @@ class ChunkedArraySelector : public TypeVisitor {
std::vector<std::shared_ptr<ArrayType>> chunks_holder;
chunks_holder.reserve(num_chunks);
- std::vector<PartitionResultByNullLikeness> partitions_by_chunk;
+ std::vector<NullLikePartition> partitions_by_chunk;
partitions_by_chunk.reserve(num_chunks);
std::vector<std::vector<uint64_t>> indices_by_chunk;
indices_by_chunk.reserve(num_chunks);
@@ -339,13 +302,11 @@ class ChunkedArraySelector : public TypeVisitor {
auto& indices = indices_by_chunk.emplace_back();
indices.resize(arr.length());
- uint64_t* indices_begin = indices.data();
- uint64_t* indices_end = indices_begin + indices.size();
- std::iota(indices_begin, indices_end, 0);
+ std::iota(indices.begin(), indices.end(), 0);
partitions_by_chunk.emplace_back(
- PartitionNullsAndNans<ArrayType, NonStablePartitioner>(
- indices_begin, indices_end, arr, 0, null_placement_));
+ PartitionNullsAndNans<ArrayType, NonStablePartitioner>(indices, arr,
0,
+
null_placement_));
null_count += partitions_by_chunk.back().null_range.size();
nan_count += partitions_by_chunk.back().nan_range.size();
@@ -493,12 +454,8 @@ class RecordBatchSelector {
const auto& first_remaining_sort_key =
selector_->sort_keys_[start_sort_key_index_];
const auto& arr = checked_cast<const
ArrayType&>(first_remaining_sort_key.array);
- uint64_t* input_indices_begin = input_indices_.data();
- uint64_t* input_indices_end = input_indices_.data() +
input_indices_.size();
-
const auto p = PartitionNullsAndNans<ArrayType, NonStablePartitioner>(
- input_indices_begin, input_indices_end, arr, 0,
- first_remaining_sort_key.null_placement);
+ input_indices_, arr, 0, first_remaining_sort_key.null_placement);
// From k = output_indices_.size(), calculate
// l = non_null_like elements to take from PartitionResult
@@ -523,13 +480,13 @@ class RecordBatchSelector {
}
if (output.nan_range.size() > 0) {
// We have the last sort_key, can just copy over the null values
- std::copy(p.nan_range.begin(), p.nan_range.begin() +
output.nan_range.size(),
- output.nan_range.begin());
+ std::copy(p.nan_begin(), p.nan_begin() + output.nan_range.size(),
+ output.nan_begin());
}
if (output.null_range.size() > 0) {
// We have the last sort_key, can just copy over the null values
- std::copy(p.null_range.begin(), p.null_range.begin() +
output.null_range.size(),
- output.null_range.begin());
+ std::copy(p.null_begin(), p.null_begin() + output.null_range.size(),
+ output.null_begin());
}
} else {
if (!output.non_null_like_range.empty()) {
diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc
b/cpp/src/arrow/compute/kernels/vector_sort.cc
index 06eed16023..de5b7e1bca 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort.cc
@@ -19,6 +19,7 @@
#include <unordered_set>
#include "arrow/compute/function.h"
+#include "arrow/compute/kernels/chunked_internal.h"
#include "arrow/compute/kernels/vector_sort_internal.h"
#include "arrow/compute/registry.h"
#include "arrow/compute/registry_internal.h"
@@ -45,13 +46,12 @@ Result<RecordBatchVector> BatchesFromTable(const Table&
table) {
// then merging the sorted chunks recursively.
class ChunkedArraySorter : public TypeVisitor {
public:
- ChunkedArraySorter(ExecContext* ctx, uint64_t* indices_begin, uint64_t*
indices_end,
+ ChunkedArraySorter(ExecContext* ctx, std::span<uint64_t> indices,
const std::shared_ptr<DataType>& physical_type,
const ArrayVector& physical_chunks, const SortOrder order,
- const NullPlacement null_placement, NullPartitionResult*
output)
+ const NullPlacement null_placement, NullLikePartition*
output)
: TypeVisitor(),
- indices_begin_(indices_begin),
- indices_end_(indices_end),
+ indices_(indices),
physical_type_(physical_type),
physical_chunks_(physical_chunks),
order_(order),
@@ -72,7 +72,7 @@ class ChunkedArraySorter : public TypeVisitor {
#undef VISIT
Status Visit(const NullType&) override {
- std::iota(indices_begin_, indices_end_, 0);
+ std::iota(indices_.begin(), indices_.end(), 0);
return Status::OK();
}
@@ -83,63 +83,58 @@ class ChunkedArraySorter : public TypeVisitor {
ArraySortOptions options(order_, null_placement_);
const auto num_chunks = static_cast<int>(physical_chunks_.size());
if (num_chunks == 0) {
- *output_ = {indices_end_, indices_end_, indices_end_, indices_end_};
+ DCHECK_EQ(static_cast<int64_t>(indices_.size()), 0);
+ *output_ = NullLikePartition::FromCounts(indices_, 0, 0, 0,
null_placement_);
return Status::OK();
}
- const int64_t num_indices = static_cast<int64_t>(indices_end_ -
indices_begin_);
+ const int64_t num_indices = static_cast<int64_t>(indices_.size());
const auto arrays = GetArrayPointers(physical_chunks_);
// Sort each chunk independently and merge to sorted indices.
// This is a serial implementation.
- std::vector<NullPartitionResult> sorted(num_chunks);
+ std::vector<NullLikePartition> sorted(num_chunks);
// First sort all individual chunks
int64_t begin_offset = 0;
- int64_t end_offset = 0;
- int64_t null_count = 0;
+ int64_t non_null_like_count = 0;
for (int i = 0; i < num_chunks; ++i) {
const auto array = checked_cast<const ArrayType*>(arrays[i]);
- end_offset += array->length();
- null_count += array->null_count();
- ARROW_ASSIGN_OR_RAISE(sorted[i], array_sorter_(indices_begin_ +
begin_offset,
- indices_begin_ +
end_offset, *array,
- begin_offset, options,
ctx_));
- begin_offset = end_offset;
+ const auto array_length = array->length();
+ ARROW_ASSIGN_OR_RAISE(
+ sorted[i], array_sorter_(indices_.subspan(begin_offset,
array_length), *array,
+ begin_offset, options, ctx_));
+ non_null_like_count += sorted[i].non_null_like_range.size();
+ begin_offset += array_length;
}
- DCHECK_EQ(end_offset, num_indices);
+ DCHECK_EQ(begin_offset, num_indices);
// Then merge them by pairs, recursively
if (sorted.size() > 1) {
- ChunkedIndexMapper chunked_mapper(arrays, indices_begin_, indices_end_);
- ARROW_ASSIGN_OR_RAISE(auto chunked_indices_pair,
- chunked_mapper.LogicalToPhysical());
- auto [chunked_indices_begin, chunked_indices_end] = chunked_indices_pair;
+ ChunkedIndexMapper chunked_mapper(arrays, indices_);
+ ARROW_ASSIGN_OR_RAISE(auto chunked_indices,
chunked_mapper.LogicalToPhysical());
- std::vector<ChunkedNullPartitionResult> chunk_sorted(num_chunks);
+ std::vector<ChunkedNullLikePartition> chunk_sorted(num_chunks);
for (int i = 0; i < num_chunks; ++i) {
- chunk_sorted[i] = sorted[i].TranslateTo(indices_begin_,
chunked_indices_begin);
+ chunk_sorted[i] = sorted[i].TranslateTo(indices_.data(),
chunked_indices.data());
}
- auto merge_nulls = [&](CompressedChunkLocation* nulls_begin,
- CompressedChunkLocation* nulls_middle,
- CompressedChunkLocation* nulls_end,
- CompressedChunkLocation* temp_indices, int64_t
null_count) {
- if (has_null_like_values<typename ArrayType::TypeClass>()) {
- PartitionNullsOnly<StablePartitioner>(nulls_begin, nulls_end, arrays,
- null_count, null_placement_);
- }
+ // merge function for merging ranges where the first sort key is equal
+ // -> no-op because we have only one sort key
+ auto merge_remaining_sort_keys = [&](std::span<CompressedChunkLocation>
range,
+ size_t middle,
+ CompressedChunkLocation*
temp_indices) {};
+
+ auto merge_non_nulls = [&](std::span<CompressedChunkLocation> range,
size_t middle,
+ CompressedChunkLocation* temp_indices) {
+ MergeNonNulls<ArrayType>(range.subspan(0, middle),
range.subspan(middle), arrays,
+ {temp_indices, range.size()});
};
- auto merge_non_nulls =
- [&](CompressedChunkLocation* range_begin, CompressedChunkLocation*
range_middle,
- CompressedChunkLocation* range_end, CompressedChunkLocation*
temp_indices) {
- MergeNonNulls<ArrayType>(range_begin, range_middle, range_end,
arrays,
- temp_indices);
- };
-
- ChunkedMergeImpl merge_impl{null_placement_, std::move(merge_nulls),
+
+ ChunkedMergeImpl merge_impl{null_placement_,
std::move(merge_remaining_sort_keys),
std::move(merge_non_nulls)};
- // std::merge is only called on non-null values, so size temp indices
accordingly
- RETURN_NOT_OK(merge_impl.Init(ctx_, num_indices - null_count));
+ // std::merge is only called on non-null-like values, so size temp
indices
+ // accordingly
+ RETURN_NOT_OK(merge_impl.Init(ctx_, non_null_like_count));
while (chunk_sorted.size() > 1) {
// Merge all pairs of chunks
@@ -149,7 +144,7 @@ class ChunkedArraySorter : public TypeVisitor {
const auto& left = *it++;
const auto& right = *it++;
DCHECK_EQ(left.overall_end(), right.overall_begin());
- const auto merged = merge_impl.Merge(left, right, null_count);
+ const auto merged = merge_impl.Merge(left, right);
*out_it++ = merged;
}
if (it < chunk_sorted.end()) {
@@ -160,37 +155,38 @@ class ChunkedArraySorter : public TypeVisitor {
// Reverse everything
sorted.resize(1);
- sorted[0] = chunk_sorted[0].TranslateTo(chunked_indices_begin,
indices_begin_);
+ sorted[0] = chunk_sorted[0].TranslateTo(chunked_indices.data(),
indices_.data());
RETURN_NOT_OK(chunked_mapper.PhysicalToLogical());
}
DCHECK_EQ(sorted.size(), 1);
- DCHECK_EQ(sorted[0].overall_begin(), indices_begin_);
- DCHECK_EQ(sorted[0].overall_end(), indices_end_);
- // Note that "nulls" can also include NaNs, hence the >= check
- DCHECK_GE(sorted[0].null_count(), null_count);
+ DCHECK_EQ(sorted[0].overall_begin(), indices_.data());
+ DCHECK_EQ(sorted[0].overall_end(), indices_.data() + indices_.size());
+ DCHECK_EQ(static_cast<int64_t>(sorted[0].non_null_like_range.size()),
+ non_null_like_count);
*output_ = sorted[0];
return Status::OK();
}
template <typename ArrayType>
- void MergeNonNulls(CompressedChunkLocation* range_begin,
- CompressedChunkLocation* range_middle,
- CompressedChunkLocation* range_end,
+ void MergeNonNulls(std::span<CompressedChunkLocation> left,
+ std::span<CompressedChunkLocation> right,
std::span<const Array* const> arrays,
- CompressedChunkLocation* temp_indices) {
+ std::span<CompressedChunkLocation> temp_indices) {
using ArrowType = typename ArrayType::TypeClass;
if (order_ == SortOrder::Ascending) {
- std::merge(range_begin, range_middle, range_middle, range_end,
temp_indices,
+ std::merge(left.data(), left.data() + left.size(), right.data(),
+ right.data() + right.size(), temp_indices.data(),
[&](CompressedChunkLocation left, CompressedChunkLocation
right) {
return ChunkValue<ArrowType>(arrays, left) <
ChunkValue<ArrowType>(arrays, right);
});
} else {
- std::merge(range_begin, range_middle, range_middle, range_end,
temp_indices,
+ std::merge(left.data(), left.data() + left.size(), right.data(),
+ right.data() + right.size(), temp_indices.data(),
[&](CompressedChunkLocation left, CompressedChunkLocation
right) {
// We don't use 'left > right' here to reduce required
// operator. If we use 'right < left' here, '<' is only
@@ -200,7 +196,8 @@ class ChunkedArraySorter : public TypeVisitor {
});
}
// Copy back temp area into main buffer
- std::copy(temp_indices, temp_indices + (range_end - range_begin),
range_begin);
+ std::copy(temp_indices.data(), temp_indices.data() + left.size() +
right.size(),
+ left.data());
}
template <typename ArrowType>
@@ -211,15 +208,14 @@ class ChunkedArraySorter : public TypeVisitor {
.template Value<ArrowType>();
}
- uint64_t* indices_begin_;
- uint64_t* indices_end_;
+ std::span<uint64_t> indices_;
const std::shared_ptr<DataType>& physical_type_;
const ArrayVector& physical_chunks_;
const SortOrder order_;
const NullPlacement null_placement_;
ArraySortFunc array_sorter_;
ExecContext* ctx_;
- NullPartitionResult* output_;
+ NullLikePartition* output_;
};
// ----------------------------------------------------------------------
@@ -228,26 +224,26 @@ class ChunkedArraySorter : public TypeVisitor {
// Visit contiguous ranges of equal values. All entries are assumed
// to be non-null.
template <typename ArrayType, typename Visitor>
-void VisitConstantRanges(const ArrayType& array, uint64_t* indices_begin,
- uint64_t* indices_end, int64_t offset, Visitor&&
visit) {
+void VisitConstantRanges(const ArrayType& array, std::span<uint64_t> indices,
+ int64_t offset, Visitor&& visit) {
using GetView = GetViewType<typename ArrayType::TypeClass>;
- if (indices_begin == indices_end) {
+ if (indices.empty()) {
return;
}
- auto range_start = indices_begin;
+ auto range_start = indices.begin();
auto range_cur = range_start;
auto last_value = GetView::LogicalValue(array.GetView(*range_cur - offset));
- while (++range_cur != indices_end) {
+ while (++range_cur != indices.end()) {
auto v = GetView::LogicalValue(array.GetView(*range_cur - offset));
if (v != last_value) {
- visit(range_start, range_cur);
+ visit({range_start, range_cur});
range_start = range_cur;
last_value = v;
}
}
if (range_start != range_cur) {
- visit(range_start, range_cur);
+ visit({range_start, range_cur});
}
}
@@ -259,8 +255,7 @@ class RecordBatchColumnSorter {
: next_column_(next_column) {}
virtual ~RecordBatchColumnSorter() {}
- virtual NullPartitionResult SortRange(uint64_t* indices_begin, uint64_t*
indices_end,
- int64_t offset) = 0;
+ virtual NullLikePartition SortRange(std::span<uint64_t> indices, int64_t
offset) = 0;
protected:
RecordBatchColumnSorter* next_column_;
@@ -281,36 +276,37 @@ class ConcreteRecordBatchColumnSorter : public
RecordBatchColumnSorter {
null_placement_(null_placement),
null_count_(array_.null_count()) {}
- NullPartitionResult SortRange(uint64_t* indices_begin, uint64_t* indices_end,
- int64_t offset) override {
+ NullLikePartition SortRange(std::span<uint64_t> indices, int64_t offset)
override {
using GetView = GetViewType<Type>;
- NullPartitionResult p;
+ NullLikePartition partitions;
if (null_count_ == 0) {
- p = NullPartitionResult::NoNulls(indices_begin, indices_end,
null_placement_);
+ partitions = PartitionNansOnly<ArrayType, StablePartitioner>(
+ indices, array_, offset, null_placement_);
+
} else {
// NOTE that null_count_ is merely an upper bound on the number of nulls
// in this particular range.
- p = PartitionNullsOnly<StablePartitioner>(indices_begin, indices_end,
array_,
- offset, null_placement_);
- DCHECK_LE(p.nulls_end - p.nulls_begin, null_count_);
+ partitions = PartitionNullsAndNans<ArrayType, StablePartitioner>(
+ indices, array_, offset, null_placement_);
+ DCHECK_LE(static_cast<int64_t>(partitions.null_range.size()),
null_count_);
}
- const NullPartitionResult q = PartitionNullLikes<ArrayType,
StablePartitioner>(
- p.non_nulls_begin, p.non_nulls_end, array_, offset, null_placement_);
// TODO This is roughly the same as ArrayCompareSorter.
// Also, we would like to use a counting sort if possible. This requires
// a counting sort compatible with indirect indexing.
if (order_ == SortOrder::Ascending) {
std::stable_sort(
- q.non_nulls_begin, q.non_nulls_end, [&](uint64_t left, uint64_t
right) {
+ partitions.non_null_like_begin(), partitions.non_null_like_end(),
+ [&](uint64_t left, uint64_t right) {
const auto lhs = GetView::LogicalValue(array_.GetView(left -
offset));
const auto rhs = GetView::LogicalValue(array_.GetView(right -
offset));
return lhs < rhs;
});
} else {
std::stable_sort(
- q.non_nulls_begin, q.non_nulls_end, [&](uint64_t left, uint64_t
right) {
+ partitions.non_null_like_begin(), partitions.non_null_like_end(),
+ [&](uint64_t left, uint64_t right) {
// We don't use 'left > right' here to reduce required operator.
// If we use 'right < left' here, '<' is only required.
const auto lhs = GetView::LogicalValue(array_.GetView(left -
offset));
@@ -322,22 +318,19 @@ class ConcreteRecordBatchColumnSorter : public
RecordBatchColumnSorter {
if (next_column_ != nullptr) {
// Visit all ranges of equal values in this column and sort them on
// the next column.
- SortNextColumn(q.nulls_begin, q.nulls_end, offset);
- SortNextColumn(p.nulls_begin, p.nulls_end, offset);
- VisitConstantRanges(array_, q.non_nulls_begin, q.non_nulls_end, offset,
- [&](uint64_t* range_start, uint64_t* range_end) {
- SortNextColumn(range_start, range_end, offset);
- });
+ SortNextColumn(partitions.null_range, offset);
+ SortNextColumn(partitions.nan_range, offset);
+ VisitConstantRanges(
+ array_, partitions.non_null_like_range, offset,
+ [&](std::span<uint64_t> indices) { SortNextColumn(indices, offset);
});
}
- return NullPartitionResult{q.non_nulls_begin, q.non_nulls_end,
- std::min(q.nulls_begin, p.nulls_begin),
- std::max(q.nulls_end, p.nulls_end)};
+ return partitions;
}
- void SortNextColumn(uint64_t* indices_begin, uint64_t* indices_end, int64_t
offset) {
+ void SortNextColumn(std::span<uint64_t> indices, int64_t offset) {
// Avoid the cost of a virtual method call in trivial cases
- if (indices_end - indices_begin > 1) {
- next_column_->SortRange(indices_begin, indices_end, offset);
+ if (indices.size() > 1) {
+ next_column_->SortRange(indices, offset);
}
}
@@ -357,12 +350,11 @@ class ConcreteRecordBatchColumnSorter<NullType> : public
RecordBatchColumnSorter
RecordBatchColumnSorter* next_column =
nullptr)
: RecordBatchColumnSorter(next_column), null_placement_(null_placement)
{}
- NullPartitionResult SortRange(uint64_t* indices_begin, uint64_t* indices_end,
- int64_t offset) {
+ NullLikePartition SortRange(std::span<uint64_t> indices, int64_t offset) {
if (next_column_ != nullptr) {
- next_column_->SortRange(indices_begin, indices_end, offset);
+ next_column_->SortRange(indices, offset);
}
- return NullPartitionResult::NullsOnly(indices_begin, indices_end,
null_placement_);
+ return NullLikePartition::FromCounts(indices, 0, 0, indices.size(),
null_placement_);
}
protected:
@@ -395,20 +387,17 @@ class RadixRecordBatchSorter {
public:
using ResolvedSortKey = ResolvedRecordBatchSortKey;
- RadixRecordBatchSorter(uint64_t* indices_begin, uint64_t* indices_end,
+ RadixRecordBatchSorter(std::span<uint64_t> indices,
std::vector<ResolvedSortKey> sort_keys)
- : sort_keys_(std::move(sort_keys)),
- indices_begin_(indices_begin),
- indices_end_(indices_end) {}
+ : sort_keys_(std::move(sort_keys)), indices_(indices) {}
- RadixRecordBatchSorter(uint64_t* indices_begin, uint64_t* indices_end,
- const RecordBatch& batch, const SortOptions& options)
+ RadixRecordBatchSorter(std::span<uint64_t> indices, const RecordBatch& batch,
+ const SortOptions& options)
: sort_keys_(ResolveRecordBatchSortKeys(batch, options.GetSortKeys(),
&status_)),
- indices_begin_(indices_begin),
- indices_end_(indices_end) {}
+ indices_(indices) {}
// Offset is for table sorting
- Result<NullPartitionResult> Sort(int64_t offset = 0) {
+ Result<NullLikePartition> Sort(int64_t offset = 0) {
ARROW_RETURN_NOT_OK(status_);
// Create column sorters from right to left
@@ -421,7 +410,7 @@ class RadixRecordBatchSorter {
}
// Sort from left to right
- return column_sorts.front()->SortRange(indices_begin_, indices_end_,
offset);
+ return column_sorts.front()->SortRange(indices_, offset);
}
protected:
@@ -474,8 +463,7 @@ class RadixRecordBatchSorter {
}
const std::vector<ResolvedSortKey> sort_keys_;
- uint64_t* indices_begin_;
- uint64_t* indices_end_;
+ std::span<uint64_t> indices_;
Status status_;
};
@@ -484,17 +472,13 @@ class MultipleKeyRecordBatchSorter : public TypeVisitor {
public:
using ResolvedSortKey = ResolvedRecordBatchSortKey;
- MultipleKeyRecordBatchSorter(uint64_t* indices_begin, uint64_t* indices_end,
+ MultipleKeyRecordBatchSorter(std::span<uint64_t> indices,
std::vector<ResolvedSortKey> sort_keys)
- : indices_begin_(indices_begin),
- indices_end_(indices_end),
- sort_keys_(std::move(sort_keys)),
- comparator_(sort_keys_) {}
+ : indices_(indices), sort_keys_(std::move(sort_keys)),
comparator_(sort_keys_) {}
- MultipleKeyRecordBatchSorter(uint64_t* indices_begin, uint64_t* indices_end,
- const RecordBatch& batch, const SortOptions&
options)
- : indices_begin_(indices_begin),
- indices_end_(indices_end),
+ MultipleKeyRecordBatchSorter(std::span<uint64_t> indices, const RecordBatch&
batch,
+ const SortOptions& options)
+ : indices_(indices),
sort_keys_(ResolveSortKeys(batch, options.GetSortKeys(), &status_)),
comparator_(sort_keys_) {}
@@ -540,72 +524,71 @@ class MultipleKeyRecordBatchSorter : public TypeVisitor {
const auto p = PartitionNullsInternal<Type>(first_sort_key);
// Sort first-key non-nulls
- std::stable_sort(
- p.non_nulls_begin, p.non_nulls_end, [&](uint64_t left, uint64_t right)
{
- // Both values are never null nor NaN
- // (otherwise they've been partitioned away above).
- const auto value_left = GetView::LogicalValue(array.GetView(left));
- const auto value_right = GetView::LogicalValue(array.GetView(right));
- if (value_left != value_right) {
- bool compared = value_left < value_right;
- if (first_sort_key.order == SortOrder::Ascending) {
- return compared;
- } else {
- return !compared;
- }
- }
- // If the left value equals to the right value,
- // we need to compare the second and following
- // sort keys.
- return comparator.Compare(left, right, 1);
- });
+ std::stable_sort(p.non_null_like_begin(), p.non_null_like_end(),
+ [&](uint64_t left, uint64_t right) {
+ // Both values are never null nor NaN
+ // (otherwise they've been partitioned away above).
+ const auto value_left =
GetView::LogicalValue(array.GetView(left));
+ const auto value_right =
+ GetView::LogicalValue(array.GetView(right));
+ if (value_left != value_right) {
+ bool compared = value_left < value_right;
+ if (first_sort_key.order == SortOrder::Ascending) {
+ return compared;
+ } else {
+ return !compared;
+ }
+ }
+ // If the left value equals to the right value,
+ // we need to compare the second and following
+ // sort keys.
+ return comparator.Compare(left, right, 1);
+ });
return comparator_.status();
}
template <typename Type>
enable_if_null<Type, Status> SortInternal() {
- std::stable_sort(indices_begin_, indices_end_, [&](uint64_t left, uint64_t
right) {
- return comparator_.Compare(left, right, 1);
- });
+ std::stable_sort(indices_.data(), indices_.data() + indices_.size(),
+ [&](uint64_t left, uint64_t right) {
+ return comparator_.Compare(left, right, 1);
+ });
return comparator_.status();
}
// Behaves like PartitionNulls() but this supports multiple sort keys.
template <typename Type>
- NullPartitionResult PartitionNullsInternal(const ResolvedSortKey&
first_sort_key) {
+ NullLikePartition PartitionNullsInternal(const ResolvedSortKey&
first_sort_key) {
using ArrayType = typename TypeTraits<Type>::ArrayType;
const ArrayType& array =
::arrow::internal::checked_cast<const
ArrayType&>(first_sort_key.array);
- const auto p = PartitionNullsOnly<StablePartitioner>(
- indices_begin_, indices_end_, array, 0, first_sort_key.null_placement);
- const auto q = PartitionNullLikes<ArrayType, StablePartitioner>(
- p.non_nulls_begin, p.non_nulls_end, array, 0,
first_sort_key.null_placement);
+ const auto p = PartitionNullsAndNans<ArrayType, StablePartitioner>(
+ indices_, array, 0, first_sort_key.null_placement);
auto& comparator = comparator_;
- if (q.nulls_begin != q.nulls_end) {
+ if (!p.nan_range.empty()) {
// Sort all NaNs by the second and following sort keys.
// TODO: could we instead run an independent sort from the second key on
// this slice?
- std::stable_sort(q.nulls_begin, q.nulls_end,
+ std::stable_sort(p.nan_begin(), p.nan_end(),
[&comparator](uint64_t left, uint64_t right) {
return comparator.Compare(left, right, 1);
});
}
- if (p.nulls_begin != p.nulls_end) {
+ if (!p.null_range.empty()) {
// Sort all nulls by the second and following sort keys.
// TODO: could we instead run an independent sort from the second key on
// this slice?
- std::stable_sort(p.nulls_begin, p.nulls_end,
+ std::stable_sort(p.null_begin(), p.null_end(),
[&comparator](uint64_t left, uint64_t right) {
return comparator.Compare(left, right, 1);
});
}
- return q;
+ return p;
}
- uint64_t* indices_begin_;
- uint64_t* indices_end_;
+ std::span<uint64_t> indices_;
Status status_;
std::vector<ResolvedSortKey> sort_keys_;
Comparator comparator_;
@@ -625,15 +608,14 @@ class TableSorter {
using Comparator = MultipleKeyComparator<ResolvedSortKey>;
public:
- TableSorter(ExecContext* ctx, uint64_t* indices_begin, uint64_t* indices_end,
- const Table& table, const SortOptions& options)
+ TableSorter(ExecContext* ctx, std::span<uint64_t> indices, const Table&
table,
+ const SortOptions& options)
: ctx_(ctx),
table_(table),
batches_(MakeBatches(table, &status_)),
options_(options),
sort_keys_(ResolveSortKeys(table, batches_, options.GetSortKeys(),
&status_)),
- indices_begin_(indices_begin),
- indices_end_(indices_end),
+ indices_(indices),
comparator_(sort_keys_) {}
// This is optimized for null partitioning and merging along the first sort
key.
@@ -670,48 +652,44 @@ class TableSorter {
if (num_batches == 0) {
return Status::OK();
}
- std::vector<NullPartitionResult> sorted(num_batches);
+ std::vector<NullLikePartition> sorted(num_batches);
// First sort all individual batches
int64_t begin_offset = 0;
int64_t end_offset = 0;
- int64_t null_count = 0;
for (int64_t i = 0; i < num_batches; ++i) {
const auto& batch = *batches_[i];
end_offset += batch.num_rows();
- RadixRecordBatchSorter sorter(indices_begin_ + begin_offset,
- indices_begin_ + end_offset, batch,
options_);
+ RadixRecordBatchSorter sorter(
+ {indices_.data() + begin_offset, indices_.data() + end_offset},
batch,
+ options_);
ARROW_ASSIGN_OR_RAISE(sorted[i], sorter.Sort(begin_offset));
- DCHECK_EQ(sorted[i].overall_begin(), indices_begin_ + begin_offset);
- DCHECK_EQ(sorted[i].overall_end(), indices_begin_ + end_offset);
- DCHECK_EQ(sorted[i].non_null_count() + sorted[i].null_count(),
batch.num_rows());
+ DCHECK_EQ(sorted[i].overall_begin(), indices_.data() + begin_offset);
+ DCHECK_EQ(sorted[i].overall_end(), indices_.data() + end_offset);
+ DCHECK_EQ(
+ static_cast<int64_t>(sorted[i].non_null_like_range.size() +
+ sorted[i].nan_range.size() +
sorted[i].null_range.size()),
+ batch.num_rows());
begin_offset = end_offset;
- // XXX this is an upper bound on the true null count
- null_count += sorted[i].null_count();
}
- DCHECK_EQ(end_offset, indices_end_ - indices_begin_);
+ DCHECK_EQ(end_offset, static_cast<int64_t>(indices_.size()));
// Then merge them by pairs, recursively
if (sorted.size() > 1) {
- ChunkedIndexMapper chunked_mapper(batches_, indices_begin_,
indices_end_);
- ARROW_ASSIGN_OR_RAISE(auto chunked_indices_pair,
- chunked_mapper.LogicalToPhysical());
- auto [chunked_indices_begin, chunked_indices_end] = chunked_indices_pair;
+ ChunkedIndexMapper chunked_mapper(batches_, indices_);
+ ARROW_ASSIGN_OR_RAISE(auto chunked_indices,
chunked_mapper.LogicalToPhysical());
- std::vector<ChunkedNullPartitionResult> chunk_sorted(num_batches);
+ std::vector<ChunkedNullLikePartition> chunk_sorted(num_batches);
for (int64_t i = 0; i < num_batches; ++i) {
- chunk_sorted[i] = sorted[i].TranslateTo(indices_begin_,
chunked_indices_begin);
+ chunk_sorted[i] = sorted[i].TranslateTo(indices_.data(),
chunked_indices.data());
}
struct Visitor {
TableSorter* sorter;
- std::vector<ChunkedNullPartitionResult>* chunk_sorted;
- int64_t null_count;
+ std::vector<ChunkedNullLikePartition>* chunk_sorted;
-#define VISIT(TYPE) \
- Status Visit(const TYPE& type) { \
- return sorter->MergeInternal<TYPE>(chunk_sorted, null_count); \
- }
+#define VISIT(TYPE) \
+ Status Visit(const TYPE& type) { return
sorter->MergeInternal<TYPE>(chunk_sorted); }
VISIT_SORTABLE_PHYSICAL_TYPES(VISIT)
VISIT(NullType)
@@ -722,12 +700,13 @@ class TableSorter {
type.ToString());
}
};
- Visitor visitor{this, &chunk_sorted, null_count};
+ Visitor visitor{this, &chunk_sorted};
RETURN_NOT_OK(VisitTypeInline(*sort_keys_[0].type, &visitor));
DCHECK_EQ(chunk_sorted.size(), 1);
- DCHECK_EQ(chunk_sorted[0].overall_begin(), chunked_indices_begin);
- DCHECK_EQ(chunk_sorted[0].overall_end(), chunked_indices_end);
+ DCHECK_EQ(chunk_sorted[0].overall_begin(), chunked_indices.data());
+ DCHECK_EQ(chunk_sorted[0].overall_end(),
+ chunked_indices.data() + chunked_indices.size());
RETURN_NOT_OK(chunked_mapper.PhysicalToLogical());
}
@@ -736,22 +715,21 @@ class TableSorter {
// Recursive merge routine, typed on the first sort key
template <typename ArrowType>
- Status MergeInternal(std::vector<ChunkedNullPartitionResult>* sorted,
- int64_t null_count) {
- auto merge_nulls = [&](CompressedChunkLocation* nulls_begin,
- CompressedChunkLocation* nulls_middle,
- CompressedChunkLocation* nulls_end,
- CompressedChunkLocation* temp_indices, int64_t
null_count) {
- MergeNulls<ArrowType>(nulls_begin, nulls_middle, nulls_end, temp_indices,
- null_count);
+ Status MergeInternal(std::vector<ChunkedNullLikePartition>* sorted) {
+ // merge function where the ranges compare equal on the first sort-key
+ auto merge_remaining_sort_keys = [&](std::span<CompressedChunkLocation>
range,
+ size_t middle,
+ CompressedChunkLocation*
temp_indices) {
+ MergeRemainingSortKeys(range, middle, temp_indices);
};
- auto merge_non_nulls =
- [&](CompressedChunkLocation* range_begin, CompressedChunkLocation*
range_middle,
- CompressedChunkLocation* range_end, CompressedChunkLocation*
temp_indices) {
- MergeNonNulls<ArrowType>(range_begin, range_middle, range_end,
temp_indices);
- };
- ChunkedMergeImpl merge_impl(sort_keys_[0].null_placement,
std::move(merge_nulls),
+ auto merge_non_nulls = [&](std::span<CompressedChunkLocation> range,
size_t middle,
+ CompressedChunkLocation* temp_indices) {
+ MergeNonNulls<ArrowType>(range, middle, temp_indices);
+ };
+
+ ChunkedMergeImpl merge_impl(sort_keys_[0].null_placement,
+ std::move(merge_remaining_sort_keys),
std::move(merge_non_nulls));
RETURN_NOT_OK(merge_impl.Init(ctx_, table_.num_rows()));
@@ -762,7 +740,7 @@ class TableSorter {
const auto& left = *it++;
const auto& right = *it++;
DCHECK_EQ(left.overall_end(), right.overall_begin());
- *out_it++ = merge_impl.Merge(left, right, null_count);
+ *out_it++ = merge_impl.Merge(left, right);
}
if (it < sorted->end()) {
*out_it++ = *it++;
@@ -772,54 +750,22 @@ class TableSorter {
return comparator_.status();
}
- template <typename ArrowType>
- void MergeNulls(CompressedChunkLocation* nulls_begin,
- CompressedChunkLocation* nulls_middle,
- CompressedChunkLocation* nulls_end,
- CompressedChunkLocation* temp_indices, int64_t null_count) {
- if constexpr (has_null_like_values<ArrowType>()) {
- // Merge rows with a null or a null-like in the first sort key
- auto& comparator = comparator_;
- const auto& first_sort_key = sort_keys_[0];
-
- std::merge(nulls_begin, nulls_middle, nulls_middle, nulls_end,
temp_indices,
- [&](CompressedChunkLocation left, CompressedChunkLocation
right) {
- // First column is either null or nan
- const auto left_loc = ChunkLocation{left};
- const auto right_loc = ChunkLocation{right};
- const auto chunk_left = first_sort_key.GetChunk(left_loc);
- const auto chunk_right = first_sort_key.GetChunk(right_loc);
- const auto left_is_null = chunk_left.IsNull();
- const auto right_is_null = chunk_right.IsNull();
- if (left_is_null == right_is_null) {
- return comparator.Compare(left_loc, right_loc, 1);
- } else if (first_sort_key.null_placement ==
NullPlacement::AtEnd) {
- return right_is_null;
- } else {
- return left_is_null;
- }
- });
- // Copy back temp area into main buffer
- std::copy(temp_indices, temp_indices + (nulls_end - nulls_begin),
nulls_begin);
- } else {
- MergeNullsOnly(nulls_begin, nulls_middle, nulls_end, temp_indices,
null_count);
- }
- }
-
- void MergeNullsOnly(CompressedChunkLocation* nulls_begin,
- CompressedChunkLocation* nulls_middle,
- CompressedChunkLocation* nulls_end,
- CompressedChunkLocation* temp_indices, int64_t
null_count) {
+ void MergeRemainingSortKeys(std::span<CompressedChunkLocation> range, size_t
middle,
+ CompressedChunkLocation* temp_indices) {
// Untyped implementation
auto& comparator = comparator_;
- std::merge(nulls_begin, nulls_middle, nulls_middle, nulls_end,
temp_indices,
+ auto left = range.subspan(0, middle);
+ auto right = range.subspan(middle);
+ std::merge(left.data(), left.data() + left.size(), right.data(),
+ right.data() + right.size(), temp_indices,
[&](CompressedChunkLocation left, CompressedChunkLocation
right) {
- // First column is always null
+ // All rows are equal on first key so just compare by
remaining
+ // keys
return comparator.Compare(ChunkLocation{left},
ChunkLocation{right}, 1);
});
// Copy back temp area into main buffer
- std::copy(temp_indices, temp_indices + (nulls_end - nulls_begin),
nulls_begin);
+ std::copy(temp_indices, temp_indices + range.size(), range.data());
}
//
@@ -827,12 +773,15 @@ class TableSorter {
//
template <typename ArrowType>
enable_if_t<!is_null_type<ArrowType>::value> MergeNonNulls(
- CompressedChunkLocation* range_begin, CompressedChunkLocation*
range_middle,
- CompressedChunkLocation* range_end, CompressedChunkLocation*
temp_indices) {
+ std::span<CompressedChunkLocation> range, size_t middle,
+ CompressedChunkLocation* temp_indices) {
auto& comparator = comparator_;
const auto& first_sort_key = sort_keys_[0];
- std::merge(range_begin, range_middle, range_middle, range_end,
temp_indices,
+ auto left = range.subspan(0, middle);
+ auto right = range.subspan(middle);
+ std::merge(left.data(), left.data() + left.size(), right.data(),
+ right.data() + right.size(), temp_indices,
[&](CompressedChunkLocation left, CompressedChunkLocation
right) {
// Both values are never null nor NaN.
const auto left_loc = ChunkLocation{left};
@@ -859,16 +808,14 @@ class TableSorter {
});
// Copy back temp area into main buffer
- std::copy(temp_indices, temp_indices + (range_end - range_begin),
range_begin);
+ std::copy(temp_indices, temp_indices + range.size(), range.data());
}
template <typename ArrowType>
- enable_if_null<ArrowType> MergeNonNulls(CompressedChunkLocation* range_begin,
- CompressedChunkLocation*
range_middle,
- CompressedChunkLocation* range_end,
+ enable_if_null<ArrowType> MergeNonNulls(std::span<CompressedChunkLocation>
range,
+ size_t middle,
CompressedChunkLocation*
temp_indices) {
- const int64_t null_count = range_end - range_begin;
- MergeNullsOnly(range_begin, range_middle, range_end, temp_indices,
null_count);
+ MergeRemainingSortKeys(range, middle, temp_indices);
}
Status status_;
@@ -877,8 +824,7 @@ class TableSorter {
const RecordBatchVector batches_;
const SortOptions& options_;
const std::vector<ResolvedSortKey> sort_keys_;
- uint64_t* indices_begin_;
- uint64_t* indices_end_;
+ std::span<uint64_t> indices_;
Comparator comparator_;
};
@@ -999,12 +945,11 @@ class SortIndicesMetaFunction : public MetaFunction {
ARROW_ASSIGN_OR_RAISE(buffers[1],
AllocateResizableBuffer(buffer_size,
ctx->memory_pool()));
auto out = std::make_shared<ArrayData>(out_type, length, buffers, 0);
- auto out_begin = out->GetMutableValues<uint64_t>(1);
- auto out_end = out_begin + length;
- std::iota(out_begin, out_end, 0);
+ std::span<uint64_t> out_span{out->GetMutableValues<uint64_t>(1),
+ static_cast<size_t>(length)};
+ std::iota(out_span.begin(), out_span.end(), 0);
- RETURN_NOT_OK(
- SortChunkedArray(ctx, out_begin, out_end, chunked_array, order,
null_placement));
+ RETURN_NOT_OK(SortChunkedArray(ctx, out_span, chunked_array, order,
null_placement));
return Datum(out);
}
@@ -1029,15 +974,15 @@ class SortIndicesMetaFunction : public MetaFunction {
ARROW_ASSIGN_OR_RAISE(buffers[1],
AllocateResizableBuffer(buffer_size,
ctx->memory_pool()));
auto out = std::make_shared<ArrayData>(out_type, length, buffers, 0);
- auto out_begin = out->GetMutableValues<uint64_t>(1);
- auto out_end = out_begin + length;
- std::iota(out_begin, out_end, 0);
+ std::span<uint64_t> out_span{out->GetMutableValues<uint64_t>(1),
+ static_cast<size_t>(length)};
+ std::iota(out_span.begin(), out_span.end(), 0);
if (n_sort_keys <= kMaxRadixSortKeys) {
- RadixRecordBatchSorter sorter(out_begin, out_end, std::move(sort_keys));
+ RadixRecordBatchSorter sorter(out_span, std::move(sort_keys));
ARROW_RETURN_NOT_OK(sorter.Sort());
} else {
- MultipleKeyRecordBatchSorter sorter(out_begin, out_end,
std::move(sort_keys));
+ MultipleKeyRecordBatchSorter sorter(out_span, std::move(sort_keys));
ARROW_RETURN_NOT_OK(sorter.Sort());
}
return Datum(out);
@@ -1069,11 +1014,11 @@ class SortIndicesMetaFunction : public MetaFunction {
ARROW_ASSIGN_OR_RAISE(buffers[1],
AllocateResizableBuffer(buffer_size,
ctx->memory_pool()));
auto out = std::make_shared<ArrayData>(out_type, length, buffers, 0);
- auto out_begin = out->GetMutableValues<uint64_t>(1);
- auto out_end = out_begin + length;
- std::iota(out_begin, out_end, 0);
+ std::span<uint64_t> out_span{out->GetMutableValues<uint64_t>(1),
+ static_cast<size_t>(length)};
+ std::iota(out_span.begin(), out_span.end(), 0);
- TableSorter sorter(ctx, out_begin, out_end, table, options);
+ TableSorter sorter(ctx, out_span, table, options);
RETURN_NOT_OK(sorter.Sort());
return Datum(out);
@@ -1145,33 +1090,31 @@ Result<std::vector<SortField>> FindSortKeys(const
Schema& schema,
return SortFieldPopulator{}.FindSortKeys(schema, sort_keys);
}
-Result<NullPartitionResult> SortChunkedArray(ExecContext* ctx, uint64_t*
indices_begin,
- uint64_t* indices_end,
- const ChunkedArray& chunked_array,
- SortOrder sort_order,
- NullPlacement null_placement) {
+Result<NullLikePartition> SortChunkedArray(ExecContext* ctx,
std::span<uint64_t> indices,
+ const ChunkedArray& chunked_array,
+ SortOrder sort_order,
+ NullPlacement null_placement) {
auto physical_type = GetPhysicalType(chunked_array.type());
auto physical_chunks = GetPhysicalChunks(chunked_array, physical_type);
- return SortChunkedArray(ctx, indices_begin, indices_end, physical_type,
physical_chunks,
- sort_order, null_placement);
+ return SortChunkedArray(ctx, indices, physical_type, physical_chunks,
sort_order,
+ null_placement);
}
-Result<NullPartitionResult> SortChunkedArray(
- ExecContext* ctx, uint64_t* indices_begin, uint64_t* indices_end,
- const std::shared_ptr<DataType>& physical_type, const ArrayVector&
physical_chunks,
- SortOrder sort_order, NullPlacement null_placement) {
- NullPartitionResult output;
- ChunkedArraySorter sorter(ctx, indices_begin, indices_end, physical_type,
- physical_chunks, sort_order, null_placement,
&output);
+Result<NullLikePartition> SortChunkedArray(ExecContext* ctx,
std::span<uint64_t> indices,
+ const std::shared_ptr<DataType>&
physical_type,
+ const ArrayVector& physical_chunks,
+ SortOrder sort_order,
+ NullPlacement null_placement) {
+ NullLikePartition output;
+ ChunkedArraySorter sorter(ctx, indices, physical_type, physical_chunks,
sort_order,
+ null_placement, &output);
RETURN_NOT_OK(sorter.Sort());
return output;
}
-Result<NullPartitionResult> SortStructArray(ExecContext* ctx, uint64_t*
indices_begin,
- uint64_t* indices_end,
- const StructArray& array,
- SortOrder sort_order,
- NullPlacement null_placement) {
+Result<NullLikePartition> SortStructArray(ExecContext* ctx,
std::span<uint64_t> indices,
+ const StructArray& array, SortOrder
sort_order,
+ NullPlacement null_placement) {
ARROW_ASSIGN_OR_RAISE(auto columns, array.Flatten());
auto batch = RecordBatch::Make(schema(array.type()->fields()),
array.length(),
std::move(columns));
@@ -1185,10 +1128,10 @@ Result<NullPartitionResult>
SortStructArray(ExecContext* ctx, uint64_t* indices_
ARROW_ASSIGN_OR_RAISE(auto sort_keys,
ResolveRecordBatchSortKeys(*batch,
options.GetSortKeys()));
if (sort_keys.size() <= kMaxRadixSortKeys) {
- RadixRecordBatchSorter sorter(indices_begin, indices_end,
std::move(sort_keys));
+ RadixRecordBatchSorter sorter(indices, std::move(sort_keys));
return sorter.Sort();
} else {
- MultipleKeyRecordBatchSorter sorter(indices_begin, indices_end,
std::move(sort_keys));
+ MultipleKeyRecordBatchSorter sorter(indices, std::move(sort_keys));
return sorter.Sort();
}
}
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_internal.h
b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
index 9d2ad650f1..38f4ab4899 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_internal.h
+++ b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
@@ -25,7 +25,9 @@
#include "arrow/array.h"
#include "arrow/compute/api_vector.h"
+#include "arrow/compute/kernel.h"
#include "arrow/compute/kernels/chunked_internal.h"
+#include "arrow/compute/ordering.h"
#include "arrow/table.h"
#include "arrow/type.h"
#include "arrow/type_traits.h"
@@ -54,19 +56,22 @@ namespace arrow::compute::internal {
// NOTE: std::partition is usually faster than std::stable_partition.
struct NonStablePartitioner {
- template <typename Predicate, typename IndexType>
- IndexType* operator()(IndexType* indices_begin, IndexType* indices_end,
- Predicate&& pred) {
- return std::partition(indices_begin, indices_end,
std::forward<Predicate>(pred));
+ // this follows std::ranges::partition semantics and returns the right-hand
partition
+ template <typename Predicate>
+ auto operator()(std::span<uint64_t> indices, Predicate&& pred) {
+ auto middle = std::partition(indices.data(), indices.data() +
indices.size(),
+ std::forward<Predicate>(pred));
+ return std::span<uint64_t>{middle, indices.data() + indices.size()};
}
};
struct StablePartitioner {
- template <typename Predicate, typename IndexType>
- IndexType* operator()(IndexType* indices_begin, IndexType* indices_end,
- Predicate&& pred) {
- return std::stable_partition(indices_begin, indices_end,
- std::forward<Predicate>(pred));
+ // this follows std::ranges::partition semantics and returns the right-hand
partition
+ template <typename Predicate>
+ auto operator()(std::span<uint64_t> indices, Predicate&& pred) {
+ auto middle = std::stable_partition(indices.data(), indices.data() +
indices.size(),
+ std::forward<Predicate>(pred));
+ return std::span<uint64_t>{middle, indices.data() + indices.size()};
}
};
@@ -106,382 +111,341 @@ int CompareTypeValues(Value&& left, Value&& right,
SortOrder order,
}
template <typename IndexType>
-struct GenericNullPartitionResult {
- IndexType* non_nulls_begin;
- IndexType* non_nulls_end;
- IndexType* nulls_begin;
- IndexType* nulls_end;
-
- IndexType* overall_begin() const { return std::min(nulls_begin,
non_nulls_begin); }
+struct GenericNullLikePartition {
+ std::span<IndexType> non_null_like_range;
+ std::span<IndexType> nan_range;
+ std::span<IndexType> null_range;
- IndexType* overall_end() const { return std::max(nulls_end, non_nulls_end); }
+ IndexType* non_null_like_begin() const { return non_null_like_range.data(); }
+ IndexType* non_null_like_end() const {
+ return non_null_like_range.data() + non_null_like_range.size();
+ }
+ IndexType* nan_begin() const { return nan_range.data(); }
+ IndexType* nan_end() const { return nan_range.data() + nan_range.size(); }
+ IndexType* null_begin() const { return null_range.data(); }
+ IndexType* null_end() const { return null_range.data() + null_range.size(); }
- int64_t non_null_count() const { return non_nulls_end - non_nulls_begin; }
+ IndexType* overall_begin() const {
+ // nans are always in the middle
+ return std::min(non_null_like_begin(), null_begin());
+ }
- int64_t null_count() const { return nulls_end - nulls_begin; }
+ IndexType* overall_end() const {
+ // nans are always in the middle
+ return std::max(non_null_like_end(), null_end());
+ }
- static GenericNullPartitionResult NoNulls(IndexType* indices_begin,
- IndexType* indices_end,
- NullPlacement null_placement) {
- if (null_placement == NullPlacement::AtStart) {
- return {indices_begin, indices_end, indices_begin, indices_begin};
+ // Note that "_begin" is not actually the begin of the stored ranges, but
can be much
+ // smaller. I.e. this function can be and is used when the Partition object
is pointing
+ // into a larger buffer and translate it into another larger buffer at the
same offset
+ template <typename TargetIndexType>
+ GenericNullLikePartition<TargetIndexType> TranslateTo(
+ IndexType* indices_begin, TargetIndexType* target_indices_begin) const {
+ size_t non_null_offset = non_null_like_range.data() - indices_begin;
+ size_t nan_offset = nan_range.data() - indices_begin;
+ size_t null_offset = null_range.data() - indices_begin;
+ return {.non_null_like_range = {target_indices_begin + non_null_offset,
+ non_null_like_range.size()},
+ .nan_range = {target_indices_begin + nan_offset, nan_range.size()},
+ .null_range = {target_indices_begin + null_offset,
null_range.size()}};
+ }
+
+ static GenericNullLikePartition FromCounts(std::span<IndexType> indices,
+ int64_t non_null_like_count,
+ int64_t nan_count, int64_t
null_count,
+ NullPlacement null_placement) {
+ GenericNullLikePartition p;
+ ARROW_DCHECK_EQ(non_null_like_count + nan_count + null_count,
+ static_cast<int64_t>(indices.size()));
+ if (null_placement == NullPlacement::AtEnd) {
+ p.non_null_like_range = indices.subspan(0, non_null_like_count);
+ p.nan_range = indices.subspan(non_null_like_count, nan_count);
+ p.null_range = indices.subspan(non_null_like_count + nan_count,
null_count);
} else {
- return {indices_begin, indices_end, indices_end, indices_end};
+ p.null_range = indices.subspan(0, null_count);
+ p.nan_range = indices.subspan(null_count, nan_count);
+ p.non_null_like_range =
+ indices.subspan(null_count + nan_count, non_null_like_count);
}
+ return p;
}
+};
+
+using NullLikePartition = GenericNullLikePartition<uint64_t>;
+using ChunkedNullLikePartition =
GenericNullLikePartition<CompressedChunkLocation>;
+
+struct NullPartition {
+ std::span<uint64_t> non_nulls;
+ std::span<uint64_t> nulls;
- static GenericNullPartitionResult NullsOnly(IndexType* indices_begin,
- IndexType* indices_end,
- NullPlacement null_placement) {
+ static NullPartition NoNulls(std::span<uint64_t> indices,
+ NullPlacement null_placement) {
if (null_placement == NullPlacement::AtStart) {
- return {indices_end, indices_end, indices_begin, indices_end};
+ return {.non_nulls = indices, .nulls = indices.subspan(0, 0)};
} else {
- return {indices_begin, indices_begin, indices_begin, indices_end};
+ return {.non_nulls = indices, .nulls = indices.subspan(indices.size(),
0)};
}
}
- static GenericNullPartitionResult NullsAtEnd(IndexType* indices_begin,
- IndexType* indices_end,
- IndexType* midpoint) {
- ARROW_DCHECK_GE(midpoint, indices_begin);
- ARROW_DCHECK_LE(midpoint, indices_end);
- return {indices_begin, midpoint, midpoint, indices_end};
+ // Note that we always pass the tail-partition second.
+ // In this function _nulls_ are passed as a second argument.
+ static NullPartition NullsAtEnd(std::span<uint64_t> indices,
+ std::span<uint64_t> null_tail) {
+ ARROW_DCHECK_GE(null_tail.data(), indices.data());
+ ARROW_DCHECK_EQ(null_tail.data() + null_tail.size(), indices.data() +
indices.size());
+ return {.non_nulls = {indices.data(), null_tail.data()}, .nulls =
null_tail};
}
- static GenericNullPartitionResult NullsAtStart(IndexType* indices_begin,
- IndexType* indices_end,
- IndexType* midpoint) {
- ARROW_DCHECK_GE(midpoint, indices_begin);
- ARROW_DCHECK_LE(midpoint, indices_end);
- return {midpoint, indices_end, indices_begin, midpoint};
- }
-
- template <typename TargetIndexType>
- GenericNullPartitionResult<TargetIndexType> TranslateTo(
- IndexType* indices_begin, TargetIndexType* target_indices_begin) const {
- return {
- (non_nulls_begin - indices_begin) + target_indices_begin,
- (non_nulls_end - indices_begin) + target_indices_begin,
- (nulls_begin - indices_begin) + target_indices_begin,
- (nulls_end - indices_begin) + target_indices_begin,
- };
+ // Note that we always pass the tail-partition second.
+ // In this function the _non-nulls_ are passed as a second argument.
+ static NullPartition NullsAtStart(std::span<uint64_t> indices,
+ std::span<uint64_t> non_null_tail) {
+ ARROW_DCHECK_GE(non_null_tail.data(), indices.data());
+ ARROW_DCHECK_EQ(non_null_tail.data() + non_null_tail.size(),
+ indices.data() + indices.size());
+ return {.non_nulls = non_null_tail, .nulls = {indices.data(),
non_null_tail.data()}};
}
};
-using NullPartitionResult = GenericNullPartitionResult<uint64_t>;
-using ChunkedNullPartitionResult =
GenericNullPartitionResult<CompressedChunkLocation>;
-
// Move nulls (not null-like values) to end of array.
//
// `offset` is used when this is called on a chunk of a chunked array
template <typename Partitioner>
-NullPartitionResult PartitionNullsOnly(uint64_t* indices_begin, uint64_t*
indices_end,
- const Array& values, int64_t offset,
- NullPlacement null_placement) {
+NullPartition PartitionNullsOnly(std::span<uint64_t> indices, const Array&
values,
+ int64_t offset, NullPlacement null_placement)
{
if (values.null_count() == 0) {
- return NullPartitionResult::NoNulls(indices_begin, indices_end,
null_placement);
+ return NullPartition::NoNulls(indices, null_placement);
}
Partitioner partitioner;
if (null_placement == NullPlacement::AtStart) {
- auto nulls_end = partitioner(
- indices_begin, indices_end,
- [&values, &offset](uint64_t ind) { return values.IsNull(ind - offset);
});
- return NullPartitionResult::NullsAtStart(indices_begin, indices_end,
nulls_end);
+ auto non_null_tail = partitioner(indices, [&values, &offset](uint64_t ind)
{
+ return values.IsNull(static_cast<int64_t>(ind) - offset);
+ });
+ return NullPartition::NullsAtStart(indices, non_null_tail);
} else {
- auto nulls_begin = partitioner(
- indices_begin, indices_end,
- [&values, &offset](uint64_t ind) { return !values.IsNull(ind -
offset); });
- return NullPartitionResult::NullsAtEnd(indices_begin, indices_end,
nulls_begin);
+ auto null_tail = partitioner(indices, [&values, &offset](uint64_t ind) {
+ return !values.IsNull(static_cast<int64_t>(ind) - offset);
+ });
+ return NullPartition::NullsAtEnd(indices, null_tail);
}
}
+struct NanPartition {
+ std::span<uint64_t> non_null_like_range;
+ std::span<uint64_t> nan_range;
+};
+
// Move non-null null-like values to end of array.
//
// `offset` is used when this is called on a chunk of a chunked array
template <typename ArrayType, typename Partitioner>
-NullPartitionResult PartitionNullLikes(uint64_t* indices_begin, uint64_t*
indices_end,
- const ArrayType& values, int64_t offset,
- NullPlacement null_placement) {
+NanPartition PartitionNans(std::span<uint64_t> indices, const ArrayType&
values,
+ int64_t offset, NullPlacement null_placement) {
if constexpr (has_null_like_values<typename ArrayType::TypeClass>()) {
Partitioner partitioner;
if (null_placement == NullPlacement::AtStart) {
- auto null_likes_end =
- partitioner(indices_begin, indices_end, [&values, &offset](uint64_t
ind) {
- return std::isnan(values.GetView(ind - offset));
- });
- return NullPartitionResult::NullsAtStart(indices_begin, indices_end,
- null_likes_end);
+ auto non_null_like_tail = partitioner(indices, [&values,
&offset](uint64_t ind) {
+ return std::isnan(values.GetView(static_cast<int64_t>(ind) - offset));
+ });
+ return NanPartition{.non_null_like_range = non_null_like_tail,
+ .nan_range = {indices.data(),
non_null_like_tail.data()}};
} else {
- auto null_likes_begin =
- partitioner(indices_begin, indices_end, [&values, &offset](uint64_t
ind) {
- return !std::isnan(values.GetView(ind - offset));
- });
- return NullPartitionResult::NullsAtEnd(indices_begin, indices_end,
- null_likes_begin);
+ auto nan_tail = partitioner(indices, [&values, &offset](uint64_t ind) {
+ return !std::isnan(values.GetView(static_cast<int64_t>(ind) - offset));
+ });
+ return NanPartition{.non_null_like_range = {indices.data(),
nan_tail.data()},
+ .nan_range = nan_tail};
}
} else {
- return NullPartitionResult::NoNulls(indices_begin, indices_end,
null_placement);
- }
-}
-
-// Move nulls to end of array.
-//
-// `offset` is used when this is called on a chunk of a chunked array
-template <typename ArrayType, typename Partitioner>
-NullPartitionResult PartitionNulls(uint64_t* indices_begin, uint64_t*
indices_end,
- const ArrayType& values, int64_t offset,
- NullPlacement null_placement) {
- // Partition nulls at start (resp. end), and null-like values just before
(resp. after)
- NullPartitionResult p = PartitionNullsOnly<Partitioner>(indices_begin,
indices_end,
- values, offset,
null_placement);
- NullPartitionResult q = PartitionNullLikes<ArrayType, Partitioner>(
- p.non_nulls_begin, p.non_nulls_end, values, offset, null_placement);
- return NullPartitionResult{q.non_nulls_begin, q.non_nulls_end,
- std::min(q.nulls_begin, p.nulls_begin),
- std::max(q.nulls_end, p.nulls_end)};
-}
-
-//
-// Null partitioning on chunked arrays, in two flavors:
-// 1) with uint64_t indices and ChunkedArrayResolver
-// 2) with CompressedChunkLocation and span of chunks
-//
-
-template <typename Partitioner>
-NullPartitionResult PartitionNullsOnly(uint64_t* indices_begin, uint64_t*
indices_end,
- const ChunkedArrayResolver& resolver,
- int64_t null_count, NullPlacement
null_placement) {
- if (null_count == 0) {
- return NullPartitionResult::NoNulls(indices_begin, indices_end,
null_placement);
- }
- Partitioner partitioner;
- if (null_placement == NullPlacement::AtStart) {
- auto nulls_end = partitioner(indices_begin, indices_end, [&](uint64_t ind)
{
- const auto chunk = resolver.Resolve(ind);
- return chunk.IsNull();
- });
- return NullPartitionResult::NullsAtStart(indices_begin, indices_end,
nulls_end);
- } else {
- auto nulls_begin = partitioner(indices_begin, indices_end, [&](uint64_t
ind) {
- const auto chunk = resolver.Resolve(ind);
- return !chunk.IsNull();
- });
- return NullPartitionResult::NullsAtEnd(indices_begin, indices_end,
nulls_begin);
- }
-}
-
-template <typename Partitioner>
-ChunkedNullPartitionResult PartitionNullsOnly(CompressedChunkLocation*
locations_begin,
- CompressedChunkLocation*
locations_end,
- std::span<const Array* const>
chunks,
- int64_t null_count,
- NullPlacement null_placement) {
- if (null_count == 0) {
- return ChunkedNullPartitionResult::NoNulls(locations_begin, locations_end,
- null_placement);
- }
- Partitioner partitioner;
- if (null_placement == NullPlacement::AtStart) {
- auto nulls_end =
- partitioner(locations_begin, locations_end,
[&](CompressedChunkLocation loc) {
- return chunks[loc.chunk_index()]->IsNull(
- static_cast<int64_t>(loc.index_in_chunk()));
- });
- return ChunkedNullPartitionResult::NullsAtStart(locations_begin,
locations_end,
- nulls_end);
- } else {
- auto nulls_begin =
- partitioner(locations_begin, locations_end,
[&](CompressedChunkLocation loc) {
- return !chunks[loc.chunk_index()]->IsNull(
- static_cast<int64_t>(loc.index_in_chunk()));
- });
- return ChunkedNullPartitionResult::NullsAtEnd(locations_begin,
locations_end,
- nulls_begin);
- }
-}
-
-template <typename ArrayType, typename Partitioner,
- typename TypeClass = typename ArrayType::TypeClass>
-NullPartitionResult PartitionNullLikes(uint64_t* indices_begin, uint64_t*
indices_end,
- const ChunkedArrayResolver& resolver,
- NullPlacement null_placement) {
- if constexpr (has_null_like_values<typename ArrayType::TypeClass>()) {
- Partitioner partitioner;
if (null_placement == NullPlacement::AtStart) {
- auto null_likes_end = partitioner(indices_begin, indices_end,
[&](uint64_t ind) {
- const auto chunk = resolver.Resolve(ind);
- return std::isnan(chunk.Value<TypeClass>());
- });
- return NullPartitionResult::NullsAtStart(indices_begin, indices_end,
- null_likes_end);
+ return NanPartition{.non_null_like_range = indices,
+ .nan_range = {indices.data(), indices.data()}};
} else {
- auto null_likes_begin = partitioner(indices_begin, indices_end,
[&](uint64_t ind) {
- const auto chunk = resolver.Resolve(ind);
- return !std::isnan(chunk.Value<TypeClass>());
- });
- return NullPartitionResult::NullsAtEnd(indices_begin, indices_end,
- null_likes_begin);
+ return NanPartition{.non_null_like_range = indices,
+ .nan_range = indices.subspan(indices.size())};
}
- } else {
- return NullPartitionResult::NoNulls(indices_begin, indices_end,
null_placement);
}
}
template <typename ArrayType, typename Partitioner>
-NullPartitionResult PartitionNulls(uint64_t* indices_begin, uint64_t*
indices_end,
- const ChunkedArrayResolver& resolver,
- int64_t null_count, NullPlacement
null_placement) {
+NullLikePartition PartitionNullsAndNans(std::span<uint64_t> indices,
+ const ArrayType& values, int64_t
offset,
+ NullPlacement null_placement) {
// Partition nulls at start (resp. end), and null-like values just before
(resp. after)
- NullPartitionResult p = PartitionNullsOnly<Partitioner>(
- indices_begin, indices_end, resolver, null_count, null_placement);
- NullPartitionResult q = PartitionNullLikes<ArrayType, Partitioner>(
- p.non_nulls_begin, p.non_nulls_end, resolver, null_placement);
- return NullPartitionResult{q.non_nulls_begin, q.non_nulls_end,
- std::min(q.nulls_begin, p.nulls_begin),
- std::max(q.nulls_end, p.nulls_end)};
+ NullPartition p =
+ PartitionNullsOnly<Partitioner>(indices, values, offset, null_placement);
+ auto q =
+ PartitionNans<ArrayType, Partitioner>(p.non_nulls, values, offset,
null_placement);
+ return NullLikePartition{.non_null_like_range = q.non_null_like_range,
+ .nan_range = q.nan_range,
+ .null_range = p.nulls};
}
-template <typename IndexType, typename NullPartitionResultType>
-struct GenericMergeImpl {
- using MergeNullsFunc = std::function<void(IndexType* nulls_begin,
- IndexType* nulls_middle,
IndexType* nulls_end,
- IndexType* temp_indices, int64_t
null_count)>;
+template <typename ArrayType, typename Partitioner>
+NullLikePartition PartitionNansOnly(std::span<uint64_t> indices, const
ArrayType& values,
+ int64_t offset, NullPlacement
null_placement) {
+ // Partition nulls at start (resp. end), and null-like values just before
(resp. after)
+ NullPartition p = NullPartition::NoNulls(indices, null_placement);
+ auto q =
+ PartitionNans<ArrayType, Partitioner>(p.non_nulls, values, offset,
null_placement);
+ return NullLikePartition{.non_null_like_range = q.non_null_like_range,
+ .nan_range = q.nan_range,
+ .null_range = p.nulls};
+}
- using MergeNonNullsFunc =
- std::function<void(IndexType* range_begin, IndexType* range_middle,
- IndexType* range_end, IndexType* temp_indices)>;
+struct ChunkedMergeImpl {
+ using MergeFunc =
+ std::function<void(std::span<CompressedChunkLocation> merge_range,
size_t middle,
+ CompressedChunkLocation* temp_indices)>;
- GenericMergeImpl(NullPlacement null_placement, MergeNullsFunc&& merge_nulls,
- MergeNonNullsFunc&& merge_non_nulls)
+ ChunkedMergeImpl(NullPlacement null_placement, MergeFunc&&
merge_remaining_sort_keys,
+ MergeFunc&& merge_non_nulls)
: null_placement_(null_placement),
- merge_nulls_(std::move(merge_nulls)),
+ merge_remaining_sort_keys_(std::move(merge_remaining_sort_keys)),
merge_non_nulls_(std::move(merge_non_nulls)) {}
Status Init(ExecContext* ctx, int64_t temp_indices_length) {
- ARROW_ASSIGN_OR_RAISE(
- temp_buffer_,
- AllocateBuffer(sizeof(IndexType) * temp_indices_length,
ctx->memory_pool()));
- temp_indices_ = reinterpret_cast<IndexType*>(temp_buffer_->mutable_data());
+ ARROW_ASSIGN_OR_RAISE(temp_buffer_,
AllocateBuffer(sizeof(CompressedChunkLocation) *
+ temp_indices_length,
+ ctx->memory_pool()));
+ temp_indices_ =
+
reinterpret_cast<CompressedChunkLocation*>(temp_buffer_->mutable_data());
return Status::OK();
}
- NullPartitionResultType Merge(const NullPartitionResultType& left,
- const NullPartitionResultType& right,
- int64_t null_count) const {
+ ChunkedNullLikePartition Merge(const ChunkedNullLikePartition& left,
+ const ChunkedNullLikePartition& right) const {
if (null_placement_ == NullPlacement::AtStart) {
- return MergeNullsAtStart(left, right, null_count);
+ return MergeNullsAtStart(left, right);
} else {
- return MergeNullsAtEnd(left, right, null_count);
+ return MergeNullsAtEnd(left, right);
}
}
- NullPartitionResultType MergeNullsAtStart(const NullPartitionResultType&
left,
- const NullPartitionResultType&
right,
- int64_t null_count) const {
+ ChunkedNullLikePartition MergeNullsAtStart(
+ const ChunkedNullLikePartition& left, const ChunkedNullLikePartition&
right) const {
// Input layout:
- // [left nulls .... left non-nulls .... right nulls .... right non-nulls]
- ARROW_DCHECK_EQ(left.nulls_end, left.non_nulls_begin);
- ARROW_DCHECK_EQ(left.non_nulls_end, right.nulls_begin);
- ARROW_DCHECK_EQ(right.nulls_end, right.non_nulls_begin);
-
- // Mutate the input, stably, to obtain the following layout:
- // [left nulls .... right nulls .... left non-nulls .... right non-nulls]
- std::rotate(left.non_nulls_begin, right.nulls_begin, right.nulls_end);
-
- const auto p = NullPartitionResultType::NullsAtStart(
- left.nulls_begin, right.non_nulls_end,
- left.nulls_begin + left.null_count() + right.null_count());
-
- // If the type has null-like values (such as NaN), ensure those plus
regular
- // nulls are partitioned in the right order. Note this assumes that all
- // null-like values (e.g. NaN) are ordered equally.
- if (p.null_count()) {
- merge_nulls_(p.nulls_begin, p.nulls_begin + left.null_count(),
p.nulls_end,
- temp_indices_, null_count);
- }
-
- // Merge the non-null values into temp area
- ARROW_DCHECK_EQ(right.non_nulls_begin - p.non_nulls_begin,
left.non_null_count());
- ARROW_DCHECK_EQ(p.non_nulls_end - right.non_nulls_begin,
right.non_null_count());
- if (p.non_null_count()) {
- merge_non_nulls_(p.non_nulls_begin, right.non_nulls_begin,
p.non_nulls_end,
+ // [left nul .. left nan .. left non-nul .. right nul .. right nan ..
right non-nul]
+ ARROW_DCHECK_EQ(left.null_end(), left.nan_begin());
+ ARROW_DCHECK_EQ(left.nan_end(), left.non_null_like_begin());
+ ARROW_DCHECK_EQ(left.non_null_like_end(), right.null_begin());
+ ARROW_DCHECK_EQ(right.null_end(), right.nan_begin());
+ ARROW_DCHECK_EQ(right.nan_end(), right.non_null_like_begin());
+
+ // Mutate the input, stably in two steps, to obtain the following layouts:
+ // [left nul .. left nan .. left non-nul .. right nul .. right nan ..
right non-nul]
+ // ┬────────────┴┴──────────────────────┬
+ // [left nul .. left nan .. right nul .. right nan .. left non-nul ..
right non-nul]
+ std::rotate(left.non_null_like_begin(), right.null_begin(),
right.nan_end());
+
+ // only use sizes of ranges that are at a different position now
+ // [left nul .. left nan .. right nul .. right nan .. left non-nul ..
right non-nul]
+ // ┬────────┴┴─────────┬
+ // [left nul .. right nul .. left nan .. right nan .. left non-nul ..
right non-nul]
+ // this is a no-op if no nan values are present
+ std::rotate(left.nan_begin(), left.nan_end(),
+ left.nan_end() + right.null_range.size());
+
+ std::span<CompressedChunkLocation> full_span{left.overall_begin(),
+ right.overall_end()};
+ const auto p = ChunkedNullLikePartition::FromCounts(
+ full_span, left.non_null_like_range.size() +
right.non_null_like_range.size(),
+ left.nan_range.size() + right.nan_range.size(),
+ left.null_range.size() + right.null_range.size(),
NullPlacement::AtStart);
+
+ if (!p.non_null_like_range.empty()) {
+ merge_non_nulls_(p.non_null_like_range, left.non_null_like_range.size(),
temp_indices_);
}
+ if (!p.nan_range.empty()) {
+ merge_remaining_sort_keys_(p.nan_range, left.nan_range.size(),
temp_indices_);
+ }
+ if (!p.null_range.empty()) {
+ merge_remaining_sort_keys_(p.null_range, left.null_range.size(),
temp_indices_);
+ }
return p;
}
- NullPartitionResultType MergeNullsAtEnd(const NullPartitionResultType& left,
- const NullPartitionResultType& right,
- int64_t null_count) const {
+ ChunkedNullLikePartition MergeNullsAtEnd(const ChunkedNullLikePartition&
left,
+ const ChunkedNullLikePartition&
right) const {
// Input layout:
- // [left non-nulls .... left nulls .... right non-nulls .... right nulls]
- ARROW_DCHECK_EQ(left.non_nulls_end, left.nulls_begin);
- ARROW_DCHECK_EQ(left.nulls_end, right.non_nulls_begin);
- ARROW_DCHECK_EQ(right.non_nulls_end, right.nulls_begin);
-
- // Mutate the input, stably, to obtain the following layout:
- // [left non-nulls .... right non-nulls .... left nulls .... right nulls]
- std::rotate(left.nulls_begin, right.non_nulls_begin, right.non_nulls_end);
-
- const auto p = NullPartitionResultType::NullsAtEnd(
- left.non_nulls_begin, right.nulls_end,
- left.non_nulls_begin + left.non_null_count() + right.non_null_count());
-
- // If the type has null-like values (such as NaN), ensure those plus
regular
- // nulls are partitioned in the right order. Note this assumes that all
- // null-like values (e.g. NaN) are ordered equally.
- if (p.null_count()) {
- merge_nulls_(p.nulls_begin, p.nulls_begin + left.null_count(),
p.nulls_end,
- temp_indices_, null_count);
- }
+ // [left non-nul .. left nan .. left nul .. right non-nul .. right nan ..
right nul]
+ ARROW_DCHECK_EQ(left.non_null_like_end(), left.nan_begin());
+ ARROW_DCHECK_EQ(left.nan_end(), left.null_begin());
+ ARROW_DCHECK_EQ(left.null_end(), right.non_null_like_begin());
+ ARROW_DCHECK_EQ(right.non_null_like_end(), right.nan_begin());
+ ARROW_DCHECK_EQ(right.nan_end(), right.null_begin());
+
+ // Mutate the input, stably in two steps, to obtain the following layouts:
+ // [left non-nul .. left nan .. left nul .. right non-nul .. right nan ..
right nul]
+ // ┬────────────────────┴┴─────────────┬
+ // [left non-nul .. right non-nul .. left nan .. left nul .. right nan ..
right nul]
+ std::rotate(left.nan_begin(), right.non_null_like_begin(),
right.non_null_like_end());
+
+ // only use sizes of ranges that are at a different position now
+ // [left non-nul .. right non-nul .. left nan .. left nul .. right nan ..
right nul]
+ // ┬────────┴┴─────────┬
+ // [left non-nul .. right non-nul .. left nan .. right nan .. left nul ..
right nul]
+ // this is a no-op if no nan values are present
+ auto new_left_null_range_begin = left.non_null_like_end() +
+ right.non_null_like_range.size() +
+ left.nan_range.size();
+ std::rotate(
+ new_left_null_range_begin, new_left_null_range_begin +
left.null_range.size(),
+ new_left_null_range_begin + left.null_range.size() +
right.nan_range.size());
+
+ std::span<CompressedChunkLocation> full_span{left.overall_begin(),
+ right.overall_end()};
+ const auto p = ChunkedNullLikePartition::FromCounts(
+ full_span, left.non_null_like_range.size() +
right.non_null_like_range.size(),
+ left.nan_range.size() + right.nan_range.size(),
+ left.null_range.size() + right.null_range.size(),
NullPlacement::AtEnd);
// Merge the non-null values into temp area
- ARROW_DCHECK_EQ(left.non_nulls_end - p.non_nulls_begin,
left.non_null_count());
- ARROW_DCHECK_EQ(p.non_nulls_end - left.non_nulls_end,
right.non_null_count());
- if (p.non_null_count()) {
- merge_non_nulls_(p.non_nulls_begin, left.non_nulls_end, p.non_nulls_end,
+ if (!p.non_null_like_range.empty()) {
+ merge_non_nulls_(p.non_null_like_range, left.non_null_like_range.size(),
temp_indices_);
}
+ if (!p.nan_range.empty()) {
+ merge_remaining_sort_keys_(p.nan_range, left.nan_range.size(),
temp_indices_);
+ }
+ if (!p.null_range.empty()) {
+ merge_remaining_sort_keys_(p.null_range, left.null_range.size(),
temp_indices_);
+ }
return p;
}
private:
NullPlacement null_placement_;
- MergeNullsFunc merge_nulls_;
- MergeNonNullsFunc merge_non_nulls_;
+ MergeFunc merge_remaining_sort_keys_;
+ MergeFunc merge_non_nulls_;
std::unique_ptr<Buffer> temp_buffer_;
- IndexType* temp_indices_ = nullptr;
+ CompressedChunkLocation* temp_indices_ = nullptr;
};
-using MergeImpl = GenericMergeImpl<uint64_t, NullPartitionResult>;
-using ChunkedMergeImpl =
- GenericMergeImpl<CompressedChunkLocation, ChunkedNullPartitionResult>;
-
// TODO make this usable if indices are non trivial on input
// (see ConcreteRecordBatchColumnSorter)
// `offset` is used when this is called on a chunk of a chunked array
-using ArraySortFunc = std::function<Result<NullPartitionResult>(
- uint64_t* indices_begin, uint64_t* indices_end, const Array& values,
int64_t offset,
+using ArraySortFunc = std::function<Result<NullLikePartition>(
+ std::span<uint64_t> indices, const Array& values, int64_t offset,
const ArraySortOptions& options, ExecContext* ctx)>;
Result<ArraySortFunc> GetArraySorter(const DataType& type);
-Result<NullPartitionResult> SortChunkedArray(ExecContext* ctx, uint64_t*
indices_begin,
- uint64_t* indices_end,
- const ChunkedArray& chunked_array,
- SortOrder sort_order,
- NullPlacement null_placement);
-
-Result<NullPartitionResult> SortChunkedArray(
- ExecContext* ctx, uint64_t* indices_begin, uint64_t* indices_end,
- const std::shared_ptr<DataType>& physical_type, const ArrayVector&
physical_chunks,
- SortOrder sort_order, NullPlacement null_placement);
-
-Result<NullPartitionResult> SortStructArray(ExecContext* ctx, uint64_t*
indices_begin,
- uint64_t* indices_end,
- const StructArray& array,
- SortOrder sort_order,
- NullPlacement null_placement);
+Result<NullLikePartition> SortChunkedArray(ExecContext* ctx,
std::span<uint64_t> indices,
+ const ChunkedArray& chunked_array,
+ SortOrder sort_order,
+ NullPlacement null_placement);
+
+Result<NullLikePartition> SortChunkedArray(ExecContext* ctx,
std::span<uint64_t> indices,
+ const std::shared_ptr<DataType>&
physical_type,
+ const ArrayVector& physical_chunks,
+ SortOrder sort_order,
+ NullPlacement null_placement);
+
+Result<NullLikePartition> SortStructArray(ExecContext* ctx,
std::span<uint64_t> indices,
+ const StructArray& array, SortOrder
sort_order,
+ NullPlacement null_placement);
// ----------------------------------------------------------------------
// Helpers for Sort/SelectK/Rank implementations
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc
b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
index cd31c85062..9bc30f2323 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
@@ -1858,6 +1858,33 @@ TEST_F(TestTableSortIndices, HeterogenousChunking) {
AssertSortIndices(table, options, "[3, 4, 2, 5, 1, 0, 6, 7]");
}
+TEST_F(TestTableSortIndices, NullPartitionAcrossChunks) {
+ auto schema = ::arrow::schema({
+ {field("a", utf8())},
+ {field("b", utf8())},
+ {field("c", utf8())},
+ {field("d", utf8())},
+ });
+
+ // Logical rows (a, b, c, d):
+ // 0: (null, null, null, "x")
+ // 1: (null, null, null, "y")
+ // 2: ("a", "p", null, null)
+ // 3: (null, "q", null, null)
+ // 4: ("b", "r", null, null)
+ // Sorted ascending on every key (nulls last on every key)
+ auto col_a = ChunkedArrayFromJSON(utf8(), {R"([null, null])", R"(["a", null,
"b"])"});
+ auto col_b = ChunkedArrayFromJSON(utf8(), {R"([null, null])", R"(["p", "q",
"r"])"});
+ auto col_c = ChunkedArrayFromJSON(utf8(), {R"([null, null])", R"([null,
null, null])"});
+ auto col_d = ChunkedArrayFromJSON(utf8(), {R"(["x", "y"])", R"([null, null,
null])"});
+ auto table = Table::Make(schema, {col_a, col_b, col_c, col_d});
+
+ SortOptions options(
+ {SortKey("a", SortOrder::Ascending), SortKey("b", SortOrder::Ascending),
+ SortKey("c", SortOrder::Ascending), SortKey("d",
SortOrder::Ascending)});
+ AssertSortIndices(table, options, "[2, 4, 3, 0, 1]");
+}
+
// Tests for temporal types
template <typename ArrowType>
class TestTableSortIndicesForTemporal : public TestTableSortIndices {
