edponce commented on a change in pull request #11019:
URL: https://github.com/apache/arrow/pull/11019#discussion_r706269910
##########
File path: cpp/src/arrow/compute/kernels/vector_sort.cc
##########
@@ -1778,6 +1798,621 @@ class SortIndicesMetaFunction : public MetaFunction {
}
};
+// ----------------------------------------------------------------------
+// TopK/BottomK implementations
+
+const auto kDefaultSelectKOptions = SelectKOptions::Defaults();
+
+const FunctionDoc select_k_doc(
+ "Returns the first k elements ordered by `options.keys`",
+ ("This function computes the k elements of the input\n"
+ "array, record batch or table specified in the column names
(`options.sort_keys`).\n"
+ "The columns that are not specified are returned as well, but not used
for\n"
+ "ordering. Null values are considered greater than any other value and
are\n"
+ "therefore sorted at the end of the array.\n"
+ "For floating-point types, NaNs are considered greater than any\n"
+ "other non-null value, but smaller than null values."),
+ {"input"}, "SelectKOptions");
+
+Result<std::shared_ptr<ArrayData>> MakeMutableUInt64Array(
+ std::shared_ptr<DataType> out_type, int64_t length, MemoryPool*
memory_pool) {
+ auto buffer_size = length * sizeof(uint64_t);
+ ARROW_ASSIGN_OR_RAISE(auto data, AllocateBuffer(buffer_size, memory_pool));
+ return ArrayData::Make(uint64(), length, {nullptr, std::move(data)},
/*null_count=*/0);
+}
+
+template <SortOrder order>
+class SelectKComparator {
+ public:
+ template <typename Type>
+ bool operator()(const Type& lval, const Type& rval);
+};
+
+template <>
+class SelectKComparator<SortOrder::Ascending> {
+ public:
+ template <typename Type>
+ bool operator()(const Type& lval, const Type& rval) {
+ return lval < rval;
+ }
+};
+
+template <>
+class SelectKComparator<SortOrder::Descending> {
+ public:
+ template <typename Type>
+ bool operator()(const Type& lval, const Type& rval) {
+ return rval < lval;
+ }
+};
+
+template <SortOrder sort_order>
+class ArraySelecter : public TypeVisitor {
Review comment:
I am a bit queazy about the `SortOrder` template parameter. It is only
used to select the corresponding
[`SelectKComparator`](https://github.com/apache/arrow/pull/11019/files#diff-bf5bdbbf3ffa5d706a841eb38face9739c18d515eeb829e31749c7deff2e9384R1890)
and it allows setting it to a different order than what is dictated by the
`sort_keys` in `SelectKOptions`. To prevent this discrepancy, you could have
the `SelectKComparator` be a private member of the `XXXSelecter` classes and
set it based on `is_top_k()` etc. This would move the checks from
[here](https://github.com/apache/arrow/pull/11019/files#diff-bf5bdbbf3ffa5d706a841eb38face9739c18d515eeb829e31749c7deff2e9384R2347-R2368)
to the constructors.
--
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]
