Skip to content

GH-46421: [C++][Acero] Asofjoin respect PauseProducing from downstream. Version 2 #51094

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
gitmodimo wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

GH-46421: [C++][Acero] Asofjoin respect PauseProducing from downstream. Version 2 #51094

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
122 changes: 108 additions & 14 deletions cpp/src/arrow/acero/asof_join_benchmark.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,21 @@
// specific language governing permissions and limitations
// under the License.

#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#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 {
Expand All @@ -33,31 +40,71 @@ 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> table;
size_t rows;
size_t bytes;
};

static Result<TableStats> MakeTable(const TableGenerationProperties& properties) {
static Result<std::shared_ptr<Table>> WithFixedSizeStringKeys(
std::shared_ptr<Table> table, int num_ids, int64_t key_bytes) {
if (key_bytes == 0) {
return table;
}

std::vector<std::string> keys;
keys.reserve(num_ids);
for (int id = 0; id < num_ids; ++id) {
std::string suffix = std::to_string(id);
if (static_cast<int64_t>(suffix.size()) > key_bytes) {
return Status::Invalid("Key size is too small for id ", id);
}
keys.emplace_back(static_cast<size_t>(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<const Int32Array&>(*chunk);
for (int64_t row = 0; row < ids.length(); ++row) {
builder.UnsafeAppend(keys[ids.Value(row)]);
}
}

std::shared_ptr<StringArray> 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<ChunkedArray>(std::move(string_keys)));
}

static Result<TableStats> MakeTable(const TableGenerationProperties& properties,
int64_t string_key_bytes) {
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Table> 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<TableStats>({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<size_t>(util::TotalBufferSize(*table));
return Result<TableStats>({std::move(table), rows, bytes});
}

static void TableJoinOverhead(benchmark::State& state,
TableGenerationProperties left_table_properties,
TableGenerationProperties right_table_properties,
int batch_size, int num_right_tables,
std::string factory_name,
std::shared_ptr<ExecNodeOptions> options) {
std::shared_ptr<ExecNodeOptions> 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;
Expand All @@ -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));
Expand All @@ -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(
Expand All @@ -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<AsofJoinNodeOptions>(
GetRepeatedOptions(int(state.range(4) + 1), kTimeCol, {kKeyCol}, tolerance));
Expand All @@ -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<AsofJoinNodeOptions>(
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.
Expand Down Expand Up @@ -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<int, int>{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
Loading