pitrou commented on code in PR #50248:
URL: https://github.com/apache/arrow/pull/50248#discussion_r3623004530
##########
cpp/src/arrow/compute/kernels/vector_sort_internal.h:
##########
@@ -106,382 +109,337 @@ 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;
Review Comment:
Can we DCHECK that `indices_begin` falls somewhere in the overall range?
##########
cpp/src/arrow/compute/kernels/vector_sort_internal.h:
##########
@@ -54,19 +56,20 @@ 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));
+ template <typename Predicate>
Review Comment:
Add a comment that this returns the right-hand partition?
##########
cpp/src/arrow/compute/kernels/vector_sort_internal.h:
##########
@@ -106,382 +109,337 @@ 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};
+ 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,
- };
+ static NullPartition NullsAtStart(std::span<uint64_t> indices,
+ std::span<uint64_t> non_null_tail) {
Review Comment:
Note: it seems a bit weird to pass the null tail to `NullsAtEnd` and the
non-null tail to `NullsAtStart`. At least add a comment?
##########
cpp/src/arrow/compute/kernels/vector_sort_internal.h:
##########
@@ -106,382 +109,337 @@ 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};
+ 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,
- };
+ 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(ind - offset);
Review Comment:
`ind - offset` is mixed signed-unsigned arithmetic, which may error out on
some compilers.
##########
cpp/src/arrow/compute/kernels/vector_array_sort.cc:
##########
@@ -359,14 +349,14 @@ 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 {
int64_t index = offset;
Review Comment:
Can we DCHECK that the nan_range is empty?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
