From 23e9f20ebc1b0d013bd3f95a9463a9fe92a2cb44 Mon Sep 17 00:00:00 2001
From: Tor Brede Vekterli <vekterli@yahooinc.com>
Date: Mon, 5 Dec 2022 13:56:34 +0100
Subject: Add tensor short form output option to vespa-visit tool

Specified with `--tensor-short-form`. No single-char option alias,
as short form output will be the default on Vespa 9 and we're running
out of usable option characters for this tool anyway.
---
 .../com/yahoo/vespavisit/StdOutVisitorHandler.java | 13 +++++++++----
 .../main/java/com/yahoo/vespavisit/VdsVisit.java   | 22 +++++++++++++++++++++-
 2 files changed, 30 insertions(+), 5 deletions(-)

(limited to 'vespaclient-java/src/main/java/com/yahoo')

diff --git a/vespaclient-java/src/main/java/com/yahoo/vespavisit/StdOutVisitorHandler.java b/vespaclient-java/src/main/java/com/yahoo/vespavisit/StdOutVisitorHandler.java
index ee8aa903d33..2ac0510a2a3 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespavisit/StdOutVisitorHandler.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespavisit/StdOutVisitorHandler.java
@@ -40,25 +40,30 @@ public class StdOutVisitorHandler extends VdsVisitHandler {
     private int processTimeMilliSecs;
     private PrintStream out;
     private final boolean jsonOutput;
+    private final boolean tensorShortForm;
 
     private VisitorDataHandler dataHandler;
 
     public StdOutVisitorHandler(boolean printIds, boolean indentXml,
                                 boolean showProgress, boolean showStatistics, boolean doStatistics,
-                                boolean abortOnClusterDown, int processtime, boolean jsonOutput)
+                                boolean abortOnClusterDown, int processtime, boolean jsonOutput,
+                                boolean tensorShortForm)
     {
-        this(printIds, indentXml, showProgress, showStatistics, doStatistics, abortOnClusterDown, processtime, jsonOutput, createStdOutPrintStream());
+        this(printIds, indentXml, showProgress, showStatistics, doStatistics, abortOnClusterDown, processtime,
+             jsonOutput, tensorShortForm, createStdOutPrintStream());
     }
 
     StdOutVisitorHandler(boolean printIds, boolean indentXml,
                          boolean showProgress, boolean showStatistics, boolean doStatistics,
-                         boolean abortOnClusterDown, int processtime, boolean jsonOutput, PrintStream out)
+                         boolean abortOnClusterDown, int processtime, boolean jsonOutput,
+                         boolean tensorShortForm, PrintStream out)
     {
         super(showProgress, showStatistics, abortOnClusterDown);
         this.printIds = printIds;
         this.indentXml = indentXml;
         this.processTimeMilliSecs = processtime;
         this.jsonOutput = jsonOutput;
+        this.tensorShortForm = tensorShortForm;
         this.out = out;
         this.dataHandler = new DataHandler(doStatistics);
     }
@@ -169,7 +174,7 @@ public class StdOutVisitorHandler extends VdsVisitHandler {
 
         private void writeJsonDocument(Document doc) throws IOException {
             writeFeedStartOrRecordSeparator();
-            out.write(JsonWriter.toByteArray(doc));
+            out.write(JsonWriter.toByteArray(doc, tensorShortForm));
         }
 
         @Override
diff --git a/vespaclient-java/src/main/java/com/yahoo/vespavisit/VdsVisit.java b/vespaclient-java/src/main/java/com/yahoo/vespavisit/VdsVisit.java
index 9a5df96b705..c022713cbb6 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespavisit/VdsVisit.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespavisit/VdsVisit.java
@@ -348,6 +348,13 @@ public class VdsVisit {
                         FixedBucketSpaces.defaultSpace(), FixedBucketSpaces.globalSpace(), FixedBucketSpaces.defaultSpace()))
                 .build());
 
+        // TODO Vespa 9 replace with --tensor-long-form ?
+        options.addOption(Option.builder()
+                .longOpt("tensor-short-form")
+                .desc("Output JSON tensors in short form. Will be the default on Vespa 9")
+                .hasArg(false)
+                .build());
+
         return options;
     }
 
@@ -363,6 +370,7 @@ public class VdsVisit {
         private int processTime = 0;
         private int fullTimeout = 7 * 24 * 60 * 60 * 1000;
         private boolean jsonOutput = false;
+        private boolean tensorShortForm = false; // TODO Vespa 9: change default to true
 
         public VisitorParameters getVisitorParameters() {
             return visitorParameters;
@@ -431,6 +439,14 @@ public class VdsVisit {
         public void setJsonOutput(boolean jsonOutput) {
             this.jsonOutput = jsonOutput;
         }
+
+        public boolean tensorShortForm() {
+            return tensorShortForm;
+        }
+
+        public void setTensorShortForm(boolean tensorShortForm) {
+            this.tensorShortForm = tensorShortForm;
+        }
     }
 
     protected static class ArgumentParser {
@@ -553,6 +569,9 @@ public class VdsVisit {
                 throttlePolicy.setMaxPendingCount(((Number)line.getParsedOptionValue("maxpendingsuperbuckets")).intValue());
                 params.setThrottlePolicy(throttlePolicy);
             }
+            if (line.hasOption("tensor-short-form")) {
+                allParams.setTensorShortForm(true);
+            }
 
             boolean jsonOutput = line.hasOption("jsonoutput");
             boolean xmlOutput = line.hasOption("xmloutput");
@@ -723,7 +742,8 @@ public class VdsVisit {
                 params.getStatisticsParts() != null,
                 params.getAbortOnClusterDown(),
                 params.getProcessTime(),
-                params.jsonOutput);
+                params.jsonOutput,
+                params.tensorShortForm);
 
         if (visitorParameters.getResumeFileName() != null) {
             handler.setProgressFileName(visitorParameters.getResumeFileName());
-- 
cgit v1.2.3