Parametrize RDataFrame snapshot benchmarks over compression levels

root/io/io/TFile_RDFSnapshot.cxx

@@ -1,68 +1,79 @@
 #include "ROOT/RDataFrame.hxx"
 #include "ROOT/RSnapshotOptions.hxx"
+#include "TFile.h"
 #include "TRandom3.h"
+#include "TSystem.h"
 
 #include "benchmark/benchmark.h"
 #include "rootbench/RBConfig.h"
 
+#include <memory>
+#include <string>
+#include <vector>
 
-auto SetupRDF() {
-   // We create an empty data frame
-   ROOT::RDataFrame tdf(100000);
-   // We now fill it with random numbers
-   gRandom->SetSeed(1);
-   auto tdf_1 = tdf.Define("rnd", []() { return gRandom->Gaus(); });
-   return tdf_1;
-}
+// Number of doubles written per snapshot: large enough that the compression
+// dominates over the fixed overheads of the snapshot, small enough to keep
+// the total runtime of the benchmarks reasonable.
+constexpr ULong64_t kNEntries = 10000000;
 
-auto SetupRDFOptions(ROOT::RCompressionSetting::EAlgorithm::EValues alg, int level) {
-  ROOT::RDF::RSnapshotOptions options;
-  options.fCompressionAlgorithm = alg;
-  options.fCompressionLevel = level;
-  return options;
+// The input data is generated only once, so that the random number generation
+// does not contribute to the benchmarked snapshot time.
+static const std::vector<double> &GetInputData()
+{
+   static const std::vector<double> data = [] {
+      std::vector<double> v(kNEntries);
+      TRandom3 rng(1);
+      for (auto &x : v)
+         x = rng.Gaus();
+      return v;
+   }();
+   return data;
 }
 
+static void BM_TFile_RDFSnapshot(benchmark::State &state, ROOT::RCompressionSetting::EAlgorithm::EValues algo,
+                                 const std::string &algoName)
+{
+   const auto &data = GetInputData();
 
-static void BM_TFile_RDFSnapshot_ZLIB(benchmark::State &state) {
-  auto tdf = SetupRDF();
-  auto options = SetupRDFOptions(ROOT::RCompressionSetting::EAlgorithm::kZLIB, 1);
-  for (auto _ : state) {
-     //And we write out the dataset on disk
-     tdf.Snapshot("randomNumbers", "bench_data.root", {"rnd"}, options);
-  }
-}
-BENCHMARK(BM_TFile_RDFSnapshot_ZLIB)->Unit(benchmark::kMicrosecond);
+   ROOT::RDataFrame df(kNEntries);
+   auto dfWithCol = df.Define("rnd", [&data](ULong64_t entry) { return data[entry]; }, {"rdfentry_"});
 
+   ROOT::RDF::RSnapshotOptions options;
+   options.fCompressionAlgorithm = algo;
+   options.fCompressionLevel = state.range(0);
 
-static void BM_TFile_RDFSnapshot_LZ4(benchmark::State &state) {
-  auto tdf = SetupRDF();
-  auto options = SetupRDFOptions(ROOT::RCompressionSetting::EAlgorithm::kLZ4, 4);
-  for (auto _ : state) {
-     //And we write out the dataset on disk
-     tdf.Snapshot("randomNumbers", "bench_data.root", {"rnd"}, options);
-  }
-}
-BENCHMARK(BM_TFile_RDFSnapshot_LZ4)->Unit(benchmark::kMicrosecond);
+   const std::string fileName =
+      RB::GetTempFs() + "/rdfsnapshot_" + algoName + "_" + std::to_string(state.range(0)) + ".root";
 
+   for (auto _ : state) {
+      // And we write out the dataset on disk
+      dfWithCol.Snapshot("randomNumbers", fileName, {"rnd"}, options);
+   }
 
-static void BM_TFile_RDFSnapshot_LZMA (benchmark::State &state) {
-  auto tdf = SetupRDF();
-  auto options = SetupRDFOptions(ROOT::RCompressionSetting::EAlgorithm::kLZMA, 8);
-  for (auto _ : state) {
-     //And we write out the dataset on disk
-     tdf.Snapshot("randomNumbers", "bench_data.root", {"rnd"}, options);
-  }
-}
-BENCHMARK(BM_TFile_RDFSnapshot_LZMA)->Unit(benchmark::kMicrosecond);
-
-static void BM_TFile_RDFSnapshot_ZSTD (benchmark::State &state) {
-  auto tdf = SetupRDF();
-  auto options = SetupRDFOptions(ROOT::RCompressionSetting::EAlgorithm::kZSTD, 6);
-  for (auto _ : state) {
-     //And we write out the dataset on disk
-     tdf.Snapshot("randomNumbers", "bench_data.root", {"rnd"}, options);
-  }
+   {
+      std::unique_ptr<TFile> file{TFile::Open(fileName.c_str())};
+      state.counters["comp_size"] = file->GetSize();
+   }
+   // Reports the throughput of uncompressed input bytes as bytes_per_second.
+   state.SetBytesProcessed(state.iterations() * kNEntries * sizeof(double));
+
+   gSystem->Unlink(fileName.c_str());
 }
-BENCHMARK(BM_TFile_RDFSnapshot_ZSTD)->Unit(benchmark::kMicrosecond);
+
+BENCHMARK_CAPTURE(BM_TFile_RDFSnapshot, ZLIB, ROOT::RCompressionSetting::EAlgorithm::kZLIB, "zlib")
+->Arg(1)->Arg(6)->Arg(9)
+->Unit(benchmark::kMillisecond);
+
+BENCHMARK_CAPTURE(BM_TFile_RDFSnapshot, LZMA, ROOT::RCompressionSetting::EAlgorithm::kLZMA, "lzma")
+->Arg(1)->Arg(6)->Arg(9)
+->Unit(benchmark::kMillisecond);
+
+BENCHMARK_CAPTURE(BM_TFile_RDFSnapshot, LZ4, ROOT::RCompressionSetting::EAlgorithm::kLZ4, "lz4")
+->Arg(1)->Arg(6)->Arg(9)
+->Unit(benchmark::kMillisecond);
+
+BENCHMARK_CAPTURE(BM_TFile_RDFSnapshot, ZSTD, ROOT::RCompressionSetting::EAlgorithm::kZSTD, "zstd")
+->Arg(1)->Arg(6)->Arg(9)
+->Unit(benchmark::kMillisecond);
 
 BENCHMARK_MAIN();