From ce5953331363ee302f60161d02d45eb42aafc7be 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] Changes from original PR

Draft sort indices on array dictionary

Update sort indices

Update sorter

Address feedback

Address feedback

Add tests and benchmarks

Some nits re:benchmarking

Some nits

Use the faster algorithm (rank-then-take-then-sort)

Remove unused code

More TODOs
---
 cpp/src/arrow/array/array_dict.h              |  4 +
 .../compute/kernels/vector_array_sort.cc      | 83 ++++++++++++++++++-
 cpp/src/arrow/compute/kernels/vector_sort.cc  |  8 +-
 .../compute/kernels/vector_sort_benchmark.cc  | 70 ++++++++++++++++
 .../compute/kernels/vector_sort_internal.h    |  8 +-
 .../arrow/compute/kernels/vector_sort_test.cc | 71 ++++++++++++++++
 cpp/src/arrow/util/benchmark_util.h           |  2 +-
 7 files changed, 236 insertions(+), 10 deletions(-)

diff --git a/cpp/src/arrow/array/array_dict.h b/cpp/src/arrow/array/array_dict.h
index 8791eaa07db..676ab87b8f3 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 ccf691669b8..0d5104e695e 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -173,6 +173,78 @@ 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& 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.
+
+    // 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
+    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);
+  }
+
+ 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
+    //   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);
+  }
+};
+
 template <typename ArrowType>
 class ArrayCountSorter {
   using ArrayType = typename TypeTraits<ArrowType>::ArrayType;
@@ -405,7 +477,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;
 };
 
@@ -508,6 +581,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({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;
@@ -562,6 +642,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 110abd27d8f..a9247d27b4c 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort.cc
@@ -259,8 +259,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>(
@@ -517,8 +517,8 @@ class MultipleKeyRecordBatchSorter : public TypeVisitor {
     const ArrayType& array =
         ::arrow::internal::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_benchmark.cc b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc
index 67f03c1eb6c..a4c192e359e 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
 
 //
 // Array sort/rank benchmark helpers
@@ -67,6 +69,8 @@ static void ArraySortFuncBenchmark(benchmark::State& state, const Runner& runner
   for (auto _ : state) {
     ABORT_NOT_OK(runner(values));
   }
+  // This may be redundant with the SetItemsProcessed() call in RegressionArgs,
+  // if size_is_bytes was false.
   state.SetItemsProcessed(state.iterations() * values->length());
 }
 
@@ -82,6 +86,55 @@ static void ArraySortFuncInt64Benchmark(benchmark::State& state, const Runner& r
   ArraySortFuncBenchmark(state, runner, values);
 }
 
+template <typename Runner>
+static void ArraySortFuncInt64DictBenchmark(benchmark::State& state, const Runner& runner,
+                                            int64_t min, int64_t max) {
+  RegressionArgs args(state);
+
+  const int64_t array_size = args.size / sizeof(int64_t);
+
+  auto rand = random::RandomArrayGenerator(kSeed);
+  auto dict_values = rand.Int64(kDictionarySize, min, max, args.null_proportion / 2);
+  auto dict_indices =
+      rand.Int64(array_size, 0, kDictionarySize - 1, args.null_proportion / 2);
+  auto dict_array = *DictionaryArray::FromArrays(dict_indices, dict_values);
+
+  ArraySortFuncBenchmark(state, runner, dict_array);
+}
+
+template <typename Runner>
+static void ArraySortFuncStringBenchmark(benchmark::State& state, const Runner& runner) {
+  RegressionArgs args(state, /*size_is_bytes=*/false);
+  // XXX This division to make numbers comparable with ArraySortFuncInt64Benchmark
+  // (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);
+
+  ArraySortFuncBenchmark(state, runner, values);
+}
+
+template <typename Runner>
+static void ArraySortFuncStringDictBenchmark(benchmark::State& state,
+                                             const Runner& runner) {
+  RegressionArgs args(state, /*size_is_bytes=*/false);
+  // XXX Same as in ArraySortFuncStringBenchmark above
+  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);
+  auto dict_indices =
+      rand.Int64(array_size, 0, kDictionarySize - 1, args.null_proportion / 2);
+  auto dict_array = *DictionaryArray::FromArrays(dict_indices, dict_values);
+
+  ArraySortFuncBenchmark(state, runner, dict_array);
+}
+
 template <typename Runner>
 static void ChunkedArraySortFuncInt64Benchmark(benchmark::State& state,
                                                const Runner& runner, int64_t min,
@@ -129,6 +182,20 @@ static void ChunkedArraySortFuncStringBenchmark(benchmark::State& state,
   ArraySortFuncBenchmark(state, runner, std::make_shared<ChunkedArray>(chunks));
 }
 
+static void ArraySortIndicesInt64Dict(benchmark::State& state) {
+  const auto min = std::numeric_limits<int64_t>::min();
+  const auto max = std::numeric_limits<int64_t>::max();
+  ArraySortFuncInt64DictBenchmark(state, SortRunner(state), min, max);
+}
+
+static void ArraySortIndicesString(benchmark::State& state) {
+  ArraySortFuncStringBenchmark(state, SortRunner(state));
+}
+
+static void ArraySortIndicesStringDict(benchmark::State& state) {
+  ArraySortFuncStringDictBenchmark(state, SortRunner(state));
+}
+
 template <typename Runner>
 static void ArraySortFuncBoolBenchmark(benchmark::State& state, const Runner& runner) {
   RegressionArgs args(state);
@@ -389,6 +456,9 @@ BENCHMARK(ArraySortIndicesInt64Wide)->Apply(ArraySortIndicesSetArgs);
 BENCHMARK(ArraySortIndicesBool)->Apply(ArraySortIndicesSetArgs);
 BENCHMARK(ArraySortIndicesStringNarrow)->Apply(ArraySortIndicesSetArgs);
 BENCHMARK(ArraySortIndicesStringWide)->Apply(ArraySortIndicesSetArgs);
+BENCHMARK(ArraySortIndicesInt64Dict)->Apply(ArraySortIndicesSetArgs);
+BENCHMARK(ArraySortIndicesString)->Apply(ArraySortIndicesSetArgs);
+BENCHMARK(ArraySortIndicesStringDict)->Apply(ArraySortIndicesSetArgs);
 
 BENCHMARK(ChunkedArraySortIndicesInt64Narrow)->Apply(ArraySortIndicesSetArgs);
 BENCHMARK(ChunkedArraySortIndicesInt64Wide)->Apply(ArraySortIndicesSetArgs);
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_internal.h b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
index 4645c593901..a51540f9865 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_internal.h
+++ b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
@@ -192,9 +192,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);
@@ -255,8 +255,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 d6b10467631..08c8357a7c9 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
@@ -368,6 +368,77 @@ TEST(ArraySortIndicesFunction, Array) {
   AssertDatumsEqual(expected, actual, /*verbose=*/true);
 }
 
+TEST(ArraySortIndicesFunction, DictionaryArray) {
+  // 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
+      "[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
+  };
+
+  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]",
+                                      "[ \"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,
+                             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);
+    }
+  }
+
+  // TODO test with sliced dict indices and values...
+}
+
 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]"});
diff --git a/cpp/src/arrow/util/benchmark_util.h b/cpp/src/arrow/util/benchmark_util.h
index 319730f5bad..2a3dcf56f88 100644
--- a/cpp/src/arrow/util/benchmark_util.h
+++ b/cpp/src/arrow/util/benchmark_util.h
@@ -108,7 +108,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 be428879967d7155bdef1e53d2a698fc24862c2e Mon Sep 17 00:00:00 2001
From: benibus <bpharks@gmx.com>
Date: Fri, 21 Apr 2023 13:09:52 -0400
Subject: [PATCH 02/11] Address internal API FIXMEs

Alters ArraySortFunc to allow for error propagation and passing through
the caller's ExecContext to use in additional kernels
---
 .../compute/kernels/vector_array_sort.cc      | 66 ++++++++++---------
 cpp/src/arrow/compute/kernels/vector_rank.cc  |  5 +-
 cpp/src/arrow/compute/kernels/vector_sort.cc  |  6 +-
 .../compute/kernels/vector_sort_internal.h    |  4 +-
 4 files changed, 42 insertions(+), 39 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index 0d5104e695e..0f596c1fa5c 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -143,9 +143,9 @@ 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) {
+  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                         const Array& array, int64_t offset,
+                                         const ArraySortOptions& options, ExecContext*) {
     const auto& values = checked_cast<const ArrayType&>(array);
 
     const auto p = PartitionNulls<ArrayType, StablePartitioner>(
@@ -176,9 +176,10 @@ 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) {
+  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                         const Array& array, int64_t offset,
+                                         const ArraySortOptions& options,
+                                         ExecContext* ctx) {
     const auto& dict_array = checked_cast<const DictionaryArray&>(array);
     const auto& dict_values = dict_array.dictionary();
     const auto& dict_indices = dict_array.indices();
@@ -200,20 +201,19 @@ class ArrayCompareSorter<DictionaryType> {
     // the original dictionary and sorting the decoded version.
 
     // TODO special-case all-nulls arrays to avoid ranking and decoding them?
+    ARROW_ASSIGN_OR_RAISE(auto ranks, RanksWithNulls(dict_values, ctx));
 
-    // FIXME Should be able to use the caller's KernelContext for rank() and take()
+    ARROW_ASSIGN_OR_RAISE(auto decoded_ranks,
+                          Take(*ranks, *dict_indices, TakeOptions::Defaults(), ctx));
 
-    // 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);
+    ARROW_ASSIGN_OR_RAISE(auto rank_sorter,
+                          GetArraySorter(*decoded_ranks->type() /* should be uint64 */));
+    return rank_sorter(indices_begin, indices_end, *decoded_ranks, offset, options, ctx);
   }
 
  private:
-  Result<std::shared_ptr<Array>> RanksWithNulls(const std::shared_ptr<Array>& array) {
+  Result<std::shared_ptr<Array>> RanksWithNulls(const std::shared_ptr<Array>& array,
+                                                ExecContext* ctx) {
     // Notes:
     // * The order is always ascending here, since the goal is to produce
     //   an exactly-equivalent-order of the dictionary values.
@@ -233,7 +233,8 @@ class ArrayCompareSorter<DictionaryType> {
       data->null_count = 0;
       DCHECK_EQ(data->offset, 0);  // FIXME
     }
-    ARROW_ASSIGN_OR_RAISE(auto rank_datum, CallFunction("rank", {array}, &rank_options));
+    ARROW_ASSIGN_OR_RAISE(auto rank_datum,
+                          CallFunction("rank", {array}, &rank_options, ctx));
     auto rank_data = rank_datum.array();
     DCHECK_EQ(rank_data->GetNullCount(), 0);
     // If there were nulls in the input, paste them in the output
@@ -261,9 +262,10 @@ 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,
+                                         ExecContext*) const {
     const auto& values = checked_cast<const ArrayType&>(array);
 
     // 32bit counter performs much better than 64bit one
@@ -345,9 +347,9 @@ 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) {
+  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                         const Array& array, int64_t offset,
+                                         const ArraySortOptions& options, ExecContext*) {
     const auto& values = checked_cast<const BooleanArray&>(array);
 
     std::array<int64_t, 3> counts{0, 0, 0};  // false, true, null
@@ -390,9 +392,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) {
+  Result<NullPartitionResult> operator()(uint64_t* indices_begin, uint64_t* indices_end,
+                                         const Array& array, int64_t offset,
+                                         const ArraySortOptions& options,
+                                         ExecContext* ctx) {
     const auto& values = checked_cast<const ArrayType&>(array);
 
     if (values.length() >= countsort_min_len_ && values.length() > values.null_count()) {
@@ -404,11 +407,11 @@ class ArrayCountOrCompareSorter {
       if (static_cast<uint64_t>(max) - static_cast<uint64_t>(min) <=
           countsort_max_range_) {
         count_sorter_.SetMinMax(min, max);
-        return count_sorter_(indices_begin, indices_end, values, offset, options);
+        return count_sorter_(indices_begin, indices_end, values, offset, options, ctx);
       }
     }
 
-    return compare_sorter_(indices_begin, indices_end, values, offset, options);
+    return compare_sorter_(indices_begin, indices_end, values, offset, options, ctx);
   }
 
  private:
@@ -430,9 +433,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, ExecContext*) {
     return NullPartitionResult::NullsOnly(indices_begin, indices_end,
                                           options.null_placement);
   }
@@ -518,8 +521,7 @@ struct ArraySortIndices {
     ArrayType arr(batch[0].array.ToArrayData());
     ARROW_ASSIGN_OR_RAISE(auto sorter, GetArraySorter(*GetPhysicalType(arr.type())));
 
-    sorter(out_begin, out_end, arr, 0, options);
-    return Status::OK();
+    return sorter(out_begin, out_end, arr, 0, options, ctx->exec_context()).status();
   }
 };
 
diff --git a/cpp/src/arrow/compute/kernels/vector_rank.cc b/cpp/src/arrow/compute/kernels/vector_rank.cc
index 828316514f6..780ae25d963 100644
--- a/cpp/src/arrow/compute/kernels/vector_rank.cc
+++ b/cpp/src/arrow/compute/kernels/vector_rank.cc
@@ -202,8 +202,9 @@ class Ranker<Array> : public RankerMixin<Array, Ranker<Array>> {
     ARROW_ASSIGN_OR_RAISE(auto array_sorter, GetArraySorter(*physical_type_));
 
     ArrayType array(input_.data());
-    NullPartitionResult sorted = array_sorter(indices_begin_, indices_end_, array, 0,
-                                              ArraySortOptions(order_, null_placement_));
+    ARROW_ASSIGN_OR_RAISE(NullPartitionResult sorted,
+                          array_sorter(indices_begin_, indices_end_, array, 0,
+                                       ArraySortOptions(order_, null_placement_), ctx_));
 
     auto value_selector = [&array](int64_t index) {
       return GetView::LogicalValue(array.GetView(index));
diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc
index a9247d27b4c..abf8fc818e5 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort.cc
@@ -95,9 +95,9 @@ class ChunkedArraySorter : public TypeVisitor {
       const auto array = checked_cast<const ArrayType*>(arrays[i]);
       end_offset += array->length();
       null_count += array->null_count();
-      sorted[i] =
-          array_sorter_(indices_begin_ + begin_offset, indices_begin_ + end_offset,
-                        *array, begin_offset, options);
+      ARROW_ASSIGN_OR_RAISE(sorted[i], array_sorter_(indices_begin_ + begin_offset,
+                                                     indices_begin_ + end_offset, *array,
+                                                     begin_offset, options, ctx_));
       begin_offset = end_offset;
     }
     DCHECK_EQ(end_offset, indices_end_ - indices_begin_);
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_internal.h b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
index a51540f9865..1241b609e57 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_internal.h
+++ b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
@@ -447,9 +447,9 @@ 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)>;
+    const ArraySortOptions& options, ExecContext* ctx)>;
 
 Result<ArraySortFunc> GetArraySorter(const DataType& type);
 

From c35c61018d0e00f4a0bdd6310c5f31dfec04042e Mon Sep 17 00:00:00 2001
From: benibus <bpharks@gmx.com>
Date: Fri, 21 Apr 2023 17:19:21 -0400
Subject: [PATCH 03/11] Address FIXMEs for handling sliced arrays

---
 .../compute/kernels/vector_array_sort.cc      |  6 +-
 .../arrow/compute/kernels/vector_sort_test.cc | 70 +++++++++++++------
 2 files changed, 52 insertions(+), 24 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index 0f596c1fa5c..05eedb6cba1 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -229,9 +229,13 @@ class ArrayCompareSorter<DictionaryType> {
     if (array->null_count() > 0) {
       null_bitmap = array->null_bitmap();
       data = array->data()->Copy();
+      if (data->offset > 0) {
+        ARROW_ASSIGN_OR_RAISE(null_bitmap, arrow::internal::CopyBitmap(
+                                               ctx->memory_pool(), null_bitmap->data(),
+                                               data->offset, data->length));
+      }
       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, ctx));
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
index 08c8357a7c9..c8ae3ab2e47 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
@@ -410,33 +410,57 @@ TEST(ArraySortIndicesFunction, RandomDictionaryArray) {
   ::arrow::random::RandomArrayGenerator rng(/*seed=*/1234);
   constexpr int64_t kLength = 200;
   constexpr int64_t kDictLength = 20;
+  // Lengths of source arrays which each dictionary array's index and value arrays are
+  // sliced from
+  constexpr int64_t kNumIndices = kLength * 4;
+  constexpr int64_t kNumValues = kDictLength * 8;
+
+  // Ensure there are duplicates in the dictionary and the indices, and nulls in both as
+  // well.
+  auto source_values = rng.StringWithRepeats(kNumValues, /*unique=*/kNumValues / 2,
+                                             /*min_length=*/1, /*max_length=*/10,
+                                             /*null_probability=*/0.2);
+  auto source_indices =
+      rng.Int64(kNumIndices, 0, kDictLength - 1, /*null_probability = */ 0.2);
+
+  auto get_test_array = [&](int64_t indices_offset,
+                            int64_t values_offset) -> std::shared_ptr<Array> {
+    auto out = *DictionaryArray::FromArrays(
+        *source_indices->SliceSafe(indices_offset, kLength),
+        *source_values->SliceSafe(values_offset, kDictLength));
+    ARROW_CHECK_EQ(out->length(), kLength);
+    ARROW_CHECK_EQ(checked_cast<const DictionaryArray&>(*out).dictionary()->length(),
+                   kDictLength);
+    return out;
+  };
 
-  // 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);
+  // Test cases are sliced from the beginning/middle/end
+  ArrayVector dict_arrays;
+  dict_arrays.push_back(get_test_array(0, 0));
+  dict_arrays.push_back(get_test_array(kNumIndices / 3, kNumValues / 3));
+  dict_arrays.push_back(get_test_array(kNumIndices - kLength, kNumValues - kDictLength));
 
-      // 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);
+  for (size_t i = 0; i < dict_arrays.size(); ++i) {
+    ARROW_SCOPED_TRACE("i = ", i);
+    auto dict_array = dict_arrays[i];
+
+    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);
+      }
     }
   }
-
-  // TODO test with sliced dict indices and values...
 }
 
 TEST(ArraySortIndicesFunction, ChunkedArray) {

From 8b3a5bc2a4cb6903a6d3d83c37ca124d50bff9e9 Mon Sep 17 00:00:00 2001
From: benibus <bpharks@gmx.com>
Date: Mon, 24 Apr 2023 13:58:37 -0400
Subject: [PATCH 04/11] Special handling for all-nulls case

---
 cpp/src/arrow/array/array_dict.h              |  4 --
 .../compute/kernels/vector_array_sort.cc      | 53 +++++++++++++++----
 .../arrow/compute/kernels/vector_sort_test.cc | 28 ++++++++++
 3 files changed, 71 insertions(+), 14 deletions(-)

diff --git a/cpp/src/arrow/array/array_dict.h b/cpp/src/arrow/array/array_dict.h
index 676ab87b8f3..8791eaa07db 100644
--- a/cpp/src/arrow/array/array_dict.h
+++ b/cpp/src/arrow/array/array_dict.h
@@ -111,10 +111,6 @@ 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 05eedb6cba1..b43935bf5e8 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -181,8 +181,9 @@ class ArrayCompareSorter<DictionaryType> {
                                          const ArraySortOptions& options,
                                          ExecContext* ctx) {
     const auto& dict_array = checked_cast<const DictionaryArray&>(array);
-    const auto& dict_values = dict_array.dictionary();
-    const auto& dict_indices = dict_array.indices();
+    // TODO: These methods should probably return a const&? They seem capable.
+    auto dict_values = dict_array.dictionary();
+    auto dict_indices = dict_array.indices();
 
     // Algorithm:
     // 1) Use the Rank function to get an exactly-equivalent-order array
@@ -200,20 +201,52 @@ 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?
-    ARROW_ASSIGN_OR_RAISE(auto ranks, RanksWithNulls(dict_values, ctx));
+    std::shared_ptr<Array> decoded_ranks;
+    // Skip the rank/take steps for cases with only nulls
+    if (IsAllNulls(dict_array)) {
+      ARROW_ASSIGN_OR_RAISE(decoded_ranks,
+                            MakeNullUInt64Array(dict_array.length(), ctx->memory_pool()));
+    } else {
+      ARROW_ASSIGN_OR_RAISE(auto ranks, RanksWithNulls(dict_values, ctx));
+
+      ARROW_ASSIGN_OR_RAISE(decoded_ranks,
+                            Take(*ranks, *dict_indices, TakeOptions::Defaults(), ctx));
+    }
 
-    ARROW_ASSIGN_OR_RAISE(auto decoded_ranks,
-                          Take(*ranks, *dict_indices, TakeOptions::Defaults(), ctx));
+    DCHECK_EQ(decoded_ranks->type_id(), Type::UINT64);
+    DCHECK_EQ(decoded_ranks->length(), dict_array.length());
+    ARROW_ASSIGN_OR_RAISE(auto rank_sorter, GetArraySorter(*decoded_ranks->type()));
 
-    ARROW_ASSIGN_OR_RAISE(auto rank_sorter,
-                          GetArraySorter(*decoded_ranks->type() /* should be uint64 */));
     return rank_sorter(indices_begin, indices_end, *decoded_ranks, offset, options, ctx);
   }
 
  private:
-  Result<std::shared_ptr<Array>> RanksWithNulls(const std::shared_ptr<Array>& array,
-                                                ExecContext* ctx) {
+  static bool IsAllNulls(const DictionaryArray& dict_array) {
+    auto indices = dict_array.indices();
+    auto values = dict_array.dictionary();
+    if (values->null_count() == values->length() ||
+        indices->null_count() == indices->length()) {
+      return true;
+    }
+    for (int64_t i = 0; i < dict_array.length(); ++i) {
+      // TODO: Special handling for dictionary types in Array::IsValid/Null?
+      if (indices->IsValid(i) && values->IsValid(dict_array.GetValueIndex(i))) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  static Result<std::shared_ptr<Array>> MakeNullUInt64Array(int64_t length,
+                                                            MemoryPool* pool) {
+    ARROW_ASSIGN_OR_RAISE(auto bitmap, AllocateEmptyBitmap(length, pool));
+    auto data = ArrayData::Make(uint64(), length, {std::move(bitmap), nullptr},
+                                /*null_count=*/length);
+    return MakeArray(data);
+  }
+
+  static Result<std::shared_ptr<Array>> RanksWithNulls(
+      const std::shared_ptr<Array>& array, ExecContext* ctx) {
     // Notes:
     // * The order is always ascending here, since the goal is to produce
     //   an exactly-equivalent-order of the dictionary values.
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
index c8ae3ab2e47..9a90bbed49e 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
@@ -397,6 +397,34 @@ TEST(ArraySortIndicesFunction, DictionaryArray) {
       }
     }
   }
+
+  auto expected_all_nulls = ArrayFromJSON(uint64(), "[0, 1, 2, 3]");
+  auto expected_empty = ArrayFromJSON(uint64(), "[]");
+
+  for (auto index_type : all_dictionary_index_types()) {
+    ARROW_SCOPED_TRACE("index_type = ", index_type->ToString());
+
+    auto dict_type = dictionary(index_type, utf8());
+    auto dict_arr_all_nulls =
+        DictArrayFromJSON(dict_type, "[null, 3, null, 1]", "[\"b\", null, \"a\", null]");
+    auto dict_arr_empty = DictArrayFromJSON(dict_type, "[]", "[\"b\", \"a\"]");
+
+    for (auto order : AllOrders()) {
+      for (auto null_placement : AllNullPlacements()) {
+        ArraySortOptions options{order, null_placement};
+
+        Datum actual;
+        ASSERT_OK_AND_ASSIGN(
+            actual, CallFunction("array_sort_indices", {dict_arr_all_nulls}, &options));
+        ValidateOutput(actual);
+        AssertDatumsEqual(expected_all_nulls, actual, /*verbose=*/true);
+        ASSERT_OK_AND_ASSIGN(
+            actual, CallFunction("array_sort_indices", {dict_arr_empty}, &options));
+        ValidateOutput(actual);
+        AssertDatumsEqual(expected_empty, actual, /*verbose=*/true);
+      }
+    }
+  }
 }
 
 Result<std::shared_ptr<Array>> DecodeDictionary(const Array& array) {

From 7e2e456ae3dfc3dafcba6d9c87aaa3f75e1b77ae Mon Sep 17 00:00:00 2001
From: Ben Harkins <60872452+benibus@users.noreply.github.com>
Date: Mon, 1 May 2023 15:52:41 -0400
Subject: [PATCH 05/11] Apply suggestion from code review

Co-authored-by: Weston Pace <weston.pace@gmail.com>
---
 cpp/src/arrow/compute/kernels/vector_sort_test.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
index 9a90bbed49e..7adc17deca3 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
@@ -379,7 +379,7 @@ TEST(ArraySortIndicesFunction, DictionaryArray) {
       "[2, 3, 4, 1, 6, 7, 0, 5, 8]"   // SortOrder::Descending NullPlacement::AtStart
   };
 
-  for (auto index_type : all_dictionary_index_types()) {
+  for (const 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()),

From 82d56dd16b9123a8d506858dde8eb87f59d98928 Mon Sep 17 00:00:00 2001
From: benibus <bpharks@gmx.com>
Date: Wed, 3 May 2023 12:08:41 -0400
Subject: [PATCH 06/11] Address review points

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

diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index b43935bf5e8..199d2182f46 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -204,8 +204,8 @@ class ArrayCompareSorter<DictionaryType> {
     std::shared_ptr<Array> decoded_ranks;
     // Skip the rank/take steps for cases with only nulls
     if (IsAllNulls(dict_array)) {
-      ARROW_ASSIGN_OR_RAISE(decoded_ranks,
-                            MakeNullUInt64Array(dict_array.length(), ctx->memory_pool()));
+      ARROW_ASSIGN_OR_RAISE(decoded_ranks, MakeArrayOfNull(uint64(), dict_array.length(),
+                                                           ctx->memory_pool()));
     } else {
       ARROW_ASSIGN_OR_RAISE(auto ranks, RanksWithNulls(dict_values, ctx));
 
@@ -237,14 +237,6 @@ class ArrayCompareSorter<DictionaryType> {
     return true;
   }
 
-  static Result<std::shared_ptr<Array>> MakeNullUInt64Array(int64_t length,
-                                                            MemoryPool* pool) {
-    ARROW_ASSIGN_OR_RAISE(auto bitmap, AllocateEmptyBitmap(length, pool));
-    auto data = ArrayData::Make(uint64(), length, {std::move(bitmap), nullptr},
-                                /*null_count=*/length);
-    return MakeArray(data);
-  }
-
   static Result<std::shared_ptr<Array>> RanksWithNulls(
       const std::shared_ptr<Array>& array, ExecContext* ctx) {
     // Notes:
@@ -271,7 +263,7 @@ class ArrayCompareSorter<DictionaryType> {
       data->null_count = 0;
     }
     ARROW_ASSIGN_OR_RAISE(auto rank_datum,
-                          CallFunction("rank", {array}, &rank_options, ctx));
+                          CallFunction("rank", {std::move(data)}, &rank_options, ctx));
     auto rank_data = rank_datum.array();
     DCHECK_EQ(rank_data->GetNullCount(), 0);
     // If there were nulls in the input, paste them in the output

From b424a6629397d6f86c894985b58847533940fa5b Mon Sep 17 00:00:00 2001
From: benibus <bpharks@gmx.com>
Date: Wed, 3 May 2023 12:10:15 -0400
Subject: [PATCH 07/11] Refactor benchmarks

---
 .../compute/kernels/vector_sort_benchmark.cc  | 145 ++++++++----------
 1 file changed, 65 insertions(+), 80 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc
index a4c192e359e..26b6dbe6614 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc
@@ -69,11 +69,25 @@ static void ArraySortFuncBenchmark(benchmark::State& state, const Runner& runner
   for (auto _ : state) {
     ABORT_NOT_OK(runner(values));
   }
-  // This may be redundant with the SetItemsProcessed() call in RegressionArgs,
-  // if size_is_bytes was false.
   state.SetItemsProcessed(state.iterations() * values->length());
 }
 
+// Compute the number of values in a StringArray based on a target `num_bytes`
+static int64_t GetStringArraySize(int64_t num_bytes, int32_t min_length,
+                                  int32_t max_length, double null_proportion,
+                                  int num_chunks = 1) {
+  // The array_size is derived from these two equations:
+  //
+  //     num_bytes = array_size * (1.0 - null_proportion) * string_mean_length
+  //     string_mean_length = (max_length + min_length) / 2.0
+  const double valid_proportion = 1.0 - null_proportion;
+  const double array_size =
+      (valid_proportion > std::numeric_limits<double>::epsilon())
+          ? (num_bytes * 2) / (valid_proportion * (max_length + min_length))
+          : (num_bytes * 2) / (max_length + min_length);
+  return static_cast<int64_t>(std::round(array_size / num_chunks));
+}
+
 template <typename Runner>
 static void ArraySortFuncInt64Benchmark(benchmark::State& state, const Runner& runner,
                                         int64_t min, int64_t max) {
@@ -88,53 +102,62 @@ static void ArraySortFuncInt64Benchmark(benchmark::State& state, const Runner& r
 
 template <typename Runner>
 static void ArraySortFuncInt64DictBenchmark(benchmark::State& state, const Runner& runner,
-                                            int64_t min, int64_t max) {
+                                            int64_t min, int64_t max, int32_t dict_size) {
   RegressionArgs args(state);
 
+  const double null_proportion = args.null_proportion / 2;
   const int64_t array_size = args.size / sizeof(int64_t);
 
   auto rand = random::RandomArrayGenerator(kSeed);
-  auto dict_values = rand.Int64(kDictionarySize, min, max, args.null_proportion / 2);
-  auto dict_indices =
-      rand.Int64(array_size, 0, kDictionarySize - 1, args.null_proportion / 2);
+  auto dict_values = rand.Int64(dict_size, min, max, null_proportion);
+  auto dict_indices = rand.Int64(array_size, 0, dict_size - 1, null_proportion);
   auto dict_array = *DictionaryArray::FromArrays(dict_indices, dict_values);
 
   ArraySortFuncBenchmark(state, runner, dict_array);
 }
 
 template <typename Runner>
-static void ArraySortFuncStringBenchmark(benchmark::State& state, const Runner& runner) {
-  RegressionArgs args(state, /*size_is_bytes=*/false);
-  // XXX This division to make numbers comparable with ArraySortFuncInt64Benchmark
-  // (including the SetItemsProcessed() call in the RegressionArgs destructor)
-  args.size /= sizeof(int64_t);
+static void ArraySortFuncStringBenchmark(benchmark::State& state, const Runner& runner,
+                                         int32_t min_length, int32_t max_length) {
+  RegressionArgs args(state);
+
+  const int64_t array_size =
+      GetStringArraySize(args.size, min_length, max_length, args.null_proportion);
 
-  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);
+  auto values = rand.String(array_size, min_length, max_length, args.null_proportion);
 
   ArraySortFuncBenchmark(state, runner, values);
 }
 
 template <typename Runner>
 static void ArraySortFuncStringDictBenchmark(benchmark::State& state,
-                                             const Runner& runner) {
-  RegressionArgs args(state, /*size_is_bytes=*/false);
-  // XXX Same as in ArraySortFuncStringBenchmark above
-  args.size /= sizeof(int64_t);
+                                             const Runner& runner, int32_t min_length,
+                                             int32_t max_length, int32_t dict_size) {
+  RegressionArgs args(state);
+
+  const double null_proportion = args.null_proportion / 2;
+  const int64_t array_size = 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);
-  auto dict_indices =
-      rand.Int64(array_size, 0, kDictionarySize - 1, args.null_proportion / 2);
+  auto dict_values = rand.String(dict_size, min_length, max_length, null_proportion);
+  auto dict_indices = rand.Int64(array_size, 0, dict_size - 1, null_proportion);
   auto dict_array = *DictionaryArray::FromArrays(dict_indices, dict_values);
 
   ArraySortFuncBenchmark(state, runner, dict_array);
 }
 
+template <typename Runner>
+static void ArraySortFuncBoolBenchmark(benchmark::State& state, const Runner& runner) {
+  RegressionArgs args(state);
+
+  const int64_t array_size = args.size * 8;
+  auto rand = random::RandomArrayGenerator(kSeed);
+  auto values = rand.Boolean(array_size, 0.5, args.null_proportion);
+
+  ArraySortFuncBenchmark(state, runner, values);
+}
+
 template <typename Runner>
 static void ChunkedArraySortFuncInt64Benchmark(benchmark::State& state,
                                                const Runner& runner, int64_t min,
@@ -158,18 +181,9 @@ static void ChunkedArraySortFuncStringBenchmark(benchmark::State& state,
                                                 int32_t max_length) {
   RegressionArgs args(state);
 
-  // The array_size is derived from these two equations:
-  //
-  //     size = array_size * (1.0 - args.null_proportion) * string_mean_length
-  //     string_mean_length = (max_length + min_length) / 2.0
-  const double valid_proportion = 1.0 - args.null_proportion;
-  const double array_size =
-      (valid_proportion > std::numeric_limits<double>::epsilon())
-          ? (args.size * 2) / (valid_proportion * (max_length + min_length))
-          : (args.size * 2) / (max_length + min_length);
-
   const auto n_chunks = 10;
-  const auto array_chunk_size = static_cast<int64_t>(std::round(array_size / n_chunks));
+  const auto array_chunk_size = GetStringArraySize(args.size, min_length, max_length,
+                                                   args.null_proportion, n_chunks);
   auto rand = random::RandomArrayGenerator(kSeed);
 
   ArrayVector chunks;
@@ -182,49 +196,6 @@ static void ChunkedArraySortFuncStringBenchmark(benchmark::State& state,
   ArraySortFuncBenchmark(state, runner, std::make_shared<ChunkedArray>(chunks));
 }
 
-static void ArraySortIndicesInt64Dict(benchmark::State& state) {
-  const auto min = std::numeric_limits<int64_t>::min();
-  const auto max = std::numeric_limits<int64_t>::max();
-  ArraySortFuncInt64DictBenchmark(state, SortRunner(state), min, max);
-}
-
-static void ArraySortIndicesString(benchmark::State& state) {
-  ArraySortFuncStringBenchmark(state, SortRunner(state));
-}
-
-static void ArraySortIndicesStringDict(benchmark::State& state) {
-  ArraySortFuncStringDictBenchmark(state, SortRunner(state));
-}
-
-template <typename Runner>
-static void ArraySortFuncBoolBenchmark(benchmark::State& state, const Runner& runner) {
-  RegressionArgs args(state);
-
-  const int64_t array_size = args.size * 8;
-  auto rand = random::RandomArrayGenerator(kSeed);
-  auto values = rand.Boolean(array_size, 0.5, args.null_proportion);
-
-  ArraySortFuncBenchmark(state, runner, values);
-}
-
-template <typename Runner>
-static void ArraySortFuncStringBenchmark(benchmark::State& state, const Runner& runner,
-                                         int32_t min_length, int32_t max_length) {
-  RegressionArgs args(state);
-
-  const double valid_proportion = 1.0 - args.null_proportion;
-  const double array_size =
-      (valid_proportion > std::numeric_limits<double>::epsilon())
-          ? (args.size * 2) / (valid_proportion * (max_length + min_length))
-          : (args.size * 2) / (max_length + min_length);
-  const auto array_size_i = static_cast<int64_t>(std::round(array_size));
-
-  auto rand = random::RandomArrayGenerator(kSeed);
-  auto values = rand.String(array_size_i, min_length, max_length, args.null_proportion);
-
-  ArraySortFuncBenchmark(state, runner, values);
-}
-
 static void ArraySortIndicesInt64Narrow(benchmark::State& state) {
   ArraySortFuncInt64Benchmark(state, SortRunner(state), -100, 100);
 }
@@ -243,6 +214,13 @@ static void ArraySortIndicesInt64Wide(benchmark::State& state) {
   ArraySortFuncInt64Benchmark(state, SortRunner(state), min, max);
 }
 
+static void ArraySortIndicesInt64WideDict(benchmark::State& state) {
+  const auto dict_size = kDictionarySize;
+  const auto min = std::numeric_limits<int64_t>::min();
+  const auto max = std::numeric_limits<int64_t>::max();
+  ArraySortFuncInt64DictBenchmark(state, SortRunner(state), min, max, dict_size);
+}
+
 static void ArrayRankInt64Wide(benchmark::State& state) {
   const auto min = std::numeric_limits<int64_t>::min();
   const auto max = std::numeric_limits<int64_t>::max();
@@ -271,6 +249,14 @@ static void ArraySortIndicesStringWide(benchmark::State& state) {
   ArraySortFuncStringBenchmark(state, SortRunner(state), min_length, max_length);
 }
 
+static void ArraySortIndicesStringWideDict(benchmark::State& state) {
+  const auto dict_size = kDictionarySize;
+  const auto min_length = 0;
+  const auto max_length = 64;
+  ArraySortFuncStringDictBenchmark(state, SortRunner(state), min_length, max_length,
+                                   dict_size);
+}
+
 static void ArrayRankStringNarrow(benchmark::State& state) {
   const auto min_length = 0;
   const auto max_length = 16;
@@ -453,12 +439,11 @@ void ArraySortIndicesSetArgs(benchmark::internal::Benchmark* bench) {
 
 BENCHMARK(ArraySortIndicesInt64Narrow)->Apply(ArraySortIndicesSetArgs);
 BENCHMARK(ArraySortIndicesInt64Wide)->Apply(ArraySortIndicesSetArgs);
+BENCHMARK(ArraySortIndicesInt64WideDict)->Apply(ArraySortIndicesSetArgs);
 BENCHMARK(ArraySortIndicesBool)->Apply(ArraySortIndicesSetArgs);
 BENCHMARK(ArraySortIndicesStringNarrow)->Apply(ArraySortIndicesSetArgs);
 BENCHMARK(ArraySortIndicesStringWide)->Apply(ArraySortIndicesSetArgs);
-BENCHMARK(ArraySortIndicesInt64Dict)->Apply(ArraySortIndicesSetArgs);
-BENCHMARK(ArraySortIndicesString)->Apply(ArraySortIndicesSetArgs);
-BENCHMARK(ArraySortIndicesStringDict)->Apply(ArraySortIndicesSetArgs);
+BENCHMARK(ArraySortIndicesStringWideDict)->Apply(ArraySortIndicesSetArgs);
 
 BENCHMARK(ChunkedArraySortIndicesInt64Narrow)->Apply(ArraySortIndicesSetArgs);
 BENCHMARK(ChunkedArraySortIndicesInt64Wide)->Apply(ArraySortIndicesSetArgs);

From 26eb022b2f051b237f34a5a5723af0c6518f8435 Mon Sep 17 00:00:00 2001
From: benibus <bpharks@gmx.com>
Date: Thu, 4 May 2023 13:18:26 -0400
Subject: [PATCH 08/11] Add issue link to TODO

---
 cpp/src/arrow/compute/kernels/vector_array_sort.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index 199d2182f46..cb0455d371b 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -182,6 +182,7 @@ class ArrayCompareSorter<DictionaryType> {
                                          ExecContext* ctx) {
     const auto& dict_array = checked_cast<const DictionaryArray&>(array);
     // TODO: These methods should probably return a const&? They seem capable.
+    // https://github.com/apache/arrow/issues/35437
     auto dict_values = dict_array.dictionary();
     auto dict_indices = dict_array.indices();
 

From 3c64de405dc09a0f78665e2f6313527e3acb8d72 Mon Sep 17 00:00:00 2001
From: benibus <bpharks@gmx.com>
Date: Thu, 18 May 2023 12:13:27 -0400
Subject: [PATCH 09/11] Alter handling of all-null and empty cases

---
 .../compute/kernels/vector_array_sort.cc      | 22 ++-------
 .../arrow/compute/kernels/vector_sort_test.cc | 47 +++++++++++--------
 2 files changed, 32 insertions(+), 37 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index cb0455d371b..1499554a960 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -203,8 +203,10 @@ class ArrayCompareSorter<DictionaryType> {
     // the original dictionary and sorting the decoded version.
 
     std::shared_ptr<Array> decoded_ranks;
-    // Skip the rank/take steps for cases with only nulls
-    if (IsAllNulls(dict_array)) {
+    // Skip the rank/take steps for cases with only nulls or no indices
+    if (dict_indices->length() == 0 ||
+        dict_indices->null_count() == dict_indices->length() ||
+        dict_values->null_count() == dict_values->length()) {
       ARROW_ASSIGN_OR_RAISE(decoded_ranks, MakeArrayOfNull(uint64(), dict_array.length(),
                                                            ctx->memory_pool()));
     } else {
@@ -222,22 +224,6 @@ class ArrayCompareSorter<DictionaryType> {
   }
 
  private:
-  static bool IsAllNulls(const DictionaryArray& dict_array) {
-    auto indices = dict_array.indices();
-    auto values = dict_array.dictionary();
-    if (values->null_count() == values->length() ||
-        indices->null_count() == indices->length()) {
-      return true;
-    }
-    for (int64_t i = 0; i < dict_array.length(); ++i) {
-      // TODO: Special handling for dictionary types in Array::IsValid/Null?
-      if (indices->IsValid(i) && values->IsValid(dict_array.GetValueIndex(i))) {
-        return false;
-      }
-    }
-    return true;
-  }
-
   static Result<std::shared_ptr<Array>> RanksWithNulls(
       const std::shared_ptr<Array>& array, ExecContext* ctx) {
     // Notes:
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
index 7adc17deca3..3429a5a8785 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc
@@ -397,31 +397,40 @@ TEST(ArraySortIndicesFunction, DictionaryArray) {
       }
     }
   }
+}
 
-  auto expected_all_nulls = ArrayFromJSON(uint64(), "[0, 1, 2, 3]");
-  auto expected_empty = ArrayFromJSON(uint64(), "[]");
+TEST(ArraySortIndicesFunction, EmptyDictionaryArray) {
+  auto expected = ArrayFromJSON(uint64(), "[]");
+  auto dict_arr = DictArrayFromJSON(dictionary(int32(), utf8()), "[]", "[\"b\", \"a\"]");
+  auto options = ArraySortOptions::Defaults();
 
-  for (auto index_type : all_dictionary_index_types()) {
-    ARROW_SCOPED_TRACE("index_type = ", index_type->ToString());
+  ASSERT_OK_AND_ASSIGN(auto actual,
+                       CallFunction("array_sort_indices", {dict_arr}, &options));
+  ValidateOutput(actual);
+  AssertDatumsEqual(expected, actual, /*verbose=*/true);
+}
 
-    auto dict_type = dictionary(index_type, utf8());
-    auto dict_arr_all_nulls =
-        DictArrayFromJSON(dict_type, "[null, 3, null, 1]", "[\"b\", null, \"a\", null]");
-    auto dict_arr_empty = DictArrayFromJSON(dict_type, "[]", "[\"b\", \"a\"]");
+TEST(ArraySortIndicesFunction, AllNullDictionaryArray) {
+  auto expected = ArrayFromJSON(uint64(), "[0, 1, 2, 3]");
+  auto dict_type = dictionary(int32(), utf8());
+  ArrayVector dict_arrs = {
+      DictArrayFromJSON(dict_type, "[null, 3, null, 1]", "[\"b\", null, \"a\", null]"),
+      DictArrayFromJSON(dict_type, "[null, null, null, null]", "[\"c\", \"a\"]"),
+      DictArrayFromJSON(dict_type, "[0, 1, 1, 0]", "[null, null]"),
+  };
 
-    for (auto order : AllOrders()) {
-      for (auto null_placement : AllNullPlacements()) {
-        ArraySortOptions options{order, null_placement};
+  int i = 0;
+  for (auto order : AllOrders()) {
+    for (auto null_placement : AllNullPlacements()) {
+      ArraySortOptions options{order, null_placement};
 
-        Datum actual;
-        ASSERT_OK_AND_ASSIGN(
-            actual, CallFunction("array_sort_indices", {dict_arr_all_nulls}, &options));
-        ValidateOutput(actual);
-        AssertDatumsEqual(expected_all_nulls, actual, /*verbose=*/true);
-        ASSERT_OK_AND_ASSIGN(
-            actual, CallFunction("array_sort_indices", {dict_arr_empty}, &options));
+      for (const auto& dict_arr : dict_arrs) {
+        ARROW_SCOPED_TRACE("i = ", i);
+        ASSERT_OK_AND_ASSIGN(auto actual,
+                             CallFunction("array_sort_indices", {dict_arr}, &options));
         ValidateOutput(actual);
-        AssertDatumsEqual(expected_empty, actual, /*verbose=*/true);
+        AssertDatumsEqual(expected, actual, /*verbose=*/true);
+        ++i;
       }
     }
   }

From 782198bf95c6ee42cff55bd070be973bbb4a3c31 Mon Sep 17 00:00:00 2001
From: benibus <bpharks@gmx.com>
Date: Thu, 18 May 2023 12:35:31 -0400
Subject: [PATCH 10/11] Revert unneeded template param

---
 cpp/src/arrow/compute/kernels/vector_sort.cc         | 8 ++++----
 cpp/src/arrow/compute/kernels/vector_sort_internal.h | 8 ++++----
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc
index abf8fc818e5..1de90cac35b 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort.cc
@@ -259,8 +259,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<ArrayType, StablePartitioner>(
-          indices_begin, indices_end, array_, offset, null_placement_);
+      p = PartitionNullsOnly<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>(
@@ -517,8 +517,8 @@ class MultipleKeyRecordBatchSorter : public TypeVisitor {
     const ArrayType& array =
         ::arrow::internal::checked_cast<const ArrayType&>(first_sort_key.array);
 
-    const auto p = PartitionNullsOnly<ArrayType, StablePartitioner>(
-        indices_begin_, indices_end_, array, 0, null_placement_);
+    const auto p = PartitionNullsOnly<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 1241b609e57..d78e5130617 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_internal.h
+++ b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
@@ -192,9 +192,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 ArrayType, typename Partitioner>
+template <typename Partitioner>
 NullPartitionResult PartitionNullsOnly(uint64_t* indices_begin, uint64_t* indices_end,
-                                       const ArrayType& values, int64_t offset,
+                                       const Array& values, int64_t offset,
                                        NullPlacement null_placement) {
   if (values.null_count() == 0) {
     return NullPartitionResult::NoNulls(indices_begin, indices_end, null_placement);
@@ -255,8 +255,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<ArrayType, Partitioner>(
-      indices_begin, indices_end, values, offset, null_placement);
+  NullPartitionResult p = PartitionNullsOnly<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,

From b62802468ed55f7f16e2fbf64803b9519db04479 Mon Sep 17 00:00:00 2001
From: benibus <bpharks@gmx.com>
Date: Thu, 18 May 2023 17:03:02 -0400
Subject: [PATCH 11/11] Tweak benchmark calculations

---
 .../compute/kernels/vector_sort_benchmark.cc  | 30 ++++++++++++++-----
 1 file changed, 22 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 26b6dbe6614..0cefddff5a1 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc
@@ -29,7 +29,7 @@ namespace arrow {
 namespace compute {
 
 constexpr auto kSeed = 0x0ff1ce;
-constexpr int32_t kDictionarySize = 20;  // a typical dictionary size
+constexpr int32_t kDictionarySize = 24;  // a typical dictionary size
 
 //
 // Array sort/rank benchmark helpers
@@ -105,12 +105,18 @@ static void ArraySortFuncInt64DictBenchmark(benchmark::State& state, const Runne
                                             int64_t min, int64_t max, int32_t dict_size) {
   RegressionArgs args(state);
 
-  const double null_proportion = args.null_proportion / 2;
+  double null_proportion = args.null_proportion;
+  double values_null_proportion = null_proportion / 2;
+  if (1.0 - null_proportion < std::numeric_limits<double>::epsilon()) {
+    values_null_proportion = null_proportion = 1.0;
+  }
+  double indices_null_proportion = null_proportion - values_null_proportion;
+
   const int64_t array_size = args.size / sizeof(int64_t);
 
   auto rand = random::RandomArrayGenerator(kSeed);
-  auto dict_values = rand.Int64(dict_size, min, max, null_proportion);
-  auto dict_indices = rand.Int64(array_size, 0, dict_size - 1, null_proportion);
+  auto dict_values = rand.Int64(dict_size, min, max, values_null_proportion);
+  auto dict_indices = rand.Int64(array_size, 0, dict_size - 1, indices_null_proportion);
   auto dict_array = *DictionaryArray::FromArrays(dict_indices, dict_values);
 
   ArraySortFuncBenchmark(state, runner, dict_array);
@@ -136,12 +142,20 @@ static void ArraySortFuncStringDictBenchmark(benchmark::State& state,
                                              int32_t max_length, int32_t dict_size) {
   RegressionArgs args(state);
 
-  const double null_proportion = args.null_proportion / 2;
-  const int64_t array_size = args.size / sizeof(int64_t);
+  double null_proportion = args.null_proportion;
+  double values_null_proportion = null_proportion / 2;
+  if (1.0 - null_proportion < std::numeric_limits<double>::epsilon()) {
+    values_null_proportion = null_proportion = 1.0;
+  }
+  double indices_null_proportion = null_proportion - values_null_proportion;
+
+  const int64_t array_size =
+      GetStringArraySize(args.size, min_length, max_length, null_proportion);
 
   auto rand = random::RandomArrayGenerator(kSeed);
-  auto dict_values = rand.String(dict_size, min_length, max_length, null_proportion);
-  auto dict_indices = rand.Int64(array_size, 0, dict_size - 1, null_proportion);
+  auto dict_values =
+      rand.String(dict_size, min_length, max_length, values_null_proportion);
+  auto dict_indices = rand.Int64(array_size, 0, dict_size - 1, indices_null_proportion);
   auto dict_array = *DictionaryArray::FromArrays(dict_indices, dict_values);
 
   ArraySortFuncBenchmark(state, runner, dict_array);