ARROW-1565: [C++] Implement TopK/BottomK

cpp/src/arrow/compute/api_vector.cc

@@ -111,6 +111,9 @@ static auto kSortOptionsType =
     GetFunctionOptionsType<SortOptions>(DataMember("sort_keys", &SortOptions::sort_keys));
 static auto kPartitionNthOptionsType = GetFunctionOptionsType<PartitionNthOptions>(
     DataMember("pivot", &PartitionNthOptions::pivot));
+static auto kSelectKOptionsType = GetFunctionOptionsType<SelectKOptions>(
+    DataMember("k", &SelectKOptions::k),
+    DataMember("sort_keys", &SelectKOptions::sort_keys));
 }  // namespace
 }  // namespace internal
 
@@ -140,6 +143,29 @@ PartitionNthOptions::PartitionNthOptions(int64_t pivot)
     : FunctionOptions(internal::kPartitionNthOptionsType), pivot(pivot) {}
 constexpr char PartitionNthOptions::kTypeName[];
 
+SelectKOptions::SelectKOptions(int64_t k, std::vector<SortKey> sort_keys)
+    : FunctionOptions(internal::kSelectKOptionsType),
+      k(k),
+      sort_keys(std::move(sort_keys)) {}
+
+bool SelectKOptions::is_top_k() const {
+  for (const auto& k : sort_keys) {
+    if (k.order != SortOrder::Descending) {
+      return false;
+    }
+  }
+  return true;
+}
+bool SelectKOptions::is_bottom_k() const {
+  for (const auto& k : sort_keys) {
+    if (k.order != SortOrder::Ascending) {
+      return false;
+    }
+  }
+  return true;
+}
+constexpr char SelectKOptions::kTypeName[];
+
 namespace internal {
 void RegisterVectorOptions(FunctionRegistry* registry) {
   DCHECK_OK(registry->AddFunctionOptionsType(kFilterOptionsType));
@@ -148,6 +174,7 @@ void RegisterVectorOptions(FunctionRegistry* registry) {
   DCHECK_OK(registry->AddFunctionOptionsType(kArraySortOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kSortOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kPartitionNthOptionsType));
+  DCHECK_OK(registry->AddFunctionOptionsType(kSelectKOptionsType));
 }
 }  // namespace internal
 
@@ -162,6 +189,13 @@ Result<std::shared_ptr<Array>> NthToIndices(const Array& values, int64_t n,
   return result.make_array();
 }
 
+Result<std::shared_ptr<Array>> SelectKUnstable(const Datum& datum, SelectKOptions options,
+                                               ExecContext* ctx) {
+  ARROW_ASSIGN_OR_RAISE(Datum result,
+                        CallFunction("select_k_unstable", {datum}, &options, ctx));
+  return result.make_array();
+}
+
 Result<Datum> ReplaceWithMask(const Datum& values, const Datum& mask,
                               const Datum& replacements, ExecContext* ctx) {
   return CallFunction("replace_with_mask", {values, mask, replacements}, ctx);

cpp/src/arrow/compute/api_vector.h

@@ -120,6 +120,46 @@ class ARROW_EXPORT SortOptions : public FunctionOptions {
   std::vector<SortKey> sort_keys;
 };
 
+/// \brief SelectK options
+class ARROW_EXPORT SelectKOptions : public FunctionOptions {
+ public:
+  explicit SelectKOptions(int64_t k = -1, std::vector<SortKey> sort_keys = {});
+
+  constexpr static char const kTypeName[] = "SelectKOptions";
+
+  static SelectKOptions Defaults() { return SelectKOptions{-1, {}}; }
+
+  static SelectKOptions TopKDefault(int64_t k, std::vector<std::string> key_names = {}) {
+    std::vector<SortKey> keys;
+    for (const auto& name : key_names) {
+      keys.emplace_back(SortKey(name, SortOrder::Descending));
+    }
+    if (key_names.empty()) {
+      keys.emplace_back(SortKey("not-used", SortOrder::Descending));
+    }
+    return SelectKOptions{k, keys};
+  }
+  static SelectKOptions BottomKDefault(int64_t k,
+                                       std::vector<std::string> key_names = {}) {
+    std::vector<SortKey> keys;
+    for (const auto& name : key_names) {
+      keys.emplace_back(SortKey(name, SortOrder::Ascending));
+    }
+    if (key_names.empty()) {
+      keys.emplace_back(SortKey("not-used", SortOrder::Ascending));
+    }
+    return SelectKOptions{k, keys};
+  }
+  bool is_top_k() const;
+
+  bool is_bottom_k() const;
+
+  /// The number of `k` elements to keep.
+  int64_t k;
+  /// Column key(s) to order by and how to order by these sort keys.
+  std::vector<SortKey> sort_keys;
+};
+
 /// \brief Partitioning options for NthToIndices
 class ARROW_EXPORT PartitionNthOptions : public FunctionOptions {
  public:
@@ -252,6 +292,21 @@ ARROW_EXPORT
 Result<std::shared_ptr<Array>> NthToIndices(const Array& values, int64_t n,
                                             ExecContext* ctx = NULLPTR);
 
+/// \brief Returns the first k elements ordered by `options.keys`.
+///
+/// Return a sorted array with its elements rearranged in such
+/// a way that the value of the element in k-th position (options.k) is in the position it
+/// would be in a sorted datum ordered by `options.keys`. Null like values will be not
+/// part of the output. Output is not guaranteed to be stable.
+///
+/// \param[in] datum datum to be partitioned
+/// \param[in] options options
+/// \param[in] ctx the function execution context, optional
+/// \return a datum with the same schema as the input
+ARROW_EXPORT
+Result<std::shared_ptr<Array>> SelectKUnstable(const Datum& datum, SelectKOptions options,
+                                               ExecContext* ctx = NULLPTR);
+
 /// \brief Returns the indices that would sort an array in the
 /// specified order.
 ///

cpp/src/arrow/compute/function_test.cc

@@ -116,6 +116,8 @@ TEST(FunctionOptions, Equality) {
       {SortKey("key", SortOrder::Descending), SortKey("value", SortOrder::Descending)}));
   options.emplace_back(new PartitionNthOptions(/*pivot=*/0));
   options.emplace_back(new PartitionNthOptions(/*pivot=*/42));
+  options.emplace_back(new SelectKOptions(0, {}));
+  options.emplace_back(new SelectKOptions(5, {{SortKey("key", SortOrder::Ascending)}}));
 
   for (size_t i = 0; i < options.size(); i++) {
     const size_t prev_i = i == 0 ? options.size() - 1 : i - 1;

cpp/src/arrow/compute/kernels/CMakeLists.txt

@@ -50,11 +50,13 @@ add_arrow_compute_test(vector_test
                        vector_replace_test.cc
                        vector_selection_test.cc
                        vector_sort_test.cc
+                       select_k_test.cc
                        test_util.cc)
 
 add_arrow_benchmark(vector_hash_benchmark PREFIX "arrow-compute")
 add_arrow_benchmark(vector_sort_benchmark PREFIX "arrow-compute")
 add_arrow_benchmark(vector_partition_benchmark PREFIX "arrow-compute")
+add_arrow_benchmark(vector_topk_benchmark PREFIX "arrow-compute")
 add_arrow_benchmark(vector_replace_benchmark PREFIX "arrow-compute")
 add_arrow_benchmark(vector_selection_benchmark PREFIX "arrow-compute")