From 546825d61bb0f6d9391b3cbd798d18ea801b3bb5 Mon Sep 17 00:00:00 2001
From: Jedi18 <akhiljnair.188@gmail.com>
Date: Mon, 10 Oct 2022 21:24:06 +0530
Subject: [PATCH 1/5] tests and struct array sorting for arrays

---
 .../compute/kernels/vector_array_sort.cc      | 56 ++++++++++++
 .../compute/kernels/vector_sort_internal.h    | 87 +++++++++++++++++++
 python/pyarrow/array.pxi                      | 39 +++++++++
 python/pyarrow/table.pxi                      | 14 +++
 python/pyarrow/tests/test_array.py            | 71 +++++++++++++++
 5 files changed, 267 insertions(+)

diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
index 324a435441f..3d37b8f4942 100644
--- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc
@@ -173,6 +173,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;
@@ -409,6 +452,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;
 
@@ -507,6 +555,13 @@ void AddArraySortingKernels(VectorKernel base, VectorFunction* func) {
   DCHECK_OK(func->AddKernel(base));
 }
 
+template <template <typename...> class ExecTemplate>
+void AddArraySortingNestedKernels(VectorKernel base, VectorFunction* func) {
+  base.signature = KernelSignature::Make({Type::STRUCT}, uint64());
+  base.exec = ExecTemplate<UInt64Type, StructType>::Exec;
+  DCHECK_OK(func->AddKernel(base));
+}
+
 const ArraySortOptions* GetDefaultArraySortOptions() {
   static const auto kDefaultArraySortOptions = ArraySortOptions::Defaults();
   return &kDefaultArraySortOptions;
@@ -561,6 +616,7 @@ void RegisterVectorArraySort(FunctionRegistry* registry) {
   base.init = ArraySortIndicesState::Init;
   base.exec_chunked = ArraySortIndicesChunked;
   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
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_internal.h b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
index d8b024525c8..47f019b1eda 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_internal.h
+++ b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
@@ -456,6 +456,93 @@ Status SortChunkedArray(ExecContext* ctx, uint64_t* indices_begin, uint64_t* ind
                         const ChunkedArray& values, SortOrder sort_order,
                         NullPlacement null_placement);
 
+class NestedValuesComparator {
+  public:
+    // StructArray Compare overload
+    int Compare(StructArray const& array, uint64_t field_index, uint64_t offset,
+                uint64_t leftrowidx, uint64_t rightrowidx) {
+      std::shared_ptr<Array> field = array.field(static_cast<int>(field_index));
+      return comparators_[field_index]->Compare(*field, offset, leftrowidx, rightrowidx);
+    }
+
+    // StructArray Prepare overload
+    Status Prepare(StructArray const& array) {
+      auto fields = array.fields();
+      for (auto field = fields.begin(); field != fields.end(); field++) {
+        std::shared_ptr<DataType> physical_type = GetPhysicalType((*field)->type());
+        NestedComparatorFactory comparator_factory = NestedComparatorFactory();
+        std::shared_ptr<NestedValueComparator> current_field_comparator;
+        ARROW_ASSIGN_OR_RAISE(
+            current_field_comparator,
+            NestedComparatorFactory().MakeFieldComparator(*physical_type));
+        comparators_.push_back(current_field_comparator);
+      }
+      return Status::OK();
+    }
+
+  private:
+    struct NestedValueComparator {
+      virtual int Compare(Array const &array, uint64_t offset, 
+                          uint64_t leftidx, uint64_t rightidx) = 0;
+      virtual ~NestedValueComparator() = default;
+    };
+
+    template <typename Type>
+    struct ConcreteNestedValueComparator : NestedValueComparator {
+      using FieldArrayType = typename TypeTraits<Type>::ArrayType;
+      using GetView = GetViewType<Type>;
+
+      virtual int Compare(Array const &array, uint64_t offset, 
+                          uint64_t leftidx, uint64_t rightidx) {
+        const FieldArrayType& values = checked_cast<const FieldArrayType&>(array);
+        auto left_value = GetView::LogicalValue(values.GetView(leftidx - offset));
+        auto right_value = GetView::LogicalValue(values.GetView(rightidx - offset));
+        if (left_value == right_value)
+          return 0;
+        else if (left_value < right_value)
+          return -1;
+        else
+          return 1;
+      }
+    };
+
+    /** 
+     * Internal use factory whose purpose is to detect the right type
+     * of comparator that should be built to be able to compare two
+     * nested values in the same field or column
+     */
+    struct NestedComparatorFactory {
+      Result<std::shared_ptr<NestedValueComparator>> MakeFieldComparator(DataType const &physical_type) {
+          RETURN_NOT_OK(VisitTypeInline(physical_type, this));
+          DCHECK_NE(current_field_comparator_, nullptr);
+          return std::move(current_field_comparator_);
+      }
+
+    #define VISIT(TYPE) \
+      Status Visit(const TYPE& type) { return VisitGeneric(type); }
+
+        VISIT_SORTABLE_PHYSICAL_TYPES(VISIT)
+    #undef VISIT
+
+      Status Visit(const DataType& type) {
+        return Status::TypeError("Unsupported type for NestedComparatorFactory: ",
+                                type.ToString());
+      }
+
+      template <typename Type>
+      Status VisitGeneric(const Type&) {
+        current_field_comparator_ = std::shared_ptr<NestedValueComparator>(
+          new ConcreteNestedValueComparator<Type>()
+        );
+        return Status::OK();
+      }
+
+      std::shared_ptr<NestedValueComparator> current_field_comparator_;
+    };
+
+  std::vector<std::shared_ptr<NestedValueComparator>> comparators_;
+};
+
 }  // namespace internal
 }  // namespace compute
 }  // namespace arrow
diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi
index f5b3e7963f6..326f3006a3c 100644
--- a/python/pyarrow/array.pxi
+++ b/python/pyarrow/array.pxi
@@ -1402,6 +1402,22 @@ cdef class Array(_PandasConvertible):
         """
         return _pc().index(self, value, start, end, memory_pool=memory_pool)
 
+    def sort(self, order="ascending"):
+        """
+        Sort the Array
+
+        Parameters
+        ----------
+        order : "ascending" or "descending"
+            The order of the sorting.
+
+        Returns
+        -------
+        result : Array
+        """
+        indices = _pc().sort_indices(self, sort_keys=[("", order)])
+        return self.take(indices)
+
     def _to_pandas(self, options, types_mapper=None, **kwargs):
         return _array_like_to_pandas(self, options, types_mapper=types_mapper)
 
@@ -2746,6 +2762,29 @@ cdef class StructArray(Array):
         result.validate()
         return result
 
+    def sort(self, order="ascending", fieldname=None):
+        """
+        Sort the StructArray
+
+        Parameters
+        ----------
+        order : "ascending" or "descending"
+            The order of the sorting.
+        fieldname : str or None, default None
+            If to sort the array by one of its fields
+            or by the whole array.
+
+        Returns
+        -------
+        result : StructArray
+        """
+        if fieldname is not None:
+            tosort = self.field(fieldname)
+        else:
+            tosort = self
+        indices = _pc().sort_indices(tosort, sort_keys=[("", order)])
+        return self.take(indices)
+
 
 cdef class ExtensionArray(Array):
     """
diff --git a/python/pyarrow/table.pxi b/python/pyarrow/table.pxi
index 30352bf3950..3d0dbc1ef3c 100644
--- a/python/pyarrow/table.pxi
+++ b/python/pyarrow/table.pxi
@@ -433,6 +433,20 @@ cdef class ChunkedArray(_PandasConvertible):
 
         return result
 
+    def sort(self, order="ascending"):
+        """
+        Sort the ChunkedArray
+        Parameters
+        ----------
+        order : "ascending" or "descending"
+            The order of the sorting.
+        Returns
+        -------
+        result : ChunkedArray
+        """
+        indices = _pc().sort_indices(self, sort_keys=[("", order)])
+        return self.take(indices)
+
     def _to_pandas(self, options, types_mapper=None, **kwargs):
         return _array_like_to_pandas(self, options, types_mapper=types_mapper)
 
diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py
index 154ca348d2b..6ddbcfab78a 100644
--- a/python/pyarrow/tests/test_array.py
+++ b/python/pyarrow/tests/test_array.py
@@ -725,6 +725,77 @@ def test_struct_array_from_chunked():
         pa.StructArray.from_arrays([chunked_arr], ["foo"])
 
 
+def test_array_sort():
+    arr = pa.array([5, 7, 35], type=pa.int64())
+    sorted_arr = arr.sort("descending")
+    assert sorted_arr.to_pylist() == [35, 7, 5]
+
+    arr = pa.chunked_array([[1, 2, 3], [4, 5, 6]])
+    sorted_arr = arr.sort("descending")
+    assert sorted_arr.to_pylist() == [6, 5, 4, 3, 2, 1]
+
+
+def test_struct_array_sort():
+    arr = pa.StructArray.from_arrays([
+        pa.array([5, 7, 7, 35], type=pa.int64()),
+        pa.array(["foo", "car", "bar", "foobar"])
+    ], names=["a", "b"])
+
+    sorted_arr = arr.sort("descending", fieldname="a")
+    assert sorted_arr.to_pylist() == [
+        {"a": 35, "b": "foobar"},
+        {"a": 7, "b": "car"},
+        {"a": 7, "b": "bar"},
+        {"a": 5, "b": "foo"},
+    ]
+
+    sorted_arr = arr.sort("descending")
+    assert sorted_arr.to_pylist() == [
+        {"a": 35, "b": "foobar"},
+        {"a": 7, "b": "car"},
+        {"a": 7, "b": "bar"},
+        {"a": 5, "b": "foo"},
+    ]
+
+    sorted_arr = arr.sort("ascending")
+    assert sorted_arr.to_pylist() == [
+        {"a": 5, "b": "foo"},
+        {"a": 7, "b": "bar"},
+        {"a": 7, "b": "car"},
+        {"a": 35, "b": "foobar"},
+    ]
+
+
+def test_struct_chunked_array_sort():
+    arr1 = pa.StructArray.from_arrays([
+        pa.array([5, 7], type=pa.int64()),
+        pa.array(["foo", "car"])
+    ], names=["a", "b"])
+
+    arr2 = pa.StructArray.from_arrays([
+        pa.array([7, 35], type=pa.int64()),
+        pa.array(["bar", "foobar"])
+    ], names=["a", "b"])
+
+    chunked_arr = pa.chunked_array([arr1, arr2])
+
+    sorted_arr = chunked_arr.sort("descending")
+    assert sorted_arr.to_pylist() == [
+        {"a": 35, "b": "foobar"},
+        {"a": 7, "b": "car"},
+        {"a": 7, "b": "bar"},
+        {"a": 5, "b": "foo"},
+    ]
+
+    sorted_arr = chunked_arr.sort("ascending")
+    assert sorted_arr.to_pylist() == [
+        {"a": 5, "b": "foo"},
+        {"a": 7, "b": "bar"},
+        {"a": 7, "b": "car"},
+        {"a": 35, "b": "foobar"},
+    ]
+
+
 @pytest.mark.parametrize("offset", (0, 1))
 def test_dictionary_from_buffers(offset):
     a = pa.array(["one", "two", "three", "two", "one"]).dictionary_encode()

From 53fd8c73c16591d92acf9c3a4f3209549f414c3e Mon Sep 17 00:00:00 2001
From: Jedi18 <akhiljnair.188@gmail.com>
Date: Tue, 11 Oct 2022 12:04:31 +0530
Subject: [PATCH 2/5] add struct support for chunked array sorting

---
 .../arrow/compute/kernels/chunked_internal.h  | 13 ++++
 cpp/src/arrow/compute/kernels/vector_sort.cc  | 19 ++---
 .../compute/kernels/vector_sort_internal.h    | 74 +++++++++++++++++++
 3 files changed, 95 insertions(+), 11 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/chunked_internal.h b/cpp/src/arrow/compute/kernels/chunked_internal.h
index 69f439fccf0..f686e5fcf72 100644
--- a/cpp/src/arrow/compute/kernels/chunked_internal.h
+++ b/cpp/src/arrow/compute/kernels/chunked_internal.h
@@ -61,6 +61,19 @@ struct ResolvedChunk<Array> {
   bool IsNull() const { return array->IsNull(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); }
+};
+
 struct ChunkedArrayResolver : protected ::arrow::internal::ChunkResolver {
   ChunkedArrayResolver(const ChunkedArrayResolver& other)
       : ::arrow::internal::ChunkResolver(other.chunks_), chunks_(other.chunks_) {}
diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc
index 94c80b9f80d..e521eabb798 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort.cc
@@ -188,6 +188,10 @@ class ChunkedArraySorter : public TypeVisitor {
     return Status::OK();
   }
 
+  Status Visit(const StructType&) override {
+    return SortInternal<StructType>();
+  }
+
  private:
   template <typename Type>
   Status SortInternal() {
@@ -269,25 +273,18 @@ class ChunkedArraySorter : public TypeVisitor {
   template <typename ArrayType>
   void MergeNonNulls(uint64_t* range_begin, uint64_t* range_middle, uint64_t* range_end,
                      const std::vector<const Array*>& arrays, uint64_t* temp_indices) {
-    const ChunkedArrayResolver left_resolver(arrays);
-    const ChunkedArrayResolver right_resolver(arrays);
+    const ChunkedArrayResolver chunk_resolver(arrays);
+    ResolvedChunkComparator<ArrayType> resolved_chunk_comparator;
 
     if (order_ == SortOrder::Ascending) {
       std::merge(range_begin, range_middle, range_middle, range_end, temp_indices,
                  [&](uint64_t left, uint64_t right) {
-                   const auto chunk_left = left_resolver.Resolve<ArrayType>(left);
-                   const auto chunk_right = right_resolver.Resolve<ArrayType>(right);
-                   return chunk_left.Value() < chunk_right.Value();
+                   return resolved_chunk_comparator.Compare(chunk_resolver, left, right);
                  });
     } else {
       std::merge(range_begin, range_middle, range_middle, range_end, temp_indices,
                  [&](uint64_t left, uint64_t right) {
-                   const auto chunk_left = left_resolver.Resolve<ArrayType>(left);
-                   const auto chunk_right = right_resolver.Resolve<ArrayType>(right);
-                   // We don't use 'left > right' here to reduce required
-                   // operator. If we use 'right < left' here, '<' is only
-                   // required.
-                   return chunk_right.Value() < chunk_left.Value();
+                   return resolved_chunk_comparator.Compare(chunk_resolver, right, left);
                  });
     }
     // Copy back temp area into main buffer
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_internal.h b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
index 47f019b1eda..8884e061a57 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_internal.h
+++ b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
@@ -480,10 +480,31 @@ class NestedValuesComparator {
       return Status::OK();
     }
 
+    // ResolvedChunk<StructArray> Prepare overload
+    Status Prepare(const ResolvedChunk<StructArray>& chunk_left,
+                  const ResolvedChunk<StructArray>&) {
+      return Prepare(*chunk_left.array);
+    }
+
+    // ResolvedChunk<StructArray> Compare overload
+    int Compare(const ResolvedChunk<StructArray>& chunk_left,
+                const ResolvedChunk<StructArray>& chunk_right, uint64_t field_index) {
+      std::shared_ptr<Array> field_left =
+          chunk_left.array->field(static_cast<int>(field_index));
+      std::shared_ptr<Array> field_right =
+          chunk_right.array->field(static_cast<int>(field_index));
+      return comparators_[field_index]->Compare(*field_left, *field_right, chunk_left.index,
+                                                chunk_right.index);
+    }
+
   private:
     struct NestedValueComparator {
       virtual int Compare(Array const &array, uint64_t offset, 
                           uint64_t leftidx, uint64_t rightidx) = 0;
+
+      virtual int Compare(Array const& left_array, Array const& right_array,
+                        uint64_t leftidx, uint64_t rightidx) = 0;
+
       virtual ~NestedValueComparator() = default;
     };
 
@@ -504,6 +525,22 @@ class NestedValuesComparator {
         else
           return 1;
       }
+
+      virtual int Compare(Array const& left_array, Array const& right_array,
+                        uint64_t leftidx, uint64_t rightidx) {
+        const FieldArrayType& left_values = checked_cast<const FieldArrayType&>(left_array);
+        const FieldArrayType& right_values =
+            checked_cast<const FieldArrayType&>(right_array);
+        auto left_value = GetView::LogicalValue(left_values.GetView(leftidx));
+        auto right_value = GetView::LogicalValue(right_values.GetView(rightidx));
+
+        if (left_value == right_value)
+          return 0;
+        else if (left_value < right_value)
+          return -1;
+        else
+          return 1;
+      }
     };
 
     /** 
@@ -543,6 +580,43 @@ class NestedValuesComparator {
   std::vector<std::shared_ptr<NestedValueComparator>> comparators_;
 };
 
+template <typename ArrayType>
+struct ResolvedChunkComparator {
+  bool Compare(const ChunkedArrayResolver& chunk_resolver, uint64_t left,
+               uint64_t right) {
+    const auto chunk_left = chunk_resolver.Resolve<ArrayType>(left);
+    const auto chunk_right = chunk_resolver.Resolve<ArrayType>(right);
+    return chunk_left.Value() < chunk_right.Value();
+  }
+};
+
+template <>
+struct ResolvedChunkComparator<StructArray> {
+  bool Compare(const ChunkedArrayResolver& chunk_resolver, uint64_t left,
+               uint64_t right) {
+    const auto chunk_left = chunk_resolver.Resolve<StructArray>(left);
+    const auto chunk_right = chunk_resolver.Resolve<StructArray>(right);
+    NestedValuesComparator nested_values_comparator;
+    auto status = nested_values_comparator.Prepare(chunk_left, chunk_right);
+
+    if (!status.ok()) {
+      return false;
+    }
+
+    for (int i = 0; i < chunk_left.array->num_fields(); i++) {
+      int val = nested_values_comparator.Compare(chunk_left, chunk_right, i);
+
+      if (val == 0) {
+        continue;
+      }
+
+      return val == -1;
+    }
+
+    return false;
+  }
+};
+
 }  // namespace internal
 }  // namespace compute
 }  // namespace arrow

From 0fe7ed43db9430e48ad334e6640f3f8b94d770d4 Mon Sep 17 00:00:00 2001
From: Jedi18 <akhiljnair.188@gmail.com>
Date: Tue, 11 Oct 2022 12:46:52 +0530
Subject: [PATCH 3/5] minor simplification

---
 .../compute/kernels/vector_sort_internal.h    | 27 ++++++++-----------
 1 file changed, 11 insertions(+), 16 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/vector_sort_internal.h b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
index 8884e061a57..7e1ea053b18 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_internal.h
+++ b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
@@ -480,21 +480,15 @@ class NestedValuesComparator {
       return Status::OK();
     }
 
-    // ResolvedChunk<StructArray> Prepare overload
-    Status Prepare(const ResolvedChunk<StructArray>& chunk_left,
-                  const ResolvedChunk<StructArray>&) {
-      return Prepare(*chunk_left.array);
-    }
-
-    // ResolvedChunk<StructArray> Compare overload
-    int Compare(const ResolvedChunk<StructArray>& chunk_left,
-                const ResolvedChunk<StructArray>& chunk_right, uint64_t field_index) {
+    // StructArray Compare overload
+    int Compare(StructArray const& array_left, StructArray const& array_right, uint64_t field_index,
+                uint64_t leftrowidx, uint64_t rightrowidx) {
       std::shared_ptr<Array> field_left =
-          chunk_left.array->field(static_cast<int>(field_index));
+          array_left.field(static_cast<int>(field_index));
       std::shared_ptr<Array> field_right =
-          chunk_right.array->field(static_cast<int>(field_index));
-      return comparators_[field_index]->Compare(*field_left, *field_right, chunk_left.index,
-                                                chunk_right.index);
+          array_right.field(static_cast<int>(field_index));
+      return comparators_[field_index]->Compare(*field_left, *field_right, leftrowidx,
+                                                rightrowidx);
     }
 
   private:
@@ -597,14 +591,15 @@ struct ResolvedChunkComparator<StructArray> {
     const auto chunk_left = chunk_resolver.Resolve<StructArray>(left);
     const auto chunk_right = chunk_resolver.Resolve<StructArray>(right);
     NestedValuesComparator nested_values_comparator;
-    auto status = nested_values_comparator.Prepare(chunk_left, chunk_right);
+    auto status = nested_values_comparator.Prepare(*chunk_left.array);
 
     if (!status.ok()) {
       return false;
     }
 
-    for (int i = 0; i < chunk_left.array->num_fields(); i++) {
-      int val = nested_values_comparator.Compare(chunk_left, chunk_right, i);
+    for (int field_index = 0; field_index < chunk_left.array->num_fields(); field_index++) {
+      int val = nested_values_comparator.Compare(*(chunk_left.array), *(chunk_right.array), field_index, chunk_left.index,
+                                                chunk_right.index);
 
       if (val == 0) {
         continue;

From 9aa07876ffeb64968ae854232e6e9a175f9d5d94 Mon Sep 17 00:00:00 2001
From: Jedi18 <akhiljnair.188@gmail.com>
Date: Tue, 11 Oct 2022 12:57:52 +0530
Subject: [PATCH 4/5] add more struct chunked array tests

---
 python/pyarrow/tests/test_array.py | 89 ++++++++++++++++++++++++++++--
 1 file changed, 85 insertions(+), 4 deletions(-)

diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py
index 6ddbcfab78a..72fd2eb96f0 100644
--- a/python/pyarrow/tests/test_array.py
+++ b/python/pyarrow/tests/test_array.py
@@ -768,13 +768,13 @@ def test_struct_array_sort():
 
 def test_struct_chunked_array_sort():
     arr1 = pa.StructArray.from_arrays([
-        pa.array([5, 7], type=pa.int64()),
-        pa.array(["foo", "car"])
+        pa.array([5, 7, 8], type=pa.int64()),
+        pa.array(["foo", "car", "far"])
     ], names=["a", "b"])
 
     arr2 = pa.StructArray.from_arrays([
-        pa.array([7, 35], type=pa.int64()),
-        pa.array(["bar", "foobar"])
+        pa.array([7, 35, 9], type=pa.int64()),
+        pa.array(["bar", "foobar", "tar"])
     ], names=["a", "b"])
 
     chunked_arr = pa.chunked_array([arr1, arr2])
@@ -782,6 +782,8 @@ def test_struct_chunked_array_sort():
     sorted_arr = chunked_arr.sort("descending")
     assert sorted_arr.to_pylist() == [
         {"a": 35, "b": "foobar"},
+        {"a": 9, "b": "tar"},
+        {"a": 8, "b": "far"},
         {"a": 7, "b": "car"},
         {"a": 7, "b": "bar"},
         {"a": 5, "b": "foo"},
@@ -792,9 +794,88 @@ def test_struct_chunked_array_sort():
         {"a": 5, "b": "foo"},
         {"a": 7, "b": "bar"},
         {"a": 7, "b": "car"},
+        {"a": 8, "b": "far"},
+        {"a": 9, "b": "tar"},
         {"a": 35, "b": "foobar"},
     ]
 
+    arr1 = pa.StructArray.from_arrays([
+        pa.array(["foo", "car", "tar"]),
+        pa.array([5, 7, 8], type=pa.int64())
+    ], names=["a", "b"])
+
+    arr2 = pa.StructArray.from_arrays([
+        pa.array(["bar", "foobar", "far"]),
+        pa.array([7, 35, 9], type=pa.int64())
+    ], names=["a", "b"])
+
+    chunked_arr = pa.chunked_array([arr1, arr2])
+
+    sorted_arr = chunked_arr.sort("ascending")
+    assert sorted_arr.to_pylist() == [
+        {"a": "bar", "b": 7},
+        {"a": "car", "b": 7},
+        {"a": "far", "b": 9},
+        {"a": "foo", "b": 5},
+        {"a": "foobar", "b": 35},
+        {"a": "tar", "b": 8},
+    ]
+
+    sorted_arr = chunked_arr.sort("descending")
+    assert sorted_arr.to_pylist() == [
+        {"a": "tar", "b": 8},
+        {"a": "foobar", "b": 35},
+        {"a": "foo", "b": 5},
+        {"a": "far", "b": 9},
+        {"a": "car", "b": 7},
+        {"a": "bar", "b": 7},
+    ]
+
+    arr1 = pa.StructArray.from_arrays([
+        pa.array([5, 5, 5], type=pa.int64()),
+        pa.array([3, 2, 2], type=pa.int64()),
+        pa.array(["foo", "car", "tar"])
+    ], names=["a", "b", "c"])
+
+    arr2 = pa.StructArray.from_arrays([
+        pa.array([5, 5, 2], type=pa.int64()),
+        pa.array([2, 8, 3], type=pa.int64()),
+        pa.array(["bar", "foobar", "far"])
+    ], names=["a", "b", "c"])
+
+    arr3 = pa.StructArray.from_arrays([
+        pa.array([5, 5, 2], type=pa.int64()),
+        pa.array([2, 8, 1], type=pa.int64()),
+        pa.array(["cat", "dog", "mouse"])
+    ], names=["a", "b", "c"])
+
+    chunked_arr = pa.chunked_array([arr1, arr2, arr3])
+
+    sorted_arr = chunked_arr.sort("ascending")
+    assert sorted_arr.to_pylist() == [
+        {"a": 2, "b": 1, "c": "mouse"},
+        {"a": 2, "b": 3, "c": "far"},
+        {"a": 5, "b": 2, "c": "bar"},
+        {"a": 5, "b": 2, "c": "car"},
+        {"a": 5, "b": 2, "c": "cat"},
+        {"a": 5, "b": 2, "c": "tar"},
+        {"a": 5, "b": 3, "c": "foo"},
+        {"a": 5, "b": 8, "c": "dog"},
+        {"a": 5, "b": 8, "c": "foobar"}
+    ]
+
+    sorted_arr = chunked_arr.sort("descending")
+    assert sorted_arr.to_pylist() == [
+        {"a": 5, "b": 8, "c": "foobar"},
+        {"a": 5, "b": 8, "c": "dog"},
+        {"a": 5, "b": 3, "c": "foo"},
+        {"a": 5, "b": 2, "c": "tar"},
+        {"a": 5, "b": 2, "c": "cat"},
+        {"a": 5, "b": 2, "c": "car"},
+        {"a": 5, "b": 2, "c": "bar"},
+        {"a": 2, "b": 3, "c": "far"},
+        {"a": 2, "b": 1, "c": "mouse"}
+    ]
 
 @pytest.mark.parametrize("offset", (0, 1))
 def test_dictionary_from_buffers(offset):

From e27fb3ab2e84fd0f5b6005f8400b98f84f3e7327 Mon Sep 17 00:00:00 2001
From: Jedi18 <akhiljnair.188@gmail.com>
Date: Tue, 11 Oct 2022 13:12:44 +0530
Subject: [PATCH 5/5] remove redundant compare overload

---
 .../arrow/compute/kernels/chunked_internal.h  |   6 +-
 cpp/src/arrow/compute/kernels/vector_sort.cc  |   4 +-
 .../compute/kernels/vector_sort_internal.h    | 199 ++++++++----------
 python/pyarrow/table.pxi                      |   2 +
 python/pyarrow/tests/test_array.py            |   1 +
 5 files changed, 100 insertions(+), 112 deletions(-)

diff --git a/cpp/src/arrow/compute/kernels/chunked_internal.h b/cpp/src/arrow/compute/kernels/chunked_internal.h
index f686e5fcf72..89ddcbcab01 100644
--- a/cpp/src/arrow/compute/kernels/chunked_internal.h
+++ b/cpp/src/arrow/compute/kernels/chunked_internal.h
@@ -64,14 +64,14 @@ struct ResolvedChunk<Array> {
 // ResolvedChunk specialization for StructArray
 template <>
 struct ResolvedChunk<StructArray> {
-  // The target struct in chunked array.
+  // The target structarray in chunked array.
   const StructArray* array;
-  // The field index in the target struct.
+  // The index in the target array.
   const int64_t index;
 
   ResolvedChunk(const StructArray* array, int64_t index) : array(array), index(index) {}
 
-  bool IsNull() const { return array->field(0)->IsNull(index); }
+  bool IsNull() const { return array->field(static_cast<int>(index)) == nullptr; }
 };
 
 struct ChunkedArrayResolver : protected ::arrow::internal::ChunkResolver {
diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc
index e521eabb798..261691bc57a 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort.cc
+++ b/cpp/src/arrow/compute/kernels/vector_sort.cc
@@ -188,9 +188,7 @@ class ChunkedArraySorter : public TypeVisitor {
     return Status::OK();
   }
 
-  Status Visit(const StructType&) override {
-    return SortInternal<StructType>();
-  }
+  Status Visit(const StructType&) override { return SortInternal<StructType>(); }
 
  private:
   template <typename Type>
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_internal.h b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
index 7e1ea053b18..bd58b9483ea 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_internal.h
+++ b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
@@ -457,119 +457,103 @@ Status SortChunkedArray(ExecContext* ctx, uint64_t* indices_begin, uint64_t* ind
                         NullPlacement null_placement);
 
 class NestedValuesComparator {
-  public:
-    // StructArray Compare overload
-    int Compare(StructArray const& array, uint64_t field_index, uint64_t offset,
-                uint64_t leftrowidx, uint64_t rightrowidx) {
-      std::shared_ptr<Array> field = array.field(static_cast<int>(field_index));
-      return comparators_[field_index]->Compare(*field, offset, leftrowidx, rightrowidx);
-    }
+ public:
+  // StructArray Compare overload
+  int Compare(StructArray const& array, uint64_t field_index, uint64_t offset,
+              uint64_t leftrowidx, uint64_t rightrowidx) {
+    std::shared_ptr<Array> field = array.field(static_cast<int>(field_index));
+    return comparators_[field_index]->Compare(*field, *field, offset, leftrowidx,
+                                              rightrowidx);
+  }
 
-    // StructArray Prepare overload
-    Status Prepare(StructArray const& array) {
-      auto fields = array.fields();
-      for (auto field = fields.begin(); field != fields.end(); field++) {
-        std::shared_ptr<DataType> physical_type = GetPhysicalType((*field)->type());
-        NestedComparatorFactory comparator_factory = NestedComparatorFactory();
-        std::shared_ptr<NestedValueComparator> current_field_comparator;
-        ARROW_ASSIGN_OR_RAISE(
-            current_field_comparator,
-            NestedComparatorFactory().MakeFieldComparator(*physical_type));
-        comparators_.push_back(current_field_comparator);
-      }
-      return Status::OK();
+  // StructArray Compare overload for different struct arrays
+  int Compare(StructArray const& array_left, StructArray const& array_right,
+              uint64_t field_index, uint64_t offset, uint64_t leftrowidx,
+              uint64_t rightrowidx) {
+    std::shared_ptr<Array> field_left = array_left.field(static_cast<int>(field_index));
+    std::shared_ptr<Array> field_right = array_right.field(static_cast<int>(field_index));
+    return comparators_[field_index]->Compare(*field_left, *field_right, offset,
+                                              leftrowidx, rightrowidx);
+  }
+
+  // StructArray Prepare overload
+  Status Prepare(StructArray const& array) {
+    auto fields = array.fields();
+    for (auto field = fields.begin(); field != fields.end(); field++) {
+      std::shared_ptr<DataType> physical_type = GetPhysicalType((*field)->type());
+      NestedComparatorFactory comparator_factory = NestedComparatorFactory();
+      std::shared_ptr<NestedValueComparator> current_field_comparator;
+      ARROW_ASSIGN_OR_RAISE(
+          current_field_comparator,
+          NestedComparatorFactory().MakeFieldComparator(*physical_type));
+      comparators_.push_back(current_field_comparator);
     }
+    return Status::OK();
+  }
 
-    // StructArray Compare overload
-    int Compare(StructArray const& array_left, StructArray const& array_right, uint64_t field_index,
-                uint64_t leftrowidx, uint64_t rightrowidx) {
-      std::shared_ptr<Array> field_left =
-          array_left.field(static_cast<int>(field_index));
-      std::shared_ptr<Array> field_right =
-          array_right.field(static_cast<int>(field_index));
-      return comparators_[field_index]->Compare(*field_left, *field_right, leftrowidx,
-                                                rightrowidx);
+ private:
+  struct NestedValueComparator {
+    virtual int Compare(Array const& left_array, Array const& right_array,
+                        uint64_t offset, uint64_t leftidx, uint64_t rightidx) = 0;
+
+    virtual ~NestedValueComparator() = default;
+  };
+
+  template <typename Type>
+  struct ConcreteNestedValueComparator : NestedValueComparator {
+    using FieldArrayType = typename TypeTraits<Type>::ArrayType;
+    using GetView = GetViewType<Type>;
+
+    virtual int Compare(Array const& left_array, Array const& right_array,
+                        uint64_t offset, uint64_t leftidx, uint64_t rightidx) {
+      const FieldArrayType& left_values = checked_cast<const FieldArrayType&>(left_array);
+      const FieldArrayType& right_values =
+          checked_cast<const FieldArrayType&>(right_array);
+      auto left_value = GetView::LogicalValue(left_values.GetView(leftidx - offset));
+      auto right_value = GetView::LogicalValue(right_values.GetView(rightidx - offset));
+
+      if (left_value == right_value)
+        return 0;
+      else if (left_value < right_value)
+        return -1;
+      else
+        return 1;
+    }
+  };
+
+  /**
+   * Internal use factory whose purpose is to detect the right type
+   * of comparator that should be built to be able to compare two
+   * nested values in the same field or column
+   */
+  struct NestedComparatorFactory {
+    Result<std::shared_ptr<NestedValueComparator>> MakeFieldComparator(
+        DataType const& physical_type) {
+      RETURN_NOT_OK(VisitTypeInline(physical_type, this));
+      DCHECK_NE(current_field_comparator_, nullptr);
+      return std::move(current_field_comparator_);
     }
 
-  private:
-    struct NestedValueComparator {
-      virtual int Compare(Array const &array, uint64_t offset, 
-                          uint64_t leftidx, uint64_t rightidx) = 0;
+#define VISIT(TYPE) \
+  Status Visit(const TYPE& type) { return VisitGeneric(type); }
 
-      virtual int Compare(Array const& left_array, Array const& right_array,
-                        uint64_t leftidx, uint64_t rightidx) = 0;
+    VISIT_SORTABLE_PHYSICAL_TYPES(VISIT)
+#undef VISIT
 
-      virtual ~NestedValueComparator() = default;
-    };
+    Status Visit(const DataType& type) {
+      return Status::TypeError("Unsupported type for NestedComparatorFactory: ",
+                               type.ToString());
+    }
 
     template <typename Type>
-    struct ConcreteNestedValueComparator : NestedValueComparator {
-      using FieldArrayType = typename TypeTraits<Type>::ArrayType;
-      using GetView = GetViewType<Type>;
-
-      virtual int Compare(Array const &array, uint64_t offset, 
-                          uint64_t leftidx, uint64_t rightidx) {
-        const FieldArrayType& values = checked_cast<const FieldArrayType&>(array);
-        auto left_value = GetView::LogicalValue(values.GetView(leftidx - offset));
-        auto right_value = GetView::LogicalValue(values.GetView(rightidx - offset));
-        if (left_value == right_value)
-          return 0;
-        else if (left_value < right_value)
-          return -1;
-        else
-          return 1;
-      }
-
-      virtual int Compare(Array const& left_array, Array const& right_array,
-                        uint64_t leftidx, uint64_t rightidx) {
-        const FieldArrayType& left_values = checked_cast<const FieldArrayType&>(left_array);
-        const FieldArrayType& right_values =
-            checked_cast<const FieldArrayType&>(right_array);
-        auto left_value = GetView::LogicalValue(left_values.GetView(leftidx));
-        auto right_value = GetView::LogicalValue(right_values.GetView(rightidx));
-
-        if (left_value == right_value)
-          return 0;
-        else if (left_value < right_value)
-          return -1;
-        else
-          return 1;
-      }
-    };
-
-    /** 
-     * Internal use factory whose purpose is to detect the right type
-     * of comparator that should be built to be able to compare two
-     * nested values in the same field or column
-     */
-    struct NestedComparatorFactory {
-      Result<std::shared_ptr<NestedValueComparator>> MakeFieldComparator(DataType const &physical_type) {
-          RETURN_NOT_OK(VisitTypeInline(physical_type, this));
-          DCHECK_NE(current_field_comparator_, nullptr);
-          return std::move(current_field_comparator_);
-      }
-
-    #define VISIT(TYPE) \
-      Status Visit(const TYPE& type) { return VisitGeneric(type); }
-
-        VISIT_SORTABLE_PHYSICAL_TYPES(VISIT)
-    #undef VISIT
-
-      Status Visit(const DataType& type) {
-        return Status::TypeError("Unsupported type for NestedComparatorFactory: ",
-                                type.ToString());
-      }
-
-      template <typename Type>
-      Status VisitGeneric(const Type&) {
-        current_field_comparator_ = std::shared_ptr<NestedValueComparator>(
-          new ConcreteNestedValueComparator<Type>()
-        );
-        return Status::OK();
-      }
+    Status VisitGeneric(const Type&) {
+      current_field_comparator_ = std::shared_ptr<NestedValueComparator>(
+          new ConcreteNestedValueComparator<Type>());
+      return Status::OK();
+    }
 
-      std::shared_ptr<NestedValueComparator> current_field_comparator_;
-    };
+    std::shared_ptr<NestedValueComparator> current_field_comparator_;
+  };
 
   std::vector<std::shared_ptr<NestedValueComparator>> comparators_;
 };
@@ -597,11 +581,14 @@ struct ResolvedChunkComparator<StructArray> {
       return false;
     }
 
-    for (int field_index = 0; field_index < chunk_left.array->num_fields(); field_index++) {
-      int val = nested_values_comparator.Compare(*(chunk_left.array), *(chunk_right.array), field_index, chunk_left.index,
-                                                chunk_right.index);
+    for (int field_index = 0; field_index < chunk_left.array->num_fields();
+         field_index++) {
+      int val = nested_values_comparator.Compare(*(chunk_left.array),
+                                                 *(chunk_right.array), field_index, 0,
+                                                 chunk_left.index, chunk_right.index);
 
       if (val == 0) {
+        // if field values are equal, check next field
         continue;
       }
 
diff --git a/python/pyarrow/table.pxi b/python/pyarrow/table.pxi
index 3d0dbc1ef3c..32c981152fa 100644
--- a/python/pyarrow/table.pxi
+++ b/python/pyarrow/table.pxi
@@ -436,10 +436,12 @@ cdef class ChunkedArray(_PandasConvertible):
     def sort(self, order="ascending"):
         """
         Sort the ChunkedArray
+
         Parameters
         ----------
         order : "ascending" or "descending"
             The order of the sorting.
+
         Returns
         -------
         result : ChunkedArray
diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py
index 72fd2eb96f0..cc018d07731 100644
--- a/python/pyarrow/tests/test_array.py
+++ b/python/pyarrow/tests/test_array.py
@@ -877,6 +877,7 @@ def test_struct_chunked_array_sort():
         {"a": 2, "b": 1, "c": "mouse"}
     ]
 
+
 @pytest.mark.parametrize("offset", (0, 1))
 def test_dictionary_from_buffers(offset):
     a = pa.array(["one", "two", "three", "two", "one"]).dictionary_encode()