From 654668cecf1699756934dd3088abbc7c3eb8fb52 Mon Sep 17 00:00:00 2001 From: "Robert A. McLeod" Date: Mon, 19 Oct 2026 17:36:58 -0700 Subject: [PATCH] GH-38386: [C++][Python] Support dictionary sort keys in ChunkedArray, RecordBatch and Table sorting --- .../compute/kernels/vector_array_sort.cc | 220 +++++--- cpp/src/arrow/compute/kernels/vector_sort.cc | 88 ++-- .../compute/kernels/vector_sort_internal.h | 51 +- .../arrow/compute/kernels/vector_sort_test.cc | 471 ++++++++++++++++++ python/pyarrow/tests/test_compute.py | 80 +++ python/pyarrow/tests/test_table.py | 45 ++ 6 files changed, 810 insertions(+), 145 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_array_sort.cc b/cpp/src/arrow/compute/kernels/vector_array_sort.cc index 1cfb11286e46..c7e6577683aa 100644 --- a/cpp/src/arrow/compute/kernels/vector_array_sort.cc +++ b/cpp/src/arrow/compute/kernels/vector_array_sort.cc @@ -21,9 +21,12 @@ #include #include #include +#include #include +#include "arrow/array/concatenate.h" #include "arrow/array/data.h" +#include "arrow/array/util.h" #include "arrow/compute/api_vector.h" #include "arrow/compute/kernels/common_internal.h" #include "arrow/compute/kernels/util_internal.h" @@ -183,82 +186,11 @@ class ArrayCompareSorter { Result operator()(std::span indices, const Array& array, int64_t offset, const ArraySortOptions& options, ExecContext* ctx) { - const auto& dict_array = checked_cast(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. - - std::shared_ptr decoded_ranks; - // 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 { - ARROW_ASSIGN_OR_RAISE(auto ranks, RanksWithNulls(dict_values, ctx)); - - ARROW_ASSIGN_OR_RAISE(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())); - - return rank_sorter(indices, *decoded_ranks, offset, options, ctx); - } - - private: - static Result> RanksWithNulls( - const std::shared_ptr& array, ExecContext* ctx) { - // 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); - - auto data = array->data(); - std::shared_ptr null_bitmap; - 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; - } - ARROW_ASSIGN_OR_RAISE(auto rank_datum, - 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 - if (null_bitmap) { - rank_data->buffers[0] = std::move(null_bitmap); - rank_data->null_count = array->null_count(); - } - return MakeArray(rank_data); + ARROW_ASSIGN_OR_RAISE(auto sortable, + DictionaryToSortable(MakeArray(array.data()), ctx)); + DCHECK_EQ(sortable->length(), array.length()); + ARROW_ASSIGN_OR_RAISE(auto sorter, GetArraySorter(*sortable->type())); + return sorter(indices, *sortable, offset, options, ctx); } }; @@ -650,6 +582,142 @@ Result GetArraySorter(const DataType& type) { return factory.MakeSorter(type); } +namespace { + +Result> RanksWithNulls(const std::shared_ptr& array, + ExecContext* ctx) { + // Notes: + // - The order is always ascending here. + // - Nulls are kept in the output, so we consider them at the end. + RankOptions rank_options(SortOrder::Ascending, NullPlacement::AtEnd, + RankOptions::Dense); + + auto data = array->data(); + std::shared_ptr null_bitmap; + 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; + } + ARROW_ASSIGN_OR_RAISE(auto rank_datum, + 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 + if (null_bitmap) { + rank_data->buffers[0] = std::move(null_bitmap); + rank_data->null_count = array->null_count(); + } + return MakeArray(rank_data); +} + +// Ranks of the concatenated `dictionaries`. Distinct dictionaries usually overlap, +// so only their distinct values are ranked. +Result> RanksOfConcatenated(const ArrayVector& dictionaries, + ExecContext* ctx) { + ARROW_ASSIGN_OR_RAISE(auto concatenated, Concatenate(dictionaries, ctx->memory_pool())); + ARROW_ASSIGN_OR_RAISE(auto encoded_datum, + CallFunction("dictionary_encode", {concatenated}, ctx)); + const auto encoded = encoded_datum.make_array(); + const auto& encoded_dict = checked_cast(*encoded); + ARROW_ASSIGN_OR_RAISE(auto distinct_ranks, + RanksWithNulls(encoded_dict.dictionary(), ctx)); + return Take(*distinct_ranks, *encoded_dict.indices(), TakeOptions::Defaults(), ctx); +} + +} // namespace + +Result> DictionaryToSortable(const ChunkedArray& values, + ExecContext* ctx) { + const auto& dict_type = checked_cast(*values.type()); + const auto& value_type = dict_type.value_type(); + ArrayVector out(values.num_chunks()); + + // Ranks cannot represent NaN, which sorts as null-like rather than as the + // greatest value, so floating-point dictionaries are decoded instead. + if (value_type->id() == Type::FLOAT || value_type->id() == Type::DOUBLE) { + for (int i = 0; i < values.num_chunks(); ++i) { + const auto& dict_array = checked_cast(*values.chunk(i)); + ARROW_ASSIGN_OR_RAISE(out[i], Take(*dict_array.dictionary(), *dict_array.indices(), + TakeOptions::Defaults(), ctx)); + } + return std::make_shared(std::move(out), value_type); + } + + // Algorithm: + // 1) Use the Rank function to get an equivalent array + // of the dictionary values. Chunks may carry different dictionaries, so all + // distinct dictionaries are ranked together to make the ranks comparable + // across chunks. + // 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 (uint64), which should be dense, so we can use a + // fast sort for most instances. + + ArrayVector dictionaries; + std::vector dictionary_offsets; + std::unordered_map dictionary_slots; + // Index into `dictionaries` per chunk, or -1 for a chunk that is entirely null + std::vector chunk_slots(values.num_chunks(), -1); + int64_t total_dictionary_length = 0; + for (int i = 0; i < values.num_chunks(); ++i) { + const auto& dict_array = checked_cast(*values.chunk(i)); + const auto& dict_indices = dict_array.indices(); + const auto& dict_values = dict_array.dictionary(); + // Skip the rank/take steps for chunks 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()) { + continue; + } + auto [it, inserted] = + dictionary_slots.try_emplace(dict_values->data().get(), dictionaries.size()); + if (inserted) { + dictionaries.push_back(dict_values); + dictionary_offsets.push_back(total_dictionary_length); + total_dictionary_length += dict_values->length(); + } + chunk_slots[i] = static_cast(it->second); + } + + std::shared_ptr ranks; + if (dictionaries.size() == 1) { + ARROW_ASSIGN_OR_RAISE(ranks, RanksWithNulls(dictionaries[0], ctx)); + } else if (dictionaries.size() > 1) { + ARROW_ASSIGN_OR_RAISE(ranks, RanksOfConcatenated(dictionaries, ctx)); + } + + for (int i = 0; i < values.num_chunks(); ++i) { + const auto& dict_array = checked_cast(*values.chunk(i)); + const int64_t slot = chunk_slots[i]; + if (slot < 0) { + ARROW_ASSIGN_OR_RAISE( + out[i], MakeArrayOfNull(uint64(), dict_array.length(), ctx->memory_pool())); + continue; + } + const auto chunk_ranks = + ranks->Slice(dictionary_offsets[slot], dictionaries[slot]->length()); + ARROW_ASSIGN_OR_RAISE( + out[i], Take(*chunk_ranks, *dict_array.indices(), TakeOptions::Defaults(), ctx)); + DCHECK_EQ(out[i]->type_id(), Type::UINT64); + } + return std::make_shared(std::move(out), uint64()); +} + +Result> DictionaryToSortable(const std::shared_ptr& values, + ExecContext* ctx) { + ARROW_ASSIGN_OR_RAISE(auto sortable, DictionaryToSortable(ChunkedArray(values), ctx)); + DCHECK_EQ(sortable->num_chunks(), 1); + return sortable->chunk(0); +} + void RegisterVectorArraySort(FunctionRegistry* registry) { // The kernel outputs into preallocated memory and is never null VectorKernel base; diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc index 3ff888570710..275ec04d1a21 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort.cc @@ -18,6 +18,7 @@ #include #include +#include "arrow/array/concatenate.h" #include "arrow/compute/function.h" #include "arrow/compute/kernels/chunked_internal.h" #include "arrow/compute/kernels/vector_sort_internal.h" @@ -362,19 +363,9 @@ class ConcreteRecordBatchColumnSorter : public RecordBatchColumnSorter }; Result> ResolveRecordBatchSortKeys( - const RecordBatch& batch, const std::vector& sort_keys) { + const RecordBatch& batch, const std::vector& sort_keys, ExecContext* ctx) { return ::arrow::compute::internal::ResolveSortKeys( - batch, sort_keys); -} - -std::vector ResolveRecordBatchSortKeys( - const RecordBatch& batch, const std::vector& sort_keys, Status* status) { - const auto maybe_resolved = ResolveRecordBatchSortKeys(batch, sort_keys); - if (!maybe_resolved.ok()) { - *status = maybe_resolved.status(); - return {}; - } - return *std::move(maybe_resolved); + batch, sort_keys, ctx); } // Radix sorting is consistently faster except when there is a large number of sort keys, @@ -391,15 +382,8 @@ class RadixRecordBatchSorter { std::vector sort_keys) : sort_keys_(std::move(sort_keys)), indices_(indices) {} - RadixRecordBatchSorter(std::span indices, const RecordBatch& batch, - const SortOptions& options) - : sort_keys_(ResolveRecordBatchSortKeys(batch, options.GetSortKeys(), &status_)), - indices_(indices) {} - // Offset is for table sorting Result Sort(int64_t offset = 0) { - ARROW_RETURN_NOT_OK(status_); - // Create column sorters from right to left std::vector> column_sorts(sort_keys_.size()); RecordBatchColumnSorter* next_column = nullptr; @@ -457,14 +441,8 @@ class RadixRecordBatchSorter { std::unique_ptr result; }; - static Result> ResolveSortKeys( - const RecordBatch& batch, const std::vector& sort_keys) { - return ::arrow::compute::internal::ResolveSortKeys(batch, sort_keys); - } - const std::vector sort_keys_; std::span indices_; - Status status_; }; // Sort a batch using a single sort and multiple-key comparisons. @@ -476,19 +454,10 @@ class MultipleKeyRecordBatchSorter : public TypeVisitor { std::vector sort_keys) : indices_(indices), sort_keys_(std::move(sort_keys)), comparator_(sort_keys_) {} - MultipleKeyRecordBatchSorter(std::span indices, const RecordBatch& batch, - const SortOptions& options) - : indices_(indices), - sort_keys_(ResolveSortKeys(batch, options.GetSortKeys(), &status_)), - comparator_(sort_keys_) {} - // This is optimized for the first sort key. The first sort key sort // is processed in this class. The second and following sort keys // are processed in Comparator. - Status Sort() { - RETURN_NOT_OK(status_); - return sort_keys_[0].type->Accept(this); - } + Status Sort() { return sort_keys_[0].type->Accept(this); } #define VISIT(TYPE) \ Status Visit(const TYPE& type) override { return SortInternal(); } @@ -501,17 +470,6 @@ class MultipleKeyRecordBatchSorter : public TypeVisitor { private: using Comparator = MultipleKeyComparator; - static std::vector ResolveSortKeys( - const RecordBatch& batch, const std::vector& sort_keys, Status* status) { - const auto maybe_resolved = - ::arrow::compute::internal::ResolveSortKeys(batch, sort_keys); - if (!maybe_resolved.ok()) { - *status = maybe_resolved.status(); - return {}; - } - return *std::move(maybe_resolved); - } - template enable_if_t::value, Status> SortInternal() { using ArrayType = typename TypeTraits::ArrayType; @@ -589,7 +547,6 @@ class MultipleKeyRecordBatchSorter : public TypeVisitor { } std::span indices_; - Status status_; std::vector sort_keys_; Comparator comparator_; }; @@ -613,8 +570,8 @@ class TableSorter { : ctx_(ctx), table_(table), batches_(MakeBatches(table, &status_)), - options_(options), - sort_keys_(ResolveSortKeys(table, batches_, options.GetSortKeys(), &status_)), + sort_keys_( + ResolveSortKeys(table, batches_, options.GetSortKeys(), ctx, &status_)), indices_(indices), comparator_(sort_keys_) {} @@ -637,8 +594,8 @@ class TableSorter { static std::vector ResolveSortKeys( const Table& table, const RecordBatchVector& batches, - const std::vector& sort_keys, Status* status) { - const auto maybe_resolved = ResolvedSortKey::Make(table, batches, sort_keys); + const std::vector& sort_keys, ExecContext* ctx, Status* status) { + const auto maybe_resolved = ResolvedSortKey::Make(table, batches, sort_keys, ctx); if (!maybe_resolved.ok()) { *status = maybe_resolved.status(); return {}; @@ -660,9 +617,16 @@ class TableSorter { for (int64_t i = 0; i < num_batches; ++i) { const auto& batch = *batches_[i]; end_offset += batch.num_rows(); + // Reuse the table-wide resolved keys + std::vector batch_sort_keys; + batch_sort_keys.reserve(sort_keys_.size()); + for (const auto& sort_key : sort_keys_) { + batch_sort_keys.emplace_back(sort_key.owned_chunks[i], sort_key.order, + sort_key.null_placement); + } RadixRecordBatchSorter sorter( - {indices_.data() + begin_offset, indices_.data() + end_offset}, batch, - options_); + {indices_.data() + begin_offset, indices_.data() + end_offset}, + std::move(batch_sort_keys)); ARROW_ASSIGN_OR_RAISE(sorted[i], sorter.Sort(begin_offset)); DCHECK_EQ(sorted[i].overall_begin(), indices_.data() + begin_offset); DCHECK_EQ(sorted[i].overall_end(), indices_.data() + end_offset); @@ -822,7 +786,6 @@ class TableSorter { ExecContext* ctx_; const Table& table_; const RecordBatchVector batches_; - const SortOptions& options_; const std::vector sort_keys_; std::span indices_; Comparator comparator_; @@ -956,7 +919,7 @@ class SortIndicesMetaFunction : public MetaFunction { Result SortIndices(const RecordBatch& batch, const SortOptions& options, ExecContext* ctx) const { ARROW_ASSIGN_OR_RAISE(auto sort_keys, - ResolveRecordBatchSortKeys(batch, options.GetSortKeys())); + ResolveRecordBatchSortKeys(batch, options.GetSortKeys(), ctx)); auto n_sort_keys = sort_keys.size(); if (n_sort_keys == 0) { @@ -1094,6 +1057,19 @@ Result SortChunkedArray(ExecContext* ctx, std::span const ChunkedArray& chunked_array, SortOrder sort_order, NullPlacement null_placement) { + if (chunked_array.type()->id() == Type::DICTIONARY) { + ARROW_ASSIGN_OR_RAISE(auto sortable, DictionaryToSortable(chunked_array, ctx)); + if (sortable->num_chunks() <= 1) { + return SortChunkedArray(ctx, indices, *sortable, sort_order, null_placement); + } + // Sort the fixed-width converted values whole, so that a single counting sort + // can replace sorting and merging each chunk. + ARROW_ASSIGN_OR_RAISE(auto contiguous, + Concatenate(sortable->chunks(), ctx->memory_pool())); + ARROW_ASSIGN_OR_RAISE(auto sorter, GetArraySorter(*contiguous->type())); + return sorter(indices, *contiguous, /*offset=*/0, + ArraySortOptions(sort_order, null_placement), ctx); + } auto physical_type = GetPhysicalType(chunked_array.type()); auto physical_chunks = GetPhysicalChunks(chunked_array, physical_type); return SortChunkedArray(ctx, indices, physical_type, physical_chunks, sort_order, @@ -1126,7 +1102,7 @@ Result SortStructArray(ExecContext* ctx, std::span } ARROW_ASSIGN_OR_RAISE(auto sort_keys, - ResolveRecordBatchSortKeys(*batch, options.GetSortKeys())); + ResolveRecordBatchSortKeys(*batch, options.GetSortKeys(), ctx)); if (sort_keys.size() <= kMaxRadixSortKeys) { RadixRecordBatchSorter sorter(indices, std::move(sort_keys)); return sorter.Sort(); diff --git a/cpp/src/arrow/compute/kernels/vector_sort_internal.h b/cpp/src/arrow/compute/kernels/vector_sort_internal.h index 8bb10efa2d75..876c200bc9b9 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_internal.h +++ b/cpp/src/arrow/compute/kernels/vector_sort_internal.h @@ -455,6 +455,16 @@ using ArraySortFunc = std::function( Result GetArraySorter(const DataType& type); +// Convert dictionary-encoded values to an equivalent non-dictionary representation +// so they can be sorted naturally. Orders are consistent across chunks, even when +// chunks carry different dictionaries, so the result can be merged chunk-wise. +// Logical nulls are unaltered. +Result> DictionaryToSortable(const ChunkedArray& values, + ExecContext* ctx); + +Result> DictionaryToSortable(const std::shared_ptr& values, + ExecContext* ctx); + Result SortChunkedArray(ExecContext* ctx, std::span indices, const ChunkedArray& chunked_array, SortOrder sort_order, @@ -521,22 +531,25 @@ Result> ResolveSortKeys( return resolved; } -template +template Result> ResolveSortKeys( - const TableOrBatch& table_or_batch, const std::vector& sort_keys) { + const RecordBatch& batch, const std::vector& sort_keys, ExecContext* ctx) { return ResolveSortKeys( - *table_or_batch.schema(), sort_keys, - [&](const SortField& f) -> Result { + *batch.schema(), sort_keys, [&](const SortField& f) -> Result { + std::shared_ptr column; if (f.is_nested()) { // TODO: Some room for improvement here, as we potentially duplicate some of the // null-flattening work for nested sort keys. For instance, given two keys with // paths [0,0,0,0] and [0,0,0,1], we shouldn't need to flatten the first three // components more than once. - ARROW_ASSIGN_OR_RAISE(auto child, f.path.GetFlattened(table_or_batch)); - return ResolvedSortKey{std::move(child), f.order, f.null_placement}; + ARROW_ASSIGN_OR_RAISE(column, f.path.GetFlattened(batch)); + } else { + column = batch.column(f.path[0]); + } + if (column->type()->id() == Type::DICTIONARY) { + ARROW_ASSIGN_OR_RAISE(column, DictionaryToSortable(column, ctx)); } - return ResolvedSortKey{table_or_batch.column(f.path[0]), f.order, - f.null_placement}; + return ResolvedSortKey{std::move(column), f.order, f.null_placement}; }); } @@ -761,18 +774,30 @@ struct ResolvedTableSortKey { // `batches` must be a chunking of `table`. static Result> Make( const Table& table, const RecordBatchVector& batches, - const std::vector& sort_keys) { + const std::vector& sort_keys, ExecContext* ctx) { auto factory = [&](const SortField& f) -> Result { // We must expose a homogenous chunking for all ResolvedSortKey, // so we can't simply access the column from the table directly. ArrayVector chunks; chunks.reserve(batches.size()); - auto physical_type = GetPhysicalType(f.type->GetSharedPtr()); - int64_t null_count = 0; for (const auto& batch : batches) { ARROW_ASSIGN_OR_RAISE(auto child, f.path.GetFlattened(*batch)); - null_count += child->null_count(); - chunks.push_back(GetPhysicalArray(*child, physical_type)); + chunks.push_back(std::move(child)); + } + auto type = f.type->GetSharedPtr(); + if (type->id() == Type::DICTIONARY) { + // Convert all batches at once so that orders agree when batches are merged + ARROW_ASSIGN_OR_RAISE( + auto sortable, + DictionaryToSortable(ChunkedArray(std::move(chunks), type), ctx)); + type = sortable->type(); + chunks = sortable->chunks(); + } + auto physical_type = GetPhysicalType(type); + int64_t null_count = 0; + for (auto& chunk : chunks) { + null_count += chunk->null_count(); + chunk = GetPhysicalArray(*chunk, physical_type); } return ResolvedTableSortKey(physical_type, std::move(chunks), f.order, diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index 9bc30f23234b..f09645733045 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -2340,6 +2340,477 @@ TEST_F(TestNestedSortIndices, SortChunkedArray) { TestSort(GetChunkedArray()); } TEST_F(TestNestedSortIndices, SortRecordBatch) { TestSort(GetRecordBatch()); } TEST_F(TestNestedSortIndices, SortTable) { TestSort(GetTable()); } +// ---------------------------------------------------------------------- +// Tests for dictionary sort keys in chunked arrays, record batches and tables + +std::shared_ptr DecodedType(const std::shared_ptr& type) { + if (type->id() == Type::DICTIONARY) { + return checked_cast(*type).value_type(); + } + if (type->id() != Type::STRUCT) { + return type; + } + FieldVector fields; + for (const auto& f : type->fields()) { + fields.push_back(f->WithType(DecodedType(f->type()))); + } + return struct_(std::move(fields)); +} + +// Decode dictionaries, recursing into struct children +Result> DecodeDictionaries(const std::shared_ptr& array) { + if (array->type_id() == Type::DICTIONARY) { + return DecodeDictionary(*array); + } + if (array->type_id() != Type::STRUCT) { + return array; + } + auto data = array->data()->Copy(); + FieldVector fields; + for (size_t i = 0; i < data->child_data.size(); ++i) { + ARROW_ASSIGN_OR_RAISE(auto child, DecodeDictionaries(MakeArray(data->child_data[i]))); + fields.push_back(array->type()->field(static_cast(i))->WithType(child->type())); + data->child_data[i] = child->data(); + } + data->type = struct_(std::move(fields)); + return MakeArray(std::move(data)); +} + +Result> DecodeDictionaries(const ChunkedArray& chunked) { + ArrayVector chunks; + for (const auto& chunk : chunked.chunks()) { + ARROW_ASSIGN_OR_RAISE(auto decoded, DecodeDictionaries(chunk)); + chunks.push_back(std::move(decoded)); + } + return ChunkedArray::Make(std::move(chunks), DecodedType(chunked.type())); +} + +Result DecodeDictionaries(const Datum& datum) { + switch (datum.kind()) { + case Datum::ARRAY: + return DecodeDictionaries(datum.make_array()); + case Datum::CHUNKED_ARRAY: + return DecodeDictionaries(*datum.chunked_array()); + case Datum::RECORD_BATCH: { + const auto& batch = *datum.record_batch(); + ArrayVector columns; + FieldVector fields; + for (int i = 0; i < batch.num_columns(); ++i) { + ARROW_ASSIGN_OR_RAISE(auto column, DecodeDictionaries(batch.column(i))); + fields.push_back(batch.schema()->field(i)->WithType(column->type())); + columns.push_back(std::move(column)); + } + return RecordBatch::Make(schema(std::move(fields)), batch.num_rows(), + std::move(columns)); + } + case Datum::TABLE: { + const auto& table = *datum.table(); + ChunkedArrayVector columns; + FieldVector fields; + for (int i = 0; i < table.num_columns(); ++i) { + ARROW_ASSIGN_OR_RAISE(auto column, DecodeDictionaries(*table.column(i))); + fields.push_back(table.schema()->field(i)->WithType(column->type())); + columns.push_back(std::move(column)); + } + return Table::Make(schema(std::move(fields)), std::move(columns), table.num_rows()); + } + default: + return Status::NotImplemented("Cannot decode ", datum.ToString()); + } +} + +// Sort options over the given key names. All order/placement combinations are used for +// up to two keys; beyond that, the combinations are applied uniformly and alternating. +std::vector DictionarySortOptions(const std::vector& refs) { + std::vector all_options; + if (refs.size() <= 2) { + std::vector> prefixes = {{}}; + for (const auto& ref : refs) { + std::vector> extended; + for (const auto& prefix : prefixes) { + for (auto order : AllOrders()) { + for (auto null_placement : AllNullPlacements()) { + auto keys = prefix; + keys.emplace_back(ref, order, null_placement); + extended.push_back(std::move(keys)); + } + } + } + prefixes = std::move(extended); + } + for (auto& keys : prefixes) { + all_options.emplace_back(std::move(keys)); + } + return all_options; + } + for (auto order : AllOrders()) { + for (auto null_placement : AllNullPlacements()) { + for (bool alternate : {false, true}) { + std::vector keys; + for (size_t i = 0; i < refs.size(); ++i) { + auto key_order = order; + if (alternate && i % 2 == 1) { + key_order = order == SortOrder::Ascending ? SortOrder::Descending + : SortOrder::Ascending; + } + keys.emplace_back(refs[i], key_order, null_placement); + } + all_options.emplace_back(std::move(keys)); + } + } + } + return all_options; +} + +void AssertSortIndicesMatchDecoded(const Datum& input, + const std::vector& refs) { + ASSERT_OK_AND_ASSIGN(auto decoded, DecodeDictionaries(input)); + for (const auto& options : DictionarySortOptions(refs)) { + ARROW_SCOPED_TRACE("options = ", options.ToString()); + ASSERT_OK_AND_ASSIGN(auto actual, CallFunction("sort_indices", {input}, &options)); + ValidateOutput(actual); + ASSERT_OK_AND_ASSIGN(auto expected, + CallFunction("sort_indices", {decoded}, &options)); + AssertDatumsEqual(expected, actual, /*verbose=*/true); + } +} + +void AssertArraySortIndicesMatchDecoded(const Datum& input) { + ASSERT_OK_AND_ASSIGN(auto decoded, DecodeDictionaries(input)); + for (auto order : AllOrders()) { + for (auto null_placement : AllNullPlacements()) { + ArraySortOptions options(order, null_placement); + ARROW_SCOPED_TRACE("options = ", options.ToString()); + ASSERT_OK_AND_ASSIGN(auto actual, + CallFunction("array_sort_indices", {input}, &options)); + ValidateOutput(actual); + ASSERT_OK_AND_ASSIGN(auto expected, + CallFunction("array_sort_indices", {decoded}, &options)); + AssertDatumsEqual(expected, actual, /*verbose=*/true); + } + } +} + +class TestDictionarySortIndices : public ::testing::Test { + protected: + static std::shared_ptr StringDictType() { + return dictionary(int32(), utf8()); + } + + // Chunks with distinct but overlapping dictionaries, a shared dictionary, null + // dictionary values without null indices, an empty chunk and an all-null chunk. + static std::shared_ptr GetStringChunkedArray() { + auto type = StringDictType(); + auto chunk0 = DictArrayFromJSON(type, "[0, 1, null, 2, 1]", R"(["b", "a", "c"])"); + auto chunk1 = DictArrayFromJSON(type, "[3, 0, 1, 2, 3]", R"(["c", "d", "a", null])"); + auto chunk2 = DictArrayFromJSON(type, "[]", R"(["e"])"); + auto chunk3 = *DictionaryArray::FromArrays( + type, ArrayFromJSON(int32(), "[2, 0, 0]"), + checked_cast(*chunk0).dictionary()); + auto chunk4 = DictArrayFromJSON(type, "[null, null]", R"(["a", "z"])"); + auto chunked = *ChunkedArray::Make({chunk0, chunk1, chunk2, chunk3, chunk4}); + ARROW_CHECK_OK(chunked->ValidateFull()); + return chunked; + } + + // Float dictionaries carry NaNs, which must sort as null-like values rather than + // as the greatest value. + static std::shared_ptr GetFloatChunkedArray() { + auto type = dictionary(int8(), float64()); + auto chunk0 = + DictArrayFromJSON(type, "[0, 1, 2, 3, 4, 1, null]", "[1.0, NaN, 0.5, null, 2.0]"); + auto chunk1 = DictArrayFromJSON(type, "[2, 0, 1, null, 2]", "[NaN, 2.0, -1.5]"); + auto chunked = *ChunkedArray::Make({chunk0, chunk1}); + ARROW_CHECK_OK(chunked->ValidateFull()); + return chunked; + } + + static std::shared_ptr RandomStringDictionary(random::RandomArrayGenerator& rng, + int64_t length, + int64_t dict_length) { + // Short strings so that separately generated dictionaries overlap + auto values = rng.StringWithRepeats(dict_length, /*unique=*/dict_length / 2 + 1, + /*min_length=*/0, /*max_length=*/2, + /*null_probability=*/0.1); + auto indices = rng.Int32(length, 0, static_cast(dict_length - 1), + /*null_probability=*/0.1); + return *DictionaryArray::FromArrays(StringDictType(), indices, values); + } + + static std::shared_ptr RandomFloatDictionary(random::RandomArrayGenerator& rng, + int64_t length, + int64_t dict_length) { + auto values = rng.Float64(dict_length, -2, 2, /*null_probability=*/0.1, + /*nan_probability=*/0.1); + auto indices = rng.Int16(length, 0, static_cast(dict_length - 1), + /*null_probability=*/0.1); + return *DictionaryArray::FromArrays(dictionary(int16(), float64()), indices, values); + } + + // A batch with alternating dictionary and low-cardinality integer columns, so that + // ties on earlier keys are common. + static std::shared_ptr RandomBatch(random::RandomArrayGenerator& rng, + int64_t length, int num_columns) { + FieldVector fields; + ArrayVector columns; + for (int i = 0; i < num_columns; ++i) { + std::shared_ptr column; + switch (i % 3) { + case 0: + column = RandomStringDictionary(rng, length, /*dict_length=*/6); + break; + case 1: + column = rng.Int32(length, 0, 2, /*null_probability=*/0.1); + break; + default: + column = RandomFloatDictionary(rng, length, /*dict_length=*/5); + break; + } + fields.push_back(field("f" + std::to_string(i), column->type())); + columns.push_back(std::move(column)); + } + auto batch = RecordBatch::Make(schema(std::move(fields)), length, std::move(columns)); + ARROW_CHECK_OK(batch->ValidateFull()); + return batch; + } + + static std::vector ColumnRefs(const Schema& schema) { + std::vector refs; + for (const auto& f : schema.fields()) { + refs.emplace_back(f->name()); + } + return refs; + } +}; + +TEST_F(TestDictionarySortIndices, Array) { + for (const auto& chunked : {GetStringChunkedArray(), GetFloatChunkedArray()}) { + for (const auto& chunk : chunked->chunks()) { + ARROW_SCOPED_TRACE("chunk = ", chunk->ToString()); + AssertArraySortIndicesMatchDecoded(chunk); + AssertSortIndicesMatchDecoded(chunk, {}); + } + } +} + +TEST_F(TestDictionarySortIndices, ChunkedArray) { + for (const auto& chunked : {GetStringChunkedArray(), GetFloatChunkedArray()}) { + ARROW_SCOPED_TRACE("type = ", chunked->type()->ToString()); + AssertArraySortIndicesMatchDecoded(chunked); + AssertSortIndicesMatchDecoded(chunked, {}); + } +} + +TEST_F(TestDictionarySortIndices, EmptyChunkedArray) { + auto chunked = *ChunkedArray::Make({}, StringDictType()); + AssertArraySortIndicesMatchDecoded(chunked); +} + +TEST_F(TestDictionarySortIndices, ChunkedArrayExpected) { + // Spot-check against a hand-computed result. Decoded values: + // ["b", "a", null, "c", "a" | null, "c", "d", "a", null | | "c", "b", "b" | null, + // null] + auto chunked = GetStringChunkedArray(); + ArraySortOptions options(SortOrder::Ascending, NullPlacement::AtEnd); + ASSERT_OK_AND_ASSIGN(auto actual, + CallFunction("array_sort_indices", {chunked}, &options)); + AssertDatumsEqual(ChunkedArrayFromJSON( + uint64(), {"[1, 4, 8, 0, 11, 12, 3, 6, 10, 7, 2, 5, 9, 13, 14]"}), + actual, /*verbose=*/true); +} + +TEST_F(TestDictionarySortIndices, FloatNaN) { + // Regression test: NaNs used to be ranked as the greatest dictionary value + auto array = DictArrayFromJSON(dictionary(int32(), float64()), "[0, 1, 2, 3, 4]", + "[1.0, NaN, 0.5, null, 2.0]"); + AssertSortIndices(array, SortOrder::Ascending, NullPlacement::AtEnd, "[2, 0, 4, 1, 3]"); + AssertSortIndices(array, SortOrder::Ascending, NullPlacement::AtStart, + "[3, 1, 2, 0, 4]"); + AssertSortIndices(array, SortOrder::Descending, NullPlacement::AtEnd, + "[4, 0, 2, 1, 3]"); + AssertSortIndices(array, SortOrder::Descending, NullPlacement::AtStart, + "[3, 1, 4, 0, 2]"); +} + +TEST_F(TestDictionarySortIndices, RecordBatch) { + auto dict_type = StringDictType(); + auto schema = ::arrow::schema({field("d", dict_type), field("i", int32()), + field("f", dictionary(int8(), float64()))}); + // Column "d" has null dictionary values but no null indices, so its null count as a + // DictionaryArray is zero. + auto d = DictArrayFromJSON(dict_type, "[0, 1, 2, 3, 0, 1, 2, 3]", + R"(["b", null, "a", "b"])"); + auto i = ArrayFromJSON(int32(), "[1, 2, 2, 0, 0, null, 1, 1]"); + auto f = DictArrayFromJSON(dictionary(int8(), float64()), "[0, 1, 2, null, 1, 0, 2, 3]", + "[NaN, 1.5, -1.0, NaN]"); + auto batch = RecordBatch::Make(schema, 8, {d, i, f}); + ASSERT_OK(batch->ValidateFull()); + + AssertSortIndicesMatchDecoded(batch, {FieldRef("d"), FieldRef("i")}); + AssertSortIndicesMatchDecoded(batch, {FieldRef("i"), FieldRef("d")}); + AssertSortIndicesMatchDecoded(batch, {FieldRef("f"), FieldRef("d")}); + AssertSortIndicesMatchDecoded(batch, {FieldRef("d"), FieldRef("f")}); + AssertSortIndicesMatchDecoded(batch, {FieldRef("d")}); + + // Decoded "d": ["b", null, "a", "b", "b", null, "a", "b"] + SortOptions options({SortKey("d"), SortKey("i", SortOrder::Descending)}); + AssertSortIndices(batch, options, "[2, 6, 0, 7, 3, 4, 1, 5]"); +} + +TEST_F(TestDictionarySortIndices, RecordBatchRandom) { + random::RandomArrayGenerator rng(/*seed=*/0x5eed); + // More keys than kMaxRadixSortKeys exercises MultipleKeyRecordBatchSorter + for (int num_columns : {2, 3, 9}) { + ARROW_SCOPED_TRACE("num_columns = ", num_columns); + auto batch = RandomBatch(rng, /*length=*/300, num_columns); + AssertSortIndicesMatchDecoded(batch, ColumnRefs(*batch->schema())); + } +} + +TEST_F(TestDictionarySortIndices, Table) { + auto string_chunked = GetStringChunkedArray(); + // Columns have different chunk boundaries, so TableBatchReader slices them to a + // common chunking. + auto float_type = dictionary(int8(), float64()); + ASSERT_OK_AND_ASSIGN( + auto float_chunked, + ChunkedArray::Make( + {DictArrayFromJSON(float_type, "[0, 1]", "[NaN, 1.0]"), + DictArrayFromJSON(float_type, "[2, 0, 1, 3, null, 2, 1, 0, 3, 3]", + "[1.0, NaN, 0.5, null]"), + DictArrayFromJSON(float_type, "[0, 1, 0]", "[2.0, -1.0]")})); + auto int_chunked = ChunkedArrayFromJSON( + int32(), {"[1, 0, 1]", "[1, 0, 0, 1, 1, 0, 1, 0, null]", "[0, 0, 1]"}); + ASSERT_EQ(string_chunked->length(), 15); + ASSERT_EQ(float_chunked->length(), 15); + ASSERT_EQ(int_chunked->length(), 15); + auto schema = ::arrow::schema( + {field("d", string_chunked->type()), field("i", int32()), field("f", float_type)}); + auto table = Table::Make(schema, {string_chunked, int_chunked, float_chunked}); + ASSERT_OK(table->ValidateFull()); + + AssertSortIndicesMatchDecoded(table, {FieldRef("d")}); + AssertSortIndicesMatchDecoded(table, {FieldRef("f")}); + AssertSortIndicesMatchDecoded(table, {FieldRef("d"), FieldRef("i")}); + AssertSortIndicesMatchDecoded(table, {FieldRef("i"), FieldRef("d")}); + AssertSortIndicesMatchDecoded(table, {FieldRef("f"), FieldRef("d")}); + AssertSortIndicesMatchDecoded(table, {FieldRef("d"), FieldRef("f")}); +} + +TEST_F(TestDictionarySortIndices, TableRandom) { + random::RandomArrayGenerator rng(/*seed=*/0xdec0de); + for (int num_columns : {2, 9}) { + ARROW_SCOPED_TRACE("num_columns = ", num_columns); + // Every batch is generated separately, so each carries its own dictionaries + RecordBatchVector batches; + for (int64_t length : {50, 0, 120, 1, 80}) { + batches.push_back(RandomBatch(rng, length, num_columns)); + } + ASSERT_OK_AND_ASSIGN(auto table, Table::FromRecordBatches(batches)); + AssertSortIndicesMatchDecoded(table, ColumnRefs(*table->schema())); + AssertSortIndicesMatchDecoded(table, {FieldRef("f0")}); + } +} + +TEST_F(TestDictionarySortIndices, ManyOverlappingDictionaries) { + // Many chunks whose short random dictionaries repeat each other's values (and + // contain nulls), so ranks must be deduplicated across dictionaries. + random::RandomArrayGenerator rng(/*seed=*/0x0ddba11); + ArrayVector dict_chunks; + ArrayVector int_chunks; + for (int i = 0; i < 40; ++i) { + const int64_t length = (i % 7 == 0) ? 0 : 20 + 13 * (i % 5); + dict_chunks.push_back(RandomStringDictionary(rng, length, /*dict_length=*/12)); + int_chunks.push_back(rng.Int32(length, 0, 2, /*null_probability=*/0.1)); + } + ASSERT_OK_AND_ASSIGN(auto dict_chunked, + ChunkedArray::Make(dict_chunks, StringDictType())); + ASSERT_OK_AND_ASSIGN(auto int_chunked, ChunkedArray::Make(int_chunks, int32())); + ASSERT_OK(dict_chunked->ValidateFull()); + + AssertArraySortIndicesMatchDecoded(dict_chunked); + AssertSortIndicesMatchDecoded(dict_chunked, {}); + + auto table = Table::Make(schema({field("d", StringDictType()), field("i", int32())}), + {dict_chunked, int_chunked}); + AssertSortIndicesMatchDecoded(table, {FieldRef("d")}); + AssertSortIndicesMatchDecoded(table, {FieldRef("d"), FieldRef("i")}); + AssertSortIndicesMatchDecoded(table, {FieldRef("i"), FieldRef("d")}); +} + +TEST_F(TestDictionarySortIndices, ValueTypes) { + struct Case { + std::shared_ptr value_type; + std::string dict0; + std::string dict1; + }; + const std::string ints0 = "[3, 1, null, 2]"; + const std::string ints1 = "[2, 7, 1, null]"; + const std::string strs0 = R"(["c", "a", null, "b"])"; + const std::string strs1 = R"(["b", "g", "a", null])"; + const std::vector cases = { + {timestamp(TimeUnit::MILLI), ints0, ints1}, + {timestamp(TimeUnit::MICRO, "UTC"), ints0, ints1}, + {timestamp(TimeUnit::NANO, "America/Edmonton"), ints0, ints1}, + {date32(), ints0, ints1}, + {time64(TimeUnit::MICRO), ints0, ints1}, + {duration(TimeUnit::SECOND), ints0, ints1}, + {int64(), ints0, ints1}, + {decimal128(10, 2), R"(["3.00", "1.50", null, "2.25"])", + R"(["2.25", "7.00", "1.50", null])"}, + {binary(), strs0, strs1}, + {large_utf8(), strs0, strs1}, + // UUID-like keys + {fixed_size_binary(16), + R"(["cccccccccccccccc", "aaaaaaaaaaaaaaaa", null, "bbbbbbbbbbbbbbbb"])", + R"(["bbbbbbbbbbbbbbbb", "gggggggggggggggg", "aaaaaaaaaaaaaaaa", null])"}, + }; + for (const auto& c : cases) { + ARROW_SCOPED_TRACE("value type = ", c.value_type->ToString()); + auto type = dictionary(int32(), c.value_type); + auto chunk0 = DictArrayFromJSON(type, "[0, 1, 2, 3, 0, null]", c.dict0); + auto chunk1 = DictArrayFromJSON(type, "[3, 0, 1, 2, 2]", c.dict1); + ASSERT_OK_AND_ASSIGN(auto d, ChunkedArray::Make({chunk0, chunk1})); + auto i = ChunkedArrayFromJSON(int32(), {"[1, 0, 1, 0, 1, 0]", "[0, 1, 0, 1, 0]"}); + + AssertArraySortIndicesMatchDecoded(chunk0); + AssertArraySortIndicesMatchDecoded(d); + auto batch = + RecordBatch::Make(schema({field("d", type), field("i", int32())}), 6, + {chunk0, ArrayFromJSON(int32(), "[1, 0, 1, 0, 1, 0]")}); + AssertSortIndicesMatchDecoded(batch, {FieldRef("d"), FieldRef("i")}); + auto table = Table::Make(schema({field("d", type), field("i", int32())}), {d, i}); + AssertSortIndicesMatchDecoded(table, {FieldRef("d"), FieldRef("i")}); + AssertSortIndicesMatchDecoded(table, {FieldRef("i"), FieldRef("d")}); + } +} + +TEST_F(TestDictionarySortIndices, Nested) { + auto dict_type = StringDictType(); + auto d = DictArrayFromJSON(dict_type, "[0, 1, 2, null, 0, 3, 2]", + R"(["b", "a", "c", null])"); + auto x = ArrayFromJSON(int32(), "[1, 0, 1, 0, 0, 1, 0]"); + std::shared_ptr parent_bitmap; + ASSERT_OK(GetBitmapFromVector({1, 1, 0, 1, 1, 1, 1}, &parent_bitmap)); + ASSERT_OK_AND_ASSIGN( + auto s, StructArray::Make({d, x}, {field("d", dict_type), field("x", int32())}, + parent_bitmap)); + auto batch = RecordBatch::Make(schema({field("s", s->type())}), s->length(), {s}); + ASSERT_OK(batch->ValidateFull()); + + // A struct key expands to its leaf fields + AssertSortIndicesMatchDecoded(batch, {FieldRef("s")}); + AssertSortIndicesMatchDecoded(batch, {FieldRef("s", "d"), FieldRef("s", "x")}); + AssertSortIndicesMatchDecoded(batch, {FieldRef("s", "x"), FieldRef("s", "d")}); + + ASSERT_OK_AND_ASSIGN(auto table, + Table::FromRecordBatches({batch->Slice(0, 3), batch->Slice(3)})); + AssertSortIndicesMatchDecoded(table, {FieldRef("s", "d"), FieldRef("s", "x")}); + AssertSortIndicesMatchDecoded(table, {FieldRef("s", "x"), FieldRef("s", "d")}); + // Sorting a struct array goes through SortStructArray + AssertArraySortIndicesMatchDecoded(s); +} + // ---------------------------------------------------------------------- // Tests for Rank, Quantile Rank and Normal Rank diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py index d4d44e9347e0..69b75557a47a 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/python/pyarrow/tests/test_compute.py @@ -3425,6 +3425,86 @@ def test_sort_indices_table(): pc.sort_indices(table, sort_keys=[("a", "nonscending", "at_end")]) +@pytest.mark.parametrize("order", ["ascending", "descending"]) +@pytest.mark.parametrize("null_placement", ["at_start", "at_end"]) +def test_sort_indices_dictionary(order, null_placement): + # Chunks carry different, overlapping dictionaries, including null + # dictionary values and NaNs. + d_type = pa.dictionary(pa.int32(), pa.string()) + d = pa.chunked_array([ + pa.DictionaryArray.from_arrays( + pa.array([0, 1, None, 2, 1], pa.int32()), pa.array(["b", "a", "c"])), + pa.DictionaryArray.from_arrays( + pa.array([3, 0, 1, 2], pa.int32()), pa.array(["c", "d", "a", None])), + ], type=d_type) + f_type = pa.dictionary(pa.int8(), pa.float64()) + f = pa.chunked_array([ + pa.DictionaryArray.from_arrays( + pa.array([0, 1, 2, 3, 1], pa.int8()), + pa.array([1.0, float("nan"), 0.5, None])), + pa.DictionaryArray.from_arrays( + pa.array([1, 0, None, 1], pa.int8()), pa.array([float("nan"), 2.0])), + ], type=f_type) + i = pa.chunked_array([[1, 0, 1, 0, 1], [0, None, 1, 0]]) + + def decode(chunked): + return chunked.cast(chunked.type.value_type) + + for column in (d, f): + expected = pc.array_sort_indices( + decode(column), order=order, null_placement=null_placement) + result = pc.array_sort_indices( + column, order=order, null_placement=null_placement) + assert result.equals(expected) + for chunk in column.chunks: + expected = pc.array_sort_indices( + chunk.dictionary_decode(), order=order, + null_placement=null_placement) + result = pc.array_sort_indices( + chunk, order=order, null_placement=null_placement) + assert result.equals(expected) + + table = pa.table({"d": d, "f": f, "i": i}) + decoded_table = pa.table({"d": decode(d), "f": decode(f), "i": i}) + cases = [(table, decoded_table), + *zip(table.to_batches(), decoded_table.to_batches())] + for keys in (["d"], ["f"], ["d", "i"], ["i", "d"], ["f", "d"], ["d", "f", "i"]): + sort_keys = [(key, order, null_placement) for key in keys] + for data, decoded in cases: + result = pc.sort_indices(data, sort_keys=sort_keys) + expected = pc.sort_indices(decoded, sort_keys=sort_keys) + assert result.equals(expected), (type(data), keys) + + +@pytest.mark.parametrize("value_type", [ + pa.timestamp("us", tz="UTC"), + pa.date32(), + pa.binary(16), +]) +def test_sort_indices_dictionary_value_types(value_type): + raw = [3, 1, None, 2, 1, 7] + if pa.types.is_fixed_size_binary(value_type): + raw = [None if v is None else bytes([v]) * 16 for v in raw] + values = pa.array(raw, value_type) + chunks = [values.slice(0, 3), values.slice(3)] + d = pa.chunked_array([c.dictionary_encode() for c in chunks]) + table = pa.table({"d": d, "i": [1, 0, 1, 0, 1, 0]}) + decoded = pa.table({"d": pa.chunked_array(chunks), "i": table["i"]}) + for sort_keys in ([("d", "ascending")], + [("d", "descending", "at_start"), ("i", "ascending")]): + result = pc.sort_indices(table, sort_keys=sort_keys) + assert result.equals(pc.sort_indices(decoded, sort_keys=sort_keys)) + + +def test_sort_indices_dictionary_nan(): + arr = pa.DictionaryArray.from_arrays( + pa.array([0, 1, 2, 3, 4]), pa.array([1.0, float("nan"), 0.5, None, 2.0])) + result = pc.array_sort_indices(arr, order="ascending", null_placement="at_start") + assert result.to_pylist() == [3, 1, 2, 0, 4] + result = pc.array_sort_indices(arr, order="descending", null_placement="at_end") + assert result.to_pylist() == [4, 0, 2, 1, 3] + + def test_is_in(): arr = pa.array([1, 2, None, 1, 2, 3]) diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py index eb02051693f2..bcce4b1e8f81 100644 --- a/python/pyarrow/tests/test_table.py +++ b/python/pyarrow/tests/test_table.py @@ -3548,6 +3548,51 @@ def test_table_sort_by(cls): assert sorted_tab_dict["b"] == ["foo", "car", "bar", "foobar"] +@pytest.mark.parametrize( + ('cls'), + [ + (pa.Table), + (pa.RecordBatch) + ] +) +def test_table_sort_by_dictionary(cls): + table = cls.from_arrays([ + pa.array(["b", "a", None, "b", "a", "c"]).dictionary_encode(), + pa.array([3, 1, 4, 2, 5, 0]), + ], names=["keys", "values"]) + + assert table.sort_by("keys").to_pydict() == { + "keys": ["a", "a", "b", "b", "c", None], + "values": [1, 5, 3, 2, 0, 4] + } + + assert table.sort_by( + [("keys", "descending", "at_start"), ("values", "ascending", "at_end")] + ).to_pydict() == { + "keys": [None, "c", "b", "b", "a", "a"], + "values": [4, 0, 2, 3, 1, 5] + } + + assert table.sort_by( + [("values", "descending", "at_end"), ("keys", "ascending", "at_end")] + ).to_pydict()["values"] == [5, 4, 3, 2, 1, 0] + + if cls is pa.Table: + # Batches carry different dictionaries + table = pa.Table.from_batches([ + pa.record_batch({"keys": pa.array(["b", "c"]).dictionary_encode(), + "values": [1, 2]}), + pa.record_batch({"keys": pa.array(["c", "a", "b"]).dictionary_encode(), + "values": [3, 4, 5]}), + ]) + assert table.sort_by( + [("keys", "ascending"), ("values", "descending")] + ).to_pydict() == { + "keys": ["a", "b", "b", "c", "c"], + "values": [4, 5, 1, 3, 2] + } + + def test_record_batch_sort(): rb = pa.RecordBatch.from_arrays([ pa.array([7, 35, 7, 5, None, None, 6, 5], type=pa.int64()),