ARROW-14656: [Python] sort helper for Array, ChunkedArray and StructArray

cpp/src/arrow/compute/kernels/chunked_internal.h

@@ -48,6 +48,19 @@ struct ResolvedChunk {
   LogicalValueType Value() const { return V::LogicalValue(array->GetView(index)); }
 };
 
+// ResolvedChunk specialization for StructArray
+template <>
+struct ResolvedChunk<StructArray> {
+  // The target struct in chunked array.
+  const StructArray* array;
+  // The field index in the target struct.
+  const int64_t index;
+
+  ResolvedChunk(const StructArray* array, int64_t index) : array(array), index(index) {}
+
+  bool IsNull() const { return array->field(0)->IsNull(index); }
+};
+
 // ResolvedChunk specialization for untyped arrays when all is needed is null lookup
 template <>
 struct ResolvedChunk<Array> {

cpp/src/arrow/compute/kernels/vector_array_sort.cc

@@ -172,6 +172,49 @@ class ArrayCompareSorter {
   }
 };
 
+template <typename ArrowType>
+class StructArrayCompareSorter {
+  using ArrayType = typename TypeTraits<ArrowType>::ArrayType;
+
+ public:
+  // `offset` is used when this is called on a chunk of a chunked array
+  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);
+    nested_value_comparator_ = std::make_shared<NestedValuesComparator>();
+
+    if (nested_value_comparator_->Prepare(values) != Status::OK()) {
+      // TODO: Improve error handling
+      return NullPartitionResult();
+    }
+
+    const auto p = PartitionNulls<ArrayType, StablePartitioner>(
+        indices_begin, indices_end, values, offset, options.null_placement);
+
+    bool asc_order = options.order == SortOrder::Ascending;
+    std::stable_sort(p.non_nulls_begin, p.non_nulls_end,
+                     [&offset, &values, asc_order, this](uint64_t left, uint64_t right) {
+                       // is better to do values.fields.size() or
+                       // values.schema().num_fields() ?
+                       for (ArrayVector::size_type fieldidx = 0;
+                            fieldidx < values.fields().size(); ++fieldidx) {
+                         int result = nested_value_comparator_->Compare(
+                             values, fieldidx, offset, asc_order ? left : right,
+                             asc_order ? right : left);
+                         if (result == -1)
+                           return true;
+                         else if (result == 1)
+                           return false;
+                       }
+                       return false;
+                     });
+    return p;
+  }
+
+  std::shared_ptr<NestedValuesComparator> nested_value_comparator_;
+};
+
 template <typename ArrowType>
 class ArrayCountSorter {
   using ArrayType = typename TypeTraits<ArrowType>::ArrayType;
@@ -410,6 +453,11 @@ struct ArraySorter<
   ArrayCompareSorter<Type> impl;
 };
 
+template <typename Type>
+struct ArraySorter<Type, enable_if_t<is_struct_type<Type>::value>> {
+  StructArrayCompareSorter<Type> impl;
+};
+
 struct ArraySorterFactory {
   ArraySortFunc sorter;
 
@@ -511,6 +559,13 @@ const ArraySortOptions* GetDefaultArraySortOptions() {
   return &kDefaultArraySortOptions;
 }
 
+template <template <typename...> class ExecTemplate>
+void AddArraySortingNestedKernels(VectorKernel base, VectorFunction* func) {
+  base.signature = KernelSignature::Make({InputType::Array(Type::STRUCT)}, uint64());
+  base.exec = ExecTemplate<UInt64Type, StructType>::Exec;
+  DCHECK_OK(func->AddKernel(base));
+}
+
 const FunctionDoc array_sort_indices_doc(
     "Return the indices that would sort an array",
     ("This function computes an array of indices that define a stable sort\n"
@@ -559,6 +614,7 @@ void RegisterVectorArraySort(FunctionRegistry* registry) {
       GetDefaultArraySortOptions());
   base.init = ArraySortIndicesState::Init;
   AddArraySortingKernels<ArraySortIndices>(base, array_sort_indices.get());
+  AddArraySortingNestedKernels<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