feat: cli can now output JSON, TEXT and BINARY proto format

substrait-io · Dec 26, 2023 · 28b8c11 · 28b8c11
1 parent 6ae7b27
commit 28b8c11
Showing 1 changed file with 19 additions and 1 deletion.
diff --git a/isthmus/src/main/java/io/substrait/isthmus/PlanEntryPoint.java b/isthmus/src/main/java/io/substrait/isthmus/PlanEntryPoint.java
@@ -5,6 +5,7 @@
 import static picocli.CommandLine.Parameters;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.protobuf.TextFormat;
 import com.google.protobuf.util.JsonFormat;
 import io.substrait.isthmus.SubstraitRelVisitor.CrossJoinPolicy;
 import io.substrait.proto.Plan;
@@ -35,6 +36,18 @@ public class PlanEntryPoint implements Callable<Integer> {
       description = "Allow multiple statements terminated with a semicolon")
   private boolean allowMultiStatement;
 
+  @Option(
+      names = {"--outputformat"},
+      defaultValue = "JSON",
+      description = "Set the output format for the generated plan: ${COMPLETION-CANDIDATES}")
+  private OutputFormat outputFormat = OutputFormat.JSON;
+
+  enum OutputFormat {
+    JSON, // protobuf JSON format
+    TEXT, // protobuf TEXT format
+    BINARY, // protobuf BINARY format
+  }
+
   @Option(
       names = {"--sqlconformancemode"},
       description = "One of built-in Calcite SQL compatibility modes: ${COMPLETION-CANDIDATES}")
@@ -57,7 +70,12 @@ public Integer call() throws Exception {
     FeatureBoard featureBoard = buildFeatureBoard();
     SqlToSubstrait converter = new SqlToSubstrait(featureBoard);
     Plan plan = converter.execute(sql, createStatements);
-    System.out.println(JsonFormat.printer().includingDefaultValueFields().print(plan));
+    switch (outputFormat) {
+      case JSON -> System.out.println(
+          JsonFormat.printer().includingDefaultValueFields().print(plan));
+      case TEXT -> TextFormat.printer().print(plan, System.out);
+      case BINARY -> plan.writeTo(System.out);
+    }
     return 0;
   }