From 85bed4ddad2c11bdbab145de34ff74770a55c912 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Fri, 6 Nov 2020 15:05:59 +0900 Subject: [PATCH 01/37] ARROW-8199: [C++] Add support for multi-column sort indices on Table Summary: * Deprecate SortToIndices() * Add SortIndices() because we use "sort_indices" as kernel name in #7240 * Add support for sort indices in descending order on Array * Rename existing "sort_indices" kernel to "array_sort_indices" and introduce "sort_indices" meta function to support RecordBatch and Table Benchmark: Summary: * No performance regression in existing sort on array * Same performance as sort on array when the number of sort keys is 1 * 1.5x-100x slower than sort on array when the number of sort keys is 2 Before: ---------------------------------------------------------------------------------- Benchmark Time CPU ---------------------------------------------------------------------------------- SortToIndicesInt64Count/32768/10000/min_time:1.000 15685 ns 15682 ns SortToIndicesInt64Count/32768/100/min_time:1.000 15961 ns 15957 ns SortToIndicesInt64Count/32768/10/min_time:1.000 16473 ns 16469 ns SortToIndicesInt64Count/32768/2/min_time:1.000 27993 ns 27987 ns SortToIndicesInt64Count/32768/1/min_time:1.000 5609 ns 5608 ns SortToIndicesInt64Count/32768/0/min_time:1.000 13143 ns 13141 ns SortToIndicesInt64Count/1048576/1/min_time:1.000 134695 ns 134670 ns SortToIndicesInt64Count/8388608/1/min_time:1.000 1243992 ns 1243260 ns SortToIndicesInt64Compare/32768/10000/min_time:1.000 193632 ns 193595 ns SortToIndicesInt64Compare/32768/100/min_time:1.000 194876 ns 194837 ns SortToIndicesInt64Compare/32768/10/min_time:1.000 182362 ns 182324 ns SortToIndicesInt64Compare/32768/2/min_time:1.000 111607 ns 111584 ns SortToIndicesInt64Compare/32768/1/min_time:1.000 5642 ns 5641 ns SortToIndicesInt64Compare/32768/0/min_time:1.000 190302 ns 190268 ns SortToIndicesInt64Compare/1048576/1/min_time:1.000 134743 ns 134718 ns SortToIndicesInt64Compare/8388608/1/min_time:1.000 1261404 ns 1249362 ns After: ------------------------------------------------------------------------------------- Benchmark Time CPU ------------------------------------------------------------------------------------- ArraySortIndicesInt64Count/32768/10000/min_time:1.000 14769 ns 14766 ns ArraySortIndicesInt64Count/32768/100/min_time:1.000 15207 ns 15204 ns ArraySortIndicesInt64Count/32768/10/min_time:1.000 15892 ns 15889 ns ArraySortIndicesInt64Count/32768/2/min_time:1.000 30107 ns 30100 ns ArraySortIndicesInt64Count/32768/1/min_time:1.000 5509 ns 5508 ns ArraySortIndicesInt64Count/32768/0/min_time:1.000 12699 ns 12696 ns ArraySortIndicesInt64Count/1048576/1/min_time:1.000 132585 ns 132557 ns ArraySortIndicesInt64Count/8388608/1/min_time:1.000 1236596 ns 1235842 ns ArraySortIndicesInt64Compare/32768/10000/min_time:1.000 193259 ns 193217 ns ArraySortIndicesInt64Compare/32768/100/min_time:1.000 190010 ns 189973 ns ArraySortIndicesInt64Compare/32768/10/min_time:1.000 179923 ns 179879 ns ArraySortIndicesInt64Compare/32768/2/min_time:1.000 111100 ns 111074 ns ArraySortIndicesInt64Compare/32768/1/min_time:1.000 5660 ns 5659 ns ArraySortIndicesInt64Compare/32768/0/min_time:1.000 186521 ns 186476 ns ArraySortIndicesInt64Compare/1048576/1/min_time:1.000 132863 ns 132832 ns ArraySortIndicesInt64Compare/8388608/1/min_time:1.000 1266383 ns 1265765 ns TableSortIndicesInt64Count/32768/10000/min_time:1.000 21812 ns 21807 ns TableSortIndicesInt64Count/32768/100/min_time:1.000 22494 ns 22490 ns TableSortIndicesInt64Count/32768/10/min_time:1.000 17300 ns 17296 ns TableSortIndicesInt64Count/32768/2/min_time:1.000 29927 ns 29921 ns TableSortIndicesInt64Count/32768/1/min_time:1.000 5877 ns 5875 ns TableSortIndicesInt64Count/32768/0/min_time:1.000 20394 ns 20390 ns TableSortIndicesInt64Count/1048576/1/min_time:1.000 132904 ns 132871 ns TableSortIndicesInt64Count/8388608/1/min_time:1.000 1342693 ns 1341943 ns TableSortIndicesInt64Compare/32768/10000/min_time:1.000 203163 ns 203106 ns TableSortIndicesInt64Compare/32768/100/min_time:1.000 199531 ns 199477 ns TableSortIndicesInt64Compare/32768/10/min_time:1.000 185968 ns 185916 ns TableSortIndicesInt64Compare/32768/2/min_time:1.000 113571 ns 113540 ns TableSortIndicesInt64Compare/32768/1/min_time:1.000 6251 ns 6249 ns TableSortIndicesInt64Compare/32768/0/min_time:1.000 183650 ns 183609 ns TableSortIndicesInt64Compare/1048576/1/min_time:1.000 131701 ns 131674 ns TableSortIndicesInt64Compare/8388608/1/min_time:1.000 1264413 ns 1263622 ns TableSortIndicesInt64Int64/32768/10000/min_time:1.000 313368 ns 313310 ns TableSortIndicesInt64Int64/32768/100/min_time:1.000 313425 ns 313361 ns TableSortIndicesInt64Int64/32768/10/min_time:1.000 337051 ns 336987 ns TableSortIndicesInt64Int64/32768/2/min_time:1.000 402063 ns 401973 ns TableSortIndicesInt64Int64/32768/1/min_time:1.000 254660 ns 254612 ns TableSortIndicesInt64Int64/32768/0/min_time:1.000 305887 ns 305815 ns TableSortIndicesInt64Int64/1048576/1/min_time:1.000 11157920 ns 11155020 ns TableSortIndicesInt64Int64/8388608/1/min_time:1.000 106529839 ns 106501576 ns Follow-up tasks: * Improve performance when the number of sort keys is 2 or greater * Idea1: Use counting sort for small range data like on array * Idea2: Use threading and merge sort * Add support multi-column partition Nth indices on Table --- c_glib/arrow-glib/compute.cpp | 37 +- c_glib/arrow-glib/compute.h | 24 +- c_glib/arrow-glib/version.h.in | 23 + c_glib/doc/arrow-glib/arrow-glib-docs.xml | 4 + ...ort-to-indices.rb => test-sort-indices.rb} | 10 +- cpp/src/arrow/compute/api_vector.cc | 22 +- cpp/src/arrow/compute/api_vector.h | 82 +++- .../kernels/vector_selection_benchmark.cc | 2 +- cpp/src/arrow/compute/kernels/vector_sort.cc | 464 ++++++++++++++++-- .../compute/kernels/vector_sort_benchmark.cc | 101 +++- .../arrow/compute/kernels/vector_sort_test.cc | 382 +++++++++++--- cpp/src/arrow/dataset/filter.cc | 2 +- 12 files changed, 1000 insertions(+), 153 deletions(-) rename c_glib/test/{test-sort-to-indices.rb => test-sort-indices.rb} (85%) diff --git a/c_glib/arrow-glib/compute.cpp b/c_glib/arrow-glib/compute.cpp index 777adee41a5..25d950db93f 100644 --- a/c_glib/arrow-glib/compute.cpp +++ b/c_glib/arrow-glib/compute.cpp @@ -2177,29 +2177,52 @@ garrow_array_is_in_chunked_array(GArrowArray *left, } /** - * garrow_array_sort_to_indices: + * garrow_array_sort_indices: * @array: A #GArrowArray. + * @order: The order for sort. * @error: (nullable): Return location for a #GError or %NULL. * * Returns: (nullable) (transfer full): The indices that would sort - * an array on success, %NULL on error. + * an array in the specified order on success, %NULL on error. * - * Since: 0.15.0 + * Since: 3.0.0 */ GArrowUInt64Array * -garrow_array_sort_to_indices(GArrowArray *array, - GError **error) +garrow_array_sort_indices(GArrowArray *array, + GArrowSortOrder order, + GError **error) { auto arrow_array = garrow_array_get_raw(array); auto arrow_array_raw = arrow_array.get(); - auto arrow_indices_array = arrow::compute::SortToIndices(*arrow_array_raw); - if (garrow::check(error, arrow_indices_array, "[array][sort-to-indices]")) { + auto arrow_order = static_cast(order); + auto arrow_indices_array = + arrow::compute::SortIndices(*arrow_array_raw, arrow_order); + if (garrow::check(error, arrow_indices_array, "[array][sort-indices]")) { return GARROW_UINT64_ARRAY(garrow_array_new_raw(&(*arrow_indices_array))); } else { return NULL; } } +/** + * garrow_array_sort_to_indices: + * @array: A #GArrowArray. + * @error: (nullable): Return location for a #GError or %NULL. + * + * Returns: (nullable) (transfer full): The indices that would sort + * an array in ascending order on success, %NULL on error. + * + * Since: 0.15.0 + * + * Deprecated: 3.0.0: Use garrow_array_sort_indices() instead. + */ +GArrowUInt64Array * +garrow_array_sort_to_indices(GArrowArray *array, + GError **error) +{ + return garrow_array_sort_indices(array, GARROW_SORT_ORDER_ASCENDING, error); +} + /** * garrow_table_filter: * @table: A #GArrowTable. diff --git a/c_glib/arrow-glib/compute.h b/c_glib/arrow-glib/compute.h index 48fdc3a1c1e..2e304d6e5a6 100644 --- a/c_glib/arrow-glib/compute.h +++ b/c_glib/arrow-glib/compute.h @@ -88,7 +88,7 @@ GArrowCastOptions *garrow_cast_options_new(void); * @GARROW_COUNT_ALL: Count all non-null values. * @GARROW_COUNT_NULL: Count all null values. * - * They are corresponding to `arrow::compute::CountOptions::mode` values. + * They are corresponding to `arrow::compute::CountOptions::Mode` values. */ typedef enum { GARROW_COUNT_ALL, @@ -377,10 +377,32 @@ GArrowBooleanArray * garrow_array_is_in_chunked_array(GArrowArray *left, GArrowChunkedArray *right, GError **error); + +/** + * GArrowSortOrder: + * @GARROW_SORT_ORDER_ASCENDING: Sort in ascending order. + * @GARROW_SORT_ORDER_DESCENDING: Sort in descending order. + * + * They are corresponding to `arrow::compute::SortOrder` values. + * + * Since: 3.0.0 + */ +typedef enum { + GARROW_SORT_ORDER_ASCENDING, + GARROW_SORT_ORDER_DESCENDING, +} GArrowSortOrder; + +GARROW_AVAILABLE_IN_3_0 +GArrowUInt64Array * +garrow_array_sort_indices(GArrowArray *array, + GArrowSortOrder order, + GError **error); +GARROW_DEPRECATED_IN_3_0_FOR(garrow_array_sort_indices) GARROW_AVAILABLE_IN_0_15 GArrowUInt64Array * garrow_array_sort_to_indices(GArrowArray *array, GError **error); + GARROW_AVAILABLE_IN_0_16 GArrowTable * garrow_table_filter(GArrowTable *table, diff --git a/c_glib/arrow-glib/version.h.in b/c_glib/arrow-glib/version.h.in index 0d2069fa626..4ad59afd1e3 100644 --- a/c_glib/arrow-glib/version.h.in +++ b/c_glib/arrow-glib/version.h.in @@ -110,6 +110,15 @@ # define GARROW_UNAVAILABLE(major, minor) G_UNAVAILABLE(major, minor) #endif +/** + * GARROW_VERSION_3_0: + * + * You can use this macro value for compile time API version check. + * + * Since: 3.0.0 + */ +#define GARROW_VERSION_3_0 G_ENCODE_VERSION(3, 0) + /** * GARROW_VERSION_2_0: * @@ -229,6 +238,20 @@ #define GARROW_AVAILABLE_IN_ALL +#if GARROW_VERSION_MIN_REQUIRED >= GARROW_VERSION_3_0 +# define GARROW_DEPRECATED_IN_3_0 GARROW_DEPRECATED +# define GARROW_DEPRECATED_IN_3_0_FOR(function) GARROW_DEPRECATED_FOR(function) +#else +# define GARROW_DEPRECATED_IN_3_0 +# define GARROW_DEPRECATED_IN_3_0_FOR(function) +#endif + +#if GARROW_VERSION_MAX_ALLOWED < GARROW_VERSION_3_0 +# define GARROW_AVAILABLE_IN_3_0 GARROW_UNAVAILABLE(3, 0) +#else +# define GARROW_AVAILABLE_IN_3_0 +#endif + #if GARROW_VERSION_MIN_REQUIRED >= GARROW_VERSION_2_0 # define GARROW_DEPRECATED_IN_2_0 GARROW_DEPRECATED # define GARROW_DEPRECATED_IN_2_0_FOR(function) GARROW_DEPRECATED_FOR(function) diff --git a/c_glib/doc/arrow-glib/arrow-glib-docs.xml b/c_glib/doc/arrow-glib/arrow-glib-docs.xml index 72a01f50e4f..bbb858eef7d 100644 --- a/c_glib/doc/arrow-glib/arrow-glib-docs.xml +++ b/c_glib/doc/arrow-glib/arrow-glib-docs.xml @@ -179,6 +179,10 @@ Index of deprecated API + + Index of new symbols in 3.0.0 + + Index of new symbols in 2.0.0 diff --git a/c_glib/test/test-sort-to-indices.rb b/c_glib/test/test-sort-indices.rb similarity index 85% rename from c_glib/test/test-sort-to-indices.rb rename to c_glib/test/test-sort-indices.rb index 355efd7e1d9..d7745c75f25 100644 --- a/c_glib/test/test-sort-to-indices.rb +++ b/c_glib/test/test-sort-indices.rb @@ -15,20 +15,20 @@ # specific language governing permissions and limitations # under the License. -class TestSortToIndices < Test::Unit::TestCase +class TestSortIndices < Test::Unit::TestCase include Helper::Buildable sub_test_case("Integer") do def test_no_null array = build_int16_array([1, 0, 4, -3]) assert_equal(build_uint64_array([3, 1, 0, 2]), - array.sort_to_indices) + array.sort_indices(:ascending)) end def test_null array = build_int16_array([nil, 1, 0, nil, 4, 3]) assert_equal(build_uint64_array([2, 1, 5, 4, 0, 3]), - array.sort_to_indices) + array.sort_indices(:ascending)) end end @@ -36,13 +36,13 @@ def test_null def test_no_null array = build_string_array(["hello", "world", "a", "z"]) assert_equal(build_uint64_array([2, 0, 1, 3]), - array.sort_to_indices) + array.sort_indices(:ascending)) end def test_null array = build_string_array([nil, "b", "a", nil, "c", "d"]) assert_equal(build_uint64_array([2, 1, 4, 5, 0, 3]), - array.sort_to_indices) + array.sort_indices(:ascending)) end end end diff --git a/cpp/src/arrow/compute/api_vector.cc b/cpp/src/arrow/compute/api_vector.cc index 1f6972860ae..9b801b1cad3 100644 --- a/cpp/src/arrow/compute/api_vector.cc +++ b/cpp/src/arrow/compute/api_vector.cc @@ -46,8 +46,22 @@ Result> NthToIndices(const Array& values, int64_t n, return result.make_array(); } -Result> SortToIndices(const Array& values, ExecContext* ctx) { - ARROW_ASSIGN_OR_RAISE(Datum result, CallFunction("sort_indices", {Datum(values)}, ctx)); +Result> SortIndices(const Array& values, ExecContext* ctx) { + return SortIndices(values, SortOrder::ASCENDING, ctx); +} + +Result> SortIndices(const Array& values, SortOrder order, + ExecContext* ctx) { + ArraySortOptions options(order); + ARROW_ASSIGN_OR_RAISE( + Datum result, CallFunction("array_sort_indices", {Datum(values)}, &options, ctx)); + return result.make_array(); +} + +Result> SortIndices(const Table& values, + const SortOptions& options, ExecContext* ctx) { + ARROW_ASSIGN_OR_RAISE(Datum result, + CallFunction("sort_indices", {Datum(values)}, &options, ctx)); return result.make_array(); } @@ -135,5 +149,9 @@ Result> Take(const Table& table, const ChunkedArray& indi return result.table(); } +Result> SortToIndices(const Array& values, ExecContext* ctx) { + return SortIndices(values, ctx); +} + } // namespace compute } // namespace arrow diff --git a/cpp/src/arrow/compute/api_vector.h b/cpp/src/arrow/compute/api_vector.h index 2c77e8ee155..5f4d8bccdf0 100644 --- a/cpp/src/arrow/compute/api_vector.h +++ b/cpp/src/arrow/compute/api_vector.h @@ -58,6 +58,34 @@ struct ARROW_EXPORT TakeOptions : public FunctionOptions { static TakeOptions Defaults() { return BoundsCheck(); } }; +enum SortOrder { + ASCENDING, + DESCENDING, +}; + +/// \brief One sort key for PartitionNthIndices (TODO) and SortIndices +struct ARROW_EXPORT SortKey { + explicit SortKey(std::string name, SortOrder order = ASCENDING) + : name(name), order(order) {} + + /// The name of the sort key. + std::string name; + /// How to order by this sort key. + SortOrder order; +}; + +struct ARROW_EXPORT ArraySortOptions : public FunctionOptions { + explicit ArraySortOptions(SortOrder order = SortOrder::ASCENDING) : order(order) {} + + SortOrder order; +}; + +struct ARROW_EXPORT SortOptions : public FunctionOptions { + explicit SortOptions(std::vector sort_keys = {}) : sort_keys(sort_keys) {} + + std::vector sort_keys; +}; + /// \brief Partitioning options for NthToIndices struct ARROW_EXPORT PartitionNthOptions : public FunctionOptions { explicit PartitionNthOptions(int64_t pivot) : pivot(pivot) {} @@ -152,7 +180,8 @@ ARROW_EXPORT Result> NthToIndices(const Array& values, int64_t n, ExecContext* ctx = NULLPTR); -/// \brief Returns the indices that would sort an array. +/// \brief Returns the indices that would sort an array in ascending +/// order. /// /// Perform an indirect sort of array. The output array will contain /// indices that would sort an array, which would be the same length @@ -165,8 +194,50 @@ Result> NthToIndices(const Array& values, int64_t n, /// \param[in] ctx the function execution context, optional /// \return offsets indices that would sort an array ARROW_EXPORT -Result> SortToIndices(const Array& values, - ExecContext* ctx = NULLPTR); +Result> SortIndices(const Array& values, + ExecContext* ctx = NULLPTR); + +/// \brief Returns the indices that would sort an array in the +/// specified order. +/// +/// Perform an indirect sort of array. The output array will contain +/// indices that would sort an array, which would be the same length +/// as input. Nulls will be stably partitioned to the end of the output. +/// +/// For example given values = [null, 1, 3.3, null, 2, 5.3] and order +/// = SortOrder::DESCENDING, the output will be [5, 2, 4, 1, 0, +/// 3]. +/// +/// \param[in] values array to sort +/// \param[in] order ascending or descending +/// \param[in] ctx the function execution context, optional +/// \return offsets indices that would sort an array +ARROW_EXPORT +Result> SortIndices(const Array& values, SortOrder order, + ExecContext* ctx = NULLPTR); + +/// \brief Returns the indices that would sort a table in the +/// specified order. +/// +/// Perform an indirect sort of table. The output array will contain +/// indices that would sort a table, which would be the same length +/// as input. Nulls will be stably partitioned to the end of the output. +/// +/// For example given table = { +/// "column1": [null, 1, 3, null, 2, 1], +/// "column2": [ 5, 3, null, null, 5, 5], +/// } and options = { +/// {"column1", SortOrder::ASCENDING}, +/// {"column2", SortOrder::DESCENDING}, +/// }, the output will be [5, 1, 4, 2, 0, 3]. +/// +/// \param[in] table table to sort +/// \param[in] options options +/// \param[in] ctx the function execution context, optional +/// \return offsets indices that would sort a table +ARROW_EXPORT +Result> SortIndices(const Table& table, const SortOptions& options, + ExecContext* ctx = NULLPTR); /// \brief Compute unique elements from an array-like object /// @@ -254,5 +325,10 @@ Result> Take(const Table& table, const ChunkedArray& indi const TakeOptions& options = TakeOptions::Defaults(), ExecContext* context = NULLPTR); +ARROW_DEPRECATED("Deprecated in 3.0.0. Use SortIndices()") +ARROW_EXPORT +Result> SortToIndices(const Array& values, + ExecContext* ctx = NULLPTR); + } // namespace compute } // namespace arrow diff --git a/cpp/src/arrow/compute/kernels/vector_selection_benchmark.cc b/cpp/src/arrow/compute/kernels/vector_selection_benchmark.cc index c595736912d..7b181eaec6e 100644 --- a/cpp/src/arrow/compute/kernels/vector_selection_benchmark.cc +++ b/cpp/src/arrow/compute/kernels/vector_selection_benchmark.cc @@ -142,7 +142,7 @@ struct TakeBenchmark { indices_null_proportion); if (monotonic_indices) { - auto arg_sorter = *SortToIndices(*indices); + auto arg_sorter = *SortIndices(*indices); indices = *Take(*indices, *arg_sorter); } diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc index 989c75757f6..367ee8d79fd 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort.cc @@ -21,9 +21,11 @@ #include #include +#include "arrow/array/concatenate.h" #include "arrow/array/data.h" #include "arrow/compute/api_vector.h" #include "arrow/compute/kernels/common.h" +#include "arrow/table.h" #include "arrow/util/optional.h" namespace arrow { @@ -164,30 +166,38 @@ inline void VisitRawValuesInline(const ArrayType& values, } template -class CompareSorter { +class ArrayCompareSorter { using ArrayType = typename TypeTraits::ArrayType; public: - void Sort(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values) { + void Sort(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, + const ArraySortOptions& options) { std::iota(indices_begin, indices_end, 0); auto nulls_begin = PartitionNulls(indices_begin, indices_end, values); - std::stable_sort(indices_begin, nulls_begin, - [&values](uint64_t left, uint64_t right) { - return values.GetView(left) < values.GetView(right); - }); + if (options.order == SortOrder::ASCENDING) { + std::stable_sort(indices_begin, nulls_begin, + [&values](uint64_t left, uint64_t right) { + return values.GetView(left) < values.GetView(right); + }); + } else { + std::stable_sort(indices_begin, nulls_begin, + [&values](uint64_t left, uint64_t right) { + return values.GetView(left) > values.GetView(right); + }); + } } }; template -class CountSorter { +class ArrayCountSorter { using ArrayType = typename TypeTraits::ArrayType; using c_type = typename ArrowType::c_type; public: - CountSorter() = default; + ArrayCountSorter() = default; - explicit CountSorter(c_type min, c_type max) { SetMinMax(min, max); } + explicit ArrayCountSorter(c_type min, c_type max) { SetMinMax(min, max); } // Assume: max >= min && (max - min) < 4Gi void SetMinMax(c_type min, c_type max) { @@ -195,12 +205,13 @@ class CountSorter { value_range_ = static_cast(max - min) + 1; } - void Sort(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values) { + void Sort(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, + const ArraySortOptions& options) { // 32bit counter performs much better than 64bit one if (values.length() < (1LL << 32)) { - SortInternal(indices_begin, indices_end, values); + SortInternal(indices_begin, indices_end, values, options); } else { - SortInternal(indices_begin, indices_end, values); + SortInternal(indices_begin, indices_end, values, options); } } @@ -210,23 +221,35 @@ class CountSorter { template void SortInternal(uint64_t* indices_begin, uint64_t* indices_end, - const ArrayType& values) { + const ArrayType& values, const ArraySortOptions& options) { const uint32_t value_range = value_range_; // first slot reserved for prefix sum std::vector counts(1 + value_range); - VisitRawValuesInline( - values, [&](c_type v) { ++counts[v - min_ + 1]; }, []() {}); - - for (uint32_t i = 1; i <= value_range; ++i) { - counts[i] += counts[i - 1]; + if (options.order == SortOrder::ASCENDING) { + VisitRawValuesInline( + values, [&](c_type v) { ++counts[v - min_ + 1]; }, []() {}); + for (uint32_t i = 1; i <= value_range; ++i) { + counts[i] += counts[i - 1]; + } + uint32_t null_position = counts[value_range]; + int64_t index = 0; + VisitRawValuesInline( + values, [&](c_type v) { indices_begin[counts[v - min_]++] = index++; }, + [&]() { indices_begin[null_position++] = index++; }); + } else { + VisitRawValuesInline( + values, [&](c_type v) { ++counts[v - min_]; }, []() {}); + for (uint32_t i = value_range; i >= 1; --i) { + counts[i - 1] += counts[i]; + } + uint32_t null_position = counts[0]; + int64_t index = 0; + VisitRawValuesInline( + values, [&](c_type v) { indices_begin[counts[v - min_ + 1]++] = index++; }, + [&]() { indices_begin[null_position++] = index++; }); } - - int64_t index = 0; - VisitRawValuesInline( - values, [&](c_type v) { indices_begin[counts[v - min_]++] = index++; }, - [&]() { indices_begin[counts[value_range]++] = index++; }); } }; @@ -234,12 +257,13 @@ class CountSorter { // - Use O(n) counting sort if values are in a small range // - Use O(nlogn) std::stable_sort otherwise template -class CountOrCompareSorter { +class ArrayCountOrCompareSorter { using ArrayType = typename TypeTraits::ArrayType; using c_type = typename ArrowType::c_type; public: - void Sort(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values) { + void Sort(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, + const ArraySortOptions& options) { if (values.length() >= countsort_min_len_ && values.length() > values.null_count()) { c_type min{std::numeric_limits::max()}; c_type max{std::numeric_limits::min()}; @@ -257,17 +281,17 @@ class CountOrCompareSorter { if (static_cast(max) - static_cast(min) <= countsort_max_range_) { count_sorter_.SetMinMax(min, max); - count_sorter_.Sort(indices_begin, indices_end, values); + count_sorter_.Sort(indices_begin, indices_end, values, options); return; } } - compare_sorter_.Sort(indices_begin, indices_end, values); + compare_sorter_.Sort(indices_begin, indices_end, values, options); } private: - CompareSorter compare_sorter_; - CountSorter count_sorter_; + ArrayCompareSorter compare_sorter_; + ArrayCountSorter count_sorter_; // Cross point to prefer counting sort than stl::stable_sort(merge sort) // - array to be sorted is longer than "count_min_len_" @@ -283,36 +307,40 @@ class CountOrCompareSorter { }; template -struct Sorter; +struct ArraySorter; template <> -struct Sorter { - CountSorter impl; - Sorter() : impl(0, 255) {} +struct ArraySorter { + ArrayCountSorter impl; + ArraySorter() : impl(0, 255) {} }; template <> -struct Sorter { - CountSorter impl; - Sorter() : impl(-128, 127) {} +struct ArraySorter { + ArrayCountSorter impl; + ArraySorter() : impl(-128, 127) {} }; template -struct Sorter::value && - (sizeof(typename Type::c_type) > 1)>> { - CountOrCompareSorter impl; +struct ArraySorter::value && + (sizeof(typename Type::c_type) > 1)>> { + ArrayCountOrCompareSorter impl; }; template -struct Sorter::value || - is_base_binary_type::value>> { - CompareSorter impl; +struct ArraySorter::value || + is_base_binary_type::value>> { + ArrayCompareSorter impl; }; +using ArraySortIndicesState = internal::OptionsWrapper; + template -struct SortIndices { +struct ArraySortIndices { using ArrayType = typename TypeTraits::ArrayType; static void Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) { + auto& options = ArraySortIndicesState::Get(ctx); + std::shared_ptr arg0; KERNEL_RETURN_IF_ERROR( ctx, @@ -322,8 +350,8 @@ struct SortIndices { uint64_t* out_begin = out_arr->GetMutableValues(1); uint64_t* out_end = out_begin + arr.length(); - Sorter sorter; - sorter.impl.Sort(out_begin, out_end, arr); + ArraySorter sorter; + sorter.impl.Sort(out_begin, out_end, arr, options); } }; @@ -346,14 +374,347 @@ void AddSortingKernels(VectorKernel base, VectorFunction* func) { } } +class TableSorter : public TypeVisitor { + private: + struct ResolvedSortKey { + ResolvedSortKey(const ChunkedArray& chunked_array, const SortOrder order) + : order(order) { + type = chunked_array.type().get(); + has_nulls = (chunked_array.null_count() > 0); + num_chunks = chunked_array.num_chunks(); + for (const auto& chunk : chunked_array.chunks()) { + chunks.push_back(chunk.get()); + } + } + + template + ArrayType* ResolveChunk(int64_t index, int64_t& chunk_index) const { + if (num_chunks == 1) { + chunk_index = index; + return static_cast(chunks[0]); + } else { + int64_t offset = 0; + for (size_t i = 0; i < num_chunks; ++i) { + if (index < offset + chunks[i]->length()) { + chunk_index = index - offset; + return static_cast(chunks[i]); + } + offset += chunks[i]->length(); + } + return nullptr; + } + } + + SortOrder order; + DataType* type; + bool has_nulls; + size_t num_chunks; + std::vector chunks; + }; + + class Comparer : public TypeVisitor { + public: + Comparer(const Table& table, const std::vector& sort_keys) + : TypeVisitor(), status_(Status::OK()) { + for (const auto& sort_key : sort_keys) { + const auto& chunked_array = table.GetColumnByName(sort_key.name); + if (!chunked_array) { + status_ = Status::Invalid("Nonexistent sort key column: ", sort_key.name); + return; + } + sort_keys_.emplace_back(*chunked_array, sort_key.order); + } + } + + Status status() { return status_; } + + const std::vector& sort_keys() { return sort_keys_; } + + bool Compare(uint64_t left, uint64_t right, size_t start_sort_key_index) { + current_left_ = left; + current_right_ = right; + current_compared_ = 0; + auto num_sort_keys = sort_keys_.size(); + for (size_t i = start_sort_key_index; i < num_sort_keys; ++i) { + current_sort_key_index_ = i; + status_ = sort_keys_[i].type->Accept(this); + if (current_compared_ != 0) { + break; + } + } + return current_compared_ < 0; + } + +#define VISIT(TYPE) \ + Status Visit(const TYPE##Type& type) override { \ + current_compared_ = CompareType(); \ + return Status::OK(); \ + } + + VISIT(Int8) + VISIT(Int16) + VISIT(Int32) + VISIT(Int64) + VISIT(UInt8) + VISIT(UInt16) + VISIT(UInt32) + VISIT(UInt64) + VISIT(Float) + VISIT(Double) + VISIT(String) + VISIT(Binary) + VISIT(LargeString) + VISIT(LargeBinary) + +#undef VISIT + + private: + template + int32_t CompareType() { + using ArrayType = typename TypeTraits::ArrayType; + const auto& sort_key = sort_keys_[current_sort_key_index_]; + auto order = sort_key.order; + int64_t index_left = 0; + auto array_left = sort_key.ResolveChunk(current_left_, index_left); + int64_t index_right = 0; + auto array_right = sort_key.ResolveChunk(current_right_, index_right); + if (sort_key.has_nulls) { + auto is_null_left = array_left->IsNull(index_left); + auto is_null_right = array_right->IsNull(index_right); + if (is_null_left && is_null_right) { + return 0; + } else if (is_null_left) { + return 1; + } else if (is_null_right) { + return -1; + } + } + auto left = array_left->GetView(index_left); + auto right = array_right->GetView(index_right); + int32_t compared; + if (left == right) { + compared = 0; + } else if (left > right) { + compared = 1; + } else { + compared = -1; + } + if (order == SortOrder::DESCENDING) { + compared = -compared; + } + return compared; + } + + Status status_; + std::vector sort_keys_; + int64_t current_left_; + int64_t current_right_; + size_t current_sort_key_index_; + int32_t current_compared_; + }; + + public: + TableSorter(uint64_t* indices_begin, uint64_t* indices_end, const Table& table, + const SortOptions& options) + : indices_begin_(indices_begin), + indices_end_(indices_end), + table_(table), + options_(options), + comparer_(table, options.sort_keys) {} + + Status Sort() { + ARROW_RETURN_NOT_OK(comparer_.status()); + return comparer_.sort_keys()[0].type->Accept(this); + } + +#define VISIT(TYPE) \ + Status Visit(const TYPE##Type& type) override { return SortInternal(); } + + VISIT(Int8) + VISIT(Int16) + VISIT(Int32) + VISIT(Int64) + VISIT(UInt8) + VISIT(UInt16) + VISIT(UInt32) + VISIT(UInt64) + VISIT(Float) + VISIT(Double) + VISIT(String) + VISIT(Binary) + VISIT(LargeString) + VISIT(LargeBinary) + +#undef VISIT + + private: + template + Status SortInternal() { + using ArrayType = typename TypeTraits::ArrayType; + std::iota(indices_begin_, indices_end_, 0); + + auto& comparer = comparer_; + const auto& first_sort_key = comparer.sort_keys()[0]; + auto nulls_begin = indices_end_; + if (first_sort_key.has_nulls > 0) { + nulls_begin = std::stable_partition( + indices_begin_, indices_end_, [&first_sort_key](uint64_t index) { + int64_t index_chunk = 0; + auto chunk = first_sort_key.ResolveChunk(index, index_chunk); + return !chunk->IsNull(index_chunk); + }); + std::stable_sort(nulls_begin, indices_end_, + [&comparer](uint64_t left, uint64_t right) { + return comparer.Compare(left, right, 1); + }); + } + std::stable_sort( + indices_begin_, nulls_begin, + [&first_sort_key, &comparer](uint64_t left, uint64_t right) { + int64_t index_left = 0; + auto array_left = first_sort_key.ResolveChunk(left, index_left); + int64_t index_right = 0; + auto array_right = first_sort_key.ResolveChunk(right, index_right); + auto value_left = array_left->GetView(index_left); + auto value_right = array_right->GetView(index_right); + if (value_left == value_right) { + return comparer.Compare(left, right, 1); + } else { + auto compared = value_left < value_right; + if (first_sort_key.order == SortOrder::ASCENDING) { + return compared; + } else { + return !compared; + } + } + }); + return Status::OK(); + } + + uint64_t* indices_begin_; + uint64_t* indices_end_; + const Table& table_; + const SortOptions& options_; + Comparer comparer_; +}; + const FunctionDoc sort_indices_doc( + "Return the indices that would sort an array, record batch or table", + ("This function computes an array of indices that define a stable sort\n" + "of the input array, record batch or table. Null values are considered\n" + "greater than any other value and are therefore sorted at the end of the\n" + "input. For floating-point types, NaNs are considered greater than any\n" + "other non-null value, but smaller than null values."), + {"input"}, "SortOptions"); + +class SortIndicesMetaFunction : public MetaFunction { + public: + SortIndicesMetaFunction() + : MetaFunction("sort_indices", Arity::Unary(), &sort_indices_doc) {} + + Result ExecuteImpl(const std::vector& args, + const FunctionOptions* options, + ExecContext* ctx) const override { + const SortOptions& sort_options = static_cast(*options); + switch (args[0].kind()) { + case Datum::ARRAY: + return SortIndices(*args[0].make_array(), sort_options, ctx); + break; + case Datum::CHUNKED_ARRAY: + return SortIndices(*args[0].chunked_array(), sort_options, ctx); + break; + case Datum::RECORD_BATCH: { + ARROW_ASSIGN_OR_RAISE(auto table, + Table::FromRecordBatches({args[0].record_batch()})); + return SortIndices(*table, sort_options, ctx); + } break; + case Datum::TABLE: + return SortIndices(*args[0].table(), sort_options, ctx); + break; + default: + break; + } + return Status::NotImplemented( + "Unsupported types for sort_indices operation: " + "values=", + args[0].ToString()); + } + + private: + Result> SortIndices(const Array& values, + const SortOptions& options, + ExecContext* ctx) const { + SortOrder order = SortOrder::ASCENDING; + if (!options.sort_keys.empty()) { + order = options.sort_keys[0].order; + } + ArraySortOptions array_options(order); + ARROW_ASSIGN_OR_RAISE( + Datum result, CallFunction("array_sort_indices", {values}, &array_options, ctx)); + return result.make_array(); + } + + Result> SortIndices(const ChunkedArray& values, + const SortOptions& options, + ExecContext* ctx) const { + SortOrder order = SortOrder::ASCENDING; + if (!options.sort_keys.empty()) { + order = options.sort_keys[0].order; + } + ArraySortOptions array_options(order); + + std::shared_ptr array_values; + if (values.num_chunks() == 1) { + array_values = values.chunk(0); + } else { + ARROW_ASSIGN_OR_RAISE(array_values, + Concatenate(values.chunks(), ctx->memory_pool())); + } + ARROW_ASSIGN_OR_RAISE(Datum result, CallFunction("array_sort_indices", {array_values}, + &array_options, ctx)); + return result.make_array(); + } + + Result SortIndices(const Table& table, const SortOptions& options, + ExecContext* ctx) const { + auto n_sort_keys = options.sort_keys.size(); + if (n_sort_keys == 0) { + return Status::Invalid("Must specify one or more sort keys"); + } + if (n_sort_keys == 1) { + auto chunked_array = table.GetColumnByName(options.sort_keys[0].name); + if (!chunked_array) { + return Status::Invalid("Nonexistent sort key column: ", + options.sort_keys[0].name); + } + return SortIndices(*chunked_array, options, ctx); + } + + auto out_type = uint64(); + auto length = table.num_rows(); + auto buffer_size = BitUtil::BytesForBits( + length * std::static_pointer_cast(out_type)->bit_width()); + std::vector> buffers(2); + ARROW_ASSIGN_OR_RAISE(buffers[1], + AllocateResizableBuffer(buffer_size, ctx->memory_pool())); + auto out = std::make_shared(out_type, length, buffers, 0); + auto out_begin = out->GetMutableValues(1); + auto out_end = out_begin + length; + + TableSorter sorter(out_begin, out_end, table, options); + ARROW_RETURN_NOT_OK(sorter.Sort()); + return Datum(out); + } +}; + +const FunctionDoc array_sort_indices_doc( "Return the indices that would sort an array", ("This function computes an array of indices that define a stable sort\n" "of the input array. Null values are considered greater than any\n" "other value and are therefore sorted at the end of the array.\n" "For floating-point types, NaNs are considered greater than any\n" "other non-null value, but smaller than null values."), - {"array"}); + {"array"}, "ArraySortOptions"); const FunctionDoc partition_nth_indices_doc( "Return the indices that would partition an array around a pivot", @@ -380,10 +741,13 @@ void RegisterVectorSort(FunctionRegistry* registry) { base.mem_allocation = MemAllocation::PREALLOCATE; base.null_handling = NullHandling::OUTPUT_NOT_NULL; - auto sort_indices = - std::make_shared("sort_indices", Arity::Unary(), &sort_indices_doc); - AddSortingKernels(base, sort_indices.get()); - DCHECK_OK(registry->AddFunction(std::move(sort_indices))); + auto array_sort_indices = std::make_shared( + "array_sort_indices", Arity::Unary(), &array_sort_indices_doc); + base.init = ArraySortIndicesState::Init; + AddSortingKernels(base, array_sort_indices.get()); + DCHECK_OK(registry->AddFunction(std::move(array_sort_indices))); + + DCHECK_OK(registry->AddFunction(std::make_shared())); // partition_nth_indices has a parameter so needs its init function auto part_indices = std::make_shared( diff --git a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc index be541b2dad6..2e6d54707ac 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc @@ -19,6 +19,7 @@ #include "arrow/compute/api_vector.h" #include "arrow/compute/kernels/test_util.h" +#include "arrow/table.h" #include "arrow/testing/gtest_util.h" #include "arrow/testing/random.h" #include "arrow/util/benchmark_util.h" @@ -27,15 +28,15 @@ namespace arrow { namespace compute { constexpr auto kSeed = 0x0ff1ce; -static void SortToIndicesBenchmark(benchmark::State& state, - const std::shared_ptr& values) { +static void ArraySortIndicesBenchmark(benchmark::State& state, + const std::shared_ptr& values) { for (auto _ : state) { - ABORT_NOT_OK(SortToIndices(*values).status()); + ABORT_NOT_OK(SortIndices(*values).status()); } state.SetItemsProcessed(state.iterations() * values->length()); } -static void SortToIndicesInt64Count(benchmark::State& state) { +static void ArraySortIndicesInt64Count(benchmark::State& state) { RegressionArgs args(state); const int64_t array_size = args.size / sizeof(int64_t); @@ -43,10 +44,10 @@ static void SortToIndicesInt64Count(benchmark::State& state) { auto values = rand.Int64(array_size, -100, 100, args.null_proportion); - SortToIndicesBenchmark(state, values); + ArraySortIndicesBenchmark(state, values); } -static void SortToIndicesInt64Compare(benchmark::State& state) { +static void ArraySortIndicesInt64Compare(benchmark::State& state) { RegressionArgs args(state); const int64_t array_size = args.size / sizeof(int64_t); @@ -56,17 +57,99 @@ static void SortToIndicesInt64Compare(benchmark::State& state) { auto max = std::numeric_limits::max(); auto values = rand.Int64(array_size, min, max, args.null_proportion); - SortToIndicesBenchmark(state, values); + ArraySortIndicesBenchmark(state, values); } -BENCHMARK(SortToIndicesInt64Count) +static void TableSortIndicesBenchmark(benchmark::State& state, + const std::shared_ptr& table, + const SortOptions& options) { + for (auto _ : state) { + ABORT_NOT_OK(SortIndices(*table, options).status()); + } + state.SetItemsProcessed(state.iterations() * table->num_rows()); +} + +static void TableSortIndicesInt64Count(benchmark::State& state) { + RegressionArgs args(state); + + const int64_t array_size = args.size / sizeof(int64_t); + auto rand = random::RandomArrayGenerator(kSeed); + auto values = rand.Int64(array_size, -100, 100, args.null_proportion); + std::vector> fields = {{field("int64", int64())}}; + auto table = Table::Make(schema(fields), {values}, array_size); + SortOptions options({SortKey("int64", SortOrder::ASCENDING)}); + + TableSortIndicesBenchmark(state, table, options); +} + +static void TableSortIndicesInt64Compare(benchmark::State& state) { + RegressionArgs args(state); + + const int64_t array_size = args.size / sizeof(int64_t); + auto rand = random::RandomArrayGenerator(kSeed); + + auto min = std::numeric_limits::min(); + auto max = std::numeric_limits::max(); + auto values = rand.Int64(array_size, min, max, args.null_proportion); + std::vector> fields = {{field("int64", int64())}}; + auto table = Table::Make(schema(fields), {values}, array_size); + SortOptions options({SortKey("int64", SortOrder::ASCENDING)}); + + TableSortIndicesBenchmark(state, table, options); +} + +static void TableSortIndicesInt64Int64(benchmark::State& state) { + RegressionArgs args(state); + + const int64_t array_size = args.size / sizeof(int64_t); + auto rand = random::RandomArrayGenerator(kSeed); + + auto min = std::numeric_limits::min(); + auto max = std::numeric_limits::max(); + auto values1 = rand.Int64(array_size, min, max, args.null_proportion); + auto values2 = rand.Int64(array_size, min, max, args.null_proportion); + std::vector> fields = { + {field("int64-1", int64())}, + {field("int64-2", int64())}, + }; + auto table = Table::Make(schema(fields), {values1, values2}, array_size); + SortOptions options({ + SortKey("int64-1", SortOrder::ASCENDING), + SortKey("int64-2", SortOrder::ASCENDING), + }); + + TableSortIndicesBenchmark(state, table, options); +} + +BENCHMARK(ArraySortIndicesInt64Count) + ->Apply(RegressionSetArgs) + ->Args({1 << 20, 1}) + ->Args({1 << 23, 1}) + ->MinTime(1.0) + ->Unit(benchmark::TimeUnit::kNanosecond); + +BENCHMARK(ArraySortIndicesInt64Compare) + ->Apply(RegressionSetArgs) + ->Args({1 << 20, 1}) + ->Args({1 << 23, 1}) + ->MinTime(1.0) + ->Unit(benchmark::TimeUnit::kNanosecond); + +BENCHMARK(TableSortIndicesInt64Count) + ->Apply(RegressionSetArgs) + ->Args({1 << 20, 1}) + ->Args({1 << 23, 1}) + ->MinTime(1.0) + ->Unit(benchmark::TimeUnit::kNanosecond); + +BENCHMARK(TableSortIndicesInt64Compare) ->Apply(RegressionSetArgs) ->Args({1 << 20, 100}) ->Args({1 << 23, 100}) ->MinTime(1.0) ->Unit(benchmark::TimeUnit::kNanosecond); -BENCHMARK(SortToIndicesInt64Compare) +BENCHMARK(TableSortIndicesInt64Int64) ->Apply(RegressionSetArgs) ->Args({1 << 20, 100}) ->Args({1 << 23, 100}) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index 448576da8ce..df0e6394900 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -22,6 +22,7 @@ #include "arrow/compute/api_vector.h" #include "arrow/compute/kernels/test_util.h" +#include "arrow/table.h" #include "arrow/testing/gtest_common.h" #include "arrow/testing/gtest_util.h" #include "arrow/testing/random.h" @@ -51,7 +52,7 @@ class NthComparator { template class SortComparator { public: - bool operator()(const ArrayType& array, uint64_t lhs, uint64_t rhs) { + bool operator()(const ArrayType& array, SortOrder order, uint64_t lhs, uint64_t rhs) { if (array.IsNull(rhs) && array.IsNull(lhs)) return lhs < rhs; if (array.IsNull(rhs)) return true; if (array.IsNull(lhs)) return false; @@ -63,7 +64,11 @@ class SortComparator { if (lhs_isnan) return false; } if (array.GetView(lhs) == array.GetView(rhs)) return lhs < rhs; - return array.GetView(lhs) < array.GetView(rhs); + if (order == SortOrder::ASCENDING) { + return array.GetView(lhs) < array.GetView(rhs); + } else { + return array.GetView(lhs) > array.GetView(rhs); + } } }; @@ -228,124 +233,150 @@ TYPED_TEST(TestNthToIndicesRandom, RandomValues) { using arrow::internal::checked_pointer_cast; template -class TestSortToIndicesKernel : public TestBase { +class TestArraySortIndicesKernel : public TestBase { private: - void AssertSortToIndicesArrays(const std::shared_ptr values, - const std::shared_ptr expected) { - ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, SortToIndices(*values)); + void AssertArraysSortIndices(const std::shared_ptr values, SortOrder order, + const std::shared_ptr expected) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, SortIndices(*values, order)); ASSERT_OK(actual->ValidateFull()); AssertArraysEqual(*expected, *actual); } protected: - virtual void AssertSortToIndices(const std::string& values, - const std::string& expected) { + virtual void AssertSortIndices(const std::string& values, SortOrder order, + const std::string& expected) { auto type = TypeTraits::type_singleton(); - AssertSortToIndicesArrays(ArrayFromJSON(type, values), - ArrayFromJSON(uint64(), expected)); + AssertArraysSortIndices(ArrayFromJSON(type, values), order, + ArrayFromJSON(uint64(), expected)); + } + + virtual void AssertSortIndices(const std::string& values, const std::string& expected) { + AssertSortIndices(values, SortOrder::ASCENDING, expected); } }; template -class TestSortToIndicesKernelForReal : public TestSortToIndicesKernel {}; -TYPED_TEST_SUITE(TestSortToIndicesKernelForReal, RealArrowTypes); +class TestArraySortIndicesKernelForReal : public TestArraySortIndicesKernel {}; +TYPED_TEST_SUITE(TestArraySortIndicesKernelForReal, RealArrowTypes); template -class TestSortToIndicesKernelForIntegral : public TestSortToIndicesKernel {}; -TYPED_TEST_SUITE(TestSortToIndicesKernelForIntegral, IntegralArrowTypes); +class TestArraySortIndicesKernelForIntegral + : public TestArraySortIndicesKernel {}; +TYPED_TEST_SUITE(TestArraySortIndicesKernelForIntegral, IntegralArrowTypes); template -class TestSortToIndicesKernelForStrings : public TestSortToIndicesKernel {}; -TYPED_TEST_SUITE(TestSortToIndicesKernelForStrings, testing::Types); - -TYPED_TEST(TestSortToIndicesKernelForReal, SortReal) { - this->AssertSortToIndices("[]", "[]"); - - this->AssertSortToIndices("[3.4, 2.6, 6.3]", "[1, 0, 2]"); - this->AssertSortToIndices("[1.1, 2.4, 3.5, 4.3, 5.1, 6.8, 7.3]", "[0,1,2,3,4,5,6]"); - this->AssertSortToIndices("[7, 6, 5, 4, 3, 2, 1]", "[6,5,4,3,2,1,0]"); - this->AssertSortToIndices("[10.4, 12, 4.2, 50, 50.3, 32, 11]", "[2,0,6,1,5,3,4]"); - - this->AssertSortToIndices("[null, 1, 3.3, null, 2, 5.3]", "[1,4,2,5,0,3]"); - - this->AssertSortToIndices("[3, 4, NaN, 1, 2, null]", "[3,4,0,1,2,5]"); - this->AssertSortToIndices("[NaN, 2, NaN, 3, 1]", "[4,1,3,0,2]"); - this->AssertSortToIndices("[null, NaN, NaN, null]", "[1,2,0,3]"); +class TestArraySortIndicesKernelForStrings + : public TestArraySortIndicesKernel {}; +TYPED_TEST_SUITE(TestArraySortIndicesKernelForStrings, testing::Types); + +TYPED_TEST(TestArraySortIndicesKernelForReal, SortReal) { + this->AssertSortIndices("[]", "[]"); + + this->AssertSortIndices("[3.4, 2.6, 6.3]", "[1, 0, 2]"); + this->AssertSortIndices("[1.1, 2.4, 3.5, 4.3, 5.1, 6.8, 7.3]", "[0,1,2,3,4,5,6]"); + this->AssertSortIndices("[7, 6, 5, 4, 3, 2, 1]", "[6,5,4,3,2,1,0]"); + this->AssertSortIndices("[10.4, 12, 4.2, 50, 50.3, 32, 11]", "[2,0,6,1,5,3,4]"); + + this->AssertSortIndices("[null, 1, 3.3, null, 2, 5.3]", SortOrder::ASCENDING, + "[1,4,2,5,0,3]"); + this->AssertSortIndices("[null, 1, 3.3, null, 2, 5.3]", SortOrder::DESCENDING, + "[5,2,4,1,0,3]"); + + this->AssertSortIndices("[3, 4, NaN, 1, 2, null]", SortOrder::ASCENDING, + "[3,4,0,1,2,5]"); + this->AssertSortIndices("[3, 4, NaN, 1, 2, null]", SortOrder::DESCENDING, + "[1,0,4,3,2,5]"); + this->AssertSortIndices("[NaN, 2, NaN, 3, 1]", SortOrder::ASCENDING, "[4,1,3,0,2]"); + this->AssertSortIndices("[NaN, 2, NaN, 3, 1]", SortOrder::DESCENDING, "[3,1,4,0,2]"); + this->AssertSortIndices("[null, NaN, NaN, null]", SortOrder::ASCENDING, "[1,2,0,3]"); + this->AssertSortIndices("[null, NaN, NaN, null]", SortOrder::DESCENDING, "[1,2,0,3]"); } -TYPED_TEST(TestSortToIndicesKernelForIntegral, SortIntegral) { - this->AssertSortToIndices("[]", "[]"); +TYPED_TEST(TestArraySortIndicesKernelForIntegral, SortIntegral) { + this->AssertSortIndices("[]", "[]"); - this->AssertSortToIndices("[3, 2, 6]", "[1, 0, 2]"); - this->AssertSortToIndices("[1, 2, 3, 4, 5, 6, 7]", "[0,1,2,3,4,5,6]"); - this->AssertSortToIndices("[7, 6, 5, 4, 3, 2, 1]", "[6,5,4,3,2,1,0]"); - this->AssertSortToIndices("[10, 12, 4, 50, 50, 32, 11]", "[2,0,6,1,5,3,4]"); + this->AssertSortIndices("[3, 2, 6]", "[1, 0, 2]"); + this->AssertSortIndices("[1, 2, 3, 4, 5, 6, 7]", "[0,1,2,3,4,5,6]"); + this->AssertSortIndices("[7, 6, 5, 4, 3, 2, 1]", "[6,5,4,3,2,1,0]"); - this->AssertSortToIndices("[null, 1, 3, null, 2, 5]", "[1,4,2,5,0,3]"); + this->AssertSortIndices("[10, 12, 4, 50, 50, 32, 11]", SortOrder::ASCENDING, + "[2,0,6,1,5,3,4]"); + this->AssertSortIndices("[10, 12, 4, 50, 50, 32, 11]", SortOrder::DESCENDING, + "[3,4,5,1,6,0,2]"); + + this->AssertSortIndices("[null, 1, 3, null, 2, 5]", SortOrder::ASCENDING, + "[1,4,2,5,0,3]"); + this->AssertSortIndices("[null, 1, 3, null, 2, 5]", SortOrder::DESCENDING, + "[5,2,4,1,0,3]"); } -TYPED_TEST(TestSortToIndicesKernelForStrings, SortStrings) { - this->AssertSortToIndices("[]", "[]"); +TYPED_TEST(TestArraySortIndicesKernelForStrings, SortStrings) { + this->AssertSortIndices("[]", "[]"); - this->AssertSortToIndices(R"(["a", "b", "c"])", "[0, 1, 2]"); - this->AssertSortToIndices(R"(["foo", "bar", "baz"])", "[1,2,0]"); - this->AssertSortToIndices(R"(["testing", "sort", "for", "strings"])", "[2, 1, 3, 0]"); + this->AssertSortIndices(R"(["a", "b", "c"])", "[0, 1, 2]"); + this->AssertSortIndices(R"(["foo", "bar", "baz"])", "[1,2,0]"); + this->AssertSortIndices(R"(["testing", "sort", "for", "strings"])", "[2, 1, 3, 0]"); } template -class TestSortToIndicesKernelForUInt8 : public TestSortToIndicesKernel {}; -TYPED_TEST_SUITE(TestSortToIndicesKernelForUInt8, UInt8Type); +class TestArraySortIndicesKernelForUInt8 : public TestArraySortIndicesKernel { +}; +TYPED_TEST_SUITE(TestArraySortIndicesKernelForUInt8, UInt8Type); template -class TestSortToIndicesKernelForInt8 : public TestSortToIndicesKernel {}; -TYPED_TEST_SUITE(TestSortToIndicesKernelForInt8, Int8Type); +class TestArraySortIndicesKernelForInt8 : public TestArraySortIndicesKernel {}; +TYPED_TEST_SUITE(TestArraySortIndicesKernelForInt8, Int8Type); -TYPED_TEST(TestSortToIndicesKernelForUInt8, SortUInt8) { - this->AssertSortToIndices("[255, null, 0, 255, 10, null, 128, 0]", "[2,7,4,6,0,3,1,5]"); +TYPED_TEST(TestArraySortIndicesKernelForUInt8, SortUInt8) { + this->AssertSortIndices("[255, null, 0, 255, 10, null, 128, 0]", "[2,7,4,6,0,3,1,5]"); } -TYPED_TEST(TestSortToIndicesKernelForInt8, SortInt8) { - this->AssertSortToIndices("[null, 10, 127, 0, -128, -128, null]", "[4,5,3,1,2,0,6]"); +TYPED_TEST(TestArraySortIndicesKernelForInt8, SortInt8) { + this->AssertSortIndices("[null, 10, 127, 0, -128, -128, null]", "[4,5,3,1,2,0,6]"); } template -class TestSortToIndicesKernelRandom : public TestBase {}; +class TestArraySortIndicesKernelRandom : public TestBase {}; template -class TestSortToIndicesKernelRandomCount : public TestBase {}; +class TestArraySortIndicesKernelRandomCount : public TestBase {}; template -class TestSortToIndicesKernelRandomCompare : public TestBase {}; +class TestArraySortIndicesKernelRandomCompare : public TestBase {}; -using SortToIndicesableTypes = +using SortIndicesableTypes = ::testing::Types; template -void ValidateSorted(const ArrayType& array, UInt64Array& offsets) { +void ValidateSorted(const ArrayType& array, UInt64Array& offsets, SortOrder order) { ASSERT_OK(array.ValidateFull()); SortComparator compare; for (int i = 1; i < array.length(); i++) { uint64_t lhs = offsets.Value(i - 1); uint64_t rhs = offsets.Value(i); - ASSERT_TRUE(compare(array, lhs, rhs)); + ASSERT_TRUE(compare(array, order, lhs, rhs)); } } -TYPED_TEST_SUITE(TestSortToIndicesKernelRandom, SortToIndicesableTypes); +TYPED_TEST_SUITE(TestArraySortIndicesKernelRandom, SortIndicesableTypes); -TYPED_TEST(TestSortToIndicesKernelRandom, SortRandomValues) { +TYPED_TEST(TestArraySortIndicesKernelRandom, SortRandomValues) { using ArrayType = typename TypeTraits::ArrayType; Random rand(0x5487655); int times = 5; int length = 1000; + times = 1; + length = 10; for (int test = 0; test < times; test++) { for (auto null_probability : {0.0, 0.1, 0.5, 1.0}) { - auto array = rand.Generate(length, null_probability); - ASSERT_OK_AND_ASSIGN(std::shared_ptr offsets, SortToIndices(*array)); - ValidateSorted(*checked_pointer_cast(array), - *checked_pointer_cast(offsets)); + for (auto order : {SortOrder::ASCENDING, SortOrder::DESCENDING}) { + auto array = rand.Generate(length, null_probability); + ASSERT_OK_AND_ASSIGN(std::shared_ptr offsets, SortIndices(*array, order)); + ValidateSorted(*checked_pointer_cast(array), + *checked_pointer_cast(offsets), order); + } } } } @@ -353,9 +384,9 @@ TYPED_TEST(TestSortToIndicesKernelRandom, SortRandomValues) { // Long array with small value range: counting sort // - length >= 1024(CountCompareSorter::countsort_min_len_) // - range <= 4096(CountCompareSorter::countsort_max_range_) -TYPED_TEST_SUITE(TestSortToIndicesKernelRandomCount, IntegralArrowTypes); +TYPED_TEST_SUITE(TestArraySortIndicesKernelRandomCount, IntegralArrowTypes); -TYPED_TEST(TestSortToIndicesKernelRandomCount, SortRandomValuesCount) { +TYPED_TEST(TestArraySortIndicesKernelRandomCount, SortRandomValuesCount) { using ArrayType = typename TypeTraits::ArrayType; RandomRange rand(0x5487656); @@ -364,18 +395,20 @@ TYPED_TEST(TestSortToIndicesKernelRandomCount, SortRandomValuesCount) { int range = 2000; for (int test = 0; test < times; test++) { for (auto null_probability : {0.0, 0.1, 0.5, 1.0}) { - auto array = rand.Generate(length, range, null_probability); - ASSERT_OK_AND_ASSIGN(std::shared_ptr offsets, SortToIndices(*array)); - ValidateSorted(*checked_pointer_cast(array), - *checked_pointer_cast(offsets)); + for (auto order : {SortOrder::ASCENDING, SortOrder::DESCENDING}) { + auto array = rand.Generate(length, range, null_probability); + ASSERT_OK_AND_ASSIGN(std::shared_ptr offsets, SortIndices(*array, order)); + ValidateSorted(*checked_pointer_cast(array), + *checked_pointer_cast(offsets), order); + } } } } // Long array with big value range: std::stable_sort -TYPED_TEST_SUITE(TestSortToIndicesKernelRandomCompare, IntegralArrowTypes); +TYPED_TEST_SUITE(TestArraySortIndicesKernelRandomCompare, IntegralArrowTypes); -TYPED_TEST(TestSortToIndicesKernelRandomCompare, SortRandomValuesCompare) { +TYPED_TEST(TestArraySortIndicesKernelRandomCompare, SortRandomValuesCompare) { using ArrayType = typename TypeTraits::ArrayType; Random rand(0x5487657); @@ -383,13 +416,214 @@ TYPED_TEST(TestSortToIndicesKernelRandomCompare, SortRandomValuesCompare) { int length = 4000; for (int test = 0; test < times; test++) { for (auto null_probability : {0.0, 0.1, 0.5, 1.0}) { - auto array = rand.Generate(length, null_probability); - ASSERT_OK_AND_ASSIGN(std::shared_ptr offsets, SortToIndices(*array)); - ValidateSorted(*checked_pointer_cast(array), - *checked_pointer_cast(offsets)); + for (auto order : {SortOrder::ASCENDING, SortOrder::DESCENDING}) { + auto array = rand.Generate(length, null_probability); + ASSERT_OK_AND_ASSIGN(std::shared_ptr offsets, SortIndices(*array, order)); + ValidateSorted(*checked_pointer_cast(array), + *checked_pointer_cast(offsets), order); + } + } + } +} + +class TestTableSortIndices : public ::testing::Test { + protected: + void AssertSortIndices(const std::shared_ptr
table, const SortOptions& options, + const std::shared_ptr expected) { + ASSERT_OK_AND_ASSIGN(auto actual, SortIndices(*table, options)); + AssertArraysEqual(*expected, *actual); + } + + void AssertSortIndices(const std::shared_ptr
table, const SortOptions& options, + const std::string expected) { + AssertSortIndices(table, options, ArrayFromJSON(uint64(), expected)); + } +}; + +TEST_F(TestTableSortIndices, SortNull) { + auto table = TableFromJSON(schema({ + {field("a", uint8())}, + {field("b", uint8())}, + }), + {"[" + "{\"a\": null, \"b\": 5}," + "{\"a\": 1, \"b\": 3}," + "{\"a\": 3, \"b\": null}," + "{\"a\": null, \"b\": null}," + "{\"a\": 2, \"b\": 5}," + "{\"a\": 1, \"b\": 5}" + "]"}); + SortOptions options( + {SortKey("a", SortOrder::ASCENDING), SortKey("b", SortOrder::DESCENDING)}); + this->AssertSortIndices(table, options, "[5, 1, 4, 2, 0, 3]"); +} + +using RandomParam = std::tuple; +class TestTableSortIndicesRandom : public testing::TestWithParam { + class Comparator : public TypeVisitor { + public: + bool operator()(const Table& table, const SortOptions& options, uint64_t lhs, + uint64_t rhs) { + for (const auto& sort_key : options.sort_keys) { + auto chunked_array = table.GetColumnByName(sort_key.name); + left_array_ = findTargetArray(chunked_array, lhs, left_index_); + right_array_ = findTargetArray(chunked_array, rhs, right_index_); + if (right_array_->IsNull(right_index_) && left_array_->IsNull(left_index_)) + continue; + if (right_array_->IsNull(right_index_)) return true; + if (left_array_->IsNull(left_index_)) return false; + int compared = Compare(); + if (compared == 0) continue; + if (sort_key.order == SortOrder::ASCENDING) { + return compared < 0; + } else { + return compared > 0; + } + } + return lhs < rhs; + } + +#define VISIT(TYPE) \ + Status Visit(const TYPE##Type& type) override { \ + CompareType(type); \ + return Status::OK(); \ + } + + VISIT(Int8) + VISIT(Int16) + VISIT(Int32) + VISIT(Int64) + VISIT(UInt8) + VISIT(UInt16) + VISIT(UInt32) + VISIT(UInt64) + VISIT(Float) + VISIT(Double) + VISIT(String) + +#undef VISIT + + private: + std::shared_ptr findTargetArray(std::shared_ptr chunked_array, + int64_t i, int64_t& chunk_index) { + int64_t offset = 0; + for (auto& chunk : chunked_array->chunks()) { + if (i < offset + chunk->length()) { + chunk_index = i - offset; + return chunk; + } + offset += chunk->length(); + } + return nullptr; + } + + int Compare() { + left_array_->type()->Accept(this); + return compared_; + } + + template + void CompareType(const Type& type) { + using ArrayType = typename TypeTraits::ArrayType; + auto left = checked_pointer_cast(left_array_)->GetView(left_index_); + auto right = checked_pointer_cast(right_array_)->GetView(right_index_); + if (left == right) { + compared_ = 0; + } else if (left > right) { + compared_ = 1; + } else { + compared_ = -1; + } + } + + std::shared_ptr left_array_; + int64_t left_index_; + std::shared_ptr right_array_; + int64_t right_index_; + int compared_; + }; + + public: + void Validate(const Table& table, const SortOptions& options, UInt64Array& offsets) { + ASSERT_OK(offsets.ValidateFull()); + Comparator comparator; + for (int i = 1; i < table.num_rows(); i++) { + uint64_t lhs = offsets.Value(i - 1); + uint64_t rhs = offsets.Value(i); + ASSERT_TRUE(comparator(table, options, lhs, rhs)); } } +}; + +TEST_P(TestTableSortIndicesRandom, Sort) { + auto first_sort_key_name = std::get<0>(GetParam()); + auto null_probability = std::get<1>(GetParam()); + auto seed = 0x61549225; + std::vector column_names = { + "uint8", "uint16", "uint32", "uint64", "int8", "int16", + "int32", "int64", "float", "double", "string", + }; + std::vector> fields = { + {field(column_names[0], uint8())}, {field(column_names[1], uint16())}, + {field(column_names[2], uint32())}, {field(column_names[3], uint64())}, + {field(column_names[4], int8())}, {field(column_names[5], int16())}, + {field(column_names[6], int32())}, {field(column_names[7], int64())}, + {field(column_names[8], float32())}, {field(column_names[9], float64())}, + {field(column_names[10], utf8())}, + }; + auto length = 4000; + std::vector> columns = { + Random(seed).Generate(length, null_probability), + Random(seed).Generate(length, null_probability), + Random(seed).Generate(length, null_probability), + Random(seed).Generate(length, null_probability), + Random(seed).Generate(length, null_probability), + Random(seed).Generate(length, null_probability), + Random(seed).Generate(length, null_probability), + Random(seed).Generate(length, null_probability), + Random(seed).Generate(length, null_probability), + Random(seed).Generate(length, null_probability), + Random(seed).Generate(length, null_probability), + }; + auto table = Table::Make(schema(fields), columns, length); + std::default_random_engine engine(seed); + std::uniform_int_distribution<> distribution(0, column_names.size() - 1); + auto n_sort_keys = 2; + std::vector sort_keys; + auto first_sort_key_order = + (distribution(engine) % 2) == 0 ? SortOrder::ASCENDING : SortOrder::DESCENDING; + sort_keys.emplace_back(first_sort_key_name, first_sort_key_order); + for (int i = 1; i < n_sort_keys; ++i) { + auto& column_name = column_names[distribution(engine)]; + auto order = + (distribution(engine) % 2) == 0 ? SortOrder::ASCENDING : SortOrder::DESCENDING; + sort_keys.emplace_back(column_name, order); + } + SortOptions options(sort_keys); + ASSERT_OK_AND_ASSIGN(auto offsets, SortIndices(*table, options)); + Validate(*table, options, *checked_pointer_cast(offsets)); } +INSTANTIATE_TEST_SUITE_P(NoNull, TestTableSortIndicesRandom, + testing::Combine(testing::Values("uint8", "uint16", "uint32", + "uint64", "int8", "int16", + "int32", "int64", "float", + "double", "string"), + testing::Values(0.0))); + +INSTANTIATE_TEST_SUITE_P(MayNull, TestTableSortIndicesRandom, + testing::Combine(testing::Values("uint8", "uint16", "uint32", + "uint64", "int8", "int16", + "int32", "int64", "float", + "double", "string"), + testing::Values(0.1, 0.5))); + +INSTANTIATE_TEST_SUITE_P(AllNull, TestTableSortIndicesRandom, + testing::Combine(testing::Values("uint8", "uint16", "uint32", + "uint64", "int8", "int16", + "int32", "int64", "float", + "double", "string"), + testing::Values(1.0))); + } // namespace compute } // namespace arrow diff --git a/cpp/src/arrow/dataset/filter.cc b/cpp/src/arrow/dataset/filter.cc index bcfc8d745dc..caab28e3f7c 100644 --- a/cpp/src/arrow/dataset/filter.cc +++ b/cpp/src/arrow/dataset/filter.cc @@ -1718,7 +1718,7 @@ Result> MakeGroupings(const StructArray& by) { ARROW_ASSIGN_OR_RAISE(auto fused, StructDictionary::Encode(by.fields())); - ARROW_ASSIGN_OR_RAISE(auto sort_indices, compute::SortToIndices(*fused.indices)); + ARROW_ASSIGN_OR_RAISE(auto sort_indices, compute::SortIndices(*fused.indices)); ARROW_ASSIGN_OR_RAISE(Datum sorted, compute::Take(fused.indices, *sort_indices)); fused.indices = checked_pointer_cast(sorted.make_array()); From 50dd97f5f8210c661430272e943825dda372fc26 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Mon, 9 Nov 2020 14:17:59 +0900 Subject: [PATCH 02/37] Remove unused variables --- cpp/src/arrow/compute/kernels/vector_sort.cc | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc index 367ee8d79fd..72c0169ceb4 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort.cc @@ -556,9 +556,10 @@ class TableSorter : public TypeVisitor { auto& comparer = comparer_; const auto& first_sort_key = comparer.sort_keys()[0]; auto nulls_begin = indices_end_; - if (first_sort_key.has_nulls > 0) { - nulls_begin = std::stable_partition( - indices_begin_, indices_end_, [&first_sort_key](uint64_t index) { + if (first_sort_key.has_nulls) { + StablePartitioner partitioner; + nulls_begin = + partitioner(indices_begin_, indices_end_, [&first_sort_key](uint64_t index) { int64_t index_chunk = 0; auto chunk = first_sort_key.ResolveChunk(index, index_chunk); return !chunk->IsNull(index_chunk); From 121b51b8f5f12f56178e0768e97fd05529d071d0 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Mon, 9 Nov 2020 14:17:59 +0900 Subject: [PATCH 03/37] Remove unused variables --- cpp/src/arrow/compute/kernels/vector_sort.cc | 2 -- 1 file changed, 2 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc index 72c0169ceb4..af43cd0d221 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort.cc @@ -594,8 +594,6 @@ class TableSorter : public TypeVisitor { uint64_t* indices_begin_; uint64_t* indices_end_; - const Table& table_; - const SortOptions& options_; Comparer comparer_; }; From e4e5d68519043f8f0114416d6d42335d5905d2ab Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Mon, 9 Nov 2020 14:24:43 +0900 Subject: [PATCH 04/37] Remove unused assignments --- cpp/src/arrow/compute/kernels/vector_sort.cc | 2 -- 1 file changed, 2 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc index af43cd0d221..039acaf8f87 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort.cc @@ -518,8 +518,6 @@ class TableSorter : public TypeVisitor { const SortOptions& options) : indices_begin_(indices_begin), indices_end_(indices_end), - table_(table), - options_(options), comparer_(table, options.sort_keys) {} Status Sort() { From b5d234f7139c9ba179a36ee4013310cedfd9cc7b Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Mon, 9 Nov 2020 14:34:37 +0900 Subject: [PATCH 05/37] Use auto --- cpp/src/arrow/compute/kernels/vector_sort.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc index 039acaf8f87..ad74f38b393 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort.cc @@ -233,7 +233,7 @@ class ArrayCountSorter { for (uint32_t i = 1; i <= value_range; ++i) { counts[i] += counts[i - 1]; } - uint32_t null_position = counts[value_range]; + auto null_position = counts[value_range]; int64_t index = 0; VisitRawValuesInline( values, [&](c_type v) { indices_begin[counts[v - min_]++] = index++; }, @@ -244,7 +244,7 @@ class ArrayCountSorter { for (uint32_t i = value_range; i >= 1; --i) { counts[i - 1] += counts[i]; } - uint32_t null_position = counts[0]; + auto null_position = counts[0]; int64_t index = 0; VisitRawValuesInline( values, [&](c_type v) { indices_begin[counts[v - min_ + 1]++] = index++; }, From 8c9c7c00c00619e9110c4c91dd657d111fec43ad Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Mon, 9 Nov 2020 14:40:29 +0900 Subject: [PATCH 06/37] Add missing Status check --- .../arrow/compute/kernels/vector_sort_test.cc | 21 +++++++++---------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index df0e6394900..47501f72236 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -472,20 +472,22 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { continue; if (right_array_->IsNull(right_index_)) return true; if (left_array_->IsNull(left_index_)) return false; - int compared = Compare(); - if (compared == 0) continue; + status_ = left_array_->type()->Accept(this); + if (compared_ == 0) continue; if (sort_key.order == SortOrder::ASCENDING) { - return compared < 0; + return compared_ < 0; } else { - return compared > 0; + return compared_ > 0; } } return lhs < rhs; } + Status status() const { return status_; } + #define VISIT(TYPE) \ Status Visit(const TYPE##Type& type) override { \ - CompareType(type); \ + CompareType(); \ return Status::OK(); \ } @@ -517,13 +519,8 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { return nullptr; } - int Compare() { - left_array_->type()->Accept(this); - return compared_; - } - template - void CompareType(const Type& type) { + void CompareType() { using ArrayType = typename TypeTraits::ArrayType; auto left = checked_pointer_cast(left_array_)->GetView(left_index_); auto right = checked_pointer_cast(right_array_)->GetView(right_index_); @@ -541,6 +538,7 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { std::shared_ptr right_array_; int64_t right_index_; int compared_; + Status status_; }; public: @@ -551,6 +549,7 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { uint64_t lhs = offsets.Value(i - 1); uint64_t rhs = offsets.Value(i); ASSERT_TRUE(comparator(table, options, lhs, rhs)); + ASSERT_OK(comparator.status()); } } }; From 96d0e522c90106f301911e96c7468fd8f52d25c9 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Mon, 9 Nov 2020 15:01:11 +0900 Subject: [PATCH 07/37] Use random::SeedType explicitly --- cpp/src/arrow/compute/kernels/vector_sort_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index 47501f72236..500b88e958d 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -557,7 +557,7 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { TEST_P(TestTableSortIndicesRandom, Sort) { auto first_sort_key_name = std::get<0>(GetParam()); auto null_probability = std::get<1>(GetParam()); - auto seed = 0x61549225; + random::SeedType seed = 0x61549225; std::vector column_names = { "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64", "float", "double", "string", From 7b4b7ae0b303208d2367fe9128563a6b0270ba43 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Mon, 9 Nov 2020 15:05:25 +0900 Subject: [PATCH 08/37] Don't use > --- cpp/src/arrow/compute/kernels/vector_sort.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc index ad74f38b393..60db85f48f7 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort.cc @@ -183,7 +183,7 @@ class ArrayCompareSorter { } else { std::stable_sort(indices_begin, nulls_begin, [&values](uint64_t left, uint64_t right) { - return values.GetView(left) > values.GetView(right); + return values.GetView(right) < values.GetView(left); }); } } From 67a3d40b0444a77def62f655773126183e4909d7 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Mon, 9 Nov 2020 15:30:14 +0900 Subject: [PATCH 09/37] Don't use max --- cpp/src/arrow/compute/kernels/vector_sort_test.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index 500b88e958d..422bfc164c4 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -557,7 +557,7 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { TEST_P(TestTableSortIndicesRandom, Sort) { auto first_sort_key_name = std::get<0>(GetParam()); auto null_probability = std::get<1>(GetParam()); - random::SeedType seed = 0x61549225; + auto seed = 0x61549225; std::vector column_names = { "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64", "float", "double", "string", @@ -586,14 +586,14 @@ TEST_P(TestTableSortIndicesRandom, Sort) { }; auto table = Table::Make(schema(fields), columns, length); std::default_random_engine engine(seed); - std::uniform_int_distribution<> distribution(0, column_names.size() - 1); + std::uniform_int_distribution<> distribution(0); auto n_sort_keys = 2; std::vector sort_keys; auto first_sort_key_order = (distribution(engine) % 2) == 0 ? SortOrder::ASCENDING : SortOrder::DESCENDING; sort_keys.emplace_back(first_sort_key_name, first_sort_key_order); for (int i = 1; i < n_sort_keys; ++i) { - auto& column_name = column_names[distribution(engine)]; + auto& column_name = column_names[distribution(engine) % column_names.size()]; auto order = (distribution(engine) % 2) == 0 ? SortOrder::ASCENDING : SortOrder::DESCENDING; sort_keys.emplace_back(column_name, order); From 3bce7274e7d5d20db12b5f2240b6c420f954cabf Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Tue, 10 Nov 2020 04:15:57 +0900 Subject: [PATCH 10/37] Add more tests --- cpp/src/arrow/compute/kernels/vector_sort_test.cc | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index 422bfc164c4..fafe8794103 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -316,6 +316,11 @@ TYPED_TEST(TestArraySortIndicesKernelForStrings, SortStrings) { this->AssertSortIndices(R"(["a", "b", "c"])", "[0, 1, 2]"); this->AssertSortIndices(R"(["foo", "bar", "baz"])", "[1,2,0]"); this->AssertSortIndices(R"(["testing", "sort", "for", "strings"])", "[2, 1, 3, 0]"); + + this->AssertSortIndices(R"(["c", "b", "a", "b"])", SortOrder::ASCENDING, + "[2, 1, 3, 0]"); + this->AssertSortIndices(R"(["c", "b", "a", "b"])", SortOrder::DESCENDING, + "[0, 1, 3, 2]"); } template From 9289fdc5fbf8b093a4f5ed3659a6b91fb908b122 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Wed, 11 Nov 2020 09:49:53 +0900 Subject: [PATCH 11/37] Adjust benchmark parameters --- .../arrow/compute/kernels/vector_sort_benchmark.cc | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc index 2e6d54707ac..ef755704149 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc @@ -123,22 +123,22 @@ static void TableSortIndicesInt64Int64(benchmark::State& state) { BENCHMARK(ArraySortIndicesInt64Count) ->Apply(RegressionSetArgs) - ->Args({1 << 20, 1}) - ->Args({1 << 23, 1}) + ->Args({1 << 20, 100}) + ->Args({1 << 23, 100}) ->MinTime(1.0) ->Unit(benchmark::TimeUnit::kNanosecond); BENCHMARK(ArraySortIndicesInt64Compare) ->Apply(RegressionSetArgs) - ->Args({1 << 20, 1}) - ->Args({1 << 23, 1}) + ->Args({1 << 20, 100}) + ->Args({1 << 23, 100}) ->MinTime(1.0) ->Unit(benchmark::TimeUnit::kNanosecond); BENCHMARK(TableSortIndicesInt64Count) ->Apply(RegressionSetArgs) - ->Args({1 << 20, 1}) - ->Args({1 << 23, 1}) + ->Args({1 << 20, 100}) + ->Args({1 << 23, 100}) ->MinTime(1.0) ->Unit(benchmark::TimeUnit::kNanosecond); From 3285131c62251a6a4e10cd6a868cdb808ad37f17 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Wed, 11 Nov 2020 10:01:15 +0900 Subject: [PATCH 12/37] Add support for NaN --- cpp/src/arrow/compute/kernels/vector_sort.cc | 84 +++++++++++++++---- .../arrow/compute/kernels/vector_sort_test.cc | 70 +++++++++++----- 2 files changed, 118 insertions(+), 36 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc index 60db85f48f7..745a4f1e6c1 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort.cc @@ -380,7 +380,7 @@ class TableSorter : public TypeVisitor { ResolvedSortKey(const ChunkedArray& chunked_array, const SortOrder order) : order(order) { type = chunked_array.type().get(); - has_nulls = (chunked_array.null_count() > 0); + null_count = chunked_array.null_count(); num_chunks = chunked_array.num_chunks(); for (const auto& chunk : chunked_array.chunks()) { chunks.push_back(chunk.get()); @@ -407,7 +407,7 @@ class TableSorter : public TypeVisitor { SortOrder order; DataType* type; - bool has_nulls; + int64_t null_count; size_t num_chunks; std::vector chunks; }; @@ -478,7 +478,7 @@ class TableSorter : public TypeVisitor { auto array_left = sort_key.ResolveChunk(current_left_, index_left); int64_t index_right = 0; auto array_right = sort_key.ResolveChunk(current_right_, index_right); - if (sort_key.has_nulls) { + if (sort_key.null_count > 0) { auto is_null_left = array_left->IsNull(index_left); auto is_null_right = array_right->IsNull(index_right); if (is_null_left && is_null_right) { @@ -554,19 +554,7 @@ class TableSorter : public TypeVisitor { auto& comparer = comparer_; const auto& first_sort_key = comparer.sort_keys()[0]; auto nulls_begin = indices_end_; - if (first_sort_key.has_nulls) { - StablePartitioner partitioner; - nulls_begin = - partitioner(indices_begin_, indices_end_, [&first_sort_key](uint64_t index) { - int64_t index_chunk = 0; - auto chunk = first_sort_key.ResolveChunk(index, index_chunk); - return !chunk->IsNull(index_chunk); - }); - std::stable_sort(nulls_begin, indices_end_, - [&comparer](uint64_t left, uint64_t right) { - return comparer.Compare(left, right, 1); - }); - } + nulls_begin = PartitionNullsInternal(first_sort_key); std::stable_sort( indices_begin_, nulls_begin, [&first_sort_key, &comparer](uint64_t left, uint64_t right) { @@ -590,6 +578,70 @@ class TableSorter : public TypeVisitor { return Status::OK(); } + template + enable_if_t::value, uint64_t*> PartitionNullsInternal( + const ResolvedSortKey& first_sort_key) { + using ArrayType = typename TypeTraits::ArrayType; + if (first_sort_key.null_count == 0) { + return indices_end_; + } + StablePartitioner partitioner; + auto nulls_begin = + partitioner(indices_begin_, indices_end_, [&first_sort_key](uint64_t index) { + int64_t index_chunk = 0; + auto chunk = first_sort_key.ResolveChunk(index, index_chunk); + return !chunk->IsNull(index_chunk); + }); + auto& comparer = comparer_; + std::stable_sort(nulls_begin, indices_end_, + [&comparer](uint64_t left, uint64_t right) { + return comparer.Compare(left, right, 1); + }); + return nulls_begin; + } + + template + enable_if_t::value, uint64_t*> PartitionNullsInternal( + const ResolvedSortKey& first_sort_key) { + using ArrayType = typename TypeTraits::ArrayType; + StablePartitioner partitioner; + if (first_sort_key.null_count == 0) { + return partitioner(indices_begin_, indices_end_, [&first_sort_key](uint64_t index) { + int64_t index_chunk = 0; + auto chunk = first_sort_key.ResolveChunk(index, index_chunk); + return !std::isnan(chunk->GetView(index_chunk)); + }); + } + auto nans_and_nulls_begin = + partitioner(indices_begin_, indices_end_, [&first_sort_key](uint64_t index) { + int64_t index_chunk = 0; + auto chunk = first_sort_key.ResolveChunk(index, index_chunk); + return !chunk->IsNull(index_chunk) && !std::isnan(chunk->GetView(index_chunk)); + }); + auto nulls_begin = nans_and_nulls_begin; + if (first_sort_key.null_count < static_cast(indices_end_ - nulls_begin)) { + // move Nulls after NaN + nulls_begin = partitioner( + nans_and_nulls_begin, indices_end_, [&first_sort_key](uint64_t index) { + int64_t index_chunk = 0; + auto chunk = first_sort_key.ResolveChunk(index, index_chunk); + return !chunk->IsNull(index_chunk); + }); + } + auto& comparer = comparer_; + if (nans_and_nulls_begin != nulls_begin) { + std::stable_sort(nans_and_nulls_begin, nulls_begin, + [&comparer](uint64_t left, uint64_t right) { + return comparer.Compare(left, right, 1); + }); + } + std::stable_sort(nulls_begin, indices_end_, + [&comparer](uint64_t left, uint64_t right) { + return comparer.Compare(left, right, 1); + }); + return nans_and_nulls_begin; + } + uint64_t* indices_begin_; uint64_t* indices_end_; Comparer comparer_; diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index fafe8794103..7546574ff7f 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -463,21 +463,42 @@ TEST_F(TestTableSortIndices, SortNull) { this->AssertSortIndices(table, options, "[5, 1, 4, 2, 0, 3]"); } +TEST_F(TestTableSortIndices, SortNaN) { + auto table = TableFromJSON(schema({ + {field("a", float32())}, + {field("b", float32())}, + }), + {"[" + "{\"a\": null, \"b\": 5}," + "{\"a\": 1, \"b\": 3}," + "{\"a\": 3, \"b\": null}," + "{\"a\": null, \"b\": null}," + "{\"a\": NaN, \"b\": null}," + "{\"a\": NaN, \"b\": 5}," + "{\"a\": NaN, \"b\": NaN}," + "{\"a\": 1, \"b\": 5}" + "]"}); + SortOptions options( + {SortKey("a", SortOrder::ASCENDING), SortKey("b", SortOrder::DESCENDING)}); + this->AssertSortIndices(table, options, "[7, 1, 2, 5, 6, 4, 0, 3]"); +} + using RandomParam = std::tuple; class TestTableSortIndicesRandom : public testing::TestWithParam { class Comparator : public TypeVisitor { public: bool operator()(const Table& table, const SortOptions& options, uint64_t lhs, uint64_t rhs) { + lhs_ = lhs; + rhs_ = rhs; for (const auto& sort_key : options.sort_keys) { auto chunked_array = table.GetColumnByName(sort_key.name); - left_array_ = findTargetArray(chunked_array, lhs, left_index_); - right_array_ = findTargetArray(chunked_array, rhs, right_index_); - if (right_array_->IsNull(right_index_) && left_array_->IsNull(left_index_)) - continue; - if (right_array_->IsNull(right_index_)) return true; - if (left_array_->IsNull(left_index_)) return false; - status_ = left_array_->type()->Accept(this); + lhs_array_ = findTargetArray(chunked_array, lhs, lhs_index_); + rhs_array_ = findTargetArray(chunked_array, rhs, rhs_index_); + if (rhs_array_->IsNull(rhs_index_) && lhs_array_->IsNull(lhs_index_)) continue; + if (rhs_array_->IsNull(rhs_index_)) return true; + if (lhs_array_->IsNull(lhs_index_)) return false; + status_ = lhs_array_->type()->Accept(this); if (compared_ == 0) continue; if (sort_key.order == SortOrder::ASCENDING) { return compared_ < 0; @@ -492,7 +513,7 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { #define VISIT(TYPE) \ Status Visit(const TYPE##Type& type) override { \ - CompareType(); \ + compared_ = CompareType(); \ return Status::OK(); \ } @@ -525,23 +546,32 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { } template - void CompareType() { + int CompareType() { using ArrayType = typename TypeTraits::ArrayType; - auto left = checked_pointer_cast(left_array_)->GetView(left_index_); - auto right = checked_pointer_cast(right_array_)->GetView(right_index_); - if (left == right) { - compared_ = 0; - } else if (left > right) { - compared_ = 1; + auto lhs_value = checked_pointer_cast(lhs_array_)->GetView(lhs_index_); + auto rhs_value = checked_pointer_cast(rhs_array_)->GetView(rhs_index_); + if (is_floating_type::value) { + const bool lhs_isnan = lhs_value != lhs_value; + const bool rhs_isnan = rhs_value != rhs_value; + if (lhs_isnan && rhs_isnan) return 0; + if (rhs_isnan) return 1; + if (lhs_isnan) return -1; + } + if (lhs_value == rhs_value) { + return 0; + } else if (lhs_value > rhs_value) { + return 1; } else { - compared_ = -1; + return -1; } } - std::shared_ptr left_array_; - int64_t left_index_; - std::shared_ptr right_array_; - int64_t right_index_; + int64_t lhs_; + std::shared_ptr lhs_array_; + int64_t lhs_index_; + int64_t rhs_; + std::shared_ptr rhs_array_; + int64_t rhs_index_; int compared_; Status status_; }; From 4a76c39c216ce854a304af56a6862068edb2bae8 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Wed, 11 Nov 2020 12:04:03 +0900 Subject: [PATCH 13/37] Update document --- docs/source/cpp/compute.rst | 37 +++++++++++++++++++++++-------------- 1 file changed, 23 insertions(+), 14 deletions(-) diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst index c1d3ac7e61b..1a019070e66 100644 --- a/docs/source/cpp/compute.rst +++ b/docs/source/cpp/compute.rst @@ -31,7 +31,7 @@ The generic Compute API Functions and function registry ------------------------------- -Functions represent compute operations over inputs of possibly varying +Functions represent compute operations over inputs of possibly varying types. Internally, a function is implemented by one or several "kernels", depending on the concrete input types (for example, a function adding values from two inputs can have different kernels depending on @@ -621,17 +621,21 @@ In these functions, nulls are considered greater than any other value Floating-point NaN values are considered greater than any other non-null value, but smaller than nulls. -+-----------------------+------------+-------------------------+-------------------+--------------------------------+-------------+ -| Function name | Arity | Input types | Output type | Options class | Notes | -+=======================+============+=========================+===================+================================+=============+ -| partition_nth_indices | Unary | Binary- and String-like | UInt64 | :struct:`PartitionNthOptions` | \(1) \(3) | -+-----------------------+------------+-------------------------+-------------------+--------------------------------+-------------+ -| partition_nth_indices | Unary | Numeric | UInt64 | :struct:`PartitionNthOptions` | \(1) | -+-----------------------+------------+-------------------------+-------------------+--------------------------------+-------------+ -| sort_indices | Unary | Binary- and String-like | UInt64 | | \(2) \(3) | -+-----------------------+------------+-------------------------+-------------------+--------------------------------+-------------+ -| sort_indices | Unary | Numeric | UInt64 | | \(2) | -+-----------------------+------------+-------------------------+-------------------+--------------------------------+-------------+ ++-----------------------+------------+-------------------------+-------------------+--------------------------------+----------------+ +| Function name | Arity | Input types | Output type | Options class | Notes | ++=======================+============+=========================+===================+================================+================+ +| partition_nth_indices | Unary | Binary- and String-like | UInt64 | :struct:`PartitionNthOptions` | \(1) \(3) | ++-----------------------+------------+-------------------------+-------------------+--------------------------------+----------------+ +| partition_nth_indices | Unary | Numeric | UInt64 | :struct:`PartitionNthOptions` | \(1) | ++-----------------------+------------+-------------------------+-------------------+--------------------------------+----------------+ +| array_sort_indices | Unary | Binary- and String-like | UInt64 | :struct:`ArraySortptions` | \(2) \(3) \(4) | ++-----------------------+------------+-------------------------+-------------------+--------------------------------+----------------+ +| array_sort_indices | Unary | Numeric | UInt64 | :struct:`ArraySortptions` | \(2) \(4) | ++-----------------------+------------+-------------------------+-------------------+--------------------------------+----------------+ +| sort_indices | Unary | Binary- and String-like | UInt64 | :struct:`Sortptions` | \(2) \(3) \(5) | ++-----------------------+------------+-------------------------+-------------------+--------------------------------+----------------+ +| sort_indices | Unary | Numeric | UInt64 | :struct:`Sortptions` | \(2) \(5) | ++-----------------------+------------+-------------------------+-------------------+--------------------------------+----------------+ * \(1) The output is an array of indices into the input array, that define a partial non-stable sort such that the *N*'th index points to the *N*'th @@ -640,12 +644,17 @@ value, but smaller than nulls. :func:`std::nth_element`). *N* is given in :member:`PartitionNthOptions::pivot`. -* \(2) The output is an array of indices into the input array, that define - a stable sort of the input array. +* \(2) The output is an array of indices into the input, that define a + stable sort of the input. * \(3) Input values are ordered lexicographically as bytestrings (even for String arrays). +* \(4) The input must be an array. The default order is ascending. + +* \(5) The input can be an array, chunked array, record batch or + table. If the input is a record batch or table, one or more sort + keys must be specified. Structural transforms ~~~~~~~~~~~~~~~~~~~~~ From efa3ee62955872bf4f586605bb99d2acb997ce69 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Tue, 17 Nov 2020 16:10:22 +0900 Subject: [PATCH 14/37] Use enum class --- cpp/src/arrow/compute/api_vector.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/compute/api_vector.h b/cpp/src/arrow/compute/api_vector.h index 5f4d8bccdf0..821a02803fa 100644 --- a/cpp/src/arrow/compute/api_vector.h +++ b/cpp/src/arrow/compute/api_vector.h @@ -58,14 +58,14 @@ struct ARROW_EXPORT TakeOptions : public FunctionOptions { static TakeOptions Defaults() { return BoundsCheck(); } }; -enum SortOrder { +enum class SortOrder { ASCENDING, DESCENDING, }; /// \brief One sort key for PartitionNthIndices (TODO) and SortIndices struct ARROW_EXPORT SortKey { - explicit SortKey(std::string name, SortOrder order = ASCENDING) + explicit SortKey(std::string name, SortOrder order = SortOrder::ASCENDING) : name(name), order(order) {} /// The name of the sort key. From 2c4d23615513b3180f821e167f8e29314955cb19 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Tue, 17 Nov 2020 16:12:15 +0900 Subject: [PATCH 15/37] Use Capital for enum name to avoid conflict --- cpp/src/arrow/compute/api_vector.cc | 2 +- cpp/src/arrow/compute/api_vector.h | 12 ++--- cpp/src/arrow/compute/kernels/vector_sort.cc | 12 ++--- .../compute/kernels/vector_sort_benchmark.cc | 8 ++-- .../arrow/compute/kernels/vector_sort_test.cc | 48 +++++++++---------- 5 files changed, 41 insertions(+), 41 deletions(-) diff --git a/cpp/src/arrow/compute/api_vector.cc b/cpp/src/arrow/compute/api_vector.cc index 9b801b1cad3..81830d91161 100644 --- a/cpp/src/arrow/compute/api_vector.cc +++ b/cpp/src/arrow/compute/api_vector.cc @@ -47,7 +47,7 @@ Result> NthToIndices(const Array& values, int64_t n, } Result> SortIndices(const Array& values, ExecContext* ctx) { - return SortIndices(values, SortOrder::ASCENDING, ctx); + return SortIndices(values, SortOrder::Ascending, ctx); } Result> SortIndices(const Array& values, SortOrder order, diff --git a/cpp/src/arrow/compute/api_vector.h b/cpp/src/arrow/compute/api_vector.h index 821a02803fa..16895a5dc09 100644 --- a/cpp/src/arrow/compute/api_vector.h +++ b/cpp/src/arrow/compute/api_vector.h @@ -59,13 +59,13 @@ struct ARROW_EXPORT TakeOptions : public FunctionOptions { }; enum class SortOrder { - ASCENDING, - DESCENDING, + Ascending, + Descending, }; /// \brief One sort key for PartitionNthIndices (TODO) and SortIndices struct ARROW_EXPORT SortKey { - explicit SortKey(std::string name, SortOrder order = SortOrder::ASCENDING) + explicit SortKey(std::string name, SortOrder order = SortOrder::Ascending) : name(name), order(order) {} /// The name of the sort key. @@ -75,7 +75,7 @@ struct ARROW_EXPORT SortKey { }; struct ARROW_EXPORT ArraySortOptions : public FunctionOptions { - explicit ArraySortOptions(SortOrder order = SortOrder::ASCENDING) : order(order) {} + explicit ArraySortOptions(SortOrder order = SortOrder::Ascending) : order(order) {} SortOrder order; }; @@ -227,8 +227,8 @@ Result> SortIndices(const Array& values, SortOrder order, /// "column1": [null, 1, 3, null, 2, 1], /// "column2": [ 5, 3, null, null, 5, 5], /// } and options = { -/// {"column1", SortOrder::ASCENDING}, -/// {"column2", SortOrder::DESCENDING}, +/// {"column1", SortOrder::Ascending}, +/// {"column2", SortOrder::Descending}, /// }, the output will be [5, 1, 4, 2, 0, 3]. /// /// \param[in] table table to sort diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc index 745a4f1e6c1..c48a5599611 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort.cc @@ -175,7 +175,7 @@ class ArrayCompareSorter { std::iota(indices_begin, indices_end, 0); auto nulls_begin = PartitionNulls(indices_begin, indices_end, values); - if (options.order == SortOrder::ASCENDING) { + if (options.order == SortOrder::Ascending) { std::stable_sort(indices_begin, nulls_begin, [&values](uint64_t left, uint64_t right) { return values.GetView(left) < values.GetView(right); @@ -227,7 +227,7 @@ class ArrayCountSorter { // first slot reserved for prefix sum std::vector counts(1 + value_range); - if (options.order == SortOrder::ASCENDING) { + if (options.order == SortOrder::Ascending) { VisitRawValuesInline( values, [&](c_type v) { ++counts[v - min_ + 1]; }, []() {}); for (uint32_t i = 1; i <= value_range; ++i) { @@ -499,7 +499,7 @@ class TableSorter : public TypeVisitor { } else { compared = -1; } - if (order == SortOrder::DESCENDING) { + if (order == SortOrder::Descending) { compared = -compared; } return compared; @@ -568,7 +568,7 @@ class TableSorter : public TypeVisitor { return comparer.Compare(left, right, 1); } else { auto compared = value_left < value_right; - if (first_sort_key.order == SortOrder::ASCENDING) { + if (first_sort_key.order == SortOrder::Ascending) { return compared; } else { return !compared; @@ -693,7 +693,7 @@ class SortIndicesMetaFunction : public MetaFunction { Result> SortIndices(const Array& values, const SortOptions& options, ExecContext* ctx) const { - SortOrder order = SortOrder::ASCENDING; + SortOrder order = SortOrder::Ascending; if (!options.sort_keys.empty()) { order = options.sort_keys[0].order; } @@ -706,7 +706,7 @@ class SortIndicesMetaFunction : public MetaFunction { Result> SortIndices(const ChunkedArray& values, const SortOptions& options, ExecContext* ctx) const { - SortOrder order = SortOrder::ASCENDING; + SortOrder order = SortOrder::Ascending; if (!options.sort_keys.empty()) { order = options.sort_keys[0].order; } diff --git a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc index ef755704149..baa149c50b5 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc @@ -77,7 +77,7 @@ static void TableSortIndicesInt64Count(benchmark::State& state) { auto values = rand.Int64(array_size, -100, 100, args.null_proportion); std::vector> fields = {{field("int64", int64())}}; auto table = Table::Make(schema(fields), {values}, array_size); - SortOptions options({SortKey("int64", SortOrder::ASCENDING)}); + SortOptions options({SortKey("int64", SortOrder::Ascending)}); TableSortIndicesBenchmark(state, table, options); } @@ -93,7 +93,7 @@ static void TableSortIndicesInt64Compare(benchmark::State& state) { auto values = rand.Int64(array_size, min, max, args.null_proportion); std::vector> fields = {{field("int64", int64())}}; auto table = Table::Make(schema(fields), {values}, array_size); - SortOptions options({SortKey("int64", SortOrder::ASCENDING)}); + SortOptions options({SortKey("int64", SortOrder::Ascending)}); TableSortIndicesBenchmark(state, table, options); } @@ -114,8 +114,8 @@ static void TableSortIndicesInt64Int64(benchmark::State& state) { }; auto table = Table::Make(schema(fields), {values1, values2}, array_size); SortOptions options({ - SortKey("int64-1", SortOrder::ASCENDING), - SortKey("int64-2", SortOrder::ASCENDING), + SortKey("int64-1", SortOrder::Ascending), + SortKey("int64-2", SortOrder::Ascending), }); TableSortIndicesBenchmark(state, table, options); diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index 7546574ff7f..aa2267f29ed 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -64,7 +64,7 @@ class SortComparator { if (lhs_isnan) return false; } if (array.GetView(lhs) == array.GetView(rhs)) return lhs < rhs; - if (order == SortOrder::ASCENDING) { + if (order == SortOrder::Ascending) { return array.GetView(lhs) < array.GetView(rhs); } else { return array.GetView(lhs) > array.GetView(rhs); @@ -251,7 +251,7 @@ class TestArraySortIndicesKernel : public TestBase { } virtual void AssertSortIndices(const std::string& values, const std::string& expected) { - AssertSortIndices(values, SortOrder::ASCENDING, expected); + AssertSortIndices(values, SortOrder::Ascending, expected); } }; @@ -277,19 +277,19 @@ TYPED_TEST(TestArraySortIndicesKernelForReal, SortReal) { this->AssertSortIndices("[7, 6, 5, 4, 3, 2, 1]", "[6,5,4,3,2,1,0]"); this->AssertSortIndices("[10.4, 12, 4.2, 50, 50.3, 32, 11]", "[2,0,6,1,5,3,4]"); - this->AssertSortIndices("[null, 1, 3.3, null, 2, 5.3]", SortOrder::ASCENDING, + this->AssertSortIndices("[null, 1, 3.3, null, 2, 5.3]", SortOrder::Ascending, "[1,4,2,5,0,3]"); - this->AssertSortIndices("[null, 1, 3.3, null, 2, 5.3]", SortOrder::DESCENDING, + this->AssertSortIndices("[null, 1, 3.3, null, 2, 5.3]", SortOrder::Descending, "[5,2,4,1,0,3]"); - this->AssertSortIndices("[3, 4, NaN, 1, 2, null]", SortOrder::ASCENDING, + this->AssertSortIndices("[3, 4, NaN, 1, 2, null]", SortOrder::Ascending, "[3,4,0,1,2,5]"); - this->AssertSortIndices("[3, 4, NaN, 1, 2, null]", SortOrder::DESCENDING, + this->AssertSortIndices("[3, 4, NaN, 1, 2, null]", SortOrder::Descending, "[1,0,4,3,2,5]"); - this->AssertSortIndices("[NaN, 2, NaN, 3, 1]", SortOrder::ASCENDING, "[4,1,3,0,2]"); - this->AssertSortIndices("[NaN, 2, NaN, 3, 1]", SortOrder::DESCENDING, "[3,1,4,0,2]"); - this->AssertSortIndices("[null, NaN, NaN, null]", SortOrder::ASCENDING, "[1,2,0,3]"); - this->AssertSortIndices("[null, NaN, NaN, null]", SortOrder::DESCENDING, "[1,2,0,3]"); + this->AssertSortIndices("[NaN, 2, NaN, 3, 1]", SortOrder::Ascending, "[4,1,3,0,2]"); + this->AssertSortIndices("[NaN, 2, NaN, 3, 1]", SortOrder::Descending, "[3,1,4,0,2]"); + this->AssertSortIndices("[null, NaN, NaN, null]", SortOrder::Ascending, "[1,2,0,3]"); + this->AssertSortIndices("[null, NaN, NaN, null]", SortOrder::Descending, "[1,2,0,3]"); } TYPED_TEST(TestArraySortIndicesKernelForIntegral, SortIntegral) { @@ -299,14 +299,14 @@ TYPED_TEST(TestArraySortIndicesKernelForIntegral, SortIntegral) { this->AssertSortIndices("[1, 2, 3, 4, 5, 6, 7]", "[0,1,2,3,4,5,6]"); this->AssertSortIndices("[7, 6, 5, 4, 3, 2, 1]", "[6,5,4,3,2,1,0]"); - this->AssertSortIndices("[10, 12, 4, 50, 50, 32, 11]", SortOrder::ASCENDING, + this->AssertSortIndices("[10, 12, 4, 50, 50, 32, 11]", SortOrder::Ascending, "[2,0,6,1,5,3,4]"); - this->AssertSortIndices("[10, 12, 4, 50, 50, 32, 11]", SortOrder::DESCENDING, + this->AssertSortIndices("[10, 12, 4, 50, 50, 32, 11]", SortOrder::Descending, "[3,4,5,1,6,0,2]"); - this->AssertSortIndices("[null, 1, 3, null, 2, 5]", SortOrder::ASCENDING, + this->AssertSortIndices("[null, 1, 3, null, 2, 5]", SortOrder::Ascending, "[1,4,2,5,0,3]"); - this->AssertSortIndices("[null, 1, 3, null, 2, 5]", SortOrder::DESCENDING, + this->AssertSortIndices("[null, 1, 3, null, 2, 5]", SortOrder::Descending, "[5,2,4,1,0,3]"); } @@ -317,9 +317,9 @@ TYPED_TEST(TestArraySortIndicesKernelForStrings, SortStrings) { this->AssertSortIndices(R"(["foo", "bar", "baz"])", "[1,2,0]"); this->AssertSortIndices(R"(["testing", "sort", "for", "strings"])", "[2, 1, 3, 0]"); - this->AssertSortIndices(R"(["c", "b", "a", "b"])", SortOrder::ASCENDING, + this->AssertSortIndices(R"(["c", "b", "a", "b"])", SortOrder::Ascending, "[2, 1, 3, 0]"); - this->AssertSortIndices(R"(["c", "b", "a", "b"])", SortOrder::DESCENDING, + this->AssertSortIndices(R"(["c", "b", "a", "b"])", SortOrder::Descending, "[0, 1, 3, 2]"); } @@ -376,7 +376,7 @@ TYPED_TEST(TestArraySortIndicesKernelRandom, SortRandomValues) { length = 10; for (int test = 0; test < times; test++) { for (auto null_probability : {0.0, 0.1, 0.5, 1.0}) { - for (auto order : {SortOrder::ASCENDING, SortOrder::DESCENDING}) { + for (auto order : {SortOrder::Ascending, SortOrder::Descending}) { auto array = rand.Generate(length, null_probability); ASSERT_OK_AND_ASSIGN(std::shared_ptr offsets, SortIndices(*array, order)); ValidateSorted(*checked_pointer_cast(array), @@ -400,7 +400,7 @@ TYPED_TEST(TestArraySortIndicesKernelRandomCount, SortRandomValuesCount) { int range = 2000; for (int test = 0; test < times; test++) { for (auto null_probability : {0.0, 0.1, 0.5, 1.0}) { - for (auto order : {SortOrder::ASCENDING, SortOrder::DESCENDING}) { + for (auto order : {SortOrder::Ascending, SortOrder::Descending}) { auto array = rand.Generate(length, range, null_probability); ASSERT_OK_AND_ASSIGN(std::shared_ptr offsets, SortIndices(*array, order)); ValidateSorted(*checked_pointer_cast(array), @@ -421,7 +421,7 @@ TYPED_TEST(TestArraySortIndicesKernelRandomCompare, SortRandomValuesCompare) { int length = 4000; for (int test = 0; test < times; test++) { for (auto null_probability : {0.0, 0.1, 0.5, 1.0}) { - for (auto order : {SortOrder::ASCENDING, SortOrder::DESCENDING}) { + for (auto order : {SortOrder::Ascending, SortOrder::Descending}) { auto array = rand.Generate(length, null_probability); ASSERT_OK_AND_ASSIGN(std::shared_ptr offsets, SortIndices(*array, order)); ValidateSorted(*checked_pointer_cast(array), @@ -459,7 +459,7 @@ TEST_F(TestTableSortIndices, SortNull) { "{\"a\": 1, \"b\": 5}" "]"}); SortOptions options( - {SortKey("a", SortOrder::ASCENDING), SortKey("b", SortOrder::DESCENDING)}); + {SortKey("a", SortOrder::Ascending), SortKey("b", SortOrder::Descending)}); this->AssertSortIndices(table, options, "[5, 1, 4, 2, 0, 3]"); } @@ -479,7 +479,7 @@ TEST_F(TestTableSortIndices, SortNaN) { "{\"a\": 1, \"b\": 5}" "]"}); SortOptions options( - {SortKey("a", SortOrder::ASCENDING), SortKey("b", SortOrder::DESCENDING)}); + {SortKey("a", SortOrder::Ascending), SortKey("b", SortOrder::Descending)}); this->AssertSortIndices(table, options, "[7, 1, 2, 5, 6, 4, 0, 3]"); } @@ -500,7 +500,7 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { if (lhs_array_->IsNull(lhs_index_)) return false; status_ = lhs_array_->type()->Accept(this); if (compared_ == 0) continue; - if (sort_key.order == SortOrder::ASCENDING) { + if (sort_key.order == SortOrder::Ascending) { return compared_ < 0; } else { return compared_ > 0; @@ -625,12 +625,12 @@ TEST_P(TestTableSortIndicesRandom, Sort) { auto n_sort_keys = 2; std::vector sort_keys; auto first_sort_key_order = - (distribution(engine) % 2) == 0 ? SortOrder::ASCENDING : SortOrder::DESCENDING; + (distribution(engine) % 2) == 0 ? SortOrder::Ascending : SortOrder::Descending; sort_keys.emplace_back(first_sort_key_name, first_sort_key_order); for (int i = 1; i < n_sort_keys; ++i) { auto& column_name = column_names[distribution(engine) % column_names.size()]; auto order = - (distribution(engine) % 2) == 0 ? SortOrder::ASCENDING : SortOrder::DESCENDING; + (distribution(engine) % 2) == 0 ? SortOrder::Ascending : SortOrder::Descending; sort_keys.emplace_back(column_name, order); } SortOptions options(sort_keys); From 0d097db9451e99d571c40cb0da2dc9d7fb894709 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Tue, 17 Nov 2020 16:13:25 +0900 Subject: [PATCH 16/37] Fix wrong description --- cpp/src/arrow/compute/api_vector.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/arrow/compute/api_vector.h b/cpp/src/arrow/compute/api_vector.h index 16895a5dc09..32726e3fe21 100644 --- a/cpp/src/arrow/compute/api_vector.h +++ b/cpp/src/arrow/compute/api_vector.h @@ -68,7 +68,7 @@ struct ARROW_EXPORT SortKey { explicit SortKey(std::string name, SortOrder order = SortOrder::Ascending) : name(name), order(order) {} - /// The name of the sort key. + /// The name of the sort column. std::string name; /// How to order by this sort key. SortOrder order; From 493313a556a8be753931fa5446a12a2cdaa333e6 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Tue, 17 Nov 2020 16:55:49 +0900 Subject: [PATCH 17/37] Remove SortIndices(Array, ExecContext) --- cpp/src/arrow/compute/api_vector.cc | 6 +----- cpp/src/arrow/compute/api_vector.h | 20 ++------------------ 2 files changed, 3 insertions(+), 23 deletions(-) diff --git a/cpp/src/arrow/compute/api_vector.cc b/cpp/src/arrow/compute/api_vector.cc index 81830d91161..9b9e03bae24 100644 --- a/cpp/src/arrow/compute/api_vector.cc +++ b/cpp/src/arrow/compute/api_vector.cc @@ -46,10 +46,6 @@ Result> NthToIndices(const Array& values, int64_t n, return result.make_array(); } -Result> SortIndices(const Array& values, ExecContext* ctx) { - return SortIndices(values, SortOrder::Ascending, ctx); -} - Result> SortIndices(const Array& values, SortOrder order, ExecContext* ctx) { ArraySortOptions options(order); @@ -150,7 +146,7 @@ Result> Take(const Table& table, const ChunkedArray& indi } Result> SortToIndices(const Array& values, ExecContext* ctx) { - return SortIndices(values, ctx); + return SortIndices(values, SortOrder::Ascending, ctx); } } // namespace compute diff --git a/cpp/src/arrow/compute/api_vector.h b/cpp/src/arrow/compute/api_vector.h index 32726e3fe21..fa8360cab8e 100644 --- a/cpp/src/arrow/compute/api_vector.h +++ b/cpp/src/arrow/compute/api_vector.h @@ -180,23 +180,6 @@ ARROW_EXPORT Result> NthToIndices(const Array& values, int64_t n, ExecContext* ctx = NULLPTR); -/// \brief Returns the indices that would sort an array in ascending -/// order. -/// -/// Perform an indirect sort of array. The output array will contain -/// indices that would sort an array, which would be the same length -/// as input. Nulls will be stably partitioned to the end of the output. -/// -/// For example given values = [null, 1, 3.3, null, 2, 5.3], the output -/// will be [1, 4, 2, 5, 0, 3] -/// -/// \param[in] values array to sort -/// \param[in] ctx the function execution context, optional -/// \return offsets indices that would sort an array -ARROW_EXPORT -Result> SortIndices(const Array& values, - ExecContext* ctx = NULLPTR); - /// \brief Returns the indices that would sort an array in the /// specified order. /// @@ -213,7 +196,8 @@ Result> SortIndices(const Array& values, /// \param[in] ctx the function execution context, optional /// \return offsets indices that would sort an array ARROW_EXPORT -Result> SortIndices(const Array& values, SortOrder order, +Result> SortIndices(const Array& values, + SortOrder order = SortOrder::Ascending, ExecContext* ctx = NULLPTR); /// \brief Returns the indices that would sort a table in the From 628e3ecc9d3ec60a95e55fc93d50bb38b9c51b79 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Tue, 17 Nov 2020 16:56:02 +0900 Subject: [PATCH 18/37] Add more description about null and order --- cpp/src/arrow/compute/api_vector.h | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/compute/api_vector.h b/cpp/src/arrow/compute/api_vector.h index fa8360cab8e..6cf3b3ab781 100644 --- a/cpp/src/arrow/compute/api_vector.h +++ b/cpp/src/arrow/compute/api_vector.h @@ -204,8 +204,9 @@ Result> SortIndices(const Array& values, /// specified order. /// /// Perform an indirect sort of table. The output array will contain -/// indices that would sort a table, which would be the same length -/// as input. Nulls will be stably partitioned to the end of the output. +/// indices that would sort a table, which would be the same length as +/// input. Nulls will be stably partitioned to the end of the output +/// regardless of order. /// /// For example given table = { /// "column1": [null, 1, 3, null, 2, 1], From 629e2915f71af5ee5ab47ddc650e1ef50ead57a0 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Tue, 17 Nov 2020 17:00:55 +0900 Subject: [PATCH 19/37] Use std::pair to return two value --- cpp/src/arrow/compute/kernels/vector_sort.cc | 94 +++++++++++--------- 1 file changed, 52 insertions(+), 42 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc index c48a5599611..b57d8460a9b 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort.cc @@ -376,6 +376,8 @@ void AddSortingKernels(VectorKernel base, VectorFunction* func) { class TableSorter : public TypeVisitor { private: + // Finds the target chunk and index in the target chunk from an + // index in ChunkedArray. struct ResolvedSortKey { ResolvedSortKey(const ChunkedArray& chunked_array, const SortOrder order) : order(order) { @@ -387,21 +389,22 @@ class TableSorter : public TypeVisitor { } } + // Finds the target chunk and index in the target chunk from an + // index in ChunkedArray. template - ArrayType* ResolveChunk(int64_t index, int64_t& chunk_index) const { + std::pair ResolveChunk(int64_t index) const { + using Chunk = std::pair; if (num_chunks == 1) { - chunk_index = index; - return static_cast(chunks[0]); + return Chunk(static_cast(chunks[0]), index); } else { int64_t offset = 0; for (size_t i = 0; i < num_chunks; ++i) { if (index < offset + chunks[i]->length()) { - chunk_index = index - offset; - return static_cast(chunks[i]); + return Chunk(static_cast(chunks[i]), index - offset); } offset += chunks[i]->length(); } - return nullptr; + return Chunk(nullptr, 0); } } @@ -474,10 +477,12 @@ class TableSorter : public TypeVisitor { using ArrayType = typename TypeTraits::ArrayType; const auto& sort_key = sort_keys_[current_sort_key_index_]; auto order = sort_key.order; - int64_t index_left = 0; - auto array_left = sort_key.ResolveChunk(current_left_, index_left); - int64_t index_right = 0; - auto array_right = sort_key.ResolveChunk(current_right_, index_right); + auto chunk_left = sort_key.ResolveChunk(current_left_); + auto array_left = chunk_left.first; + auto index_left = chunk_left.second; + auto chunk_right = sort_key.ResolveChunk(current_right_); + auto array_right = chunk_right.first; + auto index_right = chunk_right.second; if (sort_key.null_count > 0) { auto is_null_left = array_left->IsNull(index_left); auto is_null_right = array_right->IsNull(index_right); @@ -555,26 +560,27 @@ class TableSorter : public TypeVisitor { const auto& first_sort_key = comparer.sort_keys()[0]; auto nulls_begin = indices_end_; nulls_begin = PartitionNullsInternal(first_sort_key); - std::stable_sort( - indices_begin_, nulls_begin, - [&first_sort_key, &comparer](uint64_t left, uint64_t right) { - int64_t index_left = 0; - auto array_left = first_sort_key.ResolveChunk(left, index_left); - int64_t index_right = 0; - auto array_right = first_sort_key.ResolveChunk(right, index_right); - auto value_left = array_left->GetView(index_left); - auto value_right = array_right->GetView(index_right); - if (value_left == value_right) { - return comparer.Compare(left, right, 1); - } else { - auto compared = value_left < value_right; - if (first_sort_key.order == SortOrder::Ascending) { - return compared; - } else { - return !compared; - } - } - }); + std::stable_sort(indices_begin_, nulls_begin, + [&first_sort_key, &comparer](uint64_t left, uint64_t right) { + auto chunk_left = first_sort_key.ResolveChunk(left); + auto array_left = chunk_left.first; + auto index_left = chunk_left.second; + auto chunk_right = first_sort_key.ResolveChunk(right); + auto array_right = chunk_right.first; + auto index_right = chunk_right.second; + auto value_left = array_left->GetView(index_left); + auto value_right = array_right->GetView(index_right); + if (value_left == value_right) { + return comparer.Compare(left, right, 1); + } else { + auto compared = value_left < value_right; + if (first_sort_key.order == SortOrder::Ascending) { + return compared; + } else { + return !compared; + } + } + }); return Status::OK(); } @@ -588,9 +594,10 @@ class TableSorter : public TypeVisitor { StablePartitioner partitioner; auto nulls_begin = partitioner(indices_begin_, indices_end_, [&first_sort_key](uint64_t index) { - int64_t index_chunk = 0; - auto chunk = first_sort_key.ResolveChunk(index, index_chunk); - return !chunk->IsNull(index_chunk); + auto chunk = first_sort_key.ResolveChunk(index); + auto array = chunk.first; + auto index_array = chunk.second; + return !array->IsNull(index_array); }); auto& comparer = comparer_; std::stable_sort(nulls_begin, indices_end_, @@ -607,25 +614,28 @@ class TableSorter : public TypeVisitor { StablePartitioner partitioner; if (first_sort_key.null_count == 0) { return partitioner(indices_begin_, indices_end_, [&first_sort_key](uint64_t index) { - int64_t index_chunk = 0; - auto chunk = first_sort_key.ResolveChunk(index, index_chunk); - return !std::isnan(chunk->GetView(index_chunk)); + auto chunk = first_sort_key.ResolveChunk(index); + auto array = chunk.first; + auto index_array = chunk.second; + return !std::isnan(array->GetView(index_array)); }); } auto nans_and_nulls_begin = partitioner(indices_begin_, indices_end_, [&first_sort_key](uint64_t index) { - int64_t index_chunk = 0; - auto chunk = first_sort_key.ResolveChunk(index, index_chunk); - return !chunk->IsNull(index_chunk) && !std::isnan(chunk->GetView(index_chunk)); + auto chunk = first_sort_key.ResolveChunk(index); + auto array = chunk.first; + auto index_array = chunk.second; + return !array->IsNull(index_array) && !std::isnan(array->GetView(index_array)); }); auto nulls_begin = nans_and_nulls_begin; if (first_sort_key.null_count < static_cast(indices_end_ - nulls_begin)) { // move Nulls after NaN nulls_begin = partitioner( nans_and_nulls_begin, indices_end_, [&first_sort_key](uint64_t index) { - int64_t index_chunk = 0; - auto chunk = first_sort_key.ResolveChunk(index, index_chunk); - return !chunk->IsNull(index_chunk); + auto chunk = first_sort_key.ResolveChunk(index); + auto array = chunk.first; + auto index_array = chunk.second; + return !array->IsNull(index_array); }); } auto& comparer = comparer_; From ff81915c5f510af95ee11ebbb109251b967e427f Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Tue, 17 Nov 2020 17:01:07 +0900 Subject: [PATCH 20/37] Fix type --- cpp/src/arrow/compute/kernels/vector_sort.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc index b57d8460a9b..10ab916e888 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort.cc @@ -398,7 +398,7 @@ class TableSorter : public TypeVisitor { return Chunk(static_cast(chunks[0]), index); } else { int64_t offset = 0; - for (size_t i = 0; i < num_chunks; ++i) { + for (int i = 0; i < num_chunks; ++i) { if (index < offset + chunks[i]->length()) { return Chunk(static_cast(chunks[i]), index - offset); } @@ -411,7 +411,7 @@ class TableSorter : public TypeVisitor { SortOrder order; DataType* type; int64_t null_count; - size_t num_chunks; + int num_chunks; std::vector chunks; }; From 8170331253326463153257106d9091f9833eb8d9 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Tue, 17 Nov 2020 17:11:23 +0900 Subject: [PATCH 21/37] Add a comment to Compare() --- cpp/src/arrow/compute/kernels/vector_sort.cc | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc index 10ab916e888..afcd3dd964c 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort.cc @@ -433,6 +433,9 @@ class TableSorter : public TypeVisitor { const std::vector& sort_keys() { return sort_keys_; } + // Returns true if the left-th value should be ordered before the + // right-th value, false otherwise. The start_sort_key_index-th + // sort key and subsequent sort keys are used for comparison. bool Compare(uint64_t left, uint64_t right, size_t start_sort_key_index) { current_left_ = left; current_right_ = right; From ac0f654edcf84c2bca25649f24834aab847c894a Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Wed, 18 Nov 2020 15:26:54 +0900 Subject: [PATCH 22/37] Add more benchmark patterns 2020-11-18T15:27:12+09:00 Running release/arrow-compute-vector-sort-benchmark Run on (24 X 3800 MHz CPU s) CPU Caches: L1 Data 32 KiB (x12) L1 Instruction 32 KiB (x12) L2 Unified 512 KiB (x12) L3 Unified 16384 KiB (x4) Load Average: 0.48, 0.50, 0.53 ***WARNING*** CPU scaling is enabled, the benchmark real time measurements may be noisy and will incur extra overhead. ----------------------------------------------------------------------------------------------------------------------- Benchmark Time CPU Iterations UserCounters... ----------------------------------------------------------------------------------------------------------------------- ArraySortIndicesInt64Narrow/32768/10000/min_time:1.000 15722 ns 15721 ns 82919 bytes_per_second=1.94115G/s items_per_second=260.537M/s null_percent=0.01 size=32.768k ArraySortIndicesInt64Narrow/32768/100/min_time:1.000 16007 ns 16006 ns 89081 bytes_per_second=1.9066G/s items_per_second=255.899M/s null_percent=1 size=32.768k ArraySortIndicesInt64Narrow/32768/10/min_time:1.000 15178 ns 15177 ns 91254 bytes_per_second=2.01071G/s items_per_second=269.873M/s null_percent=10 size=32.768k ArraySortIndicesInt64Narrow/32768/2/min_time:1.000 29886 ns 29885 ns 47570 bytes_per_second=1045.67M/s items_per_second=137.058M/s null_percent=50 size=32.768k ArraySortIndicesInt64Narrow/32768/1/min_time:1.000 5968 ns 5968 ns 235596 bytes_per_second=5.11372G/s items_per_second=686.352M/s null_percent=100 size=32.768k ArraySortIndicesInt64Narrow/32768/0/min_time:1.000 12681 ns 12681 ns 108821 bytes_per_second=2.40665G/s items_per_second=323.015M/s null_percent=0 size=32.768k ArraySortIndicesInt64Narrow/1048576/100/min_time:1.000 813376 ns 813331 ns 1777 bytes_per_second=1.2007G/s items_per_second=161.155M/s null_percent=1 size=1048.58k ArraySortIndicesInt64Narrow/8388608/100/min_time:1.000 6543874 ns 6543621 ns 207 bytes_per_second=1.19391G/s items_per_second=160.244M/s null_percent=1 size=8.38861M ArraySortIndicesInt64Wide/32768/10000/min_time:1.000 189957 ns 189956 ns 7344 bytes_per_second=164.511M/s items_per_second=21.5628M/s null_percent=0.01 size=32.768k ArraySortIndicesInt64Wide/32768/100/min_time:1.000 190269 ns 190267 ns 7346 bytes_per_second=164.243M/s items_per_second=21.5277M/s null_percent=1 size=32.768k ArraySortIndicesInt64Wide/32768/10/min_time:1.000 175425 ns 175422 ns 7952 bytes_per_second=178.142M/s items_per_second=23.3494M/s null_percent=10 size=32.768k ArraySortIndicesInt64Wide/32768/2/min_time:1.000 107976 ns 107975 ns 12837 bytes_per_second=289.418M/s items_per_second=37.9346M/s null_percent=50 size=32.768k ArraySortIndicesInt64Wide/32768/1/min_time:1.000 5941 ns 5941 ns 235078 bytes_per_second=5.13666G/s items_per_second=689.431M/s null_percent=100 size=32.768k ArraySortIndicesInt64Wide/32768/0/min_time:1.000 184858 ns 184857 ns 7542 bytes_per_second=169.049M/s items_per_second=22.1576M/s null_percent=0 size=32.768k ArraySortIndicesInt64Wide/1048576/100/min_time:1.000 9355194 ns 9354942 ns 150 bytes_per_second=106.895M/s items_per_second=14.011M/s null_percent=1 size=1048.58k ArraySortIndicesInt64Wide/8388608/100/min_time:1.000 94101796 ns 94099551 ns 15 bytes_per_second=85.0163M/s items_per_second=11.1433M/s null_percent=1 size=8.38861M TableSortIndicesInt64Narrow/1048576/100/16/32/min_time:1.000 1386231938 ns 1386176541 ns 1 items_per_second=756.452k/s TableSortIndicesInt64Narrow/1048576/0/16/32/min_time:1.000 1350031141 ns 1349982623 ns 1 items_per_second=776.733k/s TableSortIndicesInt64Narrow/1048576/100/8/32/min_time:1.000 1401741018 ns 1401632081 ns 1 items_per_second=748.111k/s TableSortIndicesInt64Narrow/1048576/0/8/32/min_time:1.000 1373174328 ns 1373145968 ns 1 items_per_second=763.63k/s TableSortIndicesInt64Narrow/1048576/100/2/32/min_time:1.000 1035377805 ns 1035362806 ns 1 items_per_second=1012.76k/s TableSortIndicesInt64Narrow/1048576/0/2/32/min_time:1.000 1035218085 ns 1035201824 ns 1 items_per_second=1012.92k/s TableSortIndicesInt64Narrow/1048576/100/1/32/min_time:1.000 6859319 ns 6859251 ns 201 items_per_second=152.87M/s TableSortIndicesInt64Narrow/1048576/0/1/32/min_time:1.000 6847314 ns 6847048 ns 209 items_per_second=153.143M/s TableSortIndicesInt64Narrow/1048576/100/16/4/min_time:1.000 626909516 ns 626904696 ns 2 items_per_second=1.67262M/s TableSortIndicesInt64Narrow/1048576/0/16/4/min_time:1.000 591632035 ns 591602144 ns 2 items_per_second=1.77243M/s TableSortIndicesInt64Narrow/1048576/100/8/4/min_time:1.000 613197155 ns 613182727 ns 2 items_per_second=1.71005M/s TableSortIndicesInt64Narrow/1048576/0/8/4/min_time:1.000 595568690 ns 595554829 ns 2 items_per_second=1.76067M/s TableSortIndicesInt64Narrow/1048576/100/2/4/min_time:1.000 424472984 ns 424453915 ns 3 items_per_second=2.47041M/s TableSortIndicesInt64Narrow/1048576/0/2/4/min_time:1.000 397339109 ns 397335604 ns 4 items_per_second=2.63902M/s TableSortIndicesInt64Narrow/1048576/100/1/4/min_time:1.000 7027310 ns 7027241 ns 195 items_per_second=149.216M/s TableSortIndicesInt64Narrow/1048576/0/1/4/min_time:1.000 6891364 ns 6891272 ns 207 items_per_second=152.16M/s TableSortIndicesInt64Narrow/1048576/100/16/1/min_time:1.000 516832749 ns 516823293 ns 3 items_per_second=2.02889M/s TableSortIndicesInt64Narrow/1048576/0/16/1/min_time:1.000 495273237 ns 495269975 ns 3 items_per_second=2.11718M/s TableSortIndicesInt64Narrow/1048576/100/8/1/min_time:1.000 515550360 ns 515531501 ns 3 items_per_second=2.03397M/s TableSortIndicesInt64Narrow/1048576/0/8/1/min_time:1.000 496670125 ns 496663084 ns 3 items_per_second=2.11124M/s TableSortIndicesInt64Narrow/1048576/100/2/1/min_time:1.000 340060863 ns 340053172 ns 4 items_per_second=3.08356M/s TableSortIndicesInt64Narrow/1048576/0/2/1/min_time:1.000 331281499 ns 331277004 ns 4 items_per_second=3.16525M/s TableSortIndicesInt64Narrow/1048576/100/1/1/min_time:1.000 6691062 ns 6690924 ns 212 items_per_second=156.716M/s TableSortIndicesInt64Narrow/1048576/0/1/1/min_time:1.000 6384382 ns 6384323 ns 221 items_per_second=164.242M/s TableSortIndicesInt64Wide/1048576/100/16/32/min_time:1.000 622544467 ns 622531557 ns 2 items_per_second=1.68437M/s TableSortIndicesInt64Wide/1048576/0/16/32/min_time:1.000 619193843 ns 619155597 ns 2 items_per_second=1.69356M/s TableSortIndicesInt64Wide/1048576/100/8/32/min_time:1.000 615885889 ns 615884764 ns 2 items_per_second=1.70255M/s TableSortIndicesInt64Wide/1048576/0/8/32/min_time:1.000 589917731 ns 589912005 ns 2 items_per_second=1.77751M/s TableSortIndicesInt64Wide/1048576/100/2/32/min_time:1.000 600951149 ns 600947256 ns 2 items_per_second=1.74487M/s TableSortIndicesInt64Wide/1048576/0/2/32/min_time:1.000 592304437 ns 592286953 ns 2 items_per_second=1.77039M/s TableSortIndicesInt64Wide/1048576/100/1/32/min_time:1.000 98781239 ns 98777123 ns 14 items_per_second=10.6156M/s TableSortIndicesInt64Wide/1048576/0/1/32/min_time:1.000 94136230 ns 94085863 ns 15 items_per_second=11.1449M/s TableSortIndicesInt64Wide/1048576/100/16/4/min_time:1.000 232323308 ns 232319347 ns 6 items_per_second=4.51351M/s TableSortIndicesInt64Wide/1048576/0/16/4/min_time:1.000 223708791 ns 223700834 ns 6 items_per_second=4.6874M/s TableSortIndicesInt64Wide/1048576/100/8/4/min_time:1.000 221975360 ns 221968035 ns 6 items_per_second=4.724M/s TableSortIndicesInt64Wide/1048576/0/8/4/min_time:1.000 218214843 ns 218210306 ns 6 items_per_second=4.80535M/s TableSortIndicesInt64Wide/1048576/100/2/4/min_time:1.000 224756430 ns 224745751 ns 6 items_per_second=4.66561M/s TableSortIndicesInt64Wide/1048576/0/2/4/min_time:1.000 220809969 ns 220809044 ns 6 items_per_second=4.74879M/s TableSortIndicesInt64Wide/1048576/100/1/4/min_time:1.000 96159427 ns 96156899 ns 14 items_per_second=10.9048M/s TableSortIndicesInt64Wide/1048576/0/1/4/min_time:1.000 92307749 ns 92304526 ns 15 items_per_second=11.36M/s TableSortIndicesInt64Wide/1048576/100/16/1/min_time:1.000 166390841 ns 166382427 ns 8 items_per_second=6.3022M/s TableSortIndicesInt64Wide/1048576/0/16/1/min_time:1.000 164576492 ns 164570581 ns 9 items_per_second=6.37159M/s TableSortIndicesInt64Wide/1048576/100/8/1/min_time:1.000 165724919 ns 165718584 ns 8 items_per_second=6.32745M/s TableSortIndicesInt64Wide/1048576/0/8/1/min_time:1.000 164048003 ns 164046222 ns 8 items_per_second=6.39195M/s TableSortIndicesInt64Wide/1048576/100/2/1/min_time:1.000 170131788 ns 170124950 ns 8 items_per_second=6.16356M/s TableSortIndicesInt64Wide/1048576/0/2/1/min_time:1.000 170874129 ns 170871245 ns 8 items_per_second=6.13664M/s TableSortIndicesInt64Wide/1048576/100/1/1/min_time:1.000 92314674 ns 92312953 ns 15 items_per_second=11.3589M/s TableSortIndicesInt64Wide/1048576/0/1/1/min_time:1.000 92023019 ns 92022643 ns 15 items_per_second=11.3948M/s --- cpp/cmake_modules/ThirdpartyToolchain.cmake | 7 +- .../compute/kernels/vector_sort_benchmark.cc | 144 +++++++++++------- cpp/thirdparty/versions.txt | 2 +- 3 files changed, 91 insertions(+), 62 deletions(-) diff --git a/cpp/cmake_modules/ThirdpartyToolchain.cmake b/cpp/cmake_modules/ThirdpartyToolchain.cmake index 22531fcfc57..1644d5b3b91 100644 --- a/cpp/cmake_modules/ThirdpartyToolchain.cmake +++ b/cpp/cmake_modules/ThirdpartyToolchain.cmake @@ -184,9 +184,9 @@ macro(resolve_dependency DEPENDENCY_NAME) if(${DEPENDENCY_NAME}_SOURCE STREQUAL "AUTO") if(ARG_REQUIRED_VERSION) - find_package(${DEPENDENCY_NAME} ${ARG_REQUIRED_VERSION} MODULE) + find_package(${DEPENDENCY_NAME} ${ARG_REQUIRED_VERSION}) else() - find_package(${DEPENDENCY_NAME} MODULE) + find_package(${DEPENDENCY_NAME}) endif() if(${${DEPENDENCY_NAME}_FOUND}) set(${DEPENDENCY_NAME}_SOURCE "SYSTEM") @@ -1797,7 +1797,8 @@ macro(build_benchmark) endmacro() if(ARROW_BUILD_BENCHMARKS) - resolve_dependency(benchmark) + # ArgsProduct() is available since 1.5.2 + resolve_dependency(benchmark REQUIRED_VERSION 1.5.2) # TODO: Don't use global includes but rather target_include_directories get_target_property(BENCHMARK_INCLUDE_DIR benchmark::benchmark INTERFACE_INCLUDE_DIRECTORIES) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc index baa149c50b5..235c445678d 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc @@ -36,7 +36,7 @@ static void ArraySortIndicesBenchmark(benchmark::State& state, state.SetItemsProcessed(state.iterations() * values->length()); } -static void ArraySortIndicesInt64Count(benchmark::State& state) { +static void ArraySortIndicesInt64Narrow(benchmark::State& state) { RegressionArgs args(state); const int64_t array_size = args.size / sizeof(int64_t); @@ -47,7 +47,7 @@ static void ArraySortIndicesInt64Count(benchmark::State& state) { ArraySortIndicesBenchmark(state, values); } -static void ArraySortIndicesInt64Compare(benchmark::State& state) { +static void ArraySortIndicesInt64Wide(benchmark::State& state) { RegressionArgs args(state); const int64_t array_size = args.size / sizeof(int64_t); @@ -66,93 +66,121 @@ static void TableSortIndicesBenchmark(benchmark::State& state, for (auto _ : state) { ABORT_NOT_OK(SortIndices(*table, options).status()); } - state.SetItemsProcessed(state.iterations() * table->num_rows()); } -static void TableSortIndicesInt64Count(benchmark::State& state) { - RegressionArgs args(state); +// Extract benchmark args from benchmark::State +struct TableSortIndicesArgs { + // the number of records + const int64_t num_records; - const int64_t array_size = args.size / sizeof(int64_t); - auto rand = random::RandomArrayGenerator(kSeed); - auto values = rand.Int64(array_size, -100, 100, args.null_proportion); - std::vector> fields = {{field("int64", int64())}}; - auto table = Table::Make(schema(fields), {values}, array_size); - SortOptions options({SortKey("int64", SortOrder::Ascending)}); + // proportion of nulls in generated arrays + const double null_proportion; - TableSortIndicesBenchmark(state, table, options); -} + // the number of columns + const int64_t num_columns; -static void TableSortIndicesInt64Compare(benchmark::State& state) { - RegressionArgs args(state); + // the number of chunks in each generated column + const int64_t num_chunks; - const int64_t array_size = args.size / sizeof(int64_t); - auto rand = random::RandomArrayGenerator(kSeed); + // Extract args + explicit TableSortIndicesArgs(benchmark::State& state) + : num_records(state.range(0)), + null_proportion(ComputeNullProportion(state.range(1))), + num_columns(state.range(2)), + num_chunks(state.range(3)), + state_(state) { + } - auto min = std::numeric_limits::min(); - auto max = std::numeric_limits::max(); - auto values = rand.Int64(array_size, min, max, args.null_proportion); - std::vector> fields = {{field("int64", int64())}}; - auto table = Table::Make(schema(fields), {values}, array_size); - SortOptions options({SortKey("int64", SortOrder::Ascending)}); + ~TableSortIndicesArgs() { + state_.SetItemsProcessed(state_.iterations() * num_records); + } - TableSortIndicesBenchmark(state, table, options); -} + private: + double ComputeNullProportion(int64_t inverse_null_proportion) { + if (inverse_null_proportion == 0) { + return 0.0; + } else { + return std::min(1., 1. / static_cast(inverse_null_proportion)); + } + } -static void TableSortIndicesInt64Int64(benchmark::State& state) { - RegressionArgs args(state); + benchmark::State& state_; +}; - const int64_t array_size = args.size / sizeof(int64_t); - auto rand = random::RandomArrayGenerator(kSeed); +static void TableSortIndicesInt64(benchmark::State& state, int64_t min, int64_t max) { + TableSortIndicesArgs args(state); - auto min = std::numeric_limits::min(); - auto max = std::numeric_limits::max(); - auto values1 = rand.Int64(array_size, min, max, args.null_proportion); - auto values2 = rand.Int64(array_size, min, max, args.null_proportion); - std::vector> fields = { - {field("int64-1", int64())}, - {field("int64-2", int64())}, - }; - auto table = Table::Make(schema(fields), {values1, values2}, array_size); - SortOptions options({ - SortKey("int64-1", SortOrder::Ascending), - SortKey("int64-2", SortOrder::Ascending), - }); + auto rand = random::RandomArrayGenerator(kSeed); + std::vector> fields; + std::vector sort_keys; + std::vector> columns; + for (int64_t i = 0; i < args.num_columns; ++i) { + auto name = std::to_string(i); + fields.push_back(field(name, int64())); + auto order = (i % 2) == 0 ? SortOrder::Ascending : SortOrder::Descending; + sort_keys.emplace_back(name, order); + std::vector> arrays; + if ((args.num_records % args.num_chunks) != 0) { + Status::Invalid("The number of chunks (", args.num_chunks, ") must be " + "a multiple of the number of records (", args.num_records, ")").Abort(); + } + auto num_records_in_array = args.num_records / args.num_chunks; + for (int64_t j = 0; j < args.num_chunks; ++j) { + arrays.push_back(rand.Int64(num_records_in_array, + min, + max, + args.null_proportion)); + } + ASSIGN_OR_ABORT(auto chunked_array, ChunkedArray::Make(arrays, int64())); + columns.push_back(chunked_array); + } + auto table = Table::Make(schema(fields), columns, args.num_records); + SortOptions options(sort_keys); TableSortIndicesBenchmark(state, table, options); } -BENCHMARK(ArraySortIndicesInt64Count) - ->Apply(RegressionSetArgs) - ->Args({1 << 20, 100}) - ->Args({1 << 23, 100}) - ->MinTime(1.0) - ->Unit(benchmark::TimeUnit::kNanosecond); +static void TableSortIndicesInt64Narrow(benchmark::State& state) { + TableSortIndicesInt64(state, -100, 100); +} + +static void TableSortIndicesInt64Wide(benchmark::State& state) { + TableSortIndicesInt64(state, + std::numeric_limits::min(), + std::numeric_limits::max()); +} -BENCHMARK(ArraySortIndicesInt64Compare) +BENCHMARK(ArraySortIndicesInt64Narrow) ->Apply(RegressionSetArgs) ->Args({1 << 20, 100}) ->Args({1 << 23, 100}) ->MinTime(1.0) ->Unit(benchmark::TimeUnit::kNanosecond); -BENCHMARK(TableSortIndicesInt64Count) +BENCHMARK(ArraySortIndicesInt64Wide) ->Apply(RegressionSetArgs) ->Args({1 << 20, 100}) ->Args({1 << 23, 100}) ->MinTime(1.0) ->Unit(benchmark::TimeUnit::kNanosecond); -BENCHMARK(TableSortIndicesInt64Compare) - ->Apply(RegressionSetArgs) - ->Args({1 << 20, 100}) - ->Args({1 << 23, 100}) +BENCHMARK(TableSortIndicesInt64Narrow) + ->ArgsProduct({ + {1 << 20}, // the number of records + {100, 0}, // inverse null proportion + {16, 8, 2, 1}, // the number of columns + {32, 4, 1}, // the number of chunks + }) ->MinTime(1.0) ->Unit(benchmark::TimeUnit::kNanosecond); -BENCHMARK(TableSortIndicesInt64Int64) - ->Apply(RegressionSetArgs) - ->Args({1 << 20, 100}) - ->Args({1 << 23, 100}) +BENCHMARK(TableSortIndicesInt64Wide) + ->ArgsProduct({ + {1 << 20}, // the number of records + {100, 0}, // inverse null proportion + {16, 8, 2, 1}, // the number of columns + {32, 4, 1}, // the number of chunks + }) ->MinTime(1.0) ->Unit(benchmark::TimeUnit::kNanosecond); diff --git a/cpp/thirdparty/versions.txt b/cpp/thirdparty/versions.txt index 545943d30c6..969de58ea79 100644 --- a/cpp/thirdparty/versions.txt +++ b/cpp/thirdparty/versions.txt @@ -32,7 +32,7 @@ ARROW_BOOST_BUILD_VERSION=1.71.0 ARROW_BROTLI_BUILD_VERSION=v1.0.7 ARROW_BZIP2_BUILD_VERSION=1.0.8 ARROW_CARES_BUILD_VERSION=1.16.1 -ARROW_GBENCHMARK_BUILD_VERSION=v1.5.1 +ARROW_GBENCHMARK_BUILD_VERSION=v1.5.2 ARROW_GFLAGS_BUILD_VERSION=v2.2.2 ARROW_GLOG_BUILD_VERSION=v0.4.0 ARROW_GRPC_BUILD_VERSION=v1.29.1 From f45664e10de038426aa7eed19689135392fae383 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Thu, 19 Nov 2020 14:09:34 +0900 Subject: [PATCH 23/37] Add chunked array sorter and radix sort based table sorter Chunked array sorter implementation is naive. Radix sort based table sorter is slower than the existing table sorter. So this is disabled for now. Benchmark Time CPU ------------------------------------------------------------------------------------------ ArraySortIndicesInt64Narrow/32768/10000/min_time:1.000 16481 ns 16481 ns ArraySortIndicesInt64Narrow/32768/100/min_time:1.000 16782 ns 16782 ns ArraySortIndicesInt64Narrow/32768/10/min_time:1.000 16604 ns 16604 ns ArraySortIndicesInt64Narrow/32768/2/min_time:1.000 28438 ns 28438 ns ArraySortIndicesInt64Narrow/32768/1/min_time:1.000 6064 ns 6064 ns ArraySortIndicesInt64Narrow/32768/0/min_time:1.000 13949 ns 13949 ns ArraySortIndicesInt64Narrow/1048576/100/min_time:1.000 739131 ns 739116 ns ArraySortIndicesInt64Narrow/8388608/100/min_time:1.000 7270177 ns 7270007 ns ArraySortIndicesInt64Wide/32768/10000/min_time:1.000 197500 ns 197496 ns ArraySortIndicesInt64Wide/32768/100/min_time:1.000 198485 ns 198481 ns ArraySortIndicesInt64Wide/32768/10/min_time:1.000 184462 ns 184460 ns ArraySortIndicesInt64Wide/32768/2/min_time:1.000 114283 ns 114282 ns ArraySortIndicesInt64Wide/32768/1/min_time:1.000 6117 ns 6116 ns ArraySortIndicesInt64Wide/32768/0/min_time:1.000 194275 ns 194267 ns ArraySortIndicesInt64Wide/1048576/100/min_time:1.000 9688010 ns 9687762 ns ArraySortIndicesInt64Wide/8388608/100/min_time:1.000 98482097 ns 98480482 ns TableSortIndicesInt64Narrow/1048576/100/16/32/min_time:1.000 13674075901 ns 13673756372 TableSortIndicesInt64Narrow/1048576/0/16/32/min_time:1.000 13562068073 ns 13561526021 TableSortIndicesInt64Narrow/1048576/100/8/32/min_time:1.000 6816877843 ns 6816737232 ns TableSortIndicesInt64Narrow/1048576/0/8/32/min_time:1.000 6513599980 ns 6513470119 ns TableSortIndicesInt64Narrow/1048576/100/2/32/min_time:1.000 1297496469 ns 1297401147 ns TableSortIndicesInt64Narrow/1048576/0/2/32/min_time:1.000 1198835057 ns 1198795471 ns TableSortIndicesInt64Narrow/1048576/100/1/32/min_time:1.000 512438066 ns 512427916 ns TableSortIndicesInt64Narrow/1048576/0/1/32/min_time:1.000 406987001 ns 406983020 ns TableSortIndicesInt64Narrow/1048576/100/16/4/min_time:1.000 4782555568 ns 4782501462 ns TableSortIndicesInt64Narrow/1048576/0/16/4/min_time:1.000 4627874931 ns 4627776462 ns TableSortIndicesInt64Narrow/1048576/100/8/4/min_time:1.000 2347109491 ns 2347030723 ns TableSortIndicesInt64Narrow/1048576/0/8/4/min_time:1.000 2353793768 ns 2353762906 ns TableSortIndicesInt64Narrow/1048576/100/2/4/min_time:1.000 329895435 ns 329878554 ns TableSortIndicesInt64Narrow/1048576/0/2/4/min_time:1.000 318651531 ns 318643264 ns TableSortIndicesInt64Narrow/1048576/100/1/4/min_time:1.000 27022014 ns 27020904 ns TableSortIndicesInt64Narrow/1048576/0/1/4/min_time:1.000 21620126 ns 21619471 ns TableSortIndicesInt64Narrow/1048576/100/16/1/min_time:1.000 2082161247 ns 2082123210 ns TableSortIndicesInt64Narrow/1048576/0/16/1/min_time:1.000 2125744136 ns 2125688346 ns TableSortIndicesInt64Narrow/1048576/100/8/1/min_time:1.000 1054223608 ns 1054187702 ns TableSortIndicesInt64Narrow/1048576/0/8/1/min_time:1.000 1039334371 ns 1039300540 ns TableSortIndicesInt64Narrow/1048576/100/2/1/min_time:1.000 226713390 ns 226708185 ns TableSortIndicesInt64Narrow/1048576/0/2/1/min_time:1.000 232277655 ns 232265803 ns TableSortIndicesInt64Narrow/1048576/100/1/1/min_time:1.000 7317963 ns 7317641 ns TableSortIndicesInt64Narrow/1048576/0/1/1/min_time:1.000 6936323 ns 6935778 ns TableSortIndicesInt64Wide/1048576/100/16/32/min_time:1.000 15478483751 ns 15477778551 ns TableSortIndicesInt64Wide/1048576/0/16/32/min_time:1.000 15323045709 ns 15322716370 ns TableSortIndicesInt64Wide/1048576/100/8/32/min_time:1.000 7600326817 ns 7600060885 ns TableSortIndicesInt64Wide/1048576/0/8/32/min_time:1.000 7305830182 ns 7305653322 ns TableSortIndicesInt64Wide/1048576/100/2/32/min_time:1.000 1603439764 ns 1603419271 ns TableSortIndicesInt64Wide/1048576/0/2/32/min_time:1.000 1585765386 ns 1585720827 ns TableSortIndicesInt64Wide/1048576/100/1/32/min_time:1.000 848231188 ns 848213371 ns TableSortIndicesInt64Wide/1048576/0/1/32/min_time:1.000 599883956 ns 599827161 ns TableSortIndicesInt64Wide/1048576/100/16/4/min_time:1.000 6105580749 ns 6105405315 ns TableSortIndicesInt64Wide/1048576/0/16/4/min_time:1.000 6101922415 ns 6092946808 ns TableSortIndicesInt64Wide/1048576/100/8/4/min_time:1.000 3016793628 ns 3013968137 ns TableSortIndicesInt64Wide/1048576/0/8/4/min_time:1.000 2989364923 ns 2989345424 ns TableSortIndicesInt64Wide/1048576/100/2/4/min_time:1.000 595413434 ns 595398793 ns TableSortIndicesInt64Wide/1048576/0/2/4/min_time:1.000 593904971 ns 593872005 ns TableSortIndicesInt64Wide/1048576/100/1/4/min_time:1.000 126583175 ns 126578734 ns TableSortIndicesInt64Wide/1048576/0/1/4/min_time:1.000 115004382 ns 115003229 ns TableSortIndicesInt64Wide/1048576/100/16/1/min_time:1.000 3013364645 ns 3013319384 ns TableSortIndicesInt64Wide/1048576/0/16/1/min_time:1.000 3032564737 ns 3032430757 ns TableSortIndicesInt64Wide/1048576/100/8/1/min_time:1.000 1457310735 ns 1457286492 ns TableSortIndicesInt64Wide/1048576/0/8/1/min_time:1.000 1465821557 ns 1465809354 ns TableSortIndicesInt64Wide/1048576/100/2/1/min_time:1.000 335816200 ns 335812614 ns TableSortIndicesInt64Wide/1048576/0/2/1/min_time:1.000 330472127 ns 330467667 ns TableSortIndicesInt64Wide/1048576/100/1/1/min_time:1.000 103267823 ns 103262409 ns TableSortIndicesInt64Wide/1048576/0/1/1/min_time:1.000 102048116 ns 102044125 ns Diff: Range/N records/Inverse null probability/N sort keys/N chunks Narrow/1048576/100/16/32/min_time:1.000 FromTheBeginning: 1386231938 ns Radix: 13674075901 ns Narrow/1048576/0/16/32/min_time:1.000 FromTheBeginning: 1350031141 ns Radix: 13562068073 ns Narrow/1048576/100/8/32/min_time:1.000 FromTheBeginning: 1401741018 ns Radix: 6816877843 ns Narrow/1048576/0/8/32/min_time:1.000 FromTheBeginning: 1373174328 ns Radix: 6513599980 ns Narrow/1048576/100/2/32/min_time:1.000 FromTheBeginning: 1035377805 ns Radix: 1297496469 ns Narrow/1048576/0/2/32/min_time:1.000 FromTheBeginning: 1035218085 ns Radix: 1198835057 ns Narrow/1048576/100/1/32/min_time:1.000 FromTheBeginning: 6859319 ns Radix: 512438066 ns Narrow/1048576/0/1/32/min_time:1.000 FromTheBeginning: 6847314 ns Radix: 406987001 ns Narrow/1048576/100/16/4/min_time:1.000 FromTheBeginning: 626909516 ns Radix: 4782555568 ns Narrow/1048576/0/16/4/min_time:1.000 FromTheBeginning: 591632035 ns Radix: 4627874931 ns Narrow/1048576/100/8/4/min_time:1.000 FromTheBeginning: 613197155 ns Radix: 2347109491 ns Narrow/1048576/0/8/4/min_time:1.000 FromTheBeginning: 595568690 ns Radix: 2353793768 ns Narrow/1048576/100/2/4/min_time:1.000 FromTheBeginning: 424472984 ns Radix: 329895435 ns Narrow/1048576/0/2/4/min_time:1.000 FromTheBeginning: 397339109 ns Radix: 318651531 ns Narrow/1048576/100/1/4/min_time:1.000 FromTheBeginning: 7027310 ns Radix: 27022014 ns Narrow/1048576/0/1/4/min_time:1.000 FromTheBeginning: 6891364 ns Radix: 21620126 ns Narrow/1048576/100/16/1/min_time:1.000 FromTheBeginning: 516832749 ns Radix: 2082161247 ns Narrow/1048576/0/16/1/min_time:1.000 FromTheBeginning: 495273237 ns Radix: 2125744136 ns Narrow/1048576/100/8/1/min_time:1.000 FromTheBeginning: 515550360 ns Radix: 1054223608 ns Narrow/1048576/0/8/1/min_time:1.000 FromTheBeginning: 496670125 ns Radix: 1039334371 ns Narrow/1048576/100/2/1/min_time:1.000 FromTheBeginning: 340060863 ns Radix: 226713390 ns Narrow/1048576/0/2/1/min_time:1.000 FromTheBeginning: 331281499 ns Radix: 232277655 ns Narrow/1048576/100/1/1/min_time:1.000 FromTheBeginning: 6691062 ns Radix: 7317963 ns Narrow/1048576/0/1/1/min_time:1.000 FromTheBeginning: 6384382 ns Radix: 6936323 ns Wide/1048576/100/16/32/min_time:1.000 FromTheBeginning: 622544467 ns Radix: 15478483751 ns Wide/1048576/0/16/32/min_time:1.000 FromTheBeginning: 619193843 ns Radix: 15323045709 ns Wide/1048576/100/8/32/min_time:1.000 FromTheBeginning: 615885889 ns Radix: 7600326817 ns Wide/1048576/0/8/32/min_time:1.000 FromTheBeginning: 589917731 ns Radix: 7305830182 ns Wide/1048576/100/2/32/min_time:1.000 FromTheBeginning: 600951149 ns Radix: 1603439764 ns Wide/1048576/0/2/32/min_time:1.000 FromTheBeginning: 592304437 ns Radix: 1585765386 ns Wide/1048576/100/1/32/min_time:1.000 FromTheBeginning: 98781239 ns Radix: 848231188 ns Wide/1048576/0/1/32/min_time:1.000 FromTheBeginning: 94136230 ns Radix: 599883956 ns Wide/1048576/100/16/4/min_time:1.000 FromTheBeginning: 232323308 ns Radix: 6105580749 ns Wide/1048576/0/16/4/min_time:1.000 FromTheBeginning: 223708791 ns Radix: 6101922415 ns Wide/1048576/100/8/4/min_time:1.000 FromTheBeginning: 221975360 ns Radix: 3016793628 ns Wide/1048576/0/8/4/min_time:1.000 FromTheBeginning: 218214843 ns Radix: 2989364923 ns Wide/1048576/100/2/4/min_time:1.000 FromTheBeginning: 224756430 ns Radix: 595413434 ns Wide/1048576/0/2/4/min_time:1.000 FromTheBeginning: 220809969 ns Radix: 593904971 ns Wide/1048576/100/1/4/min_time:1.000 FromTheBeginning: 96159427 ns Radix: 126583175 ns Wide/1048576/0/1/4/min_time:1.000 FromTheBeginning: 92307749 ns Radix: 115004382 ns Wide/1048576/100/16/1/min_time:1.000 FromTheBeginning: 166390841 ns Radix: 3013364645 ns Wide/1048576/0/16/1/min_time:1.000 FromTheBeginning: 164576492 ns Radix: 3032564737 ns Wide/1048576/100/8/1/min_time:1.000 FromTheBeginning: 165724919 ns Radix: 1457310735 ns Wide/1048576/0/8/1/min_time:1.000 FromTheBeginning: 164048003 ns Radix: 1465821557 ns Wide/1048576/100/2/1/min_time:1.000 FromTheBeginning: 170131788 ns Radix: 335816200 ns Wide/1048576/0/2/1/min_time:1.000 FromTheBeginning: 170874129 ns Radix: 330472127 ns Wide/1048576/100/1/1/min_time:1.000 FromTheBeginning: 92314674 ns Radix: 103267823 ns Wide/1048576/0/1/1/min_time:1.000 FromTheBeginning: 92023019 ns Radix: 102048116 ns --- cpp/src/arrow/compute/api_vector.cc | 14 +- cpp/src/arrow/compute/api_vector.h | 34 +- cpp/src/arrow/compute/kernels/vector_sort.cc | 602 ++++++++++++++---- .../compute/kernels/vector_sort_benchmark.cc | 50 +- .../arrow/compute/kernels/vector_sort_test.cc | 112 +++- 5 files changed, 649 insertions(+), 163 deletions(-) diff --git a/cpp/src/arrow/compute/api_vector.cc b/cpp/src/arrow/compute/api_vector.cc index 9b9e03bae24..5eac72c9912 100644 --- a/cpp/src/arrow/compute/api_vector.cc +++ b/cpp/src/arrow/compute/api_vector.cc @@ -54,10 +54,18 @@ Result> SortIndices(const Array& values, SortOrder order, return result.make_array(); } -Result> SortIndices(const Table& values, - const SortOptions& options, ExecContext* ctx) { +Result> SortIndices(const ChunkedArray& chunked_array, + SortOrder order, ExecContext* ctx) { + SortOptions options({SortKey("not-used", order)}); + ARROW_ASSIGN_OR_RAISE( + Datum result, CallFunction("sort_indices", {Datum(chunked_array)}, &options, ctx)); + return result.make_array(); +} + +Result> SortIndices(const Table& table, const SortOptions& options, + ExecContext* ctx) { ARROW_ASSIGN_OR_RAISE(Datum result, - CallFunction("sort_indices", {Datum(values)}, &options, ctx)); + CallFunction("sort_indices", {Datum(table)}, &options, ctx)); return result.make_array(); } diff --git a/cpp/src/arrow/compute/api_vector.h b/cpp/src/arrow/compute/api_vector.h index 6cf3b3ab781..0ed3f4d36ff 100644 --- a/cpp/src/arrow/compute/api_vector.h +++ b/cpp/src/arrow/compute/api_vector.h @@ -185,18 +185,40 @@ Result> NthToIndices(const Array& values, int64_t n, /// /// Perform an indirect sort of array. The output array will contain /// indices that would sort an array, which would be the same length -/// as input. Nulls will be stably partitioned to the end of the output. +/// as input. Nulls will be stably partitioned to the end of the output +/// regardless of order. /// -/// For example given values = [null, 1, 3.3, null, 2, 5.3] and order +/// For example given array = [null, 1, 3.3, null, 2, 5.3] and order /// = SortOrder::DESCENDING, the output will be [5, 2, 4, 1, 0, /// 3]. /// -/// \param[in] values array to sort +/// \param[in] array array to sort +/// \param[in] order ascending or descending +/// \param[in] ctx the function execution context, optional +/// \return offsets indices that would sort an array +ARROW_EXPORT +Result> SortIndices(const Array& array, + SortOrder order = SortOrder::Ascending, + ExecContext* ctx = NULLPTR); + +/// \brief Returns the indices that would sort a chunked array in the +/// specified order. +/// +/// Perform an indirect sort of chunked array. The output array will +/// contain indices that would sort a chunked array, which would be +/// the same length as input. Nulls will be stably partitioned to the +/// end of the output regardless of order. +/// +/// For example given chunked_array = [[null, 1], [3.3], [null, 2, +/// 5.3]] and order = SortOrder::DESCENDING, the output will be [5, 2, +/// 4, 1, 0, 3]. +/// +/// \param[in] chunked_array chunked array to sort /// \param[in] order ascending or descending /// \param[in] ctx the function execution context, optional /// \return offsets indices that would sort an array ARROW_EXPORT -Result> SortIndices(const Array& values, +Result> SortIndices(const ChunkedArray& chunked_array, SortOrder order = SortOrder::Ascending, ExecContext* ctx = NULLPTR); @@ -209,8 +231,8 @@ Result> SortIndices(const Array& values, /// regardless of order. /// /// For example given table = { -/// "column1": [null, 1, 3, null, 2, 1], -/// "column2": [ 5, 3, null, null, 5, 5], +/// "column1": [[null, 1], [ 3, null, 2, 1]], +/// "column2": [[ 5], [3, null, null, 5, 5]], /// } and options = { /// {"column1", SortOrder::Ascending}, /// {"column2", SortOrder::Descending}, diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc index afcd3dd964c..c5a33a85b51 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort.cc @@ -21,7 +21,6 @@ #include #include -#include "arrow/array/concatenate.h" #include "arrow/array/data.h" #include "arrow/compute/api_vector.h" #include "arrow/compute/kernels/common.h" @@ -34,6 +33,65 @@ namespace internal { namespace { +// The target chunk in chunked array. +template +struct ResolvedChunk; + +// Not for string-like data. +template +struct ResolvedChunk< + ArrayType, enable_if_t::value>> { + using c_type = typename ArrayType::TypeClass::c_type; + + // The target array in chunked array. + const ArrayType* array; + + // The index in the target array. + int64_t index; + + ResolvedChunk(const ArrayType* array, int64_t index) : array(array), index(index) {} + + bool IsNull() const { return array->IsNull(index); } + + c_type GetView() const { return array->GetView(index); } +}; + +// For string-like data. +template +struct ResolvedChunk< + ArrayType, enable_if_t::value>> { + // The target array in chunked array. + const ArrayType* array; + + // The index in the target array. + int64_t index; + + ResolvedChunk(const ArrayType* array, int64_t index) : array(array), index(index) {} + + bool IsNull() const { return array->IsNull(index); } + + util::string_view GetView() const { return array->GetView(index); } +}; + +// Finds the target chunk and index in the target chunk from an index +// in chunked array. `chunks` is not shared array of +// ChunkedArray::chunks() for performance. +template +ResolvedChunk ResolveChunk(const std::vector& chunks, + int64_t index) { + const auto num_chunks = chunks.size(); + int64_t offset = 0; + for (size_t i = 0; i < num_chunks; ++i) { + if (index < offset + chunks[i]->length()) { + return ResolvedChunk(static_cast(chunks[i]), + index - offset); + } + offset += chunks[i]->length(); + } + // Never reach here. `index` must be validated in caller. + return ResolvedChunk(nullptr, 0); +} + // NOTE: std::partition is usually faster than std::stable_partition. struct NonStablePartitioner { @@ -51,37 +109,87 @@ struct StablePartitioner { } }; -// Move Nulls to end of array. Return where Null starts. +// Move nulls to end of array. Return where null starts. +// +// `offset` is for using this to an array in a chunked array. template enable_if_t::value, uint64_t*> -PartitionNulls(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values) { +PartitionNulls(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, + int64_t offset) { if (values.null_count() == 0) { return indices_end; } Partitioner partitioner; - return partitioner(indices_begin, indices_end, - [&values](uint64_t ind) { return !values.IsNull(ind); }); + return partitioner(indices_begin, indices_end, [&values, &offset](uint64_t ind) { + return !values.IsNull(ind - offset); + }); } -// Move NaNs and Nulls to end of array, Nulls after NaN. -// Return where NaN/Null starts. +// For chunked array. +template +enable_if_t::value, uint64_t*> +PartitionNulls(uint64_t* indices_begin, uint64_t* indices_end, + std::vector& arrays, int64_t null_count) { + if (null_count == 0) { + return indices_end; + } + Partitioner partitioner; + return partitioner(indices_begin, indices_end, [&arrays](uint64_t ind) { + const auto chunk = ResolveChunk(arrays, ind); + return !chunk.IsNull(); + }); +} + +// Move NaNs and nulls to end of array, nulls after NaN. Return where +// NaN/null starts. +// +// `offset` is for using this to an array in a chunked array. template enable_if_t::value, uint64_t*> -PartitionNulls(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values) { +PartitionNulls(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, + int64_t offset) { Partitioner partitioner; if (values.null_count() == 0) { - return partitioner(indices_begin, indices_end, [&values](uint64_t ind) { - return !std::isnan(values.GetView(ind)); + return partitioner(indices_begin, indices_end, [&values, &offset](uint64_t ind) { + return !std::isnan(values.GetView(ind - offset)); }); } uint64_t* nulls_begin = - partitioner(indices_begin, indices_end, [&values](uint64_t ind) { - return !values.IsNull(ind) && !std::isnan(values.GetView(ind)); + partitioner(indices_begin, indices_end, [&values, &offset](uint64_t ind) { + return !values.IsNull(ind - offset) && !std::isnan(values.GetView(ind - offset)); }); - // move Nulls after NaN + // move nulls after NaN if (values.null_count() < static_cast(indices_end - nulls_begin)) { - partitioner(nulls_begin, indices_end, - [&values](uint64_t ind) { return !values.IsNull(ind); }); + partitioner(nulls_begin, indices_end, [&values, &offset](uint64_t ind) { + return !values.IsNull(ind - offset); + }); + } + return nulls_begin; +} + +// For chunked array. +template +enable_if_t::value, uint64_t*> +PartitionNulls(uint64_t* indices_begin, uint64_t* indices_end, + std::vector& arrays, int64_t null_count) { + Partitioner partitioner; + if (null_count == 0) { + return partitioner(indices_begin, indices_end, [&arrays](uint64_t ind) { + const auto chunk = ResolveChunk(arrays, ind); + return !std::isnan(chunk.GetView()); + }); + } + uint64_t* nulls_begin = + partitioner(indices_begin, indices_end, [&arrays](uint64_t ind) { + const auto chunk = ResolveChunk(arrays, ind); + return !chunk.IsNull() && !std::isnan(chunk.GetView()); + }); + // move nulls after NaN + if (null_count < static_cast(indices_end - nulls_begin)) { + partitioner(nulls_begin, indices_end, [&arrays](uint64_t ind) { + const auto chunk = ResolveChunk(arrays, ind); + return !chunk.IsNull(); + }); } return nulls_begin; } @@ -132,7 +240,7 @@ struct PartitionNthToIndices { return; } auto nulls_begin = - PartitionNulls(out_begin, out_end, arr); + PartitionNulls(out_begin, out_end, arr, 0); auto nth_begin = out_begin + pivot; if (nth_begin < nulls_begin) { std::nth_element(out_begin, nth_begin, nulls_begin, @@ -170,22 +278,27 @@ class ArrayCompareSorter { using ArrayType = typename TypeTraits::ArrayType; public: - void Sort(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, - const ArraySortOptions& options) { - std::iota(indices_begin, indices_end, 0); - auto nulls_begin = - PartitionNulls(indices_begin, indices_end, values); + // Returns where null starts. + // + // `offset` is for using this to an array in a chunked array. + uint64_t* Sort(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, + int64_t offset, const ArraySortOptions& options) { + auto nulls_begin = PartitionNulls( + indices_begin, indices_end, values, offset); if (options.order == SortOrder::Ascending) { - std::stable_sort(indices_begin, nulls_begin, - [&values](uint64_t left, uint64_t right) { - return values.GetView(left) < values.GetView(right); - }); + std::stable_sort( + indices_begin, nulls_begin, [&values, &offset](uint64_t left, uint64_t right) { + return values.GetView(left - offset) < values.GetView(right - offset); + }); } else { - std::stable_sort(indices_begin, nulls_begin, - [&values](uint64_t left, uint64_t right) { - return values.GetView(right) < values.GetView(left); - }); + std::stable_sort( + indices_begin, nulls_begin, [&values, &offset](uint64_t left, uint64_t right) { + // We don't use 'left > right' here to reduce required operator. + // If we use 'right < left' here, '<' is only required. + return values.GetView(right - offset) < values.GetView(left - offset); + }); } + return nulls_begin; } }; @@ -205,13 +318,14 @@ class ArrayCountSorter { value_range_ = static_cast(max - min) + 1; } - void Sort(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, - const ArraySortOptions& options) { + // Returns where null starts. + uint64_t* Sort(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, + int64_t offset, const ArraySortOptions& options) { // 32bit counter performs much better than 64bit one if (values.length() < (1LL << 32)) { - SortInternal(indices_begin, indices_end, values, options); + return SortInternal(indices_begin, indices_end, values, offset, options); } else { - SortInternal(indices_begin, indices_end, values, options); + return SortInternal(indices_begin, indices_end, values, offset, options); } } @@ -219,9 +333,13 @@ class ArrayCountSorter { c_type min_{0}; uint32_t value_range_{0}; + // Returns where null starts. + // + // `offset` is for using this to an array in a chunked array. template - void SortInternal(uint64_t* indices_begin, uint64_t* indices_end, - const ArrayType& values, const ArraySortOptions& options) { + uint64_t* SortInternal(uint64_t* indices_begin, uint64_t* indices_end, + const ArrayType& values, int64_t offset, + const ArraySortOptions& options) { const uint32_t value_range = value_range_; // first slot reserved for prefix sum @@ -234,10 +352,12 @@ class ArrayCountSorter { counts[i] += counts[i - 1]; } auto null_position = counts[value_range]; - int64_t index = 0; + auto nulls_begin = indices_begin + null_position; + int64_t index = offset; VisitRawValuesInline( values, [&](c_type v) { indices_begin[counts[v - min_]++] = index++; }, [&]() { indices_begin[null_position++] = index++; }); + return nulls_begin; } else { VisitRawValuesInline( values, [&](c_type v) { ++counts[v - min_]; }, []() {}); @@ -245,10 +365,12 @@ class ArrayCountSorter { counts[i - 1] += counts[i]; } auto null_position = counts[0]; - int64_t index = 0; + auto nulls_begin = indices_begin + null_position; + int64_t index = offset; VisitRawValuesInline( values, [&](c_type v) { indices_begin[counts[v - min_ + 1]++] = index++; }, [&]() { indices_begin[null_position++] = index++; }); + return nulls_begin; } } }; @@ -262,8 +384,11 @@ class ArrayCountOrCompareSorter { using c_type = typename ArrowType::c_type; public: - void Sort(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, - const ArraySortOptions& options) { + // Returns where null starts. + // + // `offset` is for using this to an array in a chunked array. + uint64_t* Sort(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, + int64_t offset, const ArraySortOptions& options) { if (values.length() >= countsort_min_len_ && values.length() > values.null_count()) { c_type min{std::numeric_limits::max()}; c_type max{std::numeric_limits::min()}; @@ -281,12 +406,11 @@ class ArrayCountOrCompareSorter { if (static_cast(max) - static_cast(min) <= countsort_max_range_) { count_sorter_.SetMinMax(min, max); - count_sorter_.Sort(indices_begin, indices_end, values, options); - return; + return count_sorter_.Sort(indices_begin, indices_end, values, offset, options); } } - compare_sorter_.Sort(indices_begin, indices_end, values, options); + return compare_sorter_.Sort(indices_begin, indices_end, values, offset, options); } private: @@ -349,9 +473,10 @@ struct ArraySortIndices { ArrayData* out_arr = out->mutable_array(); uint64_t* out_begin = out_arr->GetMutableValues(1); uint64_t* out_end = out_begin + arr.length(); + std::iota(out_begin, out_end, 0); ArraySorter sorter; - sorter.impl.Sort(out_begin, out_end, arr, options); + sorter.impl.Sort(out_begin, out_end, arr, 0, options); } }; @@ -374,10 +499,199 @@ void AddSortingKernels(VectorKernel base, VectorFunction* func) { } } -class TableSorter : public TypeVisitor { +// Sort a chunked array directly without sorting each array in the +// chunked array. This is used for processing the second and following +// sort keys in TableRadixSorter. +// +// This uses the same algorithm as ArrayCompareSorter. +template +class ChunkedArrayCompareSorter { + using ArrayType = typename TypeTraits::ArrayType; + + public: + // Returns where null starts. + uint64_t* Sort(uint64_t* indices_begin, uint64_t* indices_end, + std::vector& arrays, int64_t null_count, + const ArraySortOptions& options) { + auto nulls_begin = PartitionNulls( + indices_begin, indices_end, arrays, null_count); + if (options.order == SortOrder::Ascending) { + std::stable_sort(indices_begin, nulls_begin, + [&arrays](uint64_t left, uint64_t right) { + const auto chunk_left = ResolveChunk(arrays, left); + const auto chunk_right = ResolveChunk(arrays, right); + return chunk_left.GetView() < chunk_right.GetView(); + }); + } else { + std::stable_sort(indices_begin, nulls_begin, + [&arrays](uint64_t left, uint64_t right) { + const auto chunk_left = ResolveChunk(arrays, left); + const auto chunk_right = ResolveChunk(arrays, right); + // We don't use 'left > right' here to reduce required operator. + // If we use 'right < left' here, '<' is only required. + return chunk_right.GetView() < chunk_left.GetView(); + }); + } + return nulls_begin; + } +}; + +// Sort a chunked array by sorting each array in the chunked array. +// +// TODO: This is a naive implementation. We'll be able to improve +// performance of this. For example, we'll be able to use threads for +// sorting each array. +class ChunkedArraySorter : public TypeVisitor { + public: + ChunkedArraySorter(uint64_t* indices_begin, uint64_t* indices_end, + const ChunkedArray& chunked_array, const SortOrder order, + bool can_use_array_sorter = true) + : TypeVisitor(), + indices_begin_(indices_begin), + indices_end_(indices_end), + chunked_array_(chunked_array), + order_(order), + can_use_array_sorter_(can_use_array_sorter) {} + + Status Sort() { return chunked_array_.type()->Accept(this); } + +#define VISIT(TYPE) \ + Status Visit(const TYPE##Type& type) override { return SortInternal(); } + + VISIT(Int8) + VISIT(Int16) + VISIT(Int32) + VISIT(Int64) + VISIT(UInt8) + VISIT(UInt16) + VISIT(UInt32) + VISIT(UInt64) + VISIT(Float) + VISIT(Double) + VISIT(String) + VISIT(Binary) + VISIT(LargeString) + VISIT(LargeBinary) + +#undef VISIT + private: - // Finds the target chunk and index in the target chunk from an - // index in ChunkedArray. + template + Status SortInternal() { + using ArrayType = typename TypeTraits::ArrayType; + ArraySortOptions options(order_); + const auto num_chunks = chunked_array_.num_chunks(); + const auto& shared_arrays = chunked_array_.chunks(); + std::vector arrays(num_chunks); + for (int i = 0; i < num_chunks; ++i) { + const auto& array = shared_arrays[i]; + arrays[i] = array.get(); + } + if (can_use_array_sorter_) { + // Sort each chunk from the beginning and merge to sorted indices. + // This is a naive implementation. + ArraySorter sorter; + int64_t begin_offset = 0; + int64_t end_offset = 0; + int64_t null_count = 0; + uint64_t* left_nulls_begin = indices_begin_; + for (int i = 0; i < num_chunks; ++i) { + const auto array = static_cast(arrays[i]); + end_offset += array->length(); + null_count += array->null_count(); + uint64_t* right_nulls_begin; + right_nulls_begin = + sorter.impl.Sort(indices_begin_ + begin_offset, indices_begin_ + end_offset, + *array, begin_offset, options); + if (i > 0) { + left_nulls_begin = Merge( + indices_begin_, indices_begin_ + begin_offset, indices_begin_ + end_offset, + left_nulls_begin, right_nulls_begin, arrays, null_count, order_); + } else { + left_nulls_begin = right_nulls_begin; + } + begin_offset = end_offset; + } + } else { + // Sort the chunked array directory. + ChunkedArrayCompareSorter sorter; + sorter.Sort(indices_begin_, indices_end_, arrays, chunked_array_.null_count(), + options); + } + return Status::OK(); + } + + // Merges two sorted indices arrays and returns where nulls starts. + // Where nulls starts is used when the next merge to detect the + // sorted indices locations. + template + uint64_t* Merge(uint64_t* indices_begin, uint64_t* indices_middle, + uint64_t* indices_end, uint64_t* left_nulls_begin, + uint64_t* right_nulls_begin, std::vector arrays, + int64_t null_count, const SortOrder order) { + auto left_num_non_nulls = left_nulls_begin - indices_begin; + auto right_num_non_nulls = right_nulls_begin - indices_middle; + auto nulls_begin = PartitionNulls( + indices_begin, indices_end, arrays, null_count); + indices_middle = indices_begin + left_num_non_nulls; + indices_end = indices_middle + right_num_non_nulls; + if (order == SortOrder::Ascending) { + std::inplace_merge(indices_begin, indices_middle, indices_end, + [&arrays](uint64_t left, uint64_t right) { + const auto chunk_left = ResolveChunk(arrays, left); + const auto chunk_right = + ResolveChunk(arrays, right); + return chunk_left.GetView() < chunk_right.GetView(); + }); + } else { + std::inplace_merge(indices_begin, indices_middle, indices_end, + [&arrays](uint64_t left, uint64_t right) { + const auto chunk_left = ResolveChunk(arrays, left); + const auto chunk_right = + ResolveChunk(arrays, right); + // We don't use 'left > right' here to reduce required + // operator. If we use 'right < left' here, '<' is only + // required. + return chunk_right.GetView() < chunk_left.GetView(); + }); + } + return nulls_begin; + } + + uint64_t* indices_begin_; + uint64_t* indices_end_; + const ChunkedArray& chunked_array_; + const SortOrder order_; + const bool can_use_array_sorter_; +}; + +// Sort a table from the back like radix sort. +class TableRadixSorter { + public: + Status Sort(uint64_t* indices_begin, uint64_t* indices_end, const Table& table, + const SortOptions& options) { + for (auto i = options.sort_keys.size(); i > 0; --i) { + const auto& sort_key = options.sort_keys[i - 1]; + const auto& chunked_array = table.GetColumnByName(sort_key.name); + if (!chunked_array) { + return Status::Invalid("Nonexistent sort key column: ", sort_key.name); + } + // We can use ArraySorter only for the sort key that is + // processed first because ArraySorter doesn't care about + // existing indices. + const auto can_use_array_sorter = (i == 0); + ChunkedArraySorter sorter(indices_begin, indices_end, *chunked_array.get(), + sort_key.order, can_use_array_sorter); + ARROW_RETURN_NOT_OK(sorter.Sort()); + } + return Status::OK(); + } +}; + +// Sort a table from the beginning. +class TableFromTheBeginningSorter : public TypeVisitor { + private: + // Preprocessed sort key. struct ResolvedSortKey { ResolvedSortKey(const ChunkedArray& chunked_array, const SortOrder order) : order(order) { @@ -390,31 +704,20 @@ class TableSorter : public TypeVisitor { } // Finds the target chunk and index in the target chunk from an - // index in ChunkedArray. + // index in chunked array. template - std::pair ResolveChunk(int64_t index) const { - using Chunk = std::pair; - if (num_chunks == 1) { - return Chunk(static_cast(chunks[0]), index); - } else { - int64_t offset = 0; - for (int i = 0; i < num_chunks; ++i) { - if (index < offset + chunks[i]->length()) { - return Chunk(static_cast(chunks[i]), index - offset); - } - offset += chunks[i]->length(); - } - return Chunk(nullptr, 0); - } + ResolvedChunk GetChunk(int64_t index) const { + return ResolveChunk(chunks, index); } SortOrder order; DataType* type; int64_t null_count; int num_chunks; - std::vector chunks; + std::vector chunks; }; + // Compare two records in the same table. class Comparer : public TypeVisitor { public: Comparer(const Table& table, const std::vector& sort_keys) @@ -444,6 +747,8 @@ class TableSorter : public TypeVisitor { for (size_t i = start_sort_key_index; i < num_sort_keys; ++i) { current_sort_key_index_ = i; status_ = sort_keys_[i].type->Accept(this); + // If the left value equals to the right value, we need to + // continue to sort. if (current_compared_ != 0) { break; } @@ -475,20 +780,24 @@ class TableSorter : public TypeVisitor { #undef VISIT private: + // Compares two records in the same table and returns -1, 0 or 1. + // + // -1: The left is less than the right. + // 0: The left equals to the right. + // 1: The left is greater than the right. + // + // This supports null and NaN. Null is processed in this and NaN + // is processed in CompareTypeValue(). template int32_t CompareType() { using ArrayType = typename TypeTraits::ArrayType; const auto& sort_key = sort_keys_[current_sort_key_index_]; auto order = sort_key.order; - auto chunk_left = sort_key.ResolveChunk(current_left_); - auto array_left = chunk_left.first; - auto index_left = chunk_left.second; - auto chunk_right = sort_key.ResolveChunk(current_right_); - auto array_right = chunk_right.first; - auto index_right = chunk_right.second; + const auto chunk_left = sort_key.GetChunk(current_left_); + const auto chunk_right = sort_key.GetChunk(current_right_); if (sort_key.null_count > 0) { - auto is_null_left = array_left->IsNull(index_left); - auto is_null_right = array_right->IsNull(index_right); + auto is_null_left = chunk_left.IsNull(); + auto is_null_right = chunk_right.IsNull(); if (is_null_left && is_null_right) { return 0; } else if (is_null_left) { @@ -497,8 +806,48 @@ class TableSorter : public TypeVisitor { return -1; } } - auto left = array_left->GetView(index_left); - auto right = array_right->GetView(index_right); + return CompareTypeValue(chunk_left, chunk_right, order); + } + + // For non-float types. Value is never NaN. + template + enable_if_t::value, int32_t> CompareTypeValue( + const ResolvedChunk::ArrayType>& chunk_left, + const ResolvedChunk::ArrayType>& chunk_right, + const SortOrder order) { + const auto left = chunk_left.GetView(); + const auto right = chunk_right.GetView(); + int32_t compared; + if (left == right) { + compared = 0; + } else if (left > right) { + compared = 1; + } else { + compared = -1; + } + if (order == SortOrder::Descending) { + compared = -compared; + } + return compared; + } + + // For float types. Value may be NaN. + template + enable_if_t::value, int32_t> CompareTypeValue( + const ResolvedChunk::ArrayType>& chunk_left, + const ResolvedChunk::ArrayType>& chunk_right, + const SortOrder order) { + const auto left = chunk_left.GetView(); + const auto right = chunk_right.GetView(); + auto is_nan_left = std::isnan(left); + auto is_nan_right = std::isnan(right); + if (is_nan_left && is_nan_right) { + return 0; + } else if (is_nan_left) { + return 1; + } else if (is_nan_right) { + return -1; + } int32_t compared; if (left == right) { compared = 0; @@ -522,12 +871,15 @@ class TableSorter : public TypeVisitor { }; public: - TableSorter(uint64_t* indices_begin, uint64_t* indices_end, const Table& table, - const SortOptions& options) + TableFromTheBeginningSorter(uint64_t* indices_begin, uint64_t* indices_end, + const Table& table, const SortOptions& options) : indices_begin_(indices_begin), indices_end_(indices_end), comparer_(table, options.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 Comparer. Status Sort() { ARROW_RETURN_NOT_OK(comparer_.status()); return comparer_.sort_keys()[0].type->Accept(this); @@ -557,7 +909,6 @@ class TableSorter : public TypeVisitor { template Status SortInternal() { using ArrayType = typename TypeTraits::ArrayType; - std::iota(indices_begin_, indices_end_, 0); auto& comparer = comparer_; const auto& first_sort_key = comparer.sort_keys()[0]; @@ -565,15 +916,15 @@ class TableSorter : public TypeVisitor { nulls_begin = PartitionNullsInternal(first_sort_key); std::stable_sort(indices_begin_, nulls_begin, [&first_sort_key, &comparer](uint64_t left, uint64_t right) { - auto chunk_left = first_sort_key.ResolveChunk(left); - auto array_left = chunk_left.first; - auto index_left = chunk_left.second; - auto chunk_right = first_sort_key.ResolveChunk(right); - auto array_right = chunk_right.first; - auto index_right = chunk_right.second; - auto value_left = array_left->GetView(index_left); - auto value_right = array_right->GetView(index_right); + // Both values are never null nor NaN. + auto chunk_left = first_sort_key.GetChunk(left); + auto chunk_right = first_sort_key.GetChunk(right); + auto value_left = chunk_left.GetView(); + auto value_right = chunk_right.GetView(); if (value_left == value_right) { + // If the left value equals to the right value, + // we need to compare the second and following + // sort keys. return comparer.Compare(left, right, 1); } else { auto compared = value_left < value_right; @@ -587,6 +938,9 @@ class TableSorter : public TypeVisitor { return Status::OK(); } + // Behaves like PatitionNulls() but this supports multiple sort keys. + // + // For non-float types. template enable_if_t::value, uint64_t*> PartitionNullsInternal( const ResolvedSortKey& first_sort_key) { @@ -597,10 +951,8 @@ class TableSorter : public TypeVisitor { StablePartitioner partitioner; auto nulls_begin = partitioner(indices_begin_, indices_end_, [&first_sort_key](uint64_t index) { - auto chunk = first_sort_key.ResolveChunk(index); - auto array = chunk.first; - auto index_array = chunk.second; - return !array->IsNull(index_array); + const auto chunk = first_sort_key.GetChunk(index); + return !chunk.IsNull(); }); auto& comparer = comparer_; std::stable_sort(nulls_begin, indices_end_, @@ -610,6 +962,9 @@ class TableSorter : public TypeVisitor { return nulls_begin; } + // Behaves like PatitionNulls() but this supports multiple sort keys. + // + // For float types. template enable_if_t::value, uint64_t*> PartitionNullsInternal( const ResolvedSortKey& first_sort_key) { @@ -617,37 +972,33 @@ class TableSorter : public TypeVisitor { StablePartitioner partitioner; if (first_sort_key.null_count == 0) { return partitioner(indices_begin_, indices_end_, [&first_sort_key](uint64_t index) { - auto chunk = first_sort_key.ResolveChunk(index); - auto array = chunk.first; - auto index_array = chunk.second; - return !std::isnan(array->GetView(index_array)); + const auto chunk = first_sort_key.GetChunk(index); + return !std::isnan(chunk.GetView()); }); } auto nans_and_nulls_begin = partitioner(indices_begin_, indices_end_, [&first_sort_key](uint64_t index) { - auto chunk = first_sort_key.ResolveChunk(index); - auto array = chunk.first; - auto index_array = chunk.second; - return !array->IsNull(index_array) && !std::isnan(array->GetView(index_array)); + const auto chunk = first_sort_key.GetChunk(index); + return !chunk.IsNull() && !std::isnan(chunk.GetView()); }); auto nulls_begin = nans_and_nulls_begin; if (first_sort_key.null_count < static_cast(indices_end_ - nulls_begin)) { - // move Nulls after NaN + // move nulls after NaN nulls_begin = partitioner( nans_and_nulls_begin, indices_end_, [&first_sort_key](uint64_t index) { - auto chunk = first_sort_key.ResolveChunk(index); - auto array = chunk.first; - auto index_array = chunk.second; - return !array->IsNull(index_array); + const auto chunk = first_sort_key.GetChunk(index); + return !chunk.IsNull(); }); } auto& comparer = comparer_; if (nans_and_nulls_begin != nulls_begin) { + // Sort all NaNs by the second and following sort keys. std::stable_sort(nans_and_nulls_begin, nulls_begin, [&comparer](uint64_t left, uint64_t right) { return comparer.Compare(left, right, 1); }); } + // Sort all nulls by the second and following sort keys. std::stable_sort(nulls_begin, indices_end_, [&comparer](uint64_t left, uint64_t right) { return comparer.Compare(left, right, 1); @@ -703,38 +1054,38 @@ class SortIndicesMetaFunction : public MetaFunction { } private: - Result> SortIndices(const Array& values, - const SortOptions& options, - ExecContext* ctx) const { + Result SortIndices(const Array& values, const SortOptions& options, + ExecContext* ctx) const { SortOrder order = SortOrder::Ascending; if (!options.sort_keys.empty()) { order = options.sort_keys[0].order; } ArraySortOptions array_options(order); - ARROW_ASSIGN_OR_RAISE( - Datum result, CallFunction("array_sort_indices", {values}, &array_options, ctx)); - return result.make_array(); + return CallFunction("array_sort_indices", {values}, &array_options, ctx); } - Result> SortIndices(const ChunkedArray& values, - const SortOptions& options, - ExecContext* ctx) const { + Result SortIndices(const ChunkedArray& chunked_array, const SortOptions& options, + ExecContext* ctx) const { SortOrder order = SortOrder::Ascending; if (!options.sort_keys.empty()) { order = options.sort_keys[0].order; } - ArraySortOptions array_options(order); - std::shared_ptr array_values; - if (values.num_chunks() == 1) { - array_values = values.chunk(0); - } else { - ARROW_ASSIGN_OR_RAISE(array_values, - Concatenate(values.chunks(), ctx->memory_pool())); - } - ARROW_ASSIGN_OR_RAISE(Datum result, CallFunction("array_sort_indices", {array_values}, - &array_options, ctx)); - return result.make_array(); + auto out_type = uint64(); + auto length = chunked_array.length(); + auto buffer_size = BitUtil::BytesForBits( + length * std::static_pointer_cast(out_type)->bit_width()); + std::vector> buffers(2); + ARROW_ASSIGN_OR_RAISE(buffers[1], + AllocateResizableBuffer(buffer_size, ctx->memory_pool())); + auto out = std::make_shared(out_type, length, buffers, 0); + auto out_begin = out->GetMutableValues(1); + auto out_end = out_begin + length; + std::iota(out_begin, out_end, 0); + + ChunkedArraySorter sorter(out_begin, out_end, chunked_array, order); + ARROW_RETURN_NOT_OK(sorter.Sort()); + return Datum(out); } Result SortIndices(const Table& table, const SortOptions& options, @@ -762,8 +1113,17 @@ class SortIndicesMetaFunction : public MetaFunction { auto out = std::make_shared(out_type, length, buffers, 0); auto out_begin = out->GetMutableValues(1); auto out_end = out_begin + length; + std::iota(out_begin, out_end, 0); - TableSorter sorter(out_begin, out_end, table, options); + // TODO: We should choose suitable sort implementation + // automatically. The current TableRadixSorter implementation is + // faster than TableFromTheBeginningSorter only when the number of + // sort keys is 2 and counting sort is used. So we always + // TableFromTheBeginningSorter for now. + // + // TableRadixSorter sorter; + // ARROW_RETURN_NOT_OK(sorter.Sort(out_begin, out_end, table, options)); + TableFromTheBeginningSorter sorter(out_begin, out_end, table, options); ARROW_RETURN_NOT_OK(sorter.Sort()); return Datum(out); } diff --git a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc index 235c445678d..9b997f7a926 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc @@ -84,16 +84,13 @@ struct TableSortIndicesArgs { // Extract args explicit TableSortIndicesArgs(benchmark::State& state) - : num_records(state.range(0)), - null_proportion(ComputeNullProportion(state.range(1))), - num_columns(state.range(2)), - num_chunks(state.range(3)), - state_(state) { - } + : num_records(state.range(0)), + null_proportion(ComputeNullProportion(state.range(1))), + num_columns(state.range(2)), + num_chunks(state.range(3)), + state_(state) {} - ~TableSortIndicesArgs() { - state_.SetItemsProcessed(state_.iterations() * num_records); - } + ~TableSortIndicesArgs() { state_.SetItemsProcessed(state_.iterations() * num_records); } private: double ComputeNullProportion(int64_t inverse_null_proportion) { @@ -121,15 +118,15 @@ static void TableSortIndicesInt64(benchmark::State& state, int64_t min, int64_t sort_keys.emplace_back(name, order); std::vector> arrays; if ((args.num_records % args.num_chunks) != 0) { - Status::Invalid("The number of chunks (", args.num_chunks, ") must be " - "a multiple of the number of records (", args.num_records, ")").Abort(); + Status::Invalid("The number of chunks (", args.num_chunks, + ") must be " + "a multiple of the number of records (", + args.num_records, ")") + .Abort(); } auto num_records_in_array = args.num_records / args.num_chunks; for (int64_t j = 0; j < args.num_chunks; ++j) { - arrays.push_back(rand.Int64(num_records_in_array, - min, - max, - args.null_proportion)); + arrays.push_back(rand.Int64(num_records_in_array, min, max, args.null_proportion)); } ASSIGN_OR_ABORT(auto chunked_array, ChunkedArray::Make(arrays, int64())); columns.push_back(chunked_array); @@ -145,8 +142,7 @@ static void TableSortIndicesInt64Narrow(benchmark::State& state) { } static void TableSortIndicesInt64Wide(benchmark::State& state) { - TableSortIndicesInt64(state, - std::numeric_limits::min(), + TableSortIndicesInt64(state, std::numeric_limits::min(), std::numeric_limits::max()); } @@ -166,21 +162,21 @@ BENCHMARK(ArraySortIndicesInt64Wide) BENCHMARK(TableSortIndicesInt64Narrow) ->ArgsProduct({ - {1 << 20}, // the number of records - {100, 0}, // inverse null proportion - {16, 8, 2, 1}, // the number of columns - {32, 4, 1}, // the number of chunks - }) + {1 << 20}, // the number of records + {100, 0}, // inverse null proportion + {16, 8, 2, 1}, // the number of columns + {32, 4, 1}, // the number of chunks + }) ->MinTime(1.0) ->Unit(benchmark::TimeUnit::kNanosecond); BENCHMARK(TableSortIndicesInt64Wide) ->ArgsProduct({ - {1 << 20}, // the number of records - {100, 0}, // inverse null proportion - {16, 8, 2, 1}, // the number of columns - {32, 4, 1}, // the number of chunks - }) + {1 << 20}, // the number of records + {100, 0}, // inverse null proportion + {16, 8, 2, 1}, // the number of columns + {32, 4, 1}, // the number of chunks + }) ->MinTime(1.0) ->Unit(benchmark::TimeUnit::kNanosecond); diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index aa2267f29ed..a0ef63c1508 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -20,8 +20,8 @@ #include #include +#include "arrow/array/concatenate.h" #include "arrow/compute/api_vector.h" -#include "arrow/compute/kernels/test_util.h" #include "arrow/table.h" #include "arrow/testing/gtest_common.h" #include "arrow/testing/gtest_util.h" @@ -372,8 +372,6 @@ TYPED_TEST(TestArraySortIndicesKernelRandom, SortRandomValues) { Random rand(0x5487655); int times = 5; int length = 1000; - times = 1; - length = 10; for (int test = 0; test < times; test++) { for (auto null_probability : {0.0, 0.1, 0.5, 1.0}) { for (auto order : {SortOrder::Ascending, SortOrder::Descending}) { @@ -431,6 +429,108 @@ TYPED_TEST(TestArraySortIndicesKernelRandomCompare, SortRandomValuesCompare) { } } +class TestChunkedArraySortIndices : public ::testing::Test { + protected: + void AssertSortIndices(const std::shared_ptr chunked_array, + SortOrder order, const std::shared_ptr expected) { + ASSERT_OK_AND_ASSIGN(auto actual, SortIndices(*chunked_array, order)); + AssertArraysEqual(*expected, *actual); + } + + void AssertSortIndices(const std::shared_ptr chunked_array, + SortOrder order, const std::string expected) { + AssertSortIndices(chunked_array, order, ArrayFromJSON(uint64(), expected)); + } +}; + +TEST_F(TestChunkedArraySortIndices, SortNull) { + auto chunked_array = ChunkedArrayFromJSON(uint8(), { + "[null, 1]", + "[3, null, 2]", + "[1]", + }); + this->AssertSortIndices(chunked_array, SortOrder::Ascending, "[1, 5, 4, 2, 0, 3]"); +} + +TEST_F(TestChunkedArraySortIndices, SortNaN) { + auto chunked_array = ChunkedArrayFromJSON(float32(), { + "[null, 1]", + "[3, null, NaN]", + "[NaN, 1]", + }); + this->AssertSortIndices(chunked_array, SortOrder::Ascending, "[1, 6, 2, 4, 5, 0, 3]"); +} + +template +class TestChunkedArrayRandomBase : public TestBase { + protected: + virtual std::shared_ptr GenerateArray(int length, double null_probability) = 0; + + void TestSortIndices(int length) { + using ArrayType = typename TypeTraits::ArrayType; + for (auto null_probability : {0.0, 0.1, 0.5, 1.0}) { + for (auto order : {SortOrder::Ascending, SortOrder::Descending}) { + for (auto num_chunks : {1, 5, 10}) { + std::vector> arrays; + for (int i = 0; i < num_chunks; ++i) { + auto array = this->GenerateArray(length, null_probability); + arrays.push_back(array); + } + ASSERT_OK_AND_ASSIGN(auto chunked_array, ChunkedArray::Make(arrays)); + ASSERT_OK_AND_ASSIGN(auto offsets, SortIndices(*chunked_array, order)); + ASSERT_OK_AND_ASSIGN(auto concatenated_array, Concatenate(arrays)); + ValidateSorted(*checked_pointer_cast(concatenated_array), + *checked_pointer_cast(offsets), order); + } + } + } + } +}; + +// Long array with big value range: std::stable_sort +template +class TestChunkedArrayRandom : public TestChunkedArrayRandomBase { + public: + void SetUp() { this->rand_.reset(new Random(0x5487655)); } + + void TearDown() { this->rand_.release(); } + + protected: + std::shared_ptr GenerateArray(int length, double null_probability) override { + return rand_->Generate(length, null_probability); + } + + private: + std::unique_ptr> rand_; +}; +TYPED_TEST_SUITE(TestChunkedArrayRandom, SortIndicesableTypes); +TYPED_TEST(TestChunkedArrayRandom, SortIndices) { this->TestSortIndices(4000); } + +// Long array with small value range: counting sort +// - length >= 1024(CountCompareSorter::countsort_min_len_) +// - range <= 4096(CountCompareSorter::countsort_max_range_) +template +class TestChunkedArrayRandomNarrow : public TestChunkedArrayRandomBase { + public: + void SetUp() { + range_ = 2000; + rand_.reset(new RandomRange(0x5487655)); + } + + void TearDown() { rand_.release(); } + + protected: + std::shared_ptr GenerateArray(int length, double null_probability) override { + return rand_->Generate(length, range_, null_probability); + } + + private: + int range_; + std::unique_ptr> rand_; +}; +TYPED_TEST_SUITE(TestChunkedArrayRandomNarrow, IntegralArrowTypes); +TYPED_TEST(TestChunkedArrayRandomNarrow, SortIndices) { this->TestSortIndices(4000); } + class TestTableSortIndices : public ::testing::Test { protected: void AssertSortIndices(const std::shared_ptr
table, const SortOptions& options, @@ -474,13 +574,13 @@ TEST_F(TestTableSortIndices, SortNaN) { "{\"a\": 3, \"b\": null}," "{\"a\": null, \"b\": null}," "{\"a\": NaN, \"b\": null}," - "{\"a\": NaN, \"b\": 5}," "{\"a\": NaN, \"b\": NaN}," + "{\"a\": NaN, \"b\": 5}," "{\"a\": 1, \"b\": 5}" "]"}); SortOptions options( {SortKey("a", SortOrder::Ascending), SortKey("b", SortOrder::Descending)}); - this->AssertSortIndices(table, options, "[7, 1, 2, 5, 6, 4, 0, 3]"); + this->AssertSortIndices(table, options, "[7, 1, 2, 6, 5, 4, 0, 3]"); } using RandomParam = std::tuple; @@ -622,7 +722,7 @@ TEST_P(TestTableSortIndicesRandom, Sort) { auto table = Table::Make(schema(fields), columns, length); std::default_random_engine engine(seed); std::uniform_int_distribution<> distribution(0); - auto n_sort_keys = 2; + auto n_sort_keys = 5; std::vector sort_keys; auto first_sort_key_order = (distribution(engine) % 2) == 0 ? SortOrder::Ascending : SortOrder::Descending; From 36317662715371b741f1588fcd48b04da0946899 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Thu, 19 Nov 2020 14:47:33 +0900 Subject: [PATCH 24/37] Format JSON --- .../arrow/compute/kernels/vector_sort_test.cc | 42 ++++++++++--------- 1 file changed, 23 insertions(+), 19 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index a0ef63c1508..76d4f61a14f 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -273,41 +273,43 @@ TYPED_TEST(TestArraySortIndicesKernelForReal, SortReal) { this->AssertSortIndices("[]", "[]"); this->AssertSortIndices("[3.4, 2.6, 6.3]", "[1, 0, 2]"); - this->AssertSortIndices("[1.1, 2.4, 3.5, 4.3, 5.1, 6.8, 7.3]", "[0,1,2,3,4,5,6]"); - this->AssertSortIndices("[7, 6, 5, 4, 3, 2, 1]", "[6,5,4,3,2,1,0]"); - this->AssertSortIndices("[10.4, 12, 4.2, 50, 50.3, 32, 11]", "[2,0,6,1,5,3,4]"); + this->AssertSortIndices("[1.1, 2.4, 3.5, 4.3, 5.1, 6.8, 7.3]", "[0, 1, 2, 3, 4, 5, 6]"); + this->AssertSortIndices("[7, 6, 5, 4, 3, 2, 1]", "[6, 5, 4, 3, 2, 1, 0]"); + this->AssertSortIndices("[10.4, 12, 4.2, 50, 50.3, 32, 11]", "[2, 0, 6, 1, 5, 3, 4]"); this->AssertSortIndices("[null, 1, 3.3, null, 2, 5.3]", SortOrder::Ascending, - "[1,4,2,5,0,3]"); + "[1, 4, 2, 5, 0, 3]"); this->AssertSortIndices("[null, 1, 3.3, null, 2, 5.3]", SortOrder::Descending, - "[5,2,4,1,0,3]"); + "[5, 2, 4, 1, 0, 3]"); this->AssertSortIndices("[3, 4, NaN, 1, 2, null]", SortOrder::Ascending, - "[3,4,0,1,2,5]"); + "[3, 4, 0, 1, 2, 5]"); this->AssertSortIndices("[3, 4, NaN, 1, 2, null]", SortOrder::Descending, - "[1,0,4,3,2,5]"); - this->AssertSortIndices("[NaN, 2, NaN, 3, 1]", SortOrder::Ascending, "[4,1,3,0,2]"); - this->AssertSortIndices("[NaN, 2, NaN, 3, 1]", SortOrder::Descending, "[3,1,4,0,2]"); - this->AssertSortIndices("[null, NaN, NaN, null]", SortOrder::Ascending, "[1,2,0,3]"); - this->AssertSortIndices("[null, NaN, NaN, null]", SortOrder::Descending, "[1,2,0,3]"); + "[1, 0, 4, 3, 2, 5]"); + this->AssertSortIndices("[NaN, 2, NaN, 3, 1]", SortOrder::Ascending, "[4, 1, 3, 0, 2]"); + this->AssertSortIndices("[NaN, 2, NaN, 3, 1]", SortOrder::Descending, + "[3, 1, 4, 0, 2]"); + this->AssertSortIndices("[null, NaN, NaN, null]", SortOrder::Ascending, "[1, 2, 0, 3]"); + this->AssertSortIndices("[null, NaN, NaN, null]", SortOrder::Descending, + "[1, 2, 0, 3]"); } TYPED_TEST(TestArraySortIndicesKernelForIntegral, SortIntegral) { this->AssertSortIndices("[]", "[]"); this->AssertSortIndices("[3, 2, 6]", "[1, 0, 2]"); - this->AssertSortIndices("[1, 2, 3, 4, 5, 6, 7]", "[0,1,2,3,4,5,6]"); - this->AssertSortIndices("[7, 6, 5, 4, 3, 2, 1]", "[6,5,4,3,2,1,0]"); + this->AssertSortIndices("[1, 2, 3, 4, 5, 6, 7]", "[0, 1, 2, 3, 4, 5, 6]"); + this->AssertSortIndices("[7, 6, 5, 4, 3, 2, 1]", "[6, 5, 4, 3, 2, 1, 0]"); this->AssertSortIndices("[10, 12, 4, 50, 50, 32, 11]", SortOrder::Ascending, - "[2,0,6,1,5,3,4]"); + "[2, 0, 6, 1, 5, 3, 4]"); this->AssertSortIndices("[10, 12, 4, 50, 50, 32, 11]", SortOrder::Descending, - "[3,4,5,1,6,0,2]"); + "[3, 4, 5, 1, 6, 0, 2]"); this->AssertSortIndices("[null, 1, 3, null, 2, 5]", SortOrder::Ascending, - "[1,4,2,5,0,3]"); + "[1, 4, 2, 5, 0, 3]"); this->AssertSortIndices("[null, 1, 3, null, 2, 5]", SortOrder::Descending, - "[5,2,4,1,0,3]"); + "[5, 2, 4, 1, 0, 3]"); } TYPED_TEST(TestArraySortIndicesKernelForStrings, SortStrings) { @@ -333,11 +335,13 @@ class TestArraySortIndicesKernelForInt8 : public TestArraySortIndicesKernelAssertSortIndices("[255, null, 0, 255, 10, null, 128, 0]", "[2,7,4,6,0,3,1,5]"); + this->AssertSortIndices("[255, null, 0, 255, 10, null, 128, 0]", + "[2, 7, 4, 6, 0, 3, 1, 5]"); } TYPED_TEST(TestArraySortIndicesKernelForInt8, SortInt8) { - this->AssertSortIndices("[null, 10, 127, 0, -128, -128, null]", "[4,5,3,1,2,0,6]"); + this->AssertSortIndices("[null, 10, 127, 0, -128, -128, null]", + "[4, 5, 3, 1, 2, 0, 6]"); } template From eca39025de49f4449ed783d3cfde49cd9cda2819 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Thu, 19 Nov 2020 14:47:44 +0900 Subject: [PATCH 25/37] Add more comments --- .../arrow/compute/kernels/vector_sort_test.cc | 27 ++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index 76d4f61a14f..df3eda5579a 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -433,6 +433,7 @@ TYPED_TEST(TestArraySortIndicesKernelRandomCompare, SortRandomValuesCompare) { } } +// Test basic cases for chunked array. class TestChunkedArraySortIndices : public ::testing::Test { protected: void AssertSortIndices(const std::shared_ptr chunked_array, @@ -465,13 +466,17 @@ TEST_F(TestChunkedArraySortIndices, SortNaN) { this->AssertSortIndices(chunked_array, SortOrder::Ascending, "[1, 6, 2, 4, 5, 0, 3]"); } +// Base class for testing against random chunked array. template class TestChunkedArrayRandomBase : public TestBase { protected: + // Generates a chunk. This should be implemented in subclasses. virtual std::shared_ptr GenerateArray(int length, double null_probability) = 0; + // All tests uses this. void TestSortIndices(int length) { using ArrayType = typename TypeTraits::ArrayType; + // We can use INSTANTIATE_TEST_SUITE_P() instead of using fors in a test. for (auto null_probability : {0.0, 0.1, 0.5, 1.0}) { for (auto order : {SortOrder::Ascending, SortOrder::Descending}) { for (auto num_chunks : {1, 5, 10}) { @@ -482,6 +487,7 @@ class TestChunkedArrayRandomBase : public TestBase { } ASSERT_OK_AND_ASSIGN(auto chunked_array, ChunkedArray::Make(arrays)); ASSERT_OK_AND_ASSIGN(auto offsets, SortIndices(*chunked_array, order)); + // Concatenates chunks to use existing ValidateSorted() for array. ASSERT_OK_AND_ASSIGN(auto concatenated_array, Concatenate(arrays)); ValidateSorted(*checked_pointer_cast(concatenated_array), *checked_pointer_cast(offsets), order); @@ -535,6 +541,7 @@ class TestChunkedArrayRandomNarrow : public TestChunkedArrayRandomBase { TYPED_TEST_SUITE(TestChunkedArrayRandomNarrow, IntegralArrowTypes); TYPED_TEST(TestChunkedArrayRandomNarrow, SortIndices) { this->TestSortIndices(4000); } +// Test basic cases for table. class TestTableSortIndices : public ::testing::Test { protected: void AssertSortIndices(const std::shared_ptr
table, const SortOptions& options, @@ -587,18 +594,24 @@ TEST_F(TestTableSortIndices, SortNaN) { this->AssertSortIndices(table, options, "[7, 1, 2, 6, 5, 4, 0, 3]"); } +// For random table tests. using RandomParam = std::tuple; class TestTableSortIndicesRandom : public testing::TestWithParam { + // Compares two records in the same table. class Comparator : public TypeVisitor { public: + // Returns true if the left record is less or equals to the right + // record, false otherwise. + // + // This supports null and NaN. bool operator()(const Table& table, const SortOptions& options, uint64_t lhs, uint64_t rhs) { lhs_ = lhs; rhs_ = rhs; for (const auto& sort_key : options.sort_keys) { auto chunked_array = table.GetColumnByName(sort_key.name); - lhs_array_ = findTargetArray(chunked_array, lhs, lhs_index_); - rhs_array_ = findTargetArray(chunked_array, rhs, rhs_index_); + lhs_array_ = FindTargetArray(chunked_array, lhs, lhs_index_); + rhs_array_ = FindTargetArray(chunked_array, rhs, rhs_index_); if (rhs_array_->IsNull(rhs_index_) && lhs_array_->IsNull(lhs_index_)) continue; if (rhs_array_->IsNull(rhs_index_)) return true; if (lhs_array_->IsNull(lhs_index_)) return false; @@ -636,7 +649,9 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { #undef VISIT private: - std::shared_ptr findTargetArray(std::shared_ptr chunked_array, + // Finds the target chunk and index in the target chunk from an + // index in chunked array. + std::shared_ptr FindTargetArray(std::shared_ptr chunked_array, int64_t i, int64_t& chunk_index) { int64_t offset = 0; for (auto& chunk : chunked_array->chunks()) { @@ -649,6 +664,11 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { return nullptr; } + // Compares two values in the same chunked array. Values are never + // null but may be NaN. + // + // Returns true if the left value is less or equals to the right + // value, false otherwise. template int CompareType() { using ArrayType = typename TypeTraits::ArrayType; @@ -681,6 +701,7 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { }; public: + // Validates the sorted indexes are really sorted. void Validate(const Table& table, const SortOptions& options, UInt64Array& offsets) { ASSERT_OK(offsets.ValidateFull()); Comparator comparator; From af9c8807c3b4bb1001d5cc6558284763b6341db5 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Thu, 19 Nov 2020 14:50:54 +0900 Subject: [PATCH 26/37] Don't use shared_ptr --- .../arrow/compute/kernels/vector_sort_test.cc | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index df3eda5579a..2fc4fa5c227 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -31,6 +31,7 @@ namespace arrow { +using internal::checked_cast; using internal::checked_pointer_cast; namespace compute { @@ -629,10 +630,10 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { Status status() const { return status_; } #define VISIT(TYPE) \ - Status Visit(const TYPE##Type& type) override { \ - compared_ = CompareType(); \ - return Status::OK(); \ - } + Status Visit(const TYPE##Type& type) override { \ + compared_ = CompareType(); \ + return Status::OK(); \ + } VISIT(Int8) VISIT(Int16) @@ -651,13 +652,13 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { private: // Finds the target chunk and index in the target chunk from an // index in chunked array. - std::shared_ptr FindTargetArray(std::shared_ptr chunked_array, + const Array* FindTargetArray(std::shared_ptr chunked_array, int64_t i, int64_t& chunk_index) { int64_t offset = 0; for (auto& chunk : chunked_array->chunks()) { if (i < offset + chunk->length()) { chunk_index = i - offset; - return chunk; + return chunk.get(); } offset += chunk->length(); } @@ -672,8 +673,8 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { template int CompareType() { using ArrayType = typename TypeTraits::ArrayType; - auto lhs_value = checked_pointer_cast(lhs_array_)->GetView(lhs_index_); - auto rhs_value = checked_pointer_cast(rhs_array_)->GetView(rhs_index_); + auto lhs_value = checked_cast(lhs_array_)->GetView(lhs_index_); + auto rhs_value = checked_cast(rhs_array_)->GetView(rhs_index_); if (is_floating_type::value) { const bool lhs_isnan = lhs_value != lhs_value; const bool rhs_isnan = rhs_value != rhs_value; @@ -691,10 +692,10 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { } int64_t lhs_; - std::shared_ptr lhs_array_; + const Array* lhs_array_; int64_t lhs_index_; int64_t rhs_; - std::shared_ptr rhs_array_; + const Array* rhs_array_; int64_t rhs_index_; int compared_; Status status_; From aa368eb851dddca758ec543858ff71a1e9f97da5 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Thu, 19 Nov 2020 15:21:12 +0900 Subject: [PATCH 27/37] Format --- cpp/src/arrow/compute/kernels/vector_sort_test.cc | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index 2fc4fa5c227..5fbd716721d 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -630,10 +630,10 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { Status status() const { return status_; } #define VISIT(TYPE) \ - Status Visit(const TYPE##Type& type) override { \ - compared_ = CompareType(); \ - return Status::OK(); \ - } + Status Visit(const TYPE##Type& type) override { \ + compared_ = CompareType(); \ + return Status::OK(); \ + } VISIT(Int8) VISIT(Int16) @@ -652,8 +652,8 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { private: // Finds the target chunk and index in the target chunk from an // index in chunked array. - const Array* FindTargetArray(std::shared_ptr chunked_array, - int64_t i, int64_t& chunk_index) { + const Array* FindTargetArray(std::shared_ptr chunked_array, int64_t i, + int64_t& chunk_index) { int64_t offset = 0; for (auto& chunk : chunked_array->chunks()) { if (i < offset + chunk->length()) { From 8deb74c67e3820ba91a543ee0e0e4bb53f3ab826 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Thu, 19 Nov 2020 15:22:46 +0900 Subject: [PATCH 28/37] Use benchmark 1.5.2 for ArgsProduct() --- ci/conda_env_cpp.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ci/conda_env_cpp.yml b/ci/conda_env_cpp.yml index 90cef3ea2d1..534ede6df1c 100644 --- a/ci/conda_env_cpp.yml +++ b/ci/conda_env_cpp.yml @@ -16,7 +16,7 @@ # under the License. aws-sdk-cpp -benchmark=1.4.1 +benchmark=1.5.2 boost-cpp>=1.68.0 brotli bzip2 From e9c03d2269c5ad021d0f7deded2183757d939f77 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Tue, 24 Nov 2020 13:38:26 +0900 Subject: [PATCH 29/37] Remove benchmark version check with conda --- cpp/cmake_modules/ThirdpartyToolchain.cmake | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/cpp/cmake_modules/ThirdpartyToolchain.cmake b/cpp/cmake_modules/ThirdpartyToolchain.cmake index 1644d5b3b91..482391fb7b7 100644 --- a/cpp/cmake_modules/ThirdpartyToolchain.cmake +++ b/cpp/cmake_modules/ThirdpartyToolchain.cmake @@ -1798,7 +1798,19 @@ endmacro() if(ARROW_BUILD_BENCHMARKS) # ArgsProduct() is available since 1.5.2 - resolve_dependency(benchmark REQUIRED_VERSION 1.5.2) + set(BENCHMARK_REQUIRED_VERSION 1.5.2) + if("${ARROW_DEPENDENCY_SOURCE}" STREQUAL "CONDA" AND "${benchmark_SOURCE}" STREQUAL "SYSTEM") + # TODO: Remove this workaround once + # https://github.com/google/benchmark/issues/1046 is resolved. + # + # benchmark doesn't set suitable version when we use released + # archive. So the benchmark package on conda-forge isn't report + # the real version. We accept all the benchmark package with + # conda. Conda users should install benchmark 1.5.2 or later by + # ci/conda_env_cpp.yml. + set(BENCHMARK_REQUIRED_VERSION 0.0.0) + endif() + resolve_dependency(benchmark REQUIRED_VERSION ${BENCHMARK_REQUIRED_VERSION}) # TODO: Don't use global includes but rather target_include_directories get_target_property(BENCHMARK_INCLUDE_DIR benchmark::benchmark INTERFACE_INCLUDE_DIRECTORIES) From f84f748ed9c0dc4735b308e7d50a78a097cd4c6a Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Tue, 24 Nov 2020 14:03:25 +0900 Subject: [PATCH 30/37] Add missing override --- cpp/src/arrow/compute/kernels/vector_sort_test.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index 5fbd716721d..09c41d23eec 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -502,9 +502,9 @@ class TestChunkedArrayRandomBase : public TestBase { template class TestChunkedArrayRandom : public TestChunkedArrayRandomBase { public: - void SetUp() { this->rand_.reset(new Random(0x5487655)); } + void SetUp() override { this->rand_.reset(new Random(0x5487655)); } - void TearDown() { this->rand_.release(); } + void TearDown() override { this->rand_.release(); } protected: std::shared_ptr GenerateArray(int length, double null_probability) override { @@ -523,12 +523,12 @@ TYPED_TEST(TestChunkedArrayRandom, SortIndices) { this->TestSortIndices(4000); } template class TestChunkedArrayRandomNarrow : public TestChunkedArrayRandomBase { public: - void SetUp() { + void SetUp() override { range_ = 2000; rand_.reset(new RandomRange(0x5487655)); } - void TearDown() { rand_.release(); } + void TearDown() override { rand_.release(); } protected: std::shared_ptr GenerateArray(int length, double null_probability) override { From a7eb62fb7ec665449ee7dc098e2d9a48fafb8d15 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Tue, 24 Nov 2020 14:03:49 +0900 Subject: [PATCH 31/37] Fix format --- cpp/cmake_modules/ThirdpartyToolchain.cmake | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cpp/cmake_modules/ThirdpartyToolchain.cmake b/cpp/cmake_modules/ThirdpartyToolchain.cmake index 482391fb7b7..131e8ad6caa 100644 --- a/cpp/cmake_modules/ThirdpartyToolchain.cmake +++ b/cpp/cmake_modules/ThirdpartyToolchain.cmake @@ -1799,7 +1799,8 @@ endmacro() if(ARROW_BUILD_BENCHMARKS) # ArgsProduct() is available since 1.5.2 set(BENCHMARK_REQUIRED_VERSION 1.5.2) - if("${ARROW_DEPENDENCY_SOURCE}" STREQUAL "CONDA" AND "${benchmark_SOURCE}" STREQUAL "SYSTEM") + if("${ARROW_DEPENDENCY_SOURCE}" STREQUAL "CONDA" + AND "${benchmark_SOURCE}" STREQUAL "SYSTEM") # TODO: Remove this workaround once # https://github.com/google/benchmark/issues/1046 is resolved. # From e9b481c8f5bfbc7650ac48d9065e1776389f8155 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Tue, 24 Nov 2020 14:58:25 +0900 Subject: [PATCH 32/37] Don't use unique_ptr --- cpp/src/arrow/compute/kernels/vector_sort_test.cc | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index 09c41d23eec..4811c518e83 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -502,9 +502,9 @@ class TestChunkedArrayRandomBase : public TestBase { template class TestChunkedArrayRandom : public TestChunkedArrayRandomBase { public: - void SetUp() override { this->rand_.reset(new Random(0x5487655)); } + void SetUp() override { rand_ = new Random(0x5487655); } - void TearDown() override { this->rand_.release(); } + void TearDown() override { delete rand_; } protected: std::shared_ptr GenerateArray(int length, double null_probability) override { @@ -512,7 +512,7 @@ class TestChunkedArrayRandom : public TestChunkedArrayRandomBase { } private: - std::unique_ptr> rand_; + Random *rand_; }; TYPED_TEST_SUITE(TestChunkedArrayRandom, SortIndicesableTypes); TYPED_TEST(TestChunkedArrayRandom, SortIndices) { this->TestSortIndices(4000); } @@ -525,10 +525,10 @@ class TestChunkedArrayRandomNarrow : public TestChunkedArrayRandomBase { public: void SetUp() override { range_ = 2000; - rand_.reset(new RandomRange(0x5487655)); + rand_ = new RandomRange(0x5487655); } - void TearDown() override { rand_.release(); } + void TearDown() override { delete rand_; } protected: std::shared_ptr GenerateArray(int length, double null_probability) override { @@ -537,7 +537,7 @@ class TestChunkedArrayRandomNarrow : public TestChunkedArrayRandomBase { private: int range_; - std::unique_ptr> rand_; + RandomRange* rand_; }; TYPED_TEST_SUITE(TestChunkedArrayRandomNarrow, IntegralArrowTypes); TYPED_TEST(TestChunkedArrayRandomNarrow, SortIndices) { this->TestSortIndices(4000); } From bbda2401ad8c12f2ee26b69ca4dedea0af64b57c Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Tue, 24 Nov 2020 15:21:04 +0900 Subject: [PATCH 33/37] Fix format --- cpp/src/arrow/compute/kernels/vector_sort_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index 4811c518e83..19b51a4458f 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -512,7 +512,7 @@ class TestChunkedArrayRandom : public TestChunkedArrayRandomBase { } private: - Random *rand_; + Random* rand_; }; TYPED_TEST_SUITE(TestChunkedArrayRandom, SortIndicesableTypes); TYPED_TEST(TestChunkedArrayRandom, SortIndices) { this->TestSortIndices(4000); } From e3ed89a49f26b79e1d0cb395832d20b086b684e3 Mon Sep 17 00:00:00 2001 From: Antoine Pitrou Date: Tue, 24 Nov 2020 15:17:33 +0100 Subject: [PATCH 34/37] Make tests faster --- .../arrow/compute/kernels/vector_sort_test.cc | 64 +++++++++++-------- 1 file changed, 37 insertions(+), 27 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index 19b51a4458f..a7ad0453fbe 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -376,7 +376,7 @@ TYPED_TEST(TestArraySortIndicesKernelRandom, SortRandomValues) { Random rand(0x5487655); int times = 5; - int length = 1000; + int length = 100; for (int test = 0; test < times; test++) { for (auto null_probability : {0.0, 0.1, 0.5, 1.0}) { for (auto order : {SortOrder::Ascending, SortOrder::Descending}) { @@ -399,7 +399,7 @@ TYPED_TEST(TestArraySortIndicesKernelRandomCount, SortRandomValuesCount) { RandomRange rand(0x5487656); int times = 5; - int length = 4000; + int length = 100; int range = 2000; for (int test = 0; test < times; test++) { for (auto null_probability : {0.0, 0.1, 0.5, 1.0}) { @@ -421,7 +421,7 @@ TYPED_TEST(TestArraySortIndicesKernelRandomCompare, SortRandomValuesCompare) { Random rand(0x5487657); int times = 5; - int length = 4000; + int length = 100; for (int test = 0; test < times; test++) { for (auto null_probability : {0.0, 0.1, 0.5, 1.0}) { for (auto order : {SortOrder::Ascending, SortOrder::Descending}) { @@ -515,7 +515,7 @@ class TestChunkedArrayRandom : public TestChunkedArrayRandomBase { Random* rand_; }; TYPED_TEST_SUITE(TestChunkedArrayRandom, SortIndicesableTypes); -TYPED_TEST(TestChunkedArrayRandom, SortIndices) { this->TestSortIndices(4000); } +TYPED_TEST(TestChunkedArrayRandom, SortIndices) { this->TestSortIndices(100); } // Long array with small value range: counting sort // - length >= 1024(CountCompareSorter::countsort_min_len_) @@ -540,7 +540,7 @@ class TestChunkedArrayRandomNarrow : public TestChunkedArrayRandomBase { RandomRange* rand_; }; TYPED_TEST_SUITE(TestChunkedArrayRandomNarrow, IntegralArrowTypes); -TYPED_TEST(TestChunkedArrayRandomNarrow, SortIndices) { this->TestSortIndices(4000); } +TYPED_TEST(TestChunkedArrayRandomNarrow, SortIndices) { this->TestSortIndices(100); } // Test basic cases for table. class TestTableSortIndices : public ::testing::Test { @@ -601,24 +601,31 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { // Compares two records in the same table. class Comparator : public TypeVisitor { public: + Comparator(const Table& table, const SortOptions& options) + : table_(table), options_(options) { + for (const auto& sort_key : options_.sort_keys) { + sort_columns_.emplace_back(table.GetColumnByName(sort_key.name).get(), + sort_key.order); + } + } + // Returns true if the left record is less or equals to the right // record, false otherwise. // // This supports null and NaN. - bool operator()(const Table& table, const SortOptions& options, uint64_t lhs, - uint64_t rhs) { + bool operator()(uint64_t lhs, uint64_t rhs) { lhs_ = lhs; rhs_ = rhs; - for (const auto& sort_key : options.sort_keys) { - auto chunked_array = table.GetColumnByName(sort_key.name); - lhs_array_ = FindTargetArray(chunked_array, lhs, lhs_index_); - rhs_array_ = FindTargetArray(chunked_array, rhs, rhs_index_); + for (const auto& pair : sort_columns_) { + const auto& chunked_array = *pair.first; + lhs_array_ = FindTargetArray(chunked_array, lhs, &lhs_index_); + rhs_array_ = FindTargetArray(chunked_array, rhs, &rhs_index_); if (rhs_array_->IsNull(rhs_index_) && lhs_array_->IsNull(lhs_index_)) continue; if (rhs_array_->IsNull(rhs_index_)) return true; if (lhs_array_->IsNull(lhs_index_)) return false; status_ = lhs_array_->type()->Accept(this); if (compared_ == 0) continue; - if (sort_key.order == SortOrder::Ascending) { + if (pair.second == SortOrder::Ascending) { return compared_ < 0; } else { return compared_ > 0; @@ -652,12 +659,12 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { private: // Finds the target chunk and index in the target chunk from an // index in chunked array. - const Array* FindTargetArray(std::shared_ptr chunked_array, int64_t i, - int64_t& chunk_index) { + const Array* FindTargetArray(const ChunkedArray& chunked_array, int64_t i, + int64_t* chunk_index) { int64_t offset = 0; - for (auto& chunk : chunked_array->chunks()) { + for (const auto& chunk : chunked_array.chunks()) { if (i < offset + chunk->length()) { - chunk_index = i - offset; + *chunk_index = i - offset; return chunk.get(); } offset += chunk->length(); @@ -691,6 +698,9 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { } } + const Table& table_; + const SortOptions& options_; + std::vector> sort_columns_; int64_t lhs_; const Array* lhs_array_; int64_t lhs_index_; @@ -705,20 +715,20 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { // Validates the sorted indexes are really sorted. void Validate(const Table& table, const SortOptions& options, UInt64Array& offsets) { ASSERT_OK(offsets.ValidateFull()); - Comparator comparator; + Comparator comparator{table, options}; for (int i = 1; i < table.num_rows(); i++) { uint64_t lhs = offsets.Value(i - 1); uint64_t rhs = offsets.Value(i); - ASSERT_TRUE(comparator(table, options, lhs, rhs)); + ASSERT_TRUE(comparator(lhs, rhs)); ASSERT_OK(comparator.status()); } } }; TEST_P(TestTableSortIndicesRandom, Sort) { - auto first_sort_key_name = std::get<0>(GetParam()); - auto null_probability = std::get<1>(GetParam()); - auto seed = 0x61549225; + const auto first_sort_key_name = std::get<0>(GetParam()); + const auto null_probability = std::get<1>(GetParam()); + const auto seed = 0x61549225; std::vector column_names = { "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64", "float", "double", "string", @@ -731,7 +741,7 @@ TEST_P(TestTableSortIndicesRandom, Sort) { {field(column_names[8], float32())}, {field(column_names[9], float64())}, {field(column_names[10], utf8())}, }; - auto length = 4000; + const auto length = 100; std::vector> columns = { Random(seed).Generate(length, null_probability), Random(seed).Generate(length, null_probability), @@ -745,17 +755,17 @@ TEST_P(TestTableSortIndicesRandom, Sort) { Random(seed).Generate(length, null_probability), Random(seed).Generate(length, null_probability), }; - auto table = Table::Make(schema(fields), columns, length); + const auto table = Table::Make(schema(fields), columns, length); std::default_random_engine engine(seed); std::uniform_int_distribution<> distribution(0); - auto n_sort_keys = 5; + const auto n_sort_keys = 5; std::vector sort_keys; - auto first_sort_key_order = + const auto first_sort_key_order = (distribution(engine) % 2) == 0 ? SortOrder::Ascending : SortOrder::Descending; sort_keys.emplace_back(first_sort_key_name, first_sort_key_order); for (int i = 1; i < n_sort_keys; ++i) { - auto& column_name = column_names[distribution(engine) % column_names.size()]; - auto order = + const auto& column_name = column_names[distribution(engine) % column_names.size()]; + const auto order = (distribution(engine) % 2) == 0 ? SortOrder::Ascending : SortOrder::Descending; sort_keys.emplace_back(column_name, order); } From 5ca9e8d2e8266dfa17d61847c6b7dce4bcf7bdaa Mon Sep 17 00:00:00 2001 From: Antoine Pitrou Date: Tue, 24 Nov 2020 15:18:38 +0100 Subject: [PATCH 35/37] Fix typos in doc --- docs/source/cpp/compute.rst | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst index 1a019070e66..9c10b95adc6 100644 --- a/docs/source/cpp/compute.rst +++ b/docs/source/cpp/compute.rst @@ -628,13 +628,13 @@ value, but smaller than nulls. +-----------------------+------------+-------------------------+-------------------+--------------------------------+----------------+ | partition_nth_indices | Unary | Numeric | UInt64 | :struct:`PartitionNthOptions` | \(1) | +-----------------------+------------+-------------------------+-------------------+--------------------------------+----------------+ -| array_sort_indices | Unary | Binary- and String-like | UInt64 | :struct:`ArraySortptions` | \(2) \(3) \(4) | +| array_sort_indices | Unary | Binary- and String-like | UInt64 | :struct:`ArraySortOptions` | \(2) \(3) \(4) | +-----------------------+------------+-------------------------+-------------------+--------------------------------+----------------+ -| array_sort_indices | Unary | Numeric | UInt64 | :struct:`ArraySortptions` | \(2) \(4) | +| array_sort_indices | Unary | Numeric | UInt64 | :struct:`ArraySortOptions` | \(2) \(4) | +-----------------------+------------+-------------------------+-------------------+--------------------------------+----------------+ -| sort_indices | Unary | Binary- and String-like | UInt64 | :struct:`Sortptions` | \(2) \(3) \(5) | +| sort_indices | Unary | Binary- and String-like | UInt64 | :struct:`SortOptions` | \(2) \(3) \(5) | +-----------------------+------------+-------------------------+-------------------+--------------------------------+----------------+ -| sort_indices | Unary | Numeric | UInt64 | :struct:`Sortptions` | \(2) \(5) | +| sort_indices | Unary | Numeric | UInt64 | :struct:`SortOptions` | \(2) \(5) | +-----------------------+------------+-------------------------+-------------------+--------------------------------+----------------+ * \(1) The output is an array of indices into the input array, that define From fecd557e365250eb8fdf22ef49f67f1fd722dac7 Mon Sep 17 00:00:00 2001 From: Antoine Pitrou Date: Tue, 24 Nov 2020 16:17:51 +0100 Subject: [PATCH 36/37] More small changes Add some chunked array sorting benchmarks --- cpp/src/arrow/compute/kernels/vector_sort.cc | 106 ++++++++---------- .../compute/kernels/vector_sort_benchmark.cc | 62 ++++++++-- .../arrow/compute/kernels/vector_sort_test.cc | 4 +- 3 files changed, 104 insertions(+), 68 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort.cc b/cpp/src/arrow/compute/kernels/vector_sort.cc index c5a33a85b51..9f7788a25de 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort.cc @@ -19,58 +19,40 @@ #include #include #include +#include #include #include "arrow/array/data.h" #include "arrow/compute/api_vector.h" #include "arrow/compute/kernels/common.h" #include "arrow/table.h" +#include "arrow/util/checked_cast.h" #include "arrow/util/optional.h" namespace arrow { + +using internal::checked_cast; + namespace compute { namespace internal { namespace { // The target chunk in chunked array. -template -struct ResolvedChunk; - -// Not for string-like data. template -struct ResolvedChunk< - ArrayType, enable_if_t::value>> { - using c_type = typename ArrayType::TypeClass::c_type; - - // The target array in chunked array. - const ArrayType* array; - - // The index in the target array. - int64_t index; +struct ResolvedChunk { + using ViewType = decltype(std::declval().GetView(0)); - ResolvedChunk(const ArrayType* array, int64_t index) : array(array), index(index) {} - - bool IsNull() const { return array->IsNull(index); } - - c_type GetView() const { return array->GetView(index); } -}; - -// For string-like data. -template -struct ResolvedChunk< - ArrayType, enable_if_t::value>> { // The target array in chunked array. const ArrayType* array; - // The index in the target array. - int64_t index; + const int64_t index; ResolvedChunk(const ArrayType* array, int64_t index) : array(array), index(index) {} bool IsNull() const { return array->IsNull(index); } - util::string_view GetView() const { return array->GetView(index); } + ViewType GetView() const { return array->GetView(index); } }; // Finds the target chunk and index in the target chunk from an index @@ -83,7 +65,7 @@ ResolvedChunk ResolveChunk(const std::vector& chunks, int64_t offset = 0; for (size_t i = 0; i < num_chunks; ++i) { if (index < offset + chunks[i]->length()) { - return ResolvedChunk(static_cast(chunks[i]), + return ResolvedChunk(checked_cast(chunks[i]), index - offset); } offset += chunks[i]->length(); @@ -111,7 +93,7 @@ struct StablePartitioner { // Move nulls to end of array. Return where null starts. // -// `offset` is for using this to an array in a chunked array. +// `offset` is used when this is called on a chunk of a chunked array template enable_if_t::value, uint64_t*> PartitionNulls(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, @@ -129,7 +111,7 @@ PartitionNulls(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& template enable_if_t::value, uint64_t*> PartitionNulls(uint64_t* indices_begin, uint64_t* indices_end, - std::vector& arrays, int64_t null_count) { + const std::vector& arrays, int64_t null_count) { if (null_count == 0) { return indices_end; } @@ -143,7 +125,7 @@ PartitionNulls(uint64_t* indices_begin, uint64_t* indices_end, // Move NaNs and nulls to end of array, nulls after NaN. Return where // NaN/null starts. // -// `offset` is for using this to an array in a chunked array. +// `offset` is used when this is called on a chunk of a chunked array template enable_if_t::value, uint64_t*> PartitionNulls(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, @@ -171,7 +153,7 @@ PartitionNulls(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& template enable_if_t::value, uint64_t*> PartitionNulls(uint64_t* indices_begin, uint64_t* indices_end, - std::vector& arrays, int64_t null_count) { + const std::vector& arrays, int64_t null_count) { Partitioner partitioner; if (null_count == 0) { return partitioner(indices_begin, indices_end, [&arrays](uint64_t ind) { @@ -280,7 +262,7 @@ class ArrayCompareSorter { public: // Returns where null starts. // - // `offset` is for using this to an array in a chunked array. + // `offset` is used when this is called on a chunk of a chunked array uint64_t* Sort(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, int64_t offset, const ArraySortOptions& options) { auto nulls_begin = PartitionNulls( @@ -335,7 +317,7 @@ class ArrayCountSorter { // Returns where null starts. // - // `offset` is for using this to an array in a chunked array. + // `offset` is used when this is called on a chunk of a chunked array template uint64_t* SortInternal(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, int64_t offset, @@ -386,7 +368,7 @@ class ArrayCountOrCompareSorter { public: // Returns where null starts. // - // `offset` is for using this to an array in a chunked array. + // `offset` is used when this is called on a chunk of a chunked array uint64_t* Sort(uint64_t* indices_begin, uint64_t* indices_end, const ArrayType& values, int64_t offset, const ArraySortOptions& options) { if (values.length() >= countsort_min_len_ && values.length() > values.null_count()) { @@ -463,7 +445,7 @@ template struct ArraySortIndices { using ArrayType = typename TypeTraits::ArrayType; static void Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) { - auto& options = ArraySortIndicesState::Get(ctx); + const auto& options = ArraySortIndicesState::Get(ctx); std::shared_ptr arg0; KERNEL_RETURN_IF_ERROR( @@ -511,7 +493,7 @@ class ChunkedArrayCompareSorter { public: // Returns where null starts. uint64_t* Sort(uint64_t* indices_begin, uint64_t* indices_end, - std::vector& arrays, int64_t null_count, + const std::vector& arrays, int64_t null_count, const ArraySortOptions& options) { auto nulls_begin = PartitionNulls( indices_begin, indices_end, arrays, null_count); @@ -596,7 +578,7 @@ class ChunkedArraySorter : public TypeVisitor { int64_t null_count = 0; uint64_t* left_nulls_begin = indices_begin_; for (int i = 0; i < num_chunks; ++i) { - const auto array = static_cast(arrays[i]); + const auto array = checked_cast(arrays[i]); end_offset += array->length(); null_count += array->null_count(); uint64_t* right_nulls_begin; @@ -627,7 +609,7 @@ class ChunkedArraySorter : public TypeVisitor { template uint64_t* Merge(uint64_t* indices_begin, uint64_t* indices_middle, uint64_t* indices_end, uint64_t* left_nulls_begin, - uint64_t* right_nulls_begin, std::vector arrays, + uint64_t* right_nulls_begin, const std::vector& arrays, int64_t null_count, const SortOrder order) { auto left_num_non_nulls = left_nulls_begin - indices_begin; auto right_num_non_nulls = right_nulls_begin - indices_middle; @@ -665,7 +647,8 @@ class ChunkedArraySorter : public TypeVisitor { const bool can_use_array_sorter_; }; -// Sort a table from the back like radix sort. +// Sort a table using a radix sort-like algorithm. +// A distinct stable sort is called for each sort key, from the last key to the first. class TableRadixSorter { public: Status Sort(uint64_t* indices_begin, uint64_t* indices_end, const Table& table, @@ -688,8 +671,8 @@ class TableRadixSorter { } }; -// Sort a table from the beginning. -class TableFromTheBeginningSorter : public TypeVisitor { +// Sort a table using a single sort and multiple-key comparisons. +class MultipleKeyTableSorter : public TypeVisitor { private: // Preprocessed sort key. struct ResolvedSortKey { @@ -721,16 +704,9 @@ class TableFromTheBeginningSorter : public TypeVisitor { class Comparer : public TypeVisitor { public: Comparer(const Table& table, const std::vector& sort_keys) - : TypeVisitor(), status_(Status::OK()) { - for (const auto& sort_key : sort_keys) { - const auto& chunked_array = table.GetColumnByName(sort_key.name); - if (!chunked_array) { - status_ = Status::Invalid("Nonexistent sort key column: ", sort_key.name); - return; - } - sort_keys_.emplace_back(*chunked_array, sort_key.order); - } - } + : TypeVisitor(), + status_(Status::OK()), + sort_keys_(ResolveSortKeys(table, sort_keys, &status_)) {} Status status() { return status_; } @@ -862,8 +838,22 @@ class TableFromTheBeginningSorter : public TypeVisitor { return compared; } + static std::vector ResolveSortKeys( + const Table& table, const std::vector& sort_keys, Status* status) { + std::vector resolved; + for (const auto& sort_key : sort_keys) { + const auto& chunked_array = table.GetColumnByName(sort_key.name); + if (!chunked_array) { + *status = Status::Invalid("Nonexistent sort key column: ", sort_key.name); + break; + } + resolved.emplace_back(*chunked_array, sort_key.order); + } + return resolved; + } + Status status_; - std::vector sort_keys_; + const std::vector sort_keys_; int64_t current_left_; int64_t current_right_; size_t current_sort_key_index_; @@ -871,8 +861,8 @@ class TableFromTheBeginningSorter : public TypeVisitor { }; public: - TableFromTheBeginningSorter(uint64_t* indices_begin, uint64_t* indices_end, - const Table& table, const SortOptions& options) + MultipleKeyTableSorter(uint64_t* indices_begin, uint64_t* indices_end, + const Table& table, const SortOptions& options) : indices_begin_(indices_begin), indices_end_(indices_end), comparer_(table, options.sort_keys) {} @@ -1117,13 +1107,13 @@ class SortIndicesMetaFunction : public MetaFunction { // TODO: We should choose suitable sort implementation // automatically. The current TableRadixSorter implementation is - // faster than TableFromTheBeginningSorter only when the number of + // faster than MultipleKeyTableSorter only when the number of // sort keys is 2 and counting sort is used. So we always - // TableFromTheBeginningSorter for now. + // MultipleKeyTableSorter for now. // // TableRadixSorter sorter; // ARROW_RETURN_NOT_OK(sorter.Sort(out_begin, out_end, table, options)); - TableFromTheBeginningSorter sorter(out_begin, out_end, table, options); + MultipleKeyTableSorter sorter(out_begin, out_end, table, options); ARROW_RETURN_NOT_OK(sorter.Sort()); return Datum(out); } diff --git a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc index 9b997f7a926..1a7eb031f20 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_benchmark.cc @@ -36,28 +36,58 @@ static void ArraySortIndicesBenchmark(benchmark::State& state, state.SetItemsProcessed(state.iterations() * values->length()); } -static void ArraySortIndicesInt64Narrow(benchmark::State& state) { +static void ChunkedArraySortIndicesBenchmark( + benchmark::State& state, const std::shared_ptr& values) { + for (auto _ : state) { + ABORT_NOT_OK(SortIndices(*values).status()); + } + state.SetItemsProcessed(state.iterations() * values->length()); +} + +static void ArraySortIndicesInt64Benchmark(benchmark::State& state, int64_t min, + int64_t max) { RegressionArgs args(state); const int64_t array_size = args.size / sizeof(int64_t); auto rand = random::RandomArrayGenerator(kSeed); - - auto values = rand.Int64(array_size, -100, 100, args.null_proportion); + auto values = rand.Int64(array_size, min, max, args.null_proportion); ArraySortIndicesBenchmark(state, values); } -static void ArraySortIndicesInt64Wide(benchmark::State& state) { +static void ChunkedArraySortIndicesInt64Benchmark(benchmark::State& state, int64_t min, + int64_t max) { RegressionArgs args(state); - const int64_t array_size = args.size / sizeof(int64_t); + const int64_t n_chunks = 10; + const int64_t array_size = args.size / n_chunks / sizeof(int64_t); auto rand = random::RandomArrayGenerator(kSeed); + ArrayVector chunks; + for (int64_t i = 0; i < n_chunks; ++i) { + chunks.push_back(rand.Int64(array_size, min, max, args.null_proportion)); + } - auto min = std::numeric_limits::min(); - auto max = std::numeric_limits::max(); - auto values = rand.Int64(array_size, min, max, args.null_proportion); + ChunkedArraySortIndicesBenchmark(state, std::make_shared(chunks)); +} - ArraySortIndicesBenchmark(state, values); +static void ArraySortIndicesInt64Narrow(benchmark::State& state) { + ArraySortIndicesInt64Benchmark(state, -100, 100); +} + +static void ArraySortIndicesInt64Wide(benchmark::State& state) { + const auto min = std::numeric_limits::min(); + const auto max = std::numeric_limits::max(); + ArraySortIndicesInt64Benchmark(state, min, max); +} + +static void ChunkedArraySortIndicesInt64Narrow(benchmark::State& state) { + ChunkedArraySortIndicesInt64Benchmark(state, -100, 100); +} + +static void ChunkedArraySortIndicesInt64Wide(benchmark::State& state) { + const auto min = std::numeric_limits::min(); + const auto max = std::numeric_limits::max(); + ChunkedArraySortIndicesInt64Benchmark(state, min, max); } static void TableSortIndicesBenchmark(benchmark::State& state, @@ -160,6 +190,20 @@ BENCHMARK(ArraySortIndicesInt64Wide) ->MinTime(1.0) ->Unit(benchmark::TimeUnit::kNanosecond); +BENCHMARK(ChunkedArraySortIndicesInt64Narrow) + ->Apply(RegressionSetArgs) + ->Args({1 << 20, 100}) + ->Args({1 << 23, 100}) + ->MinTime(1.0) + ->Unit(benchmark::TimeUnit::kNanosecond); + +BENCHMARK(ChunkedArraySortIndicesInt64Wide) + ->Apply(RegressionSetArgs) + ->Args({1 << 20, 100}) + ->Args({1 << 23, 100}) + ->MinTime(1.0) + ->Unit(benchmark::TimeUnit::kNanosecond); + BENCHMARK(TableSortIndicesInt64Narrow) ->ArgsProduct({ {1 << 20}, // the number of records diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index a7ad0453fbe..b64768962ea 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -456,6 +456,7 @@ TEST_F(TestChunkedArraySortIndices, SortNull) { "[1]", }); this->AssertSortIndices(chunked_array, SortOrder::Ascending, "[1, 5, 4, 2, 0, 3]"); + this->AssertSortIndices(chunked_array, SortOrder::Descending, "[2, 4, 1, 5, 0, 3]"); } TEST_F(TestChunkedArraySortIndices, SortNaN) { @@ -465,6 +466,7 @@ TEST_F(TestChunkedArraySortIndices, SortNaN) { "[NaN, 1]", }); this->AssertSortIndices(chunked_array, SortOrder::Ascending, "[1, 6, 2, 4, 5, 0, 3]"); + this->AssertSortIndices(chunked_array, SortOrder::Descending, "[2, 1, 6, 4, 5, 0, 3]"); } // Base class for testing against random chunked array. @@ -478,7 +480,7 @@ class TestChunkedArrayRandomBase : public TestBase { void TestSortIndices(int length) { using ArrayType = typename TypeTraits::ArrayType; // We can use INSTANTIATE_TEST_SUITE_P() instead of using fors in a test. - for (auto null_probability : {0.0, 0.1, 0.5, 1.0}) { + for (auto null_probability : {0.0, 0.1, 0.5, 0.9, 1.0}) { for (auto order : {SortOrder::Ascending, SortOrder::Descending}) { for (auto num_chunks : {1, 5, 10}) { std::vector> arrays; From 8e5026acdf1cdf084f5d9fa9e47f0d1eb97c64b3 Mon Sep 17 00:00:00 2001 From: Antoine Pitrou Date: Tue, 24 Nov 2020 16:59:23 +0100 Subject: [PATCH 37/37] Remove unused variable --- cpp/src/arrow/compute/kernels/vector_sort_test.cc | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_sort_test.cc b/cpp/src/arrow/compute/kernels/vector_sort_test.cc index b64768962ea..e70fbae2819 100644 --- a/cpp/src/arrow/compute/kernels/vector_sort_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_sort_test.cc @@ -603,8 +603,7 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { // Compares two records in the same table. class Comparator : public TypeVisitor { public: - Comparator(const Table& table, const SortOptions& options) - : table_(table), options_(options) { + Comparator(const Table& table, const SortOptions& options) : options_(options) { for (const auto& sort_key : options_.sort_keys) { sort_columns_.emplace_back(table.GetColumnByName(sort_key.name).get(), sort_key.order); @@ -700,7 +699,6 @@ class TestTableSortIndicesRandom : public testing::TestWithParam { } } - const Table& table_; const SortOptions& options_; std::vector> sort_columns_; int64_t lhs_;