Skip to content

Commit 345f5e1

Browse files
authored
apacheGH-33559: [C++] Move select-k implementation into separate module (apache#14753)
See: [ARROW-18395](https://issues.apache.org/jira/browse/ARROW-18395) Authored-by: benibus <bpharks@gmx.com> Signed-off-by: Antoine Pitrou <antoine@python.org>
1 parent df2aa38 commit 345f5e1

8 files changed

Lines changed: 1239 additions & 1162 deletions

File tree

‎cpp/src/arrow/CMakeLists.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -450,7 +450,9 @@ if(ARROW_COMPUTE)
450450
compute/kernels/vector_cumulative_ops.cc
451451
compute/kernels/vector_hash.cc
452452
compute/kernels/vector_nested.cc
453+
compute/kernels/vector_rank.cc
453454
compute/kernels/vector_replace.cc
455+
compute/kernels/vector_select_k.cc
454456
compute/kernels/vector_selection.cc
455457
compute/kernels/vector_sort.cc
456458
compute/row/encode_internal.cc
Lines changed: 248 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,248 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
#include "arrow/compute/kernels/vector_sort_internal.h"
19+
#include "arrow/compute/registry.h"
20+
21+
namespace arrow::compute::internal {
22+
23+
namespace {
24+
25+
// ----------------------------------------------------------------------
26+
// Rank implementation
27+
28+
const RankOptions* GetDefaultRankOptions() {
29+
static const auto kDefaultRankOptions = RankOptions::Defaults();
30+
return &kDefaultRankOptions;
31+
}
32+
33+
class ArrayRanker : public TypeVisitor {
34+
public:
35+
ArrayRanker(ExecContext* ctx, const Array& array, const RankOptions& options,
36+
Datum* output)
37+
: TypeVisitor(),
38+
ctx_(ctx),
39+
array_(array),
40+
options_(options),
41+
null_placement_(options.null_placement),
42+
tiebreaker_(options.tiebreaker),
43+
physical_type_(GetPhysicalType(array.type())),
44+
output_(output) {}
45+
46+
Status Run() { return physical_type_->Accept(this); }
47+
48+
#define VISIT(TYPE) \
49+
Status Visit(const TYPE& type) { return RankInternal<TYPE>(); }
50+
51+
VISIT_SORTABLE_PHYSICAL_TYPES(VISIT)
52+
53+
#undef VISIT
54+
55+
template <typename InType>
56+
Status RankInternal() {
57+
using GetView = GetViewType<InType>;
58+
using T = typename GetViewType<InType>::T;
59+
using ArrayType = typename TypeTraits<InType>::ArrayType;
60+
61+
ArrayType arr(array_.data());
62+
63+
SortOrder order = SortOrder::Ascending;
64+
if (!options_.sort_keys.empty()) {
65+
order = options_.sort_keys[0].order;
66+
}
67+
ArraySortOptions array_options(order, null_placement_);
68+
69+
auto length = array_.length();
70+
ARROW_ASSIGN_OR_RAISE(auto sort_indices,
71+
MakeMutableUInt64Array(length, ctx_->memory_pool()));
72+
auto sort_begin = sort_indices->GetMutableValues<uint64_t>(1);
73+
auto sort_end = sort_begin + length;
74+
std::iota(sort_begin, sort_end, 0);
75+
76+
ARROW_ASSIGN_OR_RAISE(auto array_sorter, GetArraySorter(*physical_type_));
77+
78+
NullPartitionResult sorted =
79+
array_sorter(sort_begin, sort_end, arr, 0, array_options);
80+
uint64_t rank;
81+
82+
ARROW_ASSIGN_OR_RAISE(auto rankings,
83+
MakeMutableUInt64Array(length, ctx_->memory_pool()));
84+
auto out_begin = rankings->GetMutableValues<uint64_t>(1);
85+
86+
switch (tiebreaker_) {
87+
case RankOptions::Dense: {
88+
T curr_value, prev_value{};
89+
rank = 0;
90+
91+
if (null_placement_ == NullPlacement::AtStart && sorted.null_count() > 0) {
92+
rank++;
93+
for (auto it = sorted.nulls_begin; it < sorted.nulls_end; it++) {
94+
out_begin[*it] = rank;
95+
}
96+
}
97+
98+
for (auto it = sorted.non_nulls_begin; it < sorted.non_nulls_end; it++) {
99+
curr_value = GetView::LogicalValue(arr.GetView(*it));
100+
if (it == sorted.non_nulls_begin || curr_value != prev_value) {
101+
rank++;
102+
}
103+
104+
out_begin[*it] = rank;
105+
prev_value = curr_value;
106+
}
107+
108+
if (null_placement_ == NullPlacement::AtEnd) {
109+
rank++;
110+
for (auto it = sorted.nulls_begin; it < sorted.nulls_end; it++) {
111+
out_begin[*it] = rank;
112+
}
113+
}
114+
break;
115+
}
116+
117+
case RankOptions::First: {
118+
rank = 0;
119+
for (auto it = sorted.overall_begin(); it < sorted.overall_end(); it++) {
120+
out_begin[*it] = ++rank;
121+
}
122+
break;
123+
}
124+
125+
case RankOptions::Min: {
126+
T curr_value, prev_value{};
127+
rank = 0;
128+
129+
if (null_placement_ == NullPlacement::AtStart) {
130+
rank++;
131+
for (auto it = sorted.nulls_begin; it < sorted.nulls_end; it++) {
132+
out_begin[*it] = rank;
133+
}
134+
}
135+
136+
for (auto it = sorted.non_nulls_begin; it < sorted.non_nulls_end; it++) {
137+
curr_value = GetView::LogicalValue(arr.GetView(*it));
138+
if (it == sorted.non_nulls_begin || curr_value != prev_value) {
139+
rank = (it - sorted.overall_begin()) + 1;
140+
}
141+
out_begin[*it] = rank;
142+
prev_value = curr_value;
143+
}
144+
145+
if (null_placement_ == NullPlacement::AtEnd) {
146+
rank = sorted.non_null_count() + 1;
147+
for (auto it = sorted.nulls_begin; it < sorted.nulls_end; it++) {
148+
out_begin[*it] = rank;
149+
}
150+
}
151+
break;
152+
}
153+
154+
case RankOptions::Max: {
155+
// The algorithm for Max is just like Min, but in reverse order.
156+
T curr_value, prev_value{};
157+
rank = length;
158+
159+
if (null_placement_ == NullPlacement::AtEnd) {
160+
for (auto it = sorted.nulls_begin; it < sorted.nulls_end; it++) {
161+
out_begin[*it] = rank;
162+
}
163+
}
164+
165+
for (auto it = sorted.non_nulls_end - 1; it >= sorted.non_nulls_begin; it--) {
166+
curr_value = GetView::LogicalValue(arr.GetView(*it));
167+
if (it == sorted.non_nulls_end - 1 || curr_value != prev_value) {
168+
rank = (it - sorted.overall_begin()) + 1;
169+
}
170+
out_begin[*it] = rank;
171+
prev_value = curr_value;
172+
}
173+
174+
if (null_placement_ == NullPlacement::AtStart) {
175+
rank = sorted.null_count();
176+
for (auto it = sorted.nulls_begin; it < sorted.nulls_end; it++) {
177+
out_begin[*it] = rank;
178+
}
179+
}
180+
181+
break;
182+
}
183+
}
184+
185+
*output_ = Datum(rankings);
186+
return Status::OK();
187+
}
188+
189+
ExecContext* ctx_;
190+
const Array& array_;
191+
const RankOptions& options_;
192+
const NullPlacement null_placement_;
193+
const RankOptions::Tiebreaker tiebreaker_;
194+
const std::shared_ptr<DataType> physical_type_;
195+
Datum* output_;
196+
};
197+
198+
const FunctionDoc rank_doc(
199+
"Compute numerical ranks of an array (1-based)",
200+
("This function computes a rank of the input array.\n"
201+
"By default, null values are considered greater than any other value and\n"
202+
"are therefore sorted at the end of the input. For floating-point types,\n"
203+
"NaNs are considered greater than any other non-null value, but smaller\n"
204+
"than null values. The default tiebreaker is to assign ranks in order of\n"
205+
"when ties appear in the input.\n"
206+
"\n"
207+
"The handling of nulls, NaNs and tiebreakers can be changed in RankOptions."),
208+
{"input"}, "RankOptions");
209+
210+
class RankMetaFunction : public MetaFunction {
211+
public:
212+
RankMetaFunction()
213+
: MetaFunction("rank", Arity::Unary(), rank_doc, GetDefaultRankOptions()) {}
214+
215+
Result<Datum> ExecuteImpl(const std::vector<Datum>& args,
216+
const FunctionOptions* options,
217+
ExecContext* ctx) const override {
218+
const auto& rank_options = checked_cast<const RankOptions&>(*options);
219+
switch (args[0].kind()) {
220+
case Datum::ARRAY: {
221+
return Rank(*args[0].make_array(), rank_options, ctx);
222+
} break;
223+
default:
224+
break;
225+
}
226+
return Status::NotImplemented(
227+
"Unsupported types for rank operation: "
228+
"values=",
229+
args[0].ToString());
230+
}
231+
232+
private:
233+
Result<Datum> Rank(const Array& array, const RankOptions& options,
234+
ExecContext* ctx) const {
235+
Datum output;
236+
ArrayRanker ranker(ctx, array, options, &output);
237+
ARROW_RETURN_NOT_OK(ranker.Run());
238+
return output;
239+
}
240+
};
241+
242+
} // namespace
243+
244+
void RegisterVectorRank(FunctionRegistry* registry) {
245+
DCHECK_OK(registry->AddFunction(std::make_shared<RankMetaFunction>()));
246+
}
247+
248+
} // namespace arrow::compute::internal

0 commit comments

Comments
 (0)