ARROW-16894: [C++] Add Benchmarks for Asof Join Node

cpp/src/arrow/compute/exec/CMakeLists.txt

@@ -62,6 +62,8 @@ add_arrow_benchmark(project_benchmark
                     benchmark_util.cc
                     project_benchmark.cc)
 
+add_arrow_benchmark(asof_join_benchmark PREFIX "arrow-compute")
+
 add_arrow_benchmark(tpch_benchmark PREFIX "arrow-compute")
 
 if(ARROW_BUILD_OPENMP_BENCHMARKS)

cpp/src/arrow/compute/exec/asof_join_benchmark.cc

@@ -0,0 +1,171 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <string>
+
+#include "benchmark/benchmark.h"
+
+#include "arrow/compute/exec/test_util.h"
+#include "arrow/dataset/file_parquet.h"
+#include "arrow/table.h"
+#include "arrow/testing/future_util.h"
+
+namespace arrow {
+namespace compute {
+
+static const char* kTimeCol = "time";
+static const char* kKeyCol = "id";
+const int kDefaultStart = 0;
+const int kDefaultEnd = 500;
+const int kDefaultMinColumnVal = -10000;
+const int kDefaultMaxColumnVal = 10000;
+
+struct TableStats {
+  std::shared_ptr<Table> table;
+  size_t total_rows;
+  size_t total_bytes;
+};
+
+static TableStats MakeTable(const TableGenerationProperties& properties) {
+  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 {table, rows, rows * row_size};
+}
+
+static ExecNode* MakeTableSourceNode(std::shared_ptr<arrow::compute::ExecPlan> plan,
+                                     std::shared_ptr<Table> table, int batch_size) {
+  return *arrow::compute::MakeExecNode(
+      "table_source", plan.get(), {},
+      arrow::compute::TableSourceNodeOptions(table, batch_size));
+}
+
+static void TableJoinOverhead(benchmark::State& state,
+                              TableGenerationProperties left_table_properties,
+                              int left_table_batch_size,
+                              TableGenerationProperties right_table_properties,
+                              int right_table_batch_size, int num_right_tables,
+                              std::string factory_name, ExecNodeOptions& options) {
+  ExecContext ctx(default_memory_pool(), nullptr);
+
+  left_table_properties.column_prefix = "lt";
+  left_table_properties.seed = 0;
+  TableStats left_table_stats = MakeTable(left_table_properties);
+
+  size_t right_hand_rows = 0;
+  size_t right_hand_bytes = 0;
+  std::vector<TableStats> right_input_tables;
+  right_input_tables.reserve(num_right_tables);
+
+  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;
+    TableStats right_table_stats = MakeTable(right_table_properties);
+    right_hand_rows += right_table_stats.total_rows;
+    right_hand_bytes += right_table_stats.total_bytes;
+    right_input_tables.push_back(right_table_stats);
+  }
+
+  for (auto _ : state) {
+    state.PauseTiming();
+    ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::compute::ExecPlan> plan,
+                         ExecPlan::Make(&ctx));
+    std::vector<ExecNode*> input_nodes = {
+        MakeTableSourceNode(plan, left_table_stats.table, left_table_batch_size)};
+    input_nodes.reserve(right_input_tables.size() + 1);
+    for (TableStats table_stats : right_input_tables) {
+      input_nodes.push_back(
+          MakeTableSourceNode(plan, table_stats.table, right_table_batch_size));
+    }
+    ASSERT_OK_AND_ASSIGN(arrow::compute::ExecNode * join_node,
+                         MakeExecNode(factory_name, plan.get(), input_nodes, options));
+    AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+    MakeExecNode("sink", plan.get(), {join_node}, SinkNodeOptions{&sink_gen});
+    state.ResumeTiming();
+    ASSERT_FINISHES_OK(StartAndCollect(plan.get(), sink_gen));
+  }
+
+  state.counters["total_rows_per_second"] = benchmark::Counter(
+      static_cast<double>(state.iterations() *
+                          (left_table_stats.total_rows + right_hand_rows)),
+      benchmark::Counter::kIsRate);
+
+  state.counters["total_bytes_per_second"] = benchmark::Counter(
+      static_cast<double>(state.iterations() *
+                          (left_table_stats.total_bytes + right_hand_bytes)),
+      benchmark::Counter::kIsRate);
+
+  state.counters["maximum_peak_memory"] =
+      benchmark::Counter(static_cast<double>(ctx.memory_pool()->max_memory()));
+}
+
+static void AsOfJoinOverhead(benchmark::State& state) {
+  int64_t tolerance = 0;
+  AsofJoinNodeOptions options = AsofJoinNodeOptions(kTimeCol, kKeyCol, tolerance);
+  TableJoinOverhead(
+      state,
+      TableGenerationProperties{int(state.range(0)), int(state.range(1)),
+                                int(state.range(2)), "", kDefaultMinColumnVal,
+                                kDefaultMaxColumnVal, 0, kDefaultStart, kDefaultEnd},
+      int(state.range(3)),
+      TableGenerationProperties{int(state.range(5)), int(state.range(6)),
+                                int(state.range(7)), "", kDefaultMinColumnVal,
+                                kDefaultMaxColumnVal, 0, kDefaultStart, kDefaultEnd},
+      int(state.range(8)), int(state.range(4)), "asofjoin", options);
+}
+
+// this generates the set of right hand tables to test on.
+void SetArgs(benchmark::internal::Benchmark* bench) {
+  bench
+      ->ArgNames({"left_freq", "left_cols", "left_ids", "left_batch_size",
+                  "num_right_tables", "right_freq", "right_cols", "right_ids",
+                  "right_batch_size"})
+      ->UseRealTime();
+  int default_freq = 5;
+  int default_cols = 20;
+  int default_ids = 500;
+  int default_num_tables = 1;
+  int default_batch_size = 100;
+
+  for (int freq : {1, 5, 10}) {
+    bench->Args({freq, default_cols, default_ids, default_batch_size, default_num_tables,
+                 freq, default_cols, default_ids, default_batch_size});
+  }
+
+  for (int cols : {10, 20, 100}) {
+    bench->Args({default_freq, cols, default_ids, default_batch_size, default_num_tables,
+                 default_freq, cols, default_ids, default_batch_size});
+  }
+  for (int ids : {100, 500, 1000}) {
+    bench->Args({default_freq, default_cols, ids, default_batch_size, default_num_tables,
+                 default_freq, default_cols, ids, default_batch_size});
+  }
+  for (int num_tables : {1, 10, 50}) {
+    bench->Args({default_freq, default_cols, default_ids, default_batch_size, num_tables,
+                 default_freq, default_cols, default_ids, default_batch_size});
+  }
+  for (int batch_size : {1, 500, 1000}) {
+    bench->Args({default_freq, default_cols, default_ids, batch_size, default_num_tables,
+                 default_freq, default_cols, default_ids, batch_size});
+  }
+}
+
+BENCHMARK(AsOfJoinOverhead)->Apply(SetArgs);
+
+}  // namespace compute
+}  // namespace arrow

cpp/src/arrow/compute/exec/test_util.cc

@@ -40,6 +40,7 @@
 #include "arrow/datum.h"
 #include "arrow/record_batch.h"
 #include "arrow/table.h"
+#include "arrow/testing/builder.h"
 #include "arrow/testing/gtest_util.h"
 #include "arrow/testing/random.h"
 #include "arrow/type.h"
@@ -459,5 +460,41 @@ void PrintTo(const Declaration& decl, std::ostream* os) {
   *os << "}";
 }
 
+std::shared_ptr<Table> MakeRandomTimeSeriesTable(
+    const TableGenerationProperties& properties) {
+  int total_columns = properties.num_columns + 2;
+  std::vector<std::shared_ptr<Array>> columns;
+  columns.reserve(total_columns);
+  arrow::FieldVector field_vector;
+  field_vector.reserve(total_columns);
+
+  field_vector.push_back(field("time", int64()));
+  field_vector.push_back(field("id", int32()));
+
+  Int64Builder time_column_builder;
+  Int32Builder id_column_builder;
+  for (int time = properties.start; time <= properties.end;
+       time += properties.time_frequency) {
+    for (int id = 0; id < properties.num_ids; id++) {
+      time_column_builder.Append(time);
+      id_column_builder.Append(id);
+    }
+  }
+
+  int num_rows = time_column_builder.length();
+  columns.push_back(time_column_builder.Finish().ValueOrDie());
+  columns.push_back(id_column_builder.Finish().ValueOrDie());
+
+  for (int i = 0; i < properties.num_columns; i++) {
+    field_vector.push_back(
+        field(properties.column_prefix + std::to_string(i), float64()));
+    random::RandomArrayGenerator rand = random::RandomArrayGenerator(properties.seed + i);
+    columns.push_back(
+        rand.Float64(num_rows, properties.min_column_value, properties.max_column_value));
+  }
+  std::shared_ptr<arrow::Schema> schema = arrow::schema(std::move(field_vector));
+  return Table::Make(schema, columns, num_rows);
+}
+
 }  // namespace compute
 }  // namespace arrow

cpp/src/arrow/compute/exec/test_util.h

@@ -145,5 +145,39 @@ class Random64Bit {
   std::uniform_int_distribution<uint64_t> dist_;
 };
 
+// Specify properties of a table to be generated.
+struct TableGenerationProperties {
+  /// Indicates the amount of time between data points that lie between
+  /// the start and end parameters.
+  int time_frequency;
+  /// The number of additional random columns in the table.
+  int num_columns;
+  /// The number of unique keys in the table.
+  int num_ids;
+  /// Specifies the prefix of each randomly generated column.
+  std::string column_prefix;
+  /// Specifies the minimum value in the randomly generated column(s).
+  int min_column_value;
+  /// Specifies the maximum value in the randomly generated column(s).
+  int max_column_value;
+  /// The random seed the random array generator is given to generate the additional
+  /// columns.
+  uint seed;
+  /// Specifies the beginning of 'time' recorded in the table, inclusive.
+  int start;
+  /// Specifies the end of 'time' recorded in the table, inclusive.
+  int end;
+};
+
+/// The table generated in accordance to the TableGenerationProperties has the following
+/// schema: time (int64) id (int32) [properties.column_prefix]0 (float64)
+///     [properties.column_prefix]1 (float64)
+///     ...
+///     [properties.column_prefix][properties.num_columns] (float64)
+/// Each id has rows corresponding to a singular data point in the time range (start, end,
+/// time_frequency). The table is sorted by time.
+std::shared_ptr<Table> MakeRandomTimeSeriesTable(
+    const TableGenerationProperties& properties);
+
 }  // namespace compute
 }  // namespace arrow