diff --git a/cpp/src/arrow/acero/asof_join_benchmark.cc b/cpp/src/arrow/acero/asof_join_benchmark.cc index ed2ac2258eb6..076a44fa4732 100644 --- a/cpp/src/arrow/acero/asof_join_benchmark.cc +++ b/cpp/src/arrow/acero/asof_join_benchmark.cc @@ -15,14 +15,21 @@ // specific language governing permissions and limitations // under the License. +#include +#include #include +#include +#include #include "benchmark/benchmark.h" #include "arrow/acero/options.h" #include "arrow/acero/test_util_internal.h" +#include "arrow/array/array_primitive.h" +#include "arrow/array/builder_binary.h" #include "arrow/table.h" #include "arrow/testing/future_util.h" +#include "arrow/util/byte_size.h" namespace arrow { namespace acero { @@ -33,6 +40,7 @@ const int kDefaultStart = 0; const int kDefaultEnd = 32000; const int kDefaultMinColumnVal = -10000; const int kDefaultMaxColumnVal = 10000; +const int64_t kLongStringKeyBytes = 128; struct TableStats { std::shared_ptr table; @@ -40,13 +48,50 @@ struct TableStats { size_t bytes; }; -static Result MakeTable(const TableGenerationProperties& properties) { +static Result> WithFixedSizeStringKeys( + std::shared_ptr
table, int num_ids, int64_t key_bytes) { + if (key_bytes == 0) { + return table; + } + + std::vector keys; + keys.reserve(num_ids); + for (int id = 0; id < num_ids; ++id) { + std::string suffix = std::to_string(id); + if (static_cast(suffix.size()) > key_bytes) { + return Status::Invalid("Key size is too small for id ", id); + } + keys.emplace_back(static_cast(key_bytes - suffix.size()), 'k'); + keys.back() += suffix; + } + + StringBuilder builder; + ARROW_RETURN_NOT_OK(builder.Resize(table->num_rows())); + ARROW_RETURN_NOT_OK(builder.ReserveData(table->num_rows() * key_bytes)); + for (const auto& chunk : table->GetColumnByName(kKeyCol)->chunks()) { + const auto& ids = static_cast(*chunk); + for (int64_t row = 0; row < ids.length(); ++row) { + builder.UnsafeAppend(keys[ids.Value(row)]); + } + } + + std::shared_ptr string_keys; + ARROW_RETURN_NOT_OK(builder.Finish(&string_keys)); + int key_index = table->schema()->GetFieldIndex(kKeyCol); + return table->SetColumn(key_index, field(kKeyCol, utf8()), + std::make_shared(std::move(string_keys))); +} + +static Result MakeTable(const TableGenerationProperties& properties, + int64_t string_key_bytes) { ARROW_ASSIGN_OR_RAISE(std::shared_ptr
table, MakeRandomTimeSeriesTable(properties)); - size_t row_size = sizeof(double) * (table.get()->schema()->num_fields() - 2) + - sizeof(int64_t) + sizeof(int32_t); - size_t rows = table.get()->num_rows(); - return Result({table, rows, rows * row_size}); + ARROW_ASSIGN_OR_RAISE( + table, + WithFixedSizeStringKeys(std::move(table), properties.num_ids, string_key_bytes)); + size_t rows = table->num_rows(); + size_t bytes = static_cast(util::TotalBufferSize(*table)); + return Result({std::move(table), rows, bytes}); } static void TableJoinOverhead(benchmark::State& state, @@ -54,10 +99,12 @@ static void TableJoinOverhead(benchmark::State& state, TableGenerationProperties right_table_properties, int batch_size, int num_right_tables, std::string factory_name, - std::shared_ptr options) { + std::shared_ptr options, bool use_threads, + int64_t string_key_bytes = 0) { left_table_properties.column_prefix = "lt"; left_table_properties.seed = 0; - ASSERT_OK_AND_ASSIGN(TableStats left_table_stats, MakeTable(left_table_properties)); + ASSERT_OK_AND_ASSIGN(TableStats left_table_stats, + MakeTable(left_table_properties, string_key_bytes)); size_t right_hand_rows = 0; size_t right_hand_bytes = 0; @@ -67,7 +114,8 @@ static void TableJoinOverhead(benchmark::State& state, for (int i = 0; i < num_right_tables; i++) { right_table_properties.column_prefix = "rt" + std::to_string(i); right_table_properties.seed = i + 1; - ASSERT_OK_AND_ASSIGN(TableStats right_table_stats, MakeTable(right_table_properties)); + ASSERT_OK_AND_ASSIGN(TableStats right_table_stats, + MakeTable(right_table_properties, string_key_bytes)); right_hand_rows += right_table_stats.rows; right_hand_bytes += right_table_stats.bytes; right_input_tables.push_back(std::move(right_table_stats)); @@ -86,9 +134,7 @@ static void TableJoinOverhead(benchmark::State& state, } Declaration join_node{factory_name, {input_nodes}, options}; state.ResumeTiming(); - // asof-join must currently be run synchronously as it relies on data arriving - // in-order - ASSERT_OK(DeclarationToStatus(std::move(join_node), /*use_threads=*/false)); + ASSERT_OK(DeclarationToStatus(std::move(join_node), use_threads)); } state.counters["rows_per_second"] = benchmark::Counter( @@ -113,7 +159,7 @@ AsofJoinNodeOptions GetRepeatedOptions(size_t repeat, FieldRef on_key, return AsofJoinNodeOptions(input_keys, tolerance); } -static void AsOfJoinOverhead(benchmark::State& state) { +static void AsOfJoinOverhead(benchmark::State& state, bool use_threads) { int64_t tolerance = 0; auto options = std::make_shared( GetRepeatedOptions(int(state.range(4) + 1), kTimeCol, {kKeyCol}, tolerance)); @@ -125,7 +171,28 @@ static void AsOfJoinOverhead(benchmark::State& state) { TableGenerationProperties{int(state.range(5)), int(state.range(6)), int(state.range(7)), "", kDefaultMinColumnVal, kDefaultMaxColumnVal, 0, kDefaultStart, kDefaultEnd}, - int(state.range(3)), int(state.range(4)), "asofjoin", std::move(options)); + int(state.range(3)), int(state.range(4)), "asofjoin", std::move(options), + use_threads); +} + +static void AsOfJoinKeyToleranceDensity(benchmark::State& state, bool use_threads, + int64_t string_key_bytes) { + constexpr int kColumns = 20; + constexpr int kIds = 500; + constexpr int kBatchSize = 4000; + constexpr int kNumRightTables = 1; + + auto options = std::make_shared( + GetRepeatedOptions(kNumRightTables + 1, kTimeCol, {kKeyCol}, state.range(2))); + TableJoinOverhead(state, + TableGenerationProperties{int(state.range(0)), kColumns, kIds, "", + kDefaultMinColumnVal, kDefaultMaxColumnVal, + 0, kDefaultStart, kDefaultEnd}, + TableGenerationProperties{int(state.range(1)), kColumns, kIds, "", + kDefaultMinColumnVal, kDefaultMaxColumnVal, + 0, kDefaultStart, kDefaultEnd}, + kBatchSize, kNumRightTables, "asofjoin", std::move(options), + use_threads, string_key_bytes); } // this generates the set of right hand tables to test on. @@ -163,7 +230,34 @@ void SetArgs(benchmark::internal::Benchmark* bench) { } } -BENCHMARK(AsOfJoinOverhead)->Apply(SetArgs); +void SetKeyToleranceDensityArgs(benchmark::internal::Benchmark* bench) { + bench->ArgNames({"left_freq", "right_freq", "tolerance"})->UseRealTime(); + + // A smaller frequency means a denser input. Include balanced inputs, dense left with + // sparse right (repeated right matches), and sparse left with dense right (many right + // candidates per left row). + for (const auto& [left_freq, right_freq] : + {std::pair{400, 400}, {200, 1000}, {1000, 200}}) { + for (int64_t tolerance : {-1000, 1000}) { + bench->Args({left_freq, right_freq, tolerance}); + } + } +} + +BENCHMARK_CAPTURE(AsOfJoinOverhead, serial_executor, false)->Apply(SetArgs); +BENCHMARK_CAPTURE(AsOfJoinKeyToleranceDensity, serial_executor_int32_keys, false, 0) + ->Apply(SetKeyToleranceDensityArgs); +BENCHMARK_CAPTURE(AsOfJoinKeyToleranceDensity, serial_executor_string128_keys, false, + kLongStringKeyBytes) + ->Apply(SetKeyToleranceDensityArgs); +#ifdef ARROW_ENABLE_THREADING +BENCHMARK_CAPTURE(AsOfJoinOverhead, threaded_executor, true)->Apply(SetArgs); +BENCHMARK_CAPTURE(AsOfJoinKeyToleranceDensity, threaded_executor_int32_keys, true, 0) + ->Apply(SetKeyToleranceDensityArgs); +BENCHMARK_CAPTURE(AsOfJoinKeyToleranceDensity, threaded_executor_string128_keys, true, + kLongStringKeyBytes) + ->Apply(SetKeyToleranceDensityArgs); +#endif } // namespace acero } // namespace arrow