pitrou commented on code in PR #50248:
URL: https://github.com/apache/arrow/pull/50248#discussion_r3629439929
##########
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:
A comment is ok, thank you!
--
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]
