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.

4 files changed, 94 insertions, 9 deletions

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());
diff --git a/vespaclient-java/src/test/java/com/yahoo/vespavisit/StdOutVisitorHandlerTest.java b/vespaclient-java/src/test/java/com/yahoo/vespavisit/StdOutVisitorHandlerTest.java
index f109b0ad56b..a2e9f91d503 100644
--- a/vespaclient-java/src/test/java/com/yahoo/vespavisit/StdOutVisitorHandlerTest.java
+++ b/vespaclient-java/src/test/java/com/yahoo/vespavisit/StdOutVisitorHandlerTest.java
@@ -1,7 +1,18 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespavisit;
 
+import com.yahoo.document.Document;
+import com.yahoo.document.DocumentPut;
+import com.yahoo.document.DocumentType;
+import com.yahoo.document.TensorDataType;
+import com.yahoo.document.datatypes.TensorFieldValue;
+import com.yahoo.documentapi.AckToken;
+import com.yahoo.documentapi.VisitorControlSession;
 import com.yahoo.documentapi.VisitorDataHandler;
+import com.yahoo.documentapi.messagebus.protocol.PutDocumentMessage;
+import com.yahoo.tensor.Tensor;
+import com.yahoo.tensor.TensorType;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 
@@ -9,6 +20,7 @@ import java.io.ByteArrayOutputStream;
 import java.io.PrintStream;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.mock;
 
 /**
  * @author bjorncs
@@ -30,7 +42,7 @@ public class StdOutVisitorHandlerTest {
         initStdOutVisitorHandlerTest(jsonOutput);
         ByteArrayOutputStream out = new ByteArrayOutputStream();
         StdOutVisitorHandler visitorHandler =
-                new StdOutVisitorHandler(/*printIds*/true, false, false, false, false, false, 0, jsonOutput, new PrintStream(out, true));
+                new StdOutVisitorHandler(/*printIds*/true, false, false, false, false, false, 0, jsonOutput, false, new PrintStream(out, true));
         VisitorDataHandler dataHandler = visitorHandler.getDataHandler();
         dataHandler.onDone();
         String output = out.toString();
@@ -43,11 +55,49 @@ public class StdOutVisitorHandlerTest {
         initStdOutVisitorHandlerTest(jsonOutput);
         ByteArrayOutputStream out = new ByteArrayOutputStream();
         StdOutVisitorHandler visitorHandler =
-                new StdOutVisitorHandler(/*printIds*/false, false, false, false, false, false, 0, jsonOutput, new PrintStream(out, true));
+                new StdOutVisitorHandler(/*printIds*/false, false, false, false, false, false, 0, jsonOutput, false, new PrintStream(out, true));
         VisitorDataHandler dataHandler = visitorHandler.getDataHandler();
         dataHandler.onDone();
         String expectedOutput = jsonOutput ? "[]" : "";
         String output = out.toString().trim();
         assertEquals(expectedOutput, output);
     }
-}
\ No newline at end of file
+
+    void do_test_json_tensor_fields_can_be_output_in_short_or_long_form(boolean tensorShortForm, String expectedOutput) {
+        var docType = new DocumentType("foo");
+        docType.addField("bar", TensorDataType.getTensor(TensorType.fromSpec("tensor(x[3])")));
+        var doc = new Document(docType, "id:baz:foo::tensor-stuff");
+        doc.setFieldValue("bar", new TensorFieldValue(Tensor.from("tensor(x[3]):[1,2,3]")));
+        var putMsg = new PutDocumentMessage(new DocumentPut(doc));
+
+        var out = new ByteArrayOutputStream();
+        var visitorHandler = new StdOutVisitorHandler(/*printIds*/false, false, false, false, false, false,
+                                                      0, true, tensorShortForm, new PrintStream(out, true));
+        var dataHandler = visitorHandler.getDataHandler();
+        var controlSession = mock(VisitorControlSession.class);
+        var ackToken = mock(AckToken.class);
+        dataHandler.setSession(controlSession);
+        dataHandler.onMessage(putMsg, ackToken);
+        dataHandler.onDone();
+
+        String output = out.toString().trim();
+        assertEquals(expectedOutput, output);
+    }
+
+    @Test
+    void json_tensor_fields_can_be_output_in_long_form() {
+        var expectedOutput = """
+        [
+        {"id":"id:baz:foo::tensor-stuff","fields":{"bar":{"cells":[{"address":{"x":"0"},"value":1.0},{"address":{"x":"1"},"value":2.0},{"address":{"x":"2"},"value":3.0}]}}}]""";
+        do_test_json_tensor_fields_can_be_output_in_short_or_long_form(false, expectedOutput);
+    }
+
+    @Test
+    void json_tensor_fields_can_be_output_in_short_form() {
+        var expectedOutput = """
+        [
+        {"id":"id:baz:foo::tensor-stuff","fields":{"bar":{"type":"tensor(x[3])","values":[1.0,2.0,3.0]}}}]""";
+        do_test_json_tensor_fields_can_be_output_in_short_or_long_form(true, expectedOutput);
+    }
+
+}
diff --git a/vespaclient-java/src/test/java/com/yahoo/vespavisit/VdsVisitTestCase.java b/vespaclient-java/src/test/java/com/yahoo/vespavisit/VdsVisitTestCase.java
index a67fee0c0d6..e0b90592c37 100644
--- a/vespaclient-java/src/test/java/com/yahoo/vespavisit/VdsVisitTestCase.java
+++ b/vespaclient-java/src/test/java/com/yahoo/vespavisit/VdsVisitTestCase.java
@@ -24,6 +24,7 @@ import java.util.List;
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -138,7 +139,8 @@ public class VdsVisitTestCase {
                 "--skipbucketsonfatalerrors",
                 "--abortonclusterdown",
                 "--visitremoves",
-                "--bucketspace", "outerspace"
+                "--bucketspace", "outerspace",
+                "--tensor-short-form"
         };
         VdsVisit.ArgumentParser parser = createMockArgumentParser();
         VdsVisit.VdsVisitParameters allParams = parser.parse(args);
@@ -158,6 +160,7 @@ public class VdsVisitTestCase {
         assertEquals(123456789, params.getTimeoutMs());
         assertEquals(7 * 24 * 60 * 60 * 1000, allParams.getFullTimeout());
         assertEquals("kittens", allParams.getCluster());
+        assertTrue(allParams.tensorShortForm());
 
         assertTrue(params.getThrottlePolicy() instanceof StaticThrottlePolicy);
         assertEquals(3, ((StaticThrottlePolicy) params.getThrottlePolicy()).getMaxPendingCount());
@@ -201,6 +204,13 @@ public class VdsVisitTestCase {
 
     private static String[] emptyArgList() { return new String[]{}; }
 
+    // TODO Vespa 9: change default from long to short
+    @Test
+    void tensor_output_format_is_long_by_default() throws Exception {
+        var allParams = createMockArgumentParser().parse(emptyArgList());
+        assertFalse(allParams.tensorShortForm());
+    }
+
     @Test
     void visitor_priority_is_low1_by_default() throws Exception {
         VdsVisit.VdsVisitParameters allParams = createMockArgumentParser().parse(emptyArgList());