From 771536add5b0f9427c248b2ddeda65dd52989589 Mon Sep 17 00:00:00 2001
From: Ariana Villegas <ariana.villegas@utec.edu.pe>
Date: Wed, 8 Jun 2022 12:48:53 -0500
Subject: [PATCH 01/11] Draft sort indices on array dictionary

---
 cpp/src/arrow/array/array_dict.h              |  4 ++
 .../compute/kernels/vector_array_sort.cc      | 57 ++++++++++++++++++-
 cpp/src/arrow/compute/kernels/vector_sort.cc  |  8 +--
 .../compute/kernels/vector_sort_internal.h    |  8 +--
 .../arrow/compute/kernels/vector_sort_test.cc | 11 ++++
 5 files changed, 79 insertions(+), 9 deletions(-)

diff --git a/cpp/src/arrow/array/array_dict.h b/cpp/src/arrow/array/array_dict.h
index 8791eaa07db..207cffe54c8 100644
--- a/cpp/src/arrow/array/array_dict.h
+++ b/cpp/src/arrow/array/array_dict.h
@@ -111,6 +111,10 @@ class ARROW_EXPORT DictionaryArray : public Array {
 
   const DictionaryType* dict_type() const { return dict_type_; }
 
+  bool IsNull(int64_t i) const {
+    return indices_->IsNull(i) || dictionary_->IsNull(GetValueIndex(i));
+  }
+
  private:
   void SetData(const std::shared_ptr<ArrayData>& data);
   const DictionaryType* dict_type_;
diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index 324a435441f..8e65308602d 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -173,6 +173,52 @@ class ArrayCompareSorter {
   }
 };
 
+template <>
+class ArrayCompareSorter<DictionaryType> {
+ public:
+  NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                 const Array& array, int64_t offset,
+                                 const ArraySortOptions& options) {
+    const auto& dict_array = checked_cast<const DictionaryArray&>(array);
+
+    const auto& indices = checked_cast<const UInt64Array&>(*(dict_array.indices()));
+    const auto& values = dict_array.dictionary();
+
+    const auto p = PartitionNulls<DictionaryArray, StablePartitioner>(
+        indices_begin, indices_end, dict_array, offset, options.null_placement);
+
+    auto indices_array =
+        CallFunction("array_sort_indices", {values}, &options).ValueOrDie().make_array();
+    const auto& indices_values = checked_cast<const UInt64Array&>(*indices_array);
+
+    std::vector<uint64_t> sort_order(indices_values.length());
+    uint64_t cur = 0;
+    auto cur_idx = GetViewType<UInt64Type>::LogicalValue(indices_values.GetView(cur));
+    auto cur_val = values->GetScalar(cur_idx);
+    for (int i = 0; i < indices_values.length(); i++) {
+      auto tmp_idx = GetViewType<UInt64Type>::LogicalValue(indices_values.GetView(i));
+      auto tmp_val = values->GetScalar(tmp_idx);
+      if (cur_val != tmp_val) {
+        cur = i;
+        cur_val = tmp_val;
+      }
+      sort_order[tmp_idx] = cur;
+    }
+
+    std::stable_sort(
+        p.non_nulls_begin, p.non_nulls_end,
+        [&indices, &sort_order, &offset](uint64_t left, uint64_t right) {
+          const auto lhs =
+              GetViewType<UInt64Type>::LogicalValue(indices.GetView(left - offset));
+          const auto rhs =
+              GetViewType<UInt64Type>::LogicalValue(indices.GetView(right - offset));
+          return sort_order[lhs] < sort_order[rhs];
+        });
+
+    return p;
+  }
+};
+
 template <typename ArrowType>
 class ArrayCountSorter {
   using ArrayType = typename TypeTraits<ArrowType>::ArrayType;
@@ -405,7 +451,8 @@ struct ArraySorter<Type, enable_if_t<is_integer_type<Type>::value &&
 template <typename Type>
 struct ArraySorter<
     Type, enable_if_t<is_floating_type<Type>::value || is_base_binary_type<Type>::value ||
-                      is_fixed_size_binary_type<Type>::value>> {
+                      is_fixed_size_binary_type<Type>::value ||
+                      is_dictionary_type<Type>::value>> {
   ArrayCompareSorter<Type> impl;
 };
 
@@ -507,6 +554,13 @@ void AddArraySortingKernels(VectorKernel base, VectorFunction* func) {
   DCHECK_OK(func->AddKernel(base));
 }
 
+template <template <typename...> class ExecTemplate>
+void AddDictArraySortingKernels(VectorKernel base, VectorFunction* func) {
+  base.signature = KernelSignature::Make({InputType::Array(Type::DICTIONARY)}, uint64());
+  base.exec = ExecTemplate<UInt64Type, DictionaryType>::Exec;
+  DCHECK_OK(func->AddKernel(base));
+}
+
 const ArraySortOptions* GetDefaultArraySortOptions() {
   static const auto kDefaultArraySortOptions = ArraySortOptions::Defaults();
   return &kDefaultArraySortOptions;
@@ -561,6 +615,7 @@ void RegisterVectorArraySort(FunctionRegistry* registry) {
   base.init = ArraySortIndicesState::Init;
   base.exec_chunked = ArraySortIndicesChunked;
   AddArraySortingKernels<ArraySortIndices>(base, array_sort_indices.get());
+  AddDictArraySortingKernels<ArraySortIndices>(base, array_sort_indices.get());
   DCHECK_OK(registry->AddFunction(std::move(array_sort_indices)));
 
   // partition_nth_indices has a parameter so needs its init function
diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc
index 28307ecca37..872f3f06404 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort.cc
@@ -374,8 +374,8 @@ class ConcreteRecordBatchColumnSorter : public RecordBatchColumnSorter {
     } else {
       // NOTE that null_count_ is merely an upper bound on the number of nulls
       // in this particular range.
-      p = PartitionNullsOnly<StablePartitioner>(indices_begin, indices_end, array_,
-                                                offset, null_placement_);
+      p = PartitionNullsOnly<ArrayType, StablePartitioner>(
+          indices_begin, indices_end, array_, offset, null_placement_);
       DCHECK_LE(p.nulls_end - p.nulls_begin, null_count_);
     }
     const NullPartitionResult q = PartitionNullLikes<ArrayType, StablePartitioner>(
@@ -797,8 +797,8 @@ class MultipleKeyRecordBatchSorter : public TypeVisitor {
     using ArrayType = typename TypeTraits<Type>::ArrayType;
     const ArrayType& array = checked_cast<const ArrayType&>(first_sort_key.array);
 
-    const auto p = PartitionNullsOnly<StablePartitioner>(indices_begin_, indices_end_,
-                                                         array, 0, null_placement_);
+    const auto p = PartitionNullsOnly<ArrayType, StablePartitioner>(
+        indices_begin_, indices_end_, array, 0, null_placement_);
     const auto q = PartitionNullLikes<ArrayType, StablePartitioner>(
         p.non_nulls_begin, p.non_nulls_end, array, 0, null_placement_);
 
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_internal.h b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
index d8b024525c8..e914564d08f 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_internal.h
+++ b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
@@ -191,9 +191,9 @@ struct NullPartitionResult {
 // 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>
+template <typename ArrayType, typename Partitioner>
 NullPartitionResult PartitionNullsOnly(uint64_t* indices_begin, uint64_t* indices_end,
-                                       const Array& values, int64_t offset,
+                                       const ArrayType& values, int64_t offset,
                                        NullPlacement null_placement) {
   if (values.null_count() == 0) {
     return NullPartitionResult::NoNulls(indices_begin, indices_end, null_placement);
@@ -254,8 +254,8 @@ NullPartitionResult PartitionNulls(uint64_t* indices_begin, uint64_t* indices_en
                                    const ArrayType& values, int64_t offset,
                                    NullPlacement null_placement) {
   // Partition nulls at start (resp. end), and null-like values just before (resp. after)
-  NullPartitionResult p = PartitionNullsOnly<Partitioner>(indices_begin, indices_end,
-                                                          values, offset, null_placement);
+  NullPartitionResult p = PartitionNullsOnly<ArrayType, Partitioner>(
+      indices_begin, indices_end, values, offset, null_placement);
   NullPartitionResult q = PartitionNullLikes<ArrayType, Partitioner>(
       p.non_nulls_begin, p.non_nulls_end, values, offset, null_placement);
   return NullPartitionResult{q.non_nulls_begin, q.non_nulls_end,
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
index c32b1a1e591..0cc3bdec2f2 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
@@ -352,6 +352,17 @@ TEST(ArraySortIndicesFunction, Array) {
   expected = ArrayFromJSON(uint64(), "[2, 4, 6, 1, 0, 3, 5]");
   ASSERT_OK_AND_ASSIGN(actual, CallFunction("array_sort_indices", {arr}, &options));
   AssertDatumsEqual(expected, actual, /*verbose=*/true);
+
+  // Dictionary Array
+  auto dict_arr = DictArrayFromJSON(dictionary(uint64(), utf8()), "[0, null, 1, 0, 2, 3]",
+                                    "[ \"b\", null, \"c\", \"a\"]");
+  expected = ArrayFromJSON(uint64(), "[5, 0, 3, 4, 1, 2]");
+  ASSERT_OK_AND_ASSIGN(actual, CallFunction("array_sort_indices", {dict_arr}));
+  AssertDatumsEqual(expected, actual, /*verbose=*/true);
+
+  expected = ArrayFromJSON(uint64(), "[1, 2, 4, 0, 3, 5]");
+  ASSERT_OK_AND_ASSIGN(actual, CallFunction("array_sort_indices", {dict_arr}, &options));
+  AssertDatumsEqual(expected, actual, /*verbose=*/true);
 }
 
 TEST(ArraySortIndicesFunction, ChunkedArray) {

From 56afbcbd5541a24d04fc2d49c79536739769ca31 Mon Sep 17 00:00:00 2001
From: Ariana Villegas <ariana.villegas@utec.edu.pe>
Date: Thu, 9 Jun 2022 11:44:04 -0500
Subject: [PATCH 02/11] Update sort indices

---
 cpp/src/arrow/compute/kernels/vector_array_sort.cc | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index 8e65308602d..3ec0d375bb1 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -181,7 +181,7 @@ class ArrayCompareSorter<DictionaryType> {
                                  const ArraySortOptions& options) {
     const auto& dict_array = checked_cast<const DictionaryArray&>(array);
 
-    const auto& indices = checked_cast<const UInt64Array&>(*(dict_array.indices()));
+    const auto& indices = dict_array.indices();
     const auto& values = dict_array.dictionary();
 
     const auto p = PartitionNulls<DictionaryArray, StablePartitioner>(
@@ -209,9 +209,9 @@ class ArrayCompareSorter<DictionaryType> {
         p.non_nulls_begin, p.non_nulls_end,
         [&indices, &sort_order, &offset](uint64_t left, uint64_t right) {
           const auto lhs =
-              GetViewType<UInt64Type>::LogicalValue(indices.GetView(left - offset));
+              std::stoull(indices->GetScalar(left - offset).ValueOrDie()->ToString());
           const auto rhs =
-              GetViewType<UInt64Type>::LogicalValue(indices.GetView(right - offset));
+              std::stoull(indices->GetScalar(right - offset).ValueOrDie()->ToString());
           return sort_order[lhs] < sort_order[rhs];
         });
 

From 895851b89f0ce2f83c226456b529336419cae415 Mon Sep 17 00:00:00 2001
From: Ariana Villegas <ariana.villegas@utec.edu.pe>
Date: Mon, 20 Jun 2022 10:05:42 -0500
Subject: [PATCH 03/11] Update sorter

---
 .../compute/kernels/vector_array_sort.cc      | 92 +++++++++++++------
 1 file changed, 63 insertions(+), 29 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index 3ec0d375bb1..08ab4a63f1a 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -17,6 +17,7 @@
 
 #include <algorithm>
 #include <cmath>
+#include <iostream>
 #include <iterator>
 #include <limits>
 #include <numeric>
@@ -175,47 +176,80 @@ class ArrayCompareSorter {
 
 template <>
 class ArrayCompareSorter<DictionaryType> {
+  struct DictionaryInternal {
+    NullPartitionResult p;
+    const std::shared_ptr<Array>& values;
+    const std::shared_ptr<Array>& indices;
+    const UInt64Array& indices_values;
+    const ArraySortOptions& options;
+    int64_t offset;
+
+    Status Visit(const DataType& index_type) {
+      return Status::TypeError("Dictionary sorting not supported for index type ",
+                               index_type.ToString());
+    }
+
+    template <typename IndexType>
+    enable_if_t<is_integer_type<IndexType>::value, Status> Visit(
+        const IndexType& index_type) {
+      return SortInternal<IndexType>();
+    }
+
+    template <typename IndexType>
+    Status SortInternal() {
+      using ArrayType = typename TypeTraits<IndexType>::ArrayType;
+      using GetView = GetViewType<IndexType>;
+      const auto& indices_array = checked_cast<const ArrayType&>(*indices);
+
+      std::vector<uint64_t> sort_order(indices_values.length());
+      uint64_t cur = 0;
+      auto cur_idx = GetViewType<UInt64Type>::LogicalValue(indices_values.GetView(cur));
+      auto cur_val = values->GetScalar(cur_idx);
+      for (int i = 0; i < indices_values.length(); i++) {
+        auto tmp_idx = GetViewType<UInt64Type>::LogicalValue(indices_values.GetView(i));
+        auto tmp_val = values->GetScalar(tmp_idx);
+        if (cur_val != tmp_val) {
+          cur = i;
+          cur_val = tmp_val;
+        }
+        sort_order[tmp_idx] = cur;
+      }
+
+      std::stable_sort(
+          this->p.non_nulls_begin, this->p.non_nulls_end,
+          [&](uint64_t left, uint64_t right) {
+            const auto lhs = GetView::LogicalValue(indices_array.GetView(left - offset));
+            const auto rhs = GetView::LogicalValue(indices_array.GetView(right - offset));
+            return sort_order[lhs] < sort_order[rhs];
+          });
+
+      return Status::OK();
+    }
+
+    Result<NullPartitionResult> Make(const std::shared_ptr<DataType>& index_type) {
+      RETURN_NOT_OK(VisitTypeInline(*index_type, this));
+      return std::move(p);
+    }
+  };
+
  public:
   NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
                                  const Array& array, int64_t offset,
                                  const ArraySortOptions& options) {
     const auto& dict_array = checked_cast<const DictionaryArray&>(array);
-
-    const auto& indices = dict_array.indices();
     const auto& values = dict_array.dictionary();
+    const auto& indices = dict_array.indices();
+    const auto& index_type = dict_array.dict_type()->index_type();
 
-    const auto p = PartitionNulls<DictionaryArray, StablePartitioner>(
+    NullPartitionResult p = PartitionNulls<DictionaryArray, StablePartitioner>(
         indices_begin, indices_end, dict_array, offset, options.null_placement);
-
     auto indices_array =
         CallFunction("array_sort_indices", {values}, &options).ValueOrDie().make_array();
-    const auto& indices_values = checked_cast<const UInt64Array&>(*indices_array);
-
-    std::vector<uint64_t> sort_order(indices_values.length());
-    uint64_t cur = 0;
-    auto cur_idx = GetViewType<UInt64Type>::LogicalValue(indices_values.GetView(cur));
-    auto cur_val = values->GetScalar(cur_idx);
-    for (int i = 0; i < indices_values.length(); i++) {
-      auto tmp_idx = GetViewType<UInt64Type>::LogicalValue(indices_values.GetView(i));
-      auto tmp_val = values->GetScalar(tmp_idx);
-      if (cur_val != tmp_val) {
-        cur = i;
-        cur_val = tmp_val;
-      }
-      sort_order[tmp_idx] = cur;
-    }
 
-    std::stable_sort(
-        p.non_nulls_begin, p.non_nulls_end,
-        [&indices, &sort_order, &offset](uint64_t left, uint64_t right) {
-          const auto lhs =
-              std::stoull(indices->GetScalar(left - offset).ValueOrDie()->ToString());
-          const auto rhs =
-              std::stoull(indices->GetScalar(right - offset).ValueOrDie()->ToString());
-          return sort_order[lhs] < sort_order[rhs];
-        });
+    const auto& indices_values = checked_cast<const UInt64Array&>(*indices_array);
+    DictionaryInternal visitor = {p, values, indices, indices_values, options, offset};
 
-    return p;
+    return visitor.Make(index_type).ValueOrDie();
   }
 };
 

From a6e314daac12a0fd41c3fa939d498f7f34a6cd08 Mon Sep 17 00:00:00 2001
From: Ariana Villegas <ariana.villegas@utec.edu.pe>
Date: Mon, 4 Jul 2022 17:43:14 -0500
Subject: [PATCH 04/11] Address feedback

---
 .../compute/kernels/vector_array_sort.cc      | 100 +++++++++---------
 cpp/src/arrow/compute/kernels/vector_sort.cc  |  18 ++--
 .../compute/kernels/vector_sort_internal.h    |   2 +-
 .../arrow/compute/kernels/vector_sort_test.cc |  33 ++++--
 4 files changed, 82 insertions(+), 71 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index 08ab4a63f1a..ef2ee4e7749 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -144,9 +144,10 @@ class ArrayCompareSorter {
   using GetView = GetViewType<ArrowType>;
 
  public:
-  NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
-                                 const Array& array, int64_t offset,
-                                 const ArraySortOptions& options) {
+  // `offset` is used when this is called on a chunk of a chunked array
+  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                         const Array& array, int64_t offset,
+                                         const ArraySortOptions& options) {
     const auto& values = checked_cast<const ArrayType&>(array);
 
     const auto p = PartitionNulls<ArrayType, StablePartitioner>(
@@ -177,10 +178,10 @@ class ArrayCompareSorter {
 template <>
 class ArrayCompareSorter<DictionaryType> {
   struct DictionaryInternal {
-    NullPartitionResult p;
+    NullPartitionResult* p;
     const std::shared_ptr<Array>& values;
     const std::shared_ptr<Array>& indices;
-    const UInt64Array& indices_values;
+    const UInt64Array& values_rank;
     const ArraySortOptions& options;
     int64_t offset;
 
@@ -201,41 +202,31 @@ class ArrayCompareSorter<DictionaryType> {
       using GetView = GetViewType<IndexType>;
       const auto& indices_array = checked_cast<const ArrayType&>(*indices);
 
-      std::vector<uint64_t> sort_order(indices_values.length());
-      uint64_t cur = 0;
-      auto cur_idx = GetViewType<UInt64Type>::LogicalValue(indices_values.GetView(cur));
-      auto cur_val = values->GetScalar(cur_idx);
-      for (int i = 0; i < indices_values.length(); i++) {
-        auto tmp_idx = GetViewType<UInt64Type>::LogicalValue(indices_values.GetView(i));
-        auto tmp_val = values->GetScalar(tmp_idx);
-        if (cur_val != tmp_val) {
-          cur = i;
-          cur_val = tmp_val;
-        }
-        sort_order[tmp_idx] = cur;
-      }
-
-      std::stable_sort(
-          this->p.non_nulls_begin, this->p.non_nulls_end,
-          [&](uint64_t left, uint64_t right) {
-            const auto lhs = GetView::LogicalValue(indices_array.GetView(left - offset));
-            const auto rhs = GetView::LogicalValue(indices_array.GetView(right - offset));
-            return sort_order[lhs] < sort_order[rhs];
-          });
+      std::stable_sort(this->p->non_nulls_begin, this->p->non_nulls_end,
+                       [&](uint64_t left, uint64_t right) {
+                         const auto lhs_idx =
+                             GetView::LogicalValue(indices_array.GetView(left - offset));
+                         const auto rhs_idx =
+                             GetView::LogicalValue(indices_array.GetView(right - offset));
+                         const auto lhs = GetViewType<UInt64Type>::LogicalValue(
+                             values_rank.GetView(lhs_idx - offset));
+                         const auto rhs = GetViewType<UInt64Type>::LogicalValue(
+                             values_rank.GetView(rhs_idx - offset));
+                         return lhs < rhs;
+                       });
 
       return Status::OK();
     }
 
-    Result<NullPartitionResult> Make(const std::shared_ptr<DataType>& index_type) {
-      RETURN_NOT_OK(VisitTypeInline(*index_type, this));
-      return std::move(p);
+    Status Make(const std::shared_ptr<DataType>& index_type) {
+      return VisitTypeInline(*index_type, this);
     }
   };
 
  public:
-  NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
-                                 const Array& array, int64_t offset,
-                                 const ArraySortOptions& options) {
+  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                         const Array& array, int64_t offset,
+                                         const ArraySortOptions& options) {
     const auto& dict_array = checked_cast<const DictionaryArray&>(array);
     const auto& values = dict_array.dictionary();
     const auto& indices = dict_array.indices();
@@ -243,13 +234,16 @@ class ArrayCompareSorter<DictionaryType> {
 
     NullPartitionResult p = PartitionNulls<DictionaryArray, StablePartitioner>(
         indices_begin, indices_end, dict_array, offset, options.null_placement);
-    auto indices_array =
-        CallFunction("array_sort_indices", {values}, &options).ValueOrDie().make_array();
 
-    const auto& indices_values = checked_cast<const UInt64Array&>(*indices_array);
-    DictionaryInternal visitor = {p, values, indices, indices_values, options, offset};
+    RankOptions rank_options(options.order, options.null_placement, RankOptions::Dense);
+    auto rank = CallFunction("rank", {values}, &rank_options).ValueOrDie().make_array();
+
+    const auto& values_rank = checked_cast<const UInt64Array&>(*rank);
 
-    return visitor.Make(index_type).ValueOrDie();
+    DictionaryInternal visitor = {&p, values, indices, values_rank, options, offset};
+    ARROW_RETURN_NOT_OK(visitor.Make(index_type));
+
+    return p;
   }
 };
 
@@ -269,9 +263,9 @@ class ArrayCountSorter {
     value_range_ = static_cast<uint32_t>(max - min) + 1;
   }
 
-  NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
-                                 const Array& array, int64_t offset,
-                                 const ArraySortOptions& options) const {
+  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                         const Array& array, int64_t offset,
+                                         const ArraySortOptions& options) const {
     const auto& values = checked_cast<const ArrayType&>(array);
 
     // 32bit counter performs much better than 64bit one
@@ -287,9 +281,9 @@ class ArrayCountSorter {
   uint32_t value_range_{0};
 
   template <typename CounterType>
-  NullPartitionResult SortInternal(uint64_t* indices_begin, uint64_t* indices_end,
-                                   const ArrayType& values, int64_t offset,
-                                   const ArraySortOptions& options) const {
+  Result<NullPartitionResult> SortInternal(uint64_t* indices_begin, uint64_t* indices_end,
+                                           const ArrayType& values, int64_t offset,
+                                           const ArraySortOptions& options) const {
     const uint32_t value_range = value_range_;
 
     // first and last slot reserved for prefix sum (depending on sort order)
@@ -353,9 +347,10 @@ class ArrayCountSorter<BooleanType> {
  public:
   ArrayCountSorter() = default;
 
-  NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
-                                 const Array& array, int64_t offset,
-                                 const ArraySortOptions& options) {
+  // `offset` is used when this is called on a chunk of a chunked array
+  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                         const Array& array, int64_t offset,
+                                         const ArraySortOptions& options) {
     const auto& values = checked_cast<const BooleanArray&>(array);
 
     std::array<int64_t, 3> counts{0, 0, 0};  // false, true, null
@@ -398,9 +393,10 @@ class ArrayCountOrCompareSorter {
   using c_type = typename ArrowType::c_type;
 
  public:
-  NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
-                                 const Array& array, int64_t offset,
-                                 const ArraySortOptions& options) {
+  // `offset` is used when this is called on a chunk of a chunked array
+  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                         const Array& array, int64_t offset,
+                                         const ArraySortOptions& options) {
     const auto& values = checked_cast<const ArrayType&>(array);
 
     if (values.length() >= countsort_min_len_ && values.length() > values.null_count()) {
@@ -438,9 +434,9 @@ class ArrayCountOrCompareSorter {
 
 class ArrayNullSorter {
  public:
-  NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
-                                 const Array& values, int64_t offset,
-                                 const ArraySortOptions& options) {
+  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                         const Array& values, int64_t offset,
+                                         const ArraySortOptions& options) {
     return NullPartitionResult::NullsOnly(indices_begin, indices_end,
                                           options.null_placement);
   }
diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc
index 872f3f06404..aea2fa8f149 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort.cc
@@ -201,7 +201,7 @@ class ChunkedArraySorter : public TypeVisitor {
 
     // Sort each chunk independently and merge to sorted indices.
     // This is a serial implementation.
-    std::vector<NullPartitionResult> sorted(num_chunks);
+    std::vector<Result<NullPartitionResult>> sorted(num_chunks);
 
     // First sort all individual chunks
     int64_t begin_offset = 0;
@@ -244,8 +244,8 @@ class ChunkedArraySorter : public TypeVisitor {
         auto out_it = sorted.begin();
         auto it = sorted.begin();
         while (it < sorted.end() - 1) {
-          const auto& left = *it++;
-          const auto& right = *it++;
+          ARROW_ASSIGN_OR_RAISE(const auto& left, *it++);
+          ARROW_ASSIGN_OR_RAISE(const auto& right, *it++);
           DCHECK_EQ(left.overall_end(), right.overall_begin());
           const auto merged = merge_impl.Merge(left, right, null_count);
           *out_it++ = merged;
@@ -258,10 +258,11 @@ class ChunkedArraySorter : public TypeVisitor {
     }
 
     DCHECK_EQ(sorted.size(), 1);
-    DCHECK_EQ(sorted[0].overall_begin(), indices_begin_);
-    DCHECK_EQ(sorted[0].overall_end(), indices_end_);
+    ARROW_ASSIGN_OR_RAISE(auto sorted_arr, sorted[0]);
+    DCHECK_EQ(sorted_arr.overall_begin(), indices_begin_);
+    DCHECK_EQ(sorted_arr.overall_end(), indices_end_);
     // Note that "nulls" can also include NaNs, hence the >= check
-    DCHECK_GE(sorted[0].null_count(), null_count);
+    DCHECK_GE(sorted_arr.null_count(), null_count);
 
     return Status::OK();
   }
@@ -1957,9 +1958,8 @@ class ArrayRanker : public TypeVisitor {
     std::iota(sort_begin, sort_end, 0);
 
     ARROW_ASSIGN_OR_RAISE(auto array_sorter, GetArraySorter(*physical_type_));
-
-    NullPartitionResult sorted =
-        array_sorter(sort_begin, sort_end, arr, 0, array_options);
+    ARROW_ASSIGN_OR_RAISE(auto sorted,
+                          array_sorter(sort_begin, sort_end, arr, 0, array_options));
     uint64_t rank;
 
     ARROW_ASSIGN_OR_RAISE(auto rankings,
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_internal.h b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
index e914564d08f..2d45ec0d6f0 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_internal.h
+++ b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
@@ -446,7 +446,7 @@ struct MergeImpl {
 // TODO make this usable if indices are non trivial on input
 // (see ConcreteRecordBatchColumnSorter)
 // `offset` is used when this is called on a chunk of a chunked array
-using ArraySortFunc = std::function<NullPartitionResult(
+using ArraySortFunc = std::function<Result<NullPartitionResult>(
     uint64_t* indices_begin, uint64_t* indices_end, const Array& values, int64_t offset,
     const ArraySortOptions& options)>;
 
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
index 0cc3bdec2f2..b3fcd617565 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
@@ -352,17 +352,32 @@ TEST(ArraySortIndicesFunction, Array) {
   expected = ArrayFromJSON(uint64(), "[2, 4, 6, 1, 0, 3, 5]");
   ASSERT_OK_AND_ASSIGN(actual, CallFunction("array_sort_indices", {arr}, &options));
   AssertDatumsEqual(expected, actual, /*verbose=*/true);
+}
 
-  // Dictionary Array
-  auto dict_arr = DictArrayFromJSON(dictionary(uint64(), utf8()), "[0, null, 1, 0, 2, 3]",
-                                    "[ \"b\", null, \"c\", \"a\"]");
-  expected = ArrayFromJSON(uint64(), "[5, 0, 3, 4, 1, 2]");
-  ASSERT_OK_AND_ASSIGN(actual, CallFunction("array_sort_indices", {dict_arr}));
-  AssertDatumsEqual(expected, actual, /*verbose=*/true);
+TEST(ArraySortIndicesFunction, DictionaryArray) {
+  // Dictionary Array (["b", null, null, null, "b", "c", "a"])
+  std::vector<std::string> expected_str = {
+      "[8, 0, 5, 1, 6, 7, 2, 3, 4]",  // SortOrder::Ascending NullPlacement::AtEnd
+      "[2, 3, 4, 8, 0, 5, 1, 6, 7]",  // SortOrder::Ascending NullPlacement::AtStart
+      "[1, 6, 7, 0, 5, 8, 2, 3, 4]",  // SortOrder::Descending NullPlacement::AtEnd
+      "[2, 3, 4, 1, 6, 7, 0, 5, 8]"   // SortOrder::Descending NullPlacement::AtStart
+  };
 
-  expected = ArrayFromJSON(uint64(), "[1, 2, 4, 0, 3, 5]");
-  ASSERT_OK_AND_ASSIGN(actual, CallFunction("array_sort_indices", {dict_arr}, &options));
-  AssertDatumsEqual(expected, actual, /*verbose=*/true);
+  for (auto index_type : {uint8(), uint16(), uint32(), uint64()}) {
+    int i = 0;
+    auto dict_arr = DictArrayFromJSON(dictionary(index_type, utf8()),
+                                      "[0, 4, null, 1, null, 0, 4, 2, 3]",
+                                      "[ \"b\", null, \"c\", \"a\", \"c\"]");
+    for (auto order : AllOrders()) {
+      for (auto null_placement : AllNullPlacements()) {
+        ArraySortOptions options{order, null_placement};
+        auto expected = ArrayFromJSON(uint64(), expected_str[i++]);
+        ASSERT_OK_AND_ASSIGN(auto actual,
+                             CallFunction("array_sort_indices", {dict_arr}, &options));
+        AssertDatumsEqual(expected, actual, /*verbose=*/true);
+      }
+    }
+  }
 }
 
 TEST(ArraySortIndicesFunction, ChunkedArray) {

From b3220104fa46ac924aa2664818cd50f5e4cdeb5c Mon Sep 17 00:00:00 2001
From: Ariana Villegas <ariana.villegas@utec.edu.pe>
Date: Fri, 15 Jul 2022 12:53:24 -0500
Subject: [PATCH 05/11] Address feedback

---
 cpp/src/arrow/array/array_dict.h              |  2 +-
 .../compute/kernels/vector_array_sort.cc      | 54 +++++++++----------
 cpp/src/arrow/compute/kernels/vector_sort.cc  | 18 +++----
 .../compute/kernels/vector_sort_internal.h    |  2 +-
 4 files changed, 36 insertions(+), 40 deletions(-)

diff --git a/cpp/src/arrow/array/array_dict.h b/cpp/src/arrow/array/array_dict.h
index 207cffe54c8..676ab87b8f3 100644
--- a/cpp/src/arrow/array/array_dict.h
+++ b/cpp/src/arrow/array/array_dict.h
@@ -112,7 +112,7 @@ class ARROW_EXPORT DictionaryArray : public Array {
   const DictionaryType* dict_type() const { return dict_type_; }
 
   bool IsNull(int64_t i) const {
-    return indices_->IsNull(i) || dictionary_->IsNull(GetValueIndex(i));
+    return indices_->IsNull(i) || dictionary()->IsNull(GetValueIndex(i));
   }
 
  private:
diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index ef2ee4e7749..3eb5afca653 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -17,7 +17,6 @@
 
 #include <algorithm>
 #include <cmath>
-#include <iostream>
 #include <iterator>
 #include <limits>
 #include <numeric>
@@ -144,10 +143,9 @@ class ArrayCompareSorter {
   using GetView = GetViewType<ArrowType>;
 
  public:
-  // `offset` is used when this is called on a chunk of a chunked array
-  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
-                                         const Array& array, int64_t offset,
-                                         const ArraySortOptions& options) {
+  NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                 const Array& array, int64_t offset,
+                                 const ArraySortOptions& options) {
     const auto& values = checked_cast<const ArrayType&>(array);
 
     const auto p = PartitionNulls<ArrayType, StablePartitioner>(
@@ -218,15 +216,15 @@ class ArrayCompareSorter<DictionaryType> {
       return Status::OK();
     }
 
-    Status Make(const std::shared_ptr<DataType>& index_type) {
-      return VisitTypeInline(*index_type, this);
+    void Make(const std::shared_ptr<DataType>& index_type) {
+      ARROW_CHECK_OK(VisitTypeInline(*index_type, this));
     }
   };
 
  public:
-  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
-                                         const Array& array, int64_t offset,
-                                         const ArraySortOptions& options) {
+  NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                 const Array& array, int64_t offset,
+                                 const ArraySortOptions& options) {
     const auto& dict_array = checked_cast<const DictionaryArray&>(array);
     const auto& values = dict_array.dictionary();
     const auto& indices = dict_array.indices();
@@ -241,7 +239,7 @@ class ArrayCompareSorter<DictionaryType> {
     const auto& values_rank = checked_cast<const UInt64Array&>(*rank);
 
     DictionaryInternal visitor = {&p, values, indices, values_rank, options, offset};
-    ARROW_RETURN_NOT_OK(visitor.Make(index_type));
+    visitor.Make(index_type);
 
     return p;
   }
@@ -263,9 +261,9 @@ class ArrayCountSorter {
     value_range_ = static_cast<uint32_t>(max - min) + 1;
   }
 
-  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
-                                         const Array& array, int64_t offset,
-                                         const ArraySortOptions& options) const {
+  NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                 const Array& array, int64_t offset,
+                                 const ArraySortOptions& options) const {
     const auto& values = checked_cast<const ArrayType&>(array);
 
     // 32bit counter performs much better than 64bit one
@@ -281,9 +279,9 @@ class ArrayCountSorter {
   uint32_t value_range_{0};
 
   template <typename CounterType>
-  Result<NullPartitionResult> SortInternal(uint64_t* indices_begin, uint64_t* indices_end,
-                                           const ArrayType& values, int64_t offset,
-                                           const ArraySortOptions& options) const {
+  NullPartitionResult SortInternal(uint64_t* indices_begin, uint64_t* indices_end,
+                                   const ArrayType& values, int64_t offset,
+                                   const ArraySortOptions& options) const {
     const uint32_t value_range = value_range_;
 
     // first and last slot reserved for prefix sum (depending on sort order)
@@ -347,10 +345,9 @@ class ArrayCountSorter<BooleanType> {
  public:
   ArrayCountSorter() = default;
 
-  // `offset` is used when this is called on a chunk of a chunked array
-  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
-                                         const Array& array, int64_t offset,
-                                         const ArraySortOptions& options) {
+  NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                 const Array& array, int64_t offset,
+                                 const ArraySortOptions& options) {
     const auto& values = checked_cast<const BooleanArray&>(array);
 
     std::array<int64_t, 3> counts{0, 0, 0};  // false, true, null
@@ -393,10 +390,9 @@ class ArrayCountOrCompareSorter {
   using c_type = typename ArrowType::c_type;
 
  public:
-  // `offset` is used when this is called on a chunk of a chunked array
-  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
-                                         const Array& array, int64_t offset,
-                                         const ArraySortOptions& options) {
+  NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                 const Array& array, int64_t offset,
+                                 const ArraySortOptions& options) {
     const auto& values = checked_cast<const ArrayType&>(array);
 
     if (values.length() >= countsort_min_len_ && values.length() > values.null_count()) {
@@ -434,9 +430,9 @@ class ArrayCountOrCompareSorter {
 
 class ArrayNullSorter {
  public:
-  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
-                                         const Array& values, int64_t offset,
-                                         const ArraySortOptions& options) {
+  NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                 const Array& values, int64_t offset,
+                                 const ArraySortOptions& options) {
     return NullPartitionResult::NullsOnly(indices_begin, indices_end,
                                           options.null_placement);
   }
@@ -586,7 +582,7 @@ void AddArraySortingKernels(VectorKernel base, VectorFunction* func) {
 
 template <template <typename...> class ExecTemplate>
 void AddDictArraySortingKernels(VectorKernel base, VectorFunction* func) {
-  base.signature = KernelSignature::Make({InputType::Array(Type::DICTIONARY)}, uint64());
+  base.signature = KernelSignature::Make({Type::DICTIONARY}, uint64());
   base.exec = ExecTemplate<UInt64Type, DictionaryType>::Exec;
   DCHECK_OK(func->AddKernel(base));
 }
diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc
index aea2fa8f149..872f3f06404 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort.cc
@@ -201,7 +201,7 @@ class ChunkedArraySorter : public TypeVisitor {
 
     // Sort each chunk independently and merge to sorted indices.
     // This is a serial implementation.
-    std::vector<Result<NullPartitionResult>> sorted(num_chunks);
+    std::vector<NullPartitionResult> sorted(num_chunks);
 
     // First sort all individual chunks
     int64_t begin_offset = 0;
@@ -244,8 +244,8 @@ class ChunkedArraySorter : public TypeVisitor {
         auto out_it = sorted.begin();
         auto it = sorted.begin();
         while (it < sorted.end() - 1) {
-          ARROW_ASSIGN_OR_RAISE(const auto& left, *it++);
-          ARROW_ASSIGN_OR_RAISE(const auto& right, *it++);
+          const auto& left = *it++;
+          const auto& right = *it++;
           DCHECK_EQ(left.overall_end(), right.overall_begin());
           const auto merged = merge_impl.Merge(left, right, null_count);
           *out_it++ = merged;
@@ -258,11 +258,10 @@ class ChunkedArraySorter : public TypeVisitor {
     }
 
     DCHECK_EQ(sorted.size(), 1);
-    ARROW_ASSIGN_OR_RAISE(auto sorted_arr, sorted[0]);
-    DCHECK_EQ(sorted_arr.overall_begin(), indices_begin_);
-    DCHECK_EQ(sorted_arr.overall_end(), indices_end_);
+    DCHECK_EQ(sorted[0].overall_begin(), indices_begin_);
+    DCHECK_EQ(sorted[0].overall_end(), indices_end_);
     // Note that "nulls" can also include NaNs, hence the >= check
-    DCHECK_GE(sorted_arr.null_count(), null_count);
+    DCHECK_GE(sorted[0].null_count(), null_count);
 
     return Status::OK();
   }
@@ -1958,8 +1957,9 @@ class ArrayRanker : public TypeVisitor {
     std::iota(sort_begin, sort_end, 0);
 
     ARROW_ASSIGN_OR_RAISE(auto array_sorter, GetArraySorter(*physical_type_));
-    ARROW_ASSIGN_OR_RAISE(auto sorted,
-                          array_sorter(sort_begin, sort_end, arr, 0, array_options));
+
+    NullPartitionResult sorted =
+        array_sorter(sort_begin, sort_end, arr, 0, array_options);
     uint64_t rank;
 
     ARROW_ASSIGN_OR_RAISE(auto rankings,
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_internal.h b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
index 2d45ec0d6f0..e914564d08f 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_internal.h
+++ b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
@@ -446,7 +446,7 @@ struct MergeImpl {
 // TODO make this usable if indices are non trivial on input
 // (see ConcreteRecordBatchColumnSorter)
 // `offset` is used when this is called on a chunk of a chunked array
-using ArraySortFunc = std::function<Result<NullPartitionResult>(
+using ArraySortFunc = std::function<NullPartitionResult(
     uint64_t* indices_begin, uint64_t* indices_end, const Array& values, int64_t offset,
     const ArraySortOptions& options)>;
 

From fc192e27eebc0329088ce8e88dbe425ec97b1fc5 Mon Sep 17 00:00:00 2001
From: Antoine Pitrou <antoine@python.org>
Date: Thu, 21 Jul 2022 16:46:54 +0200
Subject: [PATCH 06/11] Add tests and benchmarks

---
 .../compute/kernels/vector_array_sort.cc      | 25 ++++---
 .../compute/kernels/vector_sort_benchmark.cc  | 75 +++++++++++++++++++
 .../arrow/compute/kernels/vector_sort_test.cc | 46 +++++++++++-
 3 files changed, 132 insertions(+), 14 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index 3eb5afca653..4def9382e1f 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -197,20 +197,20 @@ class ArrayCompareSorter<DictionaryType> {
     template <typename IndexType>
     Status SortInternal() {
       using ArrayType = typename TypeTraits<IndexType>::ArrayType;
-      using GetView = GetViewType<IndexType>;
-      const auto& indices_array = checked_cast<const ArrayType&>(*indices);
-
+      using c_type = typename IndexType::c_type;
+      const c_type* indices_values =
+          checked_cast<const ArrayType&>(*indices).raw_values();
+      const uint64_t* ranks = values_rank.raw_values();
+
+      // TODO Instead, try to take the ranks with the dict indices,
+      // then run a sort on the results.  Since the ranks are dense,
+      // a much faster counting sort may be used.
       std::stable_sort(this->p->non_nulls_begin, this->p->non_nulls_end,
                        [&](uint64_t left, uint64_t right) {
-                         const auto lhs_idx =
-                             GetView::LogicalValue(indices_array.GetView(left - offset));
-                         const auto rhs_idx =
-                             GetView::LogicalValue(indices_array.GetView(right - offset));
-                         const auto lhs = GetViewType<UInt64Type>::LogicalValue(
-                             values_rank.GetView(lhs_idx - offset));
-                         const auto rhs = GetViewType<UInt64Type>::LogicalValue(
-                             values_rank.GetView(rhs_idx - offset));
-                         return lhs < rhs;
+                         const uint64_t left_rank = ranks[indices_values[left - offset]];
+                         const uint64_t right_rank =
+                             ranks[indices_values[right - offset]];
+                         return left_rank < right_rank;
                        });
 
       return Status::OK();
@@ -234,6 +234,7 @@ class ArrayCompareSorter<DictionaryType> {
         indices_begin, indices_end, dict_array, offset, options.null_placement);
 
     RankOptions rank_options(options.order, options.null_placement, RankOptions::Dense);
+    // FIXME propagate instead of aborting on error
     auto rank = CallFunction("rank", {values}, &rank_options).ValueOrDie().make_array();
 
     const auto& values_rank = checked_cast<const UInt64Array&>(*rank);
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc
index 6ab0bcfde97..7d60cb3aebf 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc
@@ -56,6 +56,49 @@ static void ArraySortIndicesInt64Benchmark(benchmark::State& state, int64_t min,
   ArraySortIndicesBenchmark(state, values);
 }
 
+static void ArraySortIndicesInt64DictBenchmark(benchmark::State& state, int64_t min,
+                                               int64_t max) {
+  RegressionArgs args(state);
+
+  const int64_t array_size = args.size / sizeof(int64_t);
+  const int32_t dict_values_size = 20;  // a typical dictionary size
+
+  auto rand = random::RandomArrayGenerator(kSeed);
+  auto dict_values = rand.Int64(dict_values_size, min, max, args.null_proportion / 2);
+  auto dict_indices =
+      rand.Int64(array_size, 0, dict_values_size - 1, args.null_proportion / 2);
+  auto dict_array = *DictionaryArray::FromArrays(dict_indices, dict_values);
+
+  ArraySortIndicesBenchmark(state, dict_array);
+}
+
+static void ArraySortIndicesStringsBenchmark(benchmark::State& state) {
+  RegressionArgs args(state);
+
+  const int64_t array_size = args.size / sizeof(int64_t);
+  auto rand = random::RandomArrayGenerator(kSeed);
+  auto values =
+      rand.String(array_size, /*min_length=*/3, /*max_length=*/25, args.null_proportion);
+
+  ArraySortIndicesBenchmark(state, values);
+}
+
+static void ArraySortIndicesStringsDictBenchmark(benchmark::State& state) {
+  RegressionArgs args(state);
+
+  const int64_t array_size = args.size / sizeof(int64_t);
+  const int32_t dict_values_size = 20;  // a typical dictionary size
+
+  auto rand = random::RandomArrayGenerator(kSeed);
+  auto dict_values = rand.String(dict_values_size, /*min_length=*/3, /*max_length=*/25,
+                                 args.null_proportion / 2);
+  auto dict_indices =
+      rand.Int64(array_size, 0, dict_values_size - 1, args.null_proportion / 2);
+  auto dict_array = *DictionaryArray::FromArrays(dict_indices, dict_values);
+
+  ArraySortIndicesBenchmark(state, dict_array);
+}
+
 static void ChunkedArraySortIndicesInt64Benchmark(benchmark::State& state, int64_t min,
                                                   int64_t max) {
   RegressionArgs args(state);
@@ -81,6 +124,20 @@ static void ArraySortIndicesInt64Wide(benchmark::State& state) {
   ArraySortIndicesInt64Benchmark(state, min, max);
 }
 
+static void ArraySortIndicesInt64Dict(benchmark::State& state) {
+  const auto min = std::numeric_limits<int64_t>::min();
+  const auto max = std::numeric_limits<int64_t>::max();
+  ArraySortIndicesInt64DictBenchmark(state, min, max);
+}
+
+static void ArraySortIndicesStrings(benchmark::State& state) {
+  ArraySortIndicesStringsBenchmark(state);
+}
+
+static void ArraySortIndicesStringsDict(benchmark::State& state) {
+  ArraySortIndicesStringsDictBenchmark(state);
+}
+
 static void ArraySortIndicesBool(benchmark::State& state) {
   RegressionArgs args(state);
 
@@ -249,6 +306,24 @@ BENCHMARK(ArraySortIndicesInt64Wide)
     ->Args({1 << 23, 100})
     ->Unit(benchmark::TimeUnit::kNanosecond);
 
+BENCHMARK(ArraySortIndicesInt64Dict)
+    ->Apply(RegressionSetArgs)
+    ->Args({1 << 20, 100})
+    ->Args({1 << 23, 100})
+    ->Unit(benchmark::TimeUnit::kNanosecond);
+
+BENCHMARK(ArraySortIndicesStrings)
+    ->Apply(RegressionSetArgs)
+    ->Args({1 << 20, 100})
+    ->Args({1 << 23, 100})
+    ->Unit(benchmark::TimeUnit::kNanosecond);
+
+BENCHMARK(ArraySortIndicesStringsDict)
+    ->Apply(RegressionSetArgs)
+    ->Args({1 << 20, 100})
+    ->Args({1 << 23, 100})
+    ->Unit(benchmark::TimeUnit::kNanosecond);
+
 BENCHMARK(ArraySortIndicesBool)
     ->Apply(RegressionSetArgs)
     ->Args({1 << 20, 100})
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
index b3fcd617565..f1808fae980 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
@@ -355,7 +355,9 @@ TEST(ArraySortIndicesFunction, Array) {
 }
 
 TEST(ArraySortIndicesFunction, DictionaryArray) {
-  // Dictionary Array (["b", null, null, null, "b", "c", "a"])
+  // Decoded dictionary array:
+  // (["b", "c", null, null, null, "b", "c", "c", "a"])
+
   std::vector<std::string> expected_str = {
       "[8, 0, 5, 1, 6, 7, 2, 3, 4]",  // SortOrder::Ascending NullPlacement::AtEnd
       "[2, 3, 4, 8, 0, 5, 1, 6, 7]",  // SortOrder::Ascending NullPlacement::AtStart
@@ -363,7 +365,8 @@ TEST(ArraySortIndicesFunction, DictionaryArray) {
       "[2, 3, 4, 1, 6, 7, 0, 5, 8]"   // SortOrder::Descending NullPlacement::AtStart
   };
 
-  for (auto index_type : {uint8(), uint16(), uint32(), uint64()}) {
+  for (auto index_type : all_dictionary_index_types()) {
+    ARROW_SCOPED_TRACE("index_type = ", index_type->ToString());
     int i = 0;
     auto dict_arr = DictArrayFromJSON(dictionary(index_type, utf8()),
                                       "[0, 4, null, 1, null, 0, 4, 2, 3]",
@@ -374,12 +377,51 @@ TEST(ArraySortIndicesFunction, DictionaryArray) {
         auto expected = ArrayFromJSON(uint64(), expected_str[i++]);
         ASSERT_OK_AND_ASSIGN(auto actual,
                              CallFunction("array_sort_indices", {dict_arr}, &options));
+        ValidateOutput(actual);
         AssertDatumsEqual(expected, actual, /*verbose=*/true);
       }
     }
   }
 }
 
+Result<std::shared_ptr<Array>> DecodeDictionary(const Array& array) {
+  const auto& dict_array = checked_cast<const DictionaryArray&>(array);
+  ARROW_ASSIGN_OR_RAISE(auto decoded_datum,
+                        Take(dict_array.dictionary(), dict_array.indices()));
+  return decoded_datum.make_array();
+}
+
+TEST(ArraySortIndicesFunction, RandomDictionaryArray) {
+  ::arrow::random::RandomArrayGenerator rng(/*seed=*/1234);
+  constexpr int64_t kLength = 200;
+  constexpr int64_t kDictLength = 20;
+
+  // Ensure there are duplicates in the dictionary and the indices,
+  // and nulls in both as well.
+  auto dict_values = rng.StringWithRepeats(kDictLength, /*unique=*/kDictLength / 2,
+                                           /*min_length=*/1, /*max_length=*/10,
+                                           /*null_probability=*/0.2);
+  auto dict_indices = rng.Int64(kLength, 0, kDictLength - 1, /*null_probability = */ 0.2);
+  ASSERT_OK_AND_ASSIGN(auto dict_array,
+                       DictionaryArray::FromArrays(dict_indices, dict_values));
+  ASSERT_OK_AND_ASSIGN(auto decoded, DecodeDictionary(*dict_array));
+
+  for (auto order : AllOrders()) {
+    for (auto null_placement : AllNullPlacements()) {
+      ArraySortOptions options{order, null_placement};
+      // Sorting the dictionary array...
+      ASSERT_OK_AND_ASSIGN(auto actual,
+                           CallFunction("array_sort_indices", {dict_array}, &options));
+      ValidateOutput(actual);
+
+      // should give identical results to sorting the decoded array
+      ASSERT_OK_AND_ASSIGN(auto expected,
+                           CallFunction("array_sort_indices", {decoded}, &options));
+      AssertDatumsEqual(expected, actual, /*verbose=*/true);
+    }
+  }
+}
+
 TEST(ArraySortIndicesFunction, ChunkedArray) {
   auto arr = ChunkedArrayFromJSON(int16(), {"[0, 1]", "[null, -3, null, -42, 5]"});
   auto expected = ChunkedArrayFromJSON(uint64(), {"[5, 3, 0, 1, 6, 2, 4]"});

From 9775a46e5d1b9583a6cdf81bc4ba8577c8d1cfa1 Mon Sep 17 00:00:00 2001
From: Antoine Pitrou <antoine@python.org>
Date: Thu, 21 Jul 2022 17:27:37 +0200
Subject: [PATCH 07/11] Some nits re:benchmarking

---
 .../compute/kernels/vector_sort_benchmark.cc   | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc
index 7d60cb3aebf..c0ab70a486d 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc
@@ -27,7 +27,9 @@
 
 namespace arrow {
 namespace compute {
+
 constexpr auto kSeed = 0x0ff1ce;
+constexpr int32_t kDictionarySize = 20;  // a typical dictionary size
 
 static void ArraySortIndicesBenchmark(benchmark::State& state,
                                       const std::shared_ptr<Array>& values) {
@@ -61,20 +63,20 @@ static void ArraySortIndicesInt64DictBenchmark(benchmark::State& state, int64_t
   RegressionArgs args(state);
 
   const int64_t array_size = args.size / sizeof(int64_t);
-  const int32_t dict_values_size = 20;  // a typical dictionary size
 
   auto rand = random::RandomArrayGenerator(kSeed);
-  auto dict_values = rand.Int64(dict_values_size, min, max, args.null_proportion / 2);
+  auto dict_values = rand.Int64(kDictionarySize, min, max, args.null_proportion / 2);
   auto dict_indices =
-      rand.Int64(array_size, 0, dict_values_size - 1, args.null_proportion / 2);
+      rand.Int64(array_size, 0, kDictionarySize - 1, args.null_proportion / 2);
   auto dict_array = *DictionaryArray::FromArrays(dict_indices, dict_values);
 
   ArraySortIndicesBenchmark(state, dict_array);
 }
 
 static void ArraySortIndicesStringsBenchmark(benchmark::State& state) {
-  RegressionArgs args(state);
+  RegressionArgs args(state, /*size_is_bytes=*/false);
 
+  // XXX This division to make numbers comparable with ArraySortIndicesInt64Benchmark
   const int64_t array_size = args.size / sizeof(int64_t);
   auto rand = random::RandomArrayGenerator(kSeed);
   auto values =
@@ -84,16 +86,16 @@ static void ArraySortIndicesStringsBenchmark(benchmark::State& state) {
 }
 
 static void ArraySortIndicesStringsDictBenchmark(benchmark::State& state) {
-  RegressionArgs args(state);
+  RegressionArgs args(state, /*size_is_bytes=*/false);
 
+  // XXX Same as in ArraySortIndicesStringsBenchmark above
   const int64_t array_size = args.size / sizeof(int64_t);
-  const int32_t dict_values_size = 20;  // a typical dictionary size
 
   auto rand = random::RandomArrayGenerator(kSeed);
-  auto dict_values = rand.String(dict_values_size, /*min_length=*/3, /*max_length=*/25,
+  auto dict_values = rand.String(kDictionarySize, /*min_length=*/3, /*max_length=*/25,
                                  args.null_proportion / 2);
   auto dict_indices =
-      rand.Int64(array_size, 0, dict_values_size - 1, args.null_proportion / 2);
+      rand.Int64(array_size, 0, kDictionarySize - 1, args.null_proportion / 2);
   auto dict_array = *DictionaryArray::FromArrays(dict_indices, dict_values);
 
   ArraySortIndicesBenchmark(state, dict_array);

From 429ce815bb7074fadcc6ccf42edf4ae5ec75b311 Mon Sep 17 00:00:00 2001
From: Antoine Pitrou <antoine@python.org>
Date: Thu, 21 Jul 2022 17:57:02 +0200
Subject: [PATCH 08/11] Some nits

---
 .../compute/kernels/vector_array_sort.cc      | 49 ++++++++++---------
 1 file changed, 27 insertions(+), 22 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index 4def9382e1f..02f65b615b1 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -177,9 +177,8 @@ template <>
 class ArrayCompareSorter<DictionaryType> {
   struct DictionaryInternal {
     NullPartitionResult* p;
-    const std::shared_ptr<Array>& values;
-    const std::shared_ptr<Array>& indices;
-    const UInt64Array& values_rank;
+    const std::shared_ptr<Array>& dict_indices;
+    const UInt64Array& rank_array;
     const ArraySortOptions& options;
     int64_t offset;
 
@@ -198,25 +197,20 @@ class ArrayCompareSorter<DictionaryType> {
     Status SortInternal() {
       using ArrayType = typename TypeTraits<IndexType>::ArrayType;
       using c_type = typename IndexType::c_type;
-      const c_type* indices_values =
-          checked_cast<const ArrayType&>(*indices).raw_values();
-      const uint64_t* ranks = values_rank.raw_values();
+      const c_type* indices = checked_cast<const ArrayType&>(*dict_indices).raw_values();
+      const uint64_t* ranks = rank_array.raw_values();
 
-      // TODO Instead, try to take the ranks with the dict indices,
-      // then run a sort on the results.  Since the ranks are dense,
-      // a much faster counting sort may be used.
       std::stable_sort(this->p->non_nulls_begin, this->p->non_nulls_end,
                        [&](uint64_t left, uint64_t right) {
-                         const uint64_t left_rank = ranks[indices_values[left - offset]];
-                         const uint64_t right_rank =
-                             ranks[indices_values[right - offset]];
+                         const uint64_t left_rank = ranks[indices[left - offset]];
+                         const uint64_t right_rank = ranks[indices[right - offset]];
                          return left_rank < right_rank;
                        });
 
       return Status::OK();
     }
 
-    void Make(const std::shared_ptr<DataType>& index_type) {
+    void Sort(const std::shared_ptr<DataType>& index_type) {
       ARROW_CHECK_OK(VisitTypeInline(*index_type, this));
     }
   };
@@ -226,22 +220,33 @@ class ArrayCompareSorter<DictionaryType> {
                                  const Array& array, int64_t offset,
                                  const ArraySortOptions& options) {
     const auto& dict_array = checked_cast<const DictionaryArray&>(array);
-    const auto& values = dict_array.dictionary();
-    const auto& indices = dict_array.indices();
+    const auto& dict_values = dict_array.dictionary();
+    const auto& dict_indices = dict_array.indices();
     const auto& index_type = dict_array.dict_type()->index_type();
 
     NullPartitionResult p = PartitionNulls<DictionaryArray, StablePartitioner>(
         indices_begin, indices_end, dict_array, offset, options.null_placement);
 
+    // TODO should be able to execute rank() with the caller's KernelContext
     RankOptions rank_options(options.order, options.null_placement, RankOptions::Dense);
     // FIXME propagate instead of aborting on error
-    auto rank = CallFunction("rank", {values}, &rank_options).ValueOrDie().make_array();
-
-    const auto& values_rank = checked_cast<const UInt64Array&>(*rank);
-
-    DictionaryInternal visitor = {&p, values, indices, values_rank, options, offset};
-    visitor.Make(index_type);
-
+    auto rank =
+        CallFunction("rank", {dict_values}, &rank_options).ValueOrDie().make_array();
+    DCHECK_EQ(rank->null_count(), 0);
+
+    // TODO Instead, try to take the ranks with the dict indices,
+    // then run a sort on the results.  Since the ranks are dense,
+    // a much faster counting sort may be used.
+    //
+    // However, this will first require a rank() variant that emits null
+    // for null inputs (instead of assigning them a non-null rank number).
+
+    const auto& rank_array = checked_cast<const UInt64Array&>(*rank);
+    DictionaryInternal visitor = {&p, dict_indices, rank_array, options, offset};
+    visitor.Sort(index_type);
+
+    // FIXME this is not correct, as it doesn't account for nulls in the
+    // dictionary values.
     return p;
   }
 };

From 7fa1c005540d5b393f6324309e521c5317c56afb Mon Sep 17 00:00:00 2001
From: Antoine Pitrou <antoine@python.org>
Date: Thu, 21 Jul 2022 18:42:03 +0200
Subject: [PATCH 09/11] Use the faster algorithm (rank-then-take-then-sort)

---
 .../compute/kernels/vector_array_sort.cc      | 73 +++++++++++++------
 .../compute/kernels/vector_sort_benchmark.cc  | 12 ++-
 .../arrow/compute/kernels/vector_sort_test.cc |  1 +
 cpp/src/arrow/util/benchmark_util.h           |  2 +-
 4 files changed, 62 insertions(+), 26 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index 02f65b615b1..253820073be 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -215,6 +215,37 @@ class ArrayCompareSorter<DictionaryType> {
     }
   };
 
+  Result<std::shared_ptr<Array>> RanksWithNulls(const std::shared_ptr<Array>& array) {
+    // Notes:
+    // * The order is always ascending here, since the goal is to produce
+    //   an exactly-equivalent-order of the dictionary values.
+    // * We're going to re-emit nulls in the output, so we can just always consider
+    //   them "at the end".  Note that choosing AtStart would merely shift other
+    //   ranks by 1 if there are any nulls...
+    RankOptions rank_options(SortOrder::Ascending, NullPlacement::AtEnd,
+                             RankOptions::Dense);
+
+    // XXX Should this support Type::NA?
+    auto data = array->data();
+    std::shared_ptr<Buffer> null_bitmap;
+    if (array->null_count() > 0) {
+      null_bitmap = array->null_bitmap();
+      data = array->data()->Copy();
+      data->buffers[0] = nullptr;
+      data->null_count = 0;
+      DCHECK_EQ(data->offset, 0);  // FIXME
+    }
+    ARROW_ASSIGN_OR_RAISE(auto rank_datum, CallFunction("rank", {array}, &rank_options));
+    auto rank_data = rank_datum.array();
+    DCHECK_EQ(rank_data->GetNullCount(), 0);
+    // If there were nulls in the input, paste them in the output
+    if (null_bitmap) {
+      rank_data->buffers[0] = std::move(null_bitmap);
+      rank_data->null_count = array->null_count();
+    }
+    return MakeArray(rank_data);
+  }
+
  public:
   NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
                                  const Array& array, int64_t offset,
@@ -222,32 +253,32 @@ class ArrayCompareSorter<DictionaryType> {
     const auto& dict_array = checked_cast<const DictionaryArray&>(array);
     const auto& dict_values = dict_array.dictionary();
     const auto& dict_indices = dict_array.indices();
-    const auto& index_type = dict_array.dict_type()->index_type();
 
-    NullPartitionResult p = PartitionNulls<DictionaryArray, StablePartitioner>(
-        indices_begin, indices_end, dict_array, offset, options.null_placement);
+    // Algorithm:
+    // 1) Use the Rank function to get an exactly-equivalent-order array
+    //    of the dictionary values, but with a datatype that's friendlier to
+    //    sorting (uint64).
+    // 2) Act as if we were sorting a dictionary array with the same indices,
+    //    but with the ranks as dictionary values.
+    // 2a) Dictionary-decode the ranks by calling Take.
+    // 2b) Sort the decoded ranks. Not only those are uint64, they are dense
+    //     in a [0, k) range where k is the number of unique dictionary values.
+    //     Therefore, unless the dictionary is very large, a fast counting sort
+    //     will be used.
+    //
+    // The bottom line is that performance will usually be much better
+    // (potentially an order of magnitude faster) than by naively decoding
+    // the original dictionary and sorting the decoded version.
 
-    // TODO should be able to execute rank() with the caller's KernelContext
-    RankOptions rank_options(options.order, options.null_placement, RankOptions::Dense);
-    // FIXME propagate instead of aborting on error
-    auto rank =
-        CallFunction("rank", {dict_values}, &rank_options).ValueOrDie().make_array();
-    DCHECK_EQ(rank->null_count(), 0);
+    // FIXME Should be able to use the caller's KernelContext for rank() and take()
 
-    // TODO Instead, try to take the ranks with the dict indices,
-    // then run a sort on the results.  Since the ranks are dense,
-    // a much faster counting sort may be used.
-    //
-    // However, this will first require a rank() variant that emits null
-    // for null inputs (instead of assigning them a non-null rank number).
+    // FIXME Propagate errors instead of aborting
+    auto ranks = *RanksWithNulls(dict_values);
 
-    const auto& rank_array = checked_cast<const UInt64Array&>(*rank);
-    DictionaryInternal visitor = {&p, dict_indices, rank_array, options, offset};
-    visitor.Sort(index_type);
+    auto decoded_ranks = *Take(*ranks, *dict_indices);
 
-    // FIXME this is not correct, as it doesn't account for nulls in the
-    // dictionary values.
-    return p;
+    auto rank_sorter = *GetArraySorter(*decoded_ranks->type() /* should be uint64 */);
+    return rank_sorter(indices_begin, indices_end, *decoded_ranks, offset, options);
   }
 };
 
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc
index c0ab70a486d..e3604599ee5 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc
@@ -36,6 +36,8 @@ static void ArraySortIndicesBenchmark(benchmark::State& state,
   for (auto _ : state) {
     ABORT_NOT_OK(SortIndices(*values).status());
   }
+  // This may be redundant with the SetItemsProcessed() call in RegressionArgs,
+  // if size_is_bytes was false.
   state.SetItemsProcessed(state.iterations() * values->length());
 }
 
@@ -75,9 +77,11 @@ static void ArraySortIndicesInt64DictBenchmark(benchmark::State& state, int64_t
 
 static void ArraySortIndicesStringsBenchmark(benchmark::State& state) {
   RegressionArgs args(state, /*size_is_bytes=*/false);
-
   // XXX This division to make numbers comparable with ArraySortIndicesInt64Benchmark
-  const int64_t array_size = args.size / sizeof(int64_t);
+  // (including the SetItemsProcessed() call in the RegressionArgs destructor)
+  args.size /= sizeof(int64_t);
+
+  const int64_t array_size = args.size;
   auto rand = random::RandomArrayGenerator(kSeed);
   auto values =
       rand.String(array_size, /*min_length=*/3, /*max_length=*/25, args.null_proportion);
@@ -87,10 +91,10 @@ static void ArraySortIndicesStringsBenchmark(benchmark::State& state) {
 
 static void ArraySortIndicesStringsDictBenchmark(benchmark::State& state) {
   RegressionArgs args(state, /*size_is_bytes=*/false);
-
   // XXX Same as in ArraySortIndicesStringsBenchmark above
-  const int64_t array_size = args.size / sizeof(int64_t);
+  args.size /= sizeof(int64_t);
 
+  const int64_t array_size = args.size;
   auto rand = random::RandomArrayGenerator(kSeed);
   auto dict_values = rand.String(kDictionarySize, /*min_length=*/3, /*max_length=*/25,
                                  args.null_proportion / 2);
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
index f1808fae980..cc9000e3aa0 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
@@ -373,6 +373,7 @@ TEST(ArraySortIndicesFunction, DictionaryArray) {
                                       "[ \"b\", null, \"c\", \"a\", \"c\"]");
     for (auto order : AllOrders()) {
       for (auto null_placement : AllNullPlacements()) {
+        ARROW_SCOPED_TRACE("i = ", i);
         ArraySortOptions options{order, null_placement};
         auto expected = ArrayFromJSON(uint64(), expected_str[i++]);
         ASSERT_OK_AND_ASSIGN(auto actual,
diff --git a/cpp/src/arrow/util/benchmark_util.h b/cpp/src/arrow/util/benchmark_util.h
index 79484989ac5..15f2ea94fe4 100644
--- a/cpp/src/arrow/util/benchmark_util.h
+++ b/cpp/src/arrow/util/benchmark_util.h
@@ -104,7 +104,7 @@ void RegressionSetArgs(benchmark::internal::Benchmark* bench) {
 // RAII struct to handle some of the boilerplate in regression benchmarks
 struct RegressionArgs {
   // size of memory tested (per iteration) in bytes
-  const int64_t size;
+  int64_t size;
 
   // proportion of nulls in generated arrays
   double null_proportion;

From 43ad1344bd3e246c16a02d8d2a90a1429a3cc1c9 Mon Sep 17 00:00:00 2001
From: Antoine Pitrou <antoine@python.org>
Date: Thu, 21 Jul 2022 18:56:05 +0200
Subject: [PATCH 10/11] Remove unused code

---
 .../compute/kernels/vector_array_sort.cc      | 101 ++++++------------
 1 file changed, 31 insertions(+), 70 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index 253820073be..a010ad77069 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -175,46 +175,42 @@ class ArrayCompareSorter {
 
 template <>
 class ArrayCompareSorter<DictionaryType> {
-  struct DictionaryInternal {
-    NullPartitionResult* p;
-    const std::shared_ptr<Array>& dict_indices;
-    const UInt64Array& rank_array;
-    const ArraySortOptions& options;
-    int64_t offset;
-
-    Status Visit(const DataType& index_type) {
-      return Status::TypeError("Dictionary sorting not supported for index type ",
-                               index_type.ToString());
-    }
+ public:
+  NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                 const Array& array, int64_t offset,
+                                 const ArraySortOptions& options) {
+    const auto& dict_array = checked_cast<const DictionaryArray&>(array);
+    const auto& dict_values = dict_array.dictionary();
+    const auto& dict_indices = dict_array.indices();
 
-    template <typename IndexType>
-    enable_if_t<is_integer_type<IndexType>::value, Status> Visit(
-        const IndexType& index_type) {
-      return SortInternal<IndexType>();
-    }
+    // Algorithm:
+    // 1) Use the Rank function to get an exactly-equivalent-order array
+    //    of the dictionary values, but with a datatype that's friendlier to
+    //    sorting (uint64).
+    // 2) Act as if we were sorting a dictionary array with the same indices,
+    //    but with the ranks as dictionary values.
+    // 2a) Dictionary-decode the ranks by calling Take.
+    // 2b) Sort the decoded ranks. Not only those are uint64, they are dense
+    //     in a [0, k) range where k is the number of unique dictionary values.
+    //     Therefore, unless the dictionary is very large, a fast counting sort
+    //     will be used.
+    //
+    // The bottom line is that performance will usually be much better
+    // (potentially an order of magnitude faster) than by naively decoding
+    // the original dictionary and sorting the decoded version.
 
-    template <typename IndexType>
-    Status SortInternal() {
-      using ArrayType = typename TypeTraits<IndexType>::ArrayType;
-      using c_type = typename IndexType::c_type;
-      const c_type* indices = checked_cast<const ArrayType&>(*dict_indices).raw_values();
-      const uint64_t* ranks = rank_array.raw_values();
-
-      std::stable_sort(this->p->non_nulls_begin, this->p->non_nulls_end,
-                       [&](uint64_t left, uint64_t right) {
-                         const uint64_t left_rank = ranks[indices[left - offset]];
-                         const uint64_t right_rank = ranks[indices[right - offset]];
-                         return left_rank < right_rank;
-                       });
+    // FIXME Should be able to use the caller's KernelContext for rank() and take()
 
-      return Status::OK();
-    }
+    // FIXME Propagate errors instead of aborting
+    auto ranks = *RanksWithNulls(dict_values);
 
-    void Sort(const std::shared_ptr<DataType>& index_type) {
-      ARROW_CHECK_OK(VisitTypeInline(*index_type, this));
-    }
-  };
+    auto decoded_ranks = *Take(*ranks, *dict_indices);
+
+    auto rank_sorter = *GetArraySorter(*decoded_ranks->type() /* should be uint64 */);
+    return rank_sorter(indices_begin, indices_end, *decoded_ranks, offset, options);
+  }
 
+ private:
   Result<std::shared_ptr<Array>> RanksWithNulls(const std::shared_ptr<Array>& array) {
     // Notes:
     // * The order is always ascending here, since the goal is to produce
@@ -245,41 +241,6 @@ class ArrayCompareSorter<DictionaryType> {
     }
     return MakeArray(rank_data);
   }
-
- public:
-  NullPartitionResult operator()(uint64_t* indices_begin, uint64_t* indices_end,
-                                 const Array& array, int64_t offset,
-                                 const ArraySortOptions& options) {
-    const auto& dict_array = checked_cast<const DictionaryArray&>(array);
-    const auto& dict_values = dict_array.dictionary();
-    const auto& dict_indices = dict_array.indices();
-
-    // Algorithm:
-    // 1) Use the Rank function to get an exactly-equivalent-order array
-    //    of the dictionary values, but with a datatype that's friendlier to
-    //    sorting (uint64).
-    // 2) Act as if we were sorting a dictionary array with the same indices,
-    //    but with the ranks as dictionary values.
-    // 2a) Dictionary-decode the ranks by calling Take.
-    // 2b) Sort the decoded ranks. Not only those are uint64, they are dense
-    //     in a [0, k) range where k is the number of unique dictionary values.
-    //     Therefore, unless the dictionary is very large, a fast counting sort
-    //     will be used.
-    //
-    // The bottom line is that performance will usually be much better
-    // (potentially an order of magnitude faster) than by naively decoding
-    // the original dictionary and sorting the decoded version.
-
-    // FIXME Should be able to use the caller's KernelContext for rank() and take()
-
-    // FIXME Propagate errors instead of aborting
-    auto ranks = *RanksWithNulls(dict_values);
-
-    auto decoded_ranks = *Take(*ranks, *dict_indices);
-
-    auto rank_sorter = *GetArraySorter(*decoded_ranks->type() /* should be uint64 */);
-    return rank_sorter(indices_begin, indices_end, *decoded_ranks, offset, options);
-  }
 };
 
 template <typename ArrowType>

From bae13aa284c5e46bf05707a827349dc53928e71d Mon Sep 17 00:00:00 2001
From: Antoine Pitrou <antoine@python.org>
Date: Thu, 21 Jul 2022 19:03:39 +0200
Subject: [PATCH 11/11] More TODOs

---
 cpp/src/arrow/compute/kernels/vector_array_sort.cc | 2 ++
 cpp/src/arrow/compute/kernels/vector_sort_test.cc  | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index a010ad77069..8d55c645084 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -199,6 +199,8 @@ class ArrayCompareSorter<DictionaryType> {
     // (potentially an order of magnitude faster) than by naively decoding
     // the original dictionary and sorting the decoded version.
 
+    // TODO special-case all-nulls arrays to avoid ranking and decoding them?
+
     // FIXME Should be able to use the caller's KernelContext for rank() and take()
 
     // FIXME Propagate errors instead of aborting
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
index cc9000e3aa0..7a1669c401e 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
@@ -421,6 +421,8 @@ TEST(ArraySortIndicesFunction, RandomDictionaryArray) {
       AssertDatumsEqual(expected, actual, /*verbose=*/true);
     }
   }
+
+  // TODO test with sliced dict indices and values...
 }
 
 TEST(ArraySortIndicesFunction, ChunkedArray) {