Adding support for generating data in JSON format. (apache#11778)

soumitra-st · Oct 11, 2023 · 314cc03 · 314cc03
1 parent 884dbdb
commit 314cc03
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/...r/src/main/java/org/apache/pinot/controller/recommender/data/generator/DataGenerator.java b/...r/src/main/java/org/apache/pinot/controller/recommender/data/generator/DataGenerator.java
@@ -18,6 +18,7 @@
  */
 package org.apache.pinot.controller.recommender.data.generator;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
 import java.io.File;
 import java.io.FileWriter;
 import java.io.IOException;
@@ -128,6 +129,24 @@ public void generateCsv(long totalDocs, int numFiles)
     }
   }
 
+  public void generateJson(long totalDocs, int numFiles)
+      throws IOException {
+    final int numPerFiles = (int) (totalDocs / numFiles);
+    final ObjectMapper mapper = new ObjectMapper();
+    for (int i = 0; i < numFiles; i++) {
+      try (FileWriter writer = new FileWriter(new File(_outDir, String.format("output_%d.json", i)))) {
+        for (int j = 0; j < numPerFiles; j++) {
+          Map<String, Object> row = new HashMap<>();
+          for (int k = 0; k < _genSpec.getColumns().size(); k++) {
+            String key = _genSpec.getColumns().get(k);
+            row.put(key, _generators.get(key).next());
+          }
+          writer.append(mapper.writeValueAsString(row)).append('\n');
+        }
+      }
+    }
+  }
+
   private Object serializeIfMultiValue(Object obj) {
     if (obj instanceof List) {
       return StringUtils.join((List) obj, ";");

diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/GenerateDataCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/GenerateDataCommand.java
@@ -54,6 +54,7 @@ public class GenerateDataCommand extends AbstractBaseAdminCommand implements Com
 
   private final static String FORMAT_AVRO = "avro";
   private final static String FORMAT_CSV = "csv";
+  private final static String FORMAT_JSON = "json";
 
   @CommandLine.Option(names = {"-numRecords"}, required = true, description = "Number of records to generate.")
   private int _numRecords = 0;
@@ -79,7 +80,7 @@ public class GenerateDataCommand extends AbstractBaseAdminCommand implements Com
   private boolean _help = false;
 
   @CommandLine.Option(names = {"-format"}, required = false, help = true,
-      description = "Output format ('AVRO' or 'CSV').")
+      description = "Output format ('AVRO' or 'CSV' or 'JSON').")
   private String _format = FORMAT_AVRO;
 
   @Override
@@ -150,6 +151,8 @@ public boolean execute()
       gen.generateAvro(_numRecords, _numFiles);
     } else if (FORMAT_CSV.equalsIgnoreCase(_format)) {
       gen.generateCsv(_numRecords, _numFiles);
+    } else if (FORMAT_JSON.equalsIgnoreCase(_format)) {
+      gen.generateJson(_numRecords, _numFiles);
     } else {
       throw new IllegalArgumentException(String.format("Invalid output format '%s'", _format));
     }