ARROW-18184: [C++] Improve JSON parser benchmarks (apache#14552)

benibus · pitrou · web-flow · commit e1883ae5230a · 2022-11-03T19:13:09.000+01:00
See: [ARROW-18184](https://issues.apache.org/jira/browse/ARROW-18184). For context, this is based on [pull/14100](apache#14100). This enhances the capabilities of the existing random JSON generators and adds benchmark coverage for ordered/unordered fields, missing fields, and inferred schemas over a range of field counts. Lead-authored-by: benibus <bpharks@gmx.com> Co-authored-by: Antoine Pitrou <antoine@python.org> Co-authored-by: Ben Harkins <60872452+benibus@users.noreply.github.com> Co-authored-by: Antoine Pitrou <pitrou@free.fr> Signed-off-by: Antoine Pitrou <antoine@python.org>
diff --git a/cpp/src/arrow/json/parser_benchmark.cc b/cpp/src/arrow/json/parser_benchmark.cc
@@ -17,7 +17,7 @@
 
 #include "benchmark/benchmark.h"
 
-#include <string>
+#include <unordered_set>
 
 #include "arrow/json/chunker.h"
 #include "arrow/json/options.h"
@@ -30,29 +30,62 @@
 namespace arrow {
 namespace json {
 
-std::shared_ptr<Schema> TestSchema() {
-  return schema({field("int", int32()), field("str", utf8())});
-}
-
-constexpr int seed = 0x432432;
+constexpr int kSeed = 0x432432;
 
-std::string TestJsonData(int num_rows, bool pretty = false) {
-  std::default_random_engine engine(seed);
+template <typename Input>
+std::string GenerateTestData(const Input& input, int num_rows,
+                             const GenerateOptions& options, bool pretty = false) {
+  std::default_random_engine engine(kSeed);
   std::string json;
   for (int i = 0; i < num_rows; ++i) {
     StringBuffer sb;
     Writer writer(sb);
-    ABORT_NOT_OK(Generate(TestSchema(), engine, &writer));
+    ABORT_NOT_OK(Generate(input, engine, &writer, options));
     json += pretty ? PrettyPrint(sb.GetString()) : sb.GetString();
     json += "\n";
   }
-
   return json;
 }
 
-static void BenchmarkJSONChunking(benchmark::State& state,
+template <typename Input>
+std::string GenerateTestData(const Input& input, int num_rows, bool pretty = false) {
+  return GenerateTestData(input, num_rows, GenerateOptions::Defaults(), pretty);
+}
+
+FieldVector GenerateTestFields(int num_fields, int mean_name_length) {
+  const std::shared_ptr<DataType> types[] = {boolean(), int64(), float64(), utf8()};
+
+  std::default_random_engine engine(kSeed);
+
+  std::poisson_distribution<int> length_dist(mean_name_length);
+  std::uniform_int_distribution<uint16_t> char_dist(32, 126);
+  std::uniform_int_distribution<size_t> type_dist(0, std::size(types) - 1);
+
+  std::unordered_set<std::string> names;
+  names.reserve(num_fields);
+
+  while (static_cast<int>(names.size()) < num_fields) {
+    auto length = length_dist(engine);
+    if (!length) continue;
+    std::string name(length, '\0');
+    for (auto& ch : name) ch = static_cast<char>(char_dist(engine));
+    names.emplace(std::move(name));
+  }
+
+  FieldVector fields;
+  fields.reserve(num_fields);
+  for (const auto& name : names) {
+    fields.push_back(field(name, types[type_dist(engine)]));
+  }
+
+  return fields;
+}
+
+FieldVector TestFields() { return {field("int", int32()), field("str", utf8())}; }
+
+static void BenchmarkJSONChunking(benchmark::State& state,  // NOLINT non-const reference
                                   const std::shared_ptr<Buffer>& json,
-                                  ParseOptions options) {  // NOLINT non-const reference
+                                  ParseOptions options) {
   auto chunker = MakeChunker(options);
 
   for (auto _ : state) {
@@ -61,6 +94,7 @@ static void BenchmarkJSONChunking(benchmark::State& state,
   }
 
   state.SetBytesProcessed(state.iterations() * json->size());
+  state.counters["json_size"] = static_cast<double>(json->size());
 }
 
 static void ChunkJSONPrettyPrinted(
@@ -69,9 +103,9 @@ static void ChunkJSONPrettyPrinted(
 
   auto options = ParseOptions::Defaults();
   options.newlines_in_values = true;
-  options.explicit_schema = TestSchema();
+  options.explicit_schema = schema(TestFields());
 
-  auto json = TestJsonData(num_rows, /* pretty */ true);
+  auto json = GenerateTestData(options.explicit_schema, num_rows, /*pretty=*/true);
   BenchmarkJSONChunking(state, std::make_shared<Buffer>(json), options);
 }
 
@@ -81,15 +115,15 @@ static void ChunkJSONLineDelimited(
 
   auto options = ParseOptions::Defaults();
   options.newlines_in_values = false;
-  options.explicit_schema = TestSchema();
+  options.explicit_schema = schema(TestFields());
 
-  auto json = TestJsonData(num_rows);
+  auto json = GenerateTestData(options.explicit_schema, num_rows);
   BenchmarkJSONChunking(state, std::make_shared<Buffer>(json), options);
   state.SetBytesProcessed(0);
 }
 
 static void BenchmarkJSONParsing(benchmark::State& state,  // NOLINT non-const reference
-                                 const std::shared_ptr<Buffer>& json, int32_t num_rows,
+                                 const std::shared_ptr<Buffer>& json,
                                  ParseOptions options) {
   for (auto _ : state) {
     std::unique_ptr<BlockParser> parser;
@@ -100,22 +134,23 @@ static void BenchmarkJSONParsing(benchmark::State& state,  // NOLINT non-const r
     ABORT_NOT_OK(parser->Finish(&parsed));
   }
   state.SetBytesProcessed(state.iterations() * json->size());
+  state.counters["json_size"] = static_cast<double>(json->size());
 }
 
 static void ParseJSONBlockWithSchema(
     benchmark::State& state) {  // NOLINT non-const reference
   const int32_t num_rows = 5000;
   auto options = ParseOptions::Defaults();
   options.unexpected_field_behavior = UnexpectedFieldBehavior::Error;
-  options.explicit_schema = TestSchema();
+  options.explicit_schema = schema(TestFields());
 
-  auto json = TestJsonData(num_rows);
-  BenchmarkJSONParsing(state, std::make_shared<Buffer>(json), num_rows, options);
+  auto json = GenerateTestData(options.explicit_schema, num_rows);
+  BenchmarkJSONParsing(state, std::make_shared<Buffer>(json), options);
 }
 
 static void BenchmarkJSONReading(benchmark::State& state,  // NOLINT non-const reference
-                                 const std::string& json, int32_t num_rows,
-                                 ReadOptions read_options, ParseOptions parse_options) {
+                                 const std::string& json, ReadOptions read_options,
+                                 ParseOptions parse_options) {
   for (auto _ : state) {
     std::shared_ptr<io::InputStream> input;
     ABORT_NOT_OK(MakeStream(json, &input));
@@ -127,20 +162,22 @@ static void BenchmarkJSONReading(benchmark::State& state,  // NOLINT non-const r
   }
 
   state.SetBytesProcessed(state.iterations() * json.size());
+  state.counters["json_size"] = static_cast<double>(json.size());
 }
 
 static void BenchmarkReadJSONBlockWithSchema(
-    benchmark::State& state, bool use_threads) {  // NOLINT non-const reference
+    benchmark::State& state,  // NOLINT non-const reference
+    bool use_threads) {
   const int32_t num_rows = 500000;
   auto read_options = ReadOptions::Defaults();
   read_options.use_threads = use_threads;
 
   auto parse_options = ParseOptions::Defaults();
   parse_options.unexpected_field_behavior = UnexpectedFieldBehavior::Error;
-  parse_options.explicit_schema = TestSchema();
+  parse_options.explicit_schema = schema(TestFields());
 
-  auto json = TestJsonData(num_rows);
-  BenchmarkJSONReading(state, json, num_rows, read_options, parse_options);
+  auto json = GenerateTestData(parse_options.explicit_schema, num_rows);
+  BenchmarkJSONReading(state, json, read_options, parse_options);
 }
 
 static void ReadJSONBlockWithSchemaSingleThread(
@@ -153,12 +190,47 @@ static void ReadJSONBlockWithSchemaMultiThread(
   BenchmarkReadJSONBlockWithSchema(state, true);
 }
 
+static void ParseJSONFields(benchmark::State& state) {  // NOLINT non-const reference
+  const bool ordered = !!state.range(0);
+  const bool with_schema = !!state.range(1);
+  const double sparsity = state.range(2) / 100.0;
+  const auto num_fields = static_cast<int>(state.range(3));
+
+  // This would generate approximately 400 kB of JSON data
+  int32_t num_rows = static_cast<int32_t>(2e4 / (1.0 - sparsity) / num_fields);
+  // ... however, we want enough rows to make setup/finish overhead negligible
+  num_rows = std::max<int32_t>(num_rows, 200);
+  // ... and also we want to avoid an "Exceeded maximum rows" error.
+  num_rows = std::min<int32_t>(num_rows, kMaxParserNumRows);
+  // In the end, we will empirically generate between 400 kB and 4 MB of JSON data.
+
+  auto fields = GenerateTestFields(num_fields, 10);
+
+  auto parse_options = ParseOptions::Defaults();
+  if (with_schema) {
+    parse_options.explicit_schema = schema(fields);
+    parse_options.unexpected_field_behavior = UnexpectedFieldBehavior::Error;
+  }
+
+  auto gen_options = GenerateOptions::Defaults();
+  gen_options.field_probability = 1.0 - sparsity;
+  gen_options.randomize_field_order = !ordered;
+
+  auto json = GenerateTestData(fields, num_rows, gen_options);
+  BenchmarkJSONParsing(state, std::make_shared<Buffer>(json), parse_options);
+}
+
 BENCHMARK(ChunkJSONPrettyPrinted);
 BENCHMARK(ChunkJSONLineDelimited);
 BENCHMARK(ParseJSONBlockWithSchema);
 
 BENCHMARK(ReadJSONBlockWithSchemaSingleThread);
 BENCHMARK(ReadJSONBlockWithSchemaMultiThread)->UseRealTime();
 
+BENCHMARK(ParseJSONFields)
+    // NOTE: "sparsity" is the percentage of missing fields
+    ->ArgNames({"ordered", "schema", "sparsity", "num_fields"})
+    ->ArgsProduct({{1, 0}, {1, 0}, {0, 10, 90}, {10, 100, 1000}});
+
 }  // namespace json
 }  // namespace arrow
diff --git a/cpp/src/arrow/json/test_common.h b/cpp/src/arrow/json/test_common.h
@@ -54,20 +54,34 @@ using rj::StringBuffer;
 using std::string_view;
 using Writer = rj::Writer<StringBuffer>;
 
+struct GenerateOptions {
+  // Probability of a field being written
+  double field_probability = 1.0;
+  // Probability of a value being null
+  double null_probability = 0.2;
+  // Whether to randomize the order of written fields
+  bool randomize_field_order = false;
+
+  static constexpr GenerateOptions Defaults() { return GenerateOptions{}; }
+};
+
 inline static Status OK(bool ok) { return ok ? Status::OK() : Status::Invalid(""); }
 
 template <typename Engine>
-inline static Status Generate(const std::shared_ptr<DataType>& type, Engine& e,
-                              Writer* writer);
+inline static Status Generate(
+    const std::shared_ptr<DataType>& type, Engine& e, Writer* writer,
+    const GenerateOptions& options = GenerateOptions::Defaults());
 
 template <typename Engine>
-inline static Status Generate(const std::vector<std::shared_ptr<Field>>& fields,
-                              Engine& e, Writer* writer);
+inline static Status Generate(
+    const std::vector<std::shared_ptr<Field>>& fields, Engine& e, Writer* writer,
+    const GenerateOptions& options = GenerateOptions::Defaults());
 
 template <typename Engine>
-inline static Status Generate(const std::shared_ptr<Schema>& schm, Engine& e,
-                              Writer* writer) {
-  return Generate(schm->fields(), e, writer);
+inline static Status Generate(
+    const std::shared_ptr<Schema>& schm, Engine& e, Writer* writer,
+    const GenerateOptions& options = GenerateOptions::Defaults()) {
+  return Generate(schm->fields(), e, writer, options);
 }
 
 template <typename Engine>
@@ -99,7 +113,7 @@ struct GenerateImpl {
   template <typename T>
   enable_if_base_binary<T, Status> Visit(const T&) {
     auto size = std::poisson_distribution<>{4}(e);
-    std::uniform_int_distribution<uint16_t> gen_char(32, 127);  // FIXME generate UTF8
+    std::uniform_int_distribution<uint16_t> gen_char(32, 126);  // FIXME generate UTF8
     std::string s(size, '\0');
     for (char& ch : s) ch = static_cast<char>(gen_char(e));
     return OK(writer.String(s.c_str()));
@@ -109,11 +123,13 @@ struct GenerateImpl {
   enable_if_list_like<T, Status> Visit(const T& t) {
     auto size = std::poisson_distribution<>{4}(e);
     writer.StartArray();
-    for (int i = 0; i < size; ++i) RETURN_NOT_OK(Generate(t.value_type(), e, &writer));
+    for (int i = 0; i < size; ++i) {
+      RETURN_NOT_OK(Generate(t.value_type(), e, &writer, options));
+    }
     return OK(writer.EndArray(size));
   }
 
-  Status Visit(const StructType& t) { return Generate(t.fields(), e, &writer); }
+  Status Visit(const StructType& t) { return Generate(t.fields(), e, &writer, options); }
 
   Status Visit(const DayTimeIntervalType& t) { return NotImplemented(t); }
 
@@ -135,29 +151,54 @@ struct GenerateImpl {
 
   Engine& e;
   rj::Writer<rj::StringBuffer>& writer;
+  const GenerateOptions& options;
 };
 
 template <typename Engine>
 inline static Status Generate(const std::shared_ptr<DataType>& type, Engine& e,
-                              Writer* writer) {
-  if (std::uniform_real_distribution<>{0, 1}(e) < .2) {
-    // one out of 5 chance of null, anywhere
+                              Writer* writer, const GenerateOptions& options) {
+  if (std::bernoulli_distribution(options.null_probability)(e)) {
     writer->Null();
     return Status::OK();
   }
-  GenerateImpl<Engine> visitor = {e, *writer};
+  GenerateImpl<Engine> visitor = {e, *writer, options};
   return VisitTypeInline(*type, &visitor);
 }
 
 template <typename Engine>
 inline static Status Generate(const std::vector<std::shared_ptr<Field>>& fields,
-                              Engine& e, Writer* writer) {
+                              Engine& e, Writer* writer, const GenerateOptions& options) {
   RETURN_NOT_OK(OK(writer->StartObject()));
-  for (const auto& f : fields) {
-    writer->Key(f->name().c_str());
-    RETURN_NOT_OK(Generate(f->type(), e, writer));
+
+  int num_fields = 0;
+  auto write_field = [&](const Field& f) {
+    ++num_fields;
+    writer->Key(f.name().c_str());
+    return Generate(f.type(), e, writer, options);
+  };
+
+  std::bernoulli_distribution bool_dist(options.field_probability);
+  if (options.randomize_field_order) {
+    std::vector<size_t> indices;
+    indices.reserve(static_cast<size_t>(fields.size() * options.field_probability));
+    for (size_t i = 0; i < fields.size(); ++i) {
+      if (bool_dist(e)) {
+        indices.push_back(i);
+      }
+    }
+    std::shuffle(indices.begin(), indices.end(), e);
+    for (auto i : indices) {
+      RETURN_NOT_OK(write_field(*fields[i]));
+    }
+  } else {
+    for (const auto& f : fields) {
+      if (bool_dist(e)) {
+        RETURN_NOT_OK(write_field(*f));
+      }
+    }
   }
-  return OK(writer->EndObject(static_cast<int>(fields.size())));
+
+  return OK(writer->EndObject(num_fields));
 }
 
 inline static Status MakeStream(string_view src_str,
