ARROW-14656: [Python] Add sorting for StructArray

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 structarray in chunked array.
+ const StructArray* array;
+ // 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(static_cast<int>(index)) == nullptr; }
+};
+
 struct ChunkedArrayResolver : protected ::arrow::internal::ChunkResolver {
  ChunkedArrayResolver(const ChunkedArrayResolver& other)
  : ::arrow::internal::ChunkResolver(other.chunks_), chunks_(other.chunks_) {}

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

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

@@ -188,6 +188,8 @@ class ChunkedArraySorter : public TypeVisitor {
  return Status::OK();
  }
 
+ Status Visit(const StructType&) override { return SortInternal<StructType>(); }
+
  private:
  template <typename Type>
  Status SortInternal() {
@@ -269,25 +271,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

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

@@ -456,6 +456,149 @@ 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, *field, offset, leftrowidx,
+ rightrowidx);
+ }
+
+ // 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();
+ }
+
+ 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_);
+ }
+
+#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_;
+};
+
+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.array);
+
+ if (!status.ok()) {
+ 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, 0,
+ chunk_left.index, chunk_right.index);
+
+ if (val == 0) {
+ // if field values are equal, check next field
+ continue;
+ }
+
+ return val == -1;
+ }
+
+ return false;
+ }
+};
+
 } // namespace internal
 } // namespace compute
 } // namespace arrow

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):
  """

python/pyarrow/table.pxi

@@ -433,6 +433,22 @@ 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)