Merge pull request #25549 from vespa-engine/bratseth/tensor-direct-values

Parse tensor JSON values at root

18 files changed, 197 insertions, 148 deletions

diff --git a/document/src/main/java/com/yahoo/document/json/JsonSerializationHelper.java b/document/src/main/java/com/yahoo/document/json/JsonSerializationHelper.java
index d33cc8078dd..7f6ead528fe 100644
--- a/document/src/main/java/com/yahoo/document/json/JsonSerializationHelper.java
+++ b/document/src/main/java/com/yahoo/document/json/JsonSerializationHelper.java
@@ -48,6 +48,7 @@ import java.util.Set;
  * @author Vegard Sjonfjell
  */
 public class JsonSerializationHelper {
+
     private final static Base64.Encoder base64Encoder = Base64.getEncoder(); // Important: _basic_ format
 
     static class JsonSerializationException extends RuntimeException {
@@ -99,14 +100,6 @@ public class JsonSerializationHelper {
         });
     }
 
-    private static void serializeTensorDimensions(JsonGenerator generator, Set<String> dimensions) throws IOException {
-        generator.writeArrayFieldStart(TensorReader.TENSOR_DIMENSIONS);
-        for (String dimension : dimensions) {
-            generator.writeString(dimension);
-        }
-        generator.writeEndArray();
-    }
-
     static void serializeTensorCells(JsonGenerator generator, Tensor tensor) throws IOException {
         generator.writeArrayFieldStart(TensorReader.TENSOR_CELLS);
         for (Map.Entry<TensorAddress, Double> cell : tensor.cells().entrySet()) {
diff --git a/document/src/main/java/com/yahoo/document/json/TokenBuffer.java b/document/src/main/java/com/yahoo/document/json/TokenBuffer.java
index 6b2bdbd53d8..6bee48ea2d7 100644
--- a/document/src/main/java/com/yahoo/document/json/TokenBuffer.java
+++ b/document/src/main/java/com/yahoo/document/json/TokenBuffer.java
@@ -2,9 +2,9 @@
 package com.yahoo.document.json;
 
 import java.io.IOException;
-import java.util.ArrayDeque;
-import java.util.Deque;
+import java.util.ArrayList;
 import java.util.Iterator;
+import java.util.List;
 
 import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.core.JsonToken;
@@ -17,60 +17,73 @@ import com.google.common.base.Preconditions;
  */
 public class TokenBuffer {
 
-    private final Deque<Token> buffer;
+    private final List<Token> tokens;
+
+    private int position = 0;
     private int nesting = 0;
 
     public TokenBuffer() {
-        this(new ArrayDeque<>());
+        this(new ArrayList<>());
     }
 
-    private TokenBuffer(Deque<Token> buffer) {
-        this.buffer = buffer;
-        if (buffer.size() > 0) {
-            updateNesting(buffer.peekFirst().token);
-        }
+    public TokenBuffer(List<Token> tokens) {
+        this.tokens = tokens;
+        if (tokens.size() > 0)
+            updateNesting(tokens.get(position).token);
     }
 
     /** Returns whether any tokens are available in this */
-    public boolean isEmpty() { return size() == 0; }
+    public boolean isEmpty() { return remaining() == 0; }
 
-    public JsonToken next() {
-        buffer.removeFirst();
-        Token t = buffer.peekFirst();
-        if (t == null) {
-            return null;
-        }
-        updateNesting(t.token);
-        return t.token;
+    public JsonToken previous() {
+        updateNestingGoingBackwards(current());
+        position--;
+        return current();
     }
 
     /** Returns the current token without changing position, or null if none */
-    public JsonToken currentToken() {
-        Token token = buffer.peekFirst();
+    public JsonToken current() {
+        if (isEmpty()) return null;
+        Token token = tokens.get(position);
         if (token == null) return null;
         return token.token;
     }
 
+    public JsonToken next() {
+        position++;
+        JsonToken token = current();
+        updateNesting(token);
+        return token;
+    }
+
+    /** Returns a given number of tokens ahead, or null if none */
+    public JsonToken peek(int ahead) {
+        if (tokens.size() <= position + ahead) return null;
+        return tokens.get(position + ahead).token;
+    }
+
     /** Returns the current token name without changing position, or null if none */
     public String currentName() {
-        Token token = buffer.peekFirst();
+        if (isEmpty()) return null;
+        Token token = tokens.get(position);
         if (token == null) return null;
         return token.name;
     }
 
     /** Returns the current token text without changing position, or null if none */
     public String currentText() {
-        Token token = buffer.peekFirst();
+        if (isEmpty()) return null;
+        Token token = tokens.get(position);
         if (token == null) return null;
         return token.text;
     }
 
-    public int size() {
-        return buffer.size();
+    public int remaining() {
+        return tokens.size() - position;
     }
 
     private void add(JsonToken token, String name, String text) {
-        buffer.addLast(new Token(token, name, text));
+        tokens.add(tokens.size(), new Token(token, name, text));
     }
 
     public void bufferObject(JsonToken first, JsonParser tokens) {
@@ -83,7 +96,7 @@ public class TokenBuffer {
 
         Preconditions.checkArgument(first == firstToken,
                 "Expected %s, got %s.", firstToken.name(), t);
-        if (size() == 0) {
+        if (remaining() == 0) {
             updateNesting(t);
         }
         localNesting = storeAndPeekNesting(t, localNesting, tokens);
@@ -98,10 +111,11 @@ public class TokenBuffer {
         return nesting + nestingOffset(t);
     }
 
-    private int nestingOffset(JsonToken t) {
-        if (t.isStructStart()) {
+    private int nestingOffset(JsonToken token) {
+        if (token == null) return 0;
+        if (token.isStructStart()) {
             return 1;
-        } else if (t.isStructEnd()) {
+        } else if (token.isStructEnd()) {
             return -1;
         } else {
             return 0;
@@ -112,7 +126,6 @@ public class TokenBuffer {
         try {
             add(t, tokens.getCurrentName(), tokens.getText());
         } catch (IOException e) {
-            // TODO something sane
             throw new IllegalArgumentException(e);
         }
     }
@@ -121,13 +134,16 @@ public class TokenBuffer {
         try {
             return tokens.nextValue();
         } catch (IOException e) {
-            // TODO something sane
             throw new IllegalArgumentException(e);
         }
     }
 
-    private void updateNesting(JsonToken t) {
-        nesting += nestingOffset(t);
+    private void updateNesting(JsonToken token) {
+        nesting += nestingOffset(token);
+    }
+
+    private void updateNestingGoingBackwards(JsonToken token) {
+        nesting -= nestingOffset(token);
     }
 
     public int nesting() {
@@ -140,10 +156,10 @@ public class TokenBuffer {
         Token toReturn = null;
         Iterator<Token> i;
 
-        if (name.equals(currentName()) && currentToken().isScalarValue()) {
-            toReturn = buffer.peekFirst();
+        if (name.equals(currentName()) && current().isScalarValue()) {
+            toReturn = tokens.get(position);
         } else {
-            i = buffer.iterator();
+            i = tokens.iterator();
             i.next(); // just ignore the first value, as we know it's not what
                       // we're looking for, and it's nesting effect is already
                       // included
@@ -169,8 +185,8 @@ public class TokenBuffer {
         } while ( nesting() > initialNesting + relativeNesting);
     }
 
-    public Deque<Token> rest() {
-        return buffer;
+    public List<Token> rest() {
+        return tokens.subList(position, tokens.size());
     }
 
     public static final class Token {
diff --git a/document/src/main/java/com/yahoo/document/json/document/DocumentParser.java b/document/src/main/java/com/yahoo/document/json/document/DocumentParser.java
index 5e1c1eb6ac4..b63a39f51c5 100644
--- a/document/src/main/java/com/yahoo/document/json/document/DocumentParser.java
+++ b/document/src/main/java/com/yahoo/document/json/document/DocumentParser.java
@@ -76,18 +76,10 @@ public class DocumentParser {
             throw new IllegalArgumentException("Could not read document, no document?");
         }
         switch (currentToken) {
-            case START_OBJECT:
-                indentLevel++;
-                break;
-            case END_OBJECT:
-                indentLevel--;
-                break;
-            case START_ARRAY:
-                indentLevel += 10000L;
-                break;
-            case END_ARRAY:
-                indentLevel -= 10000L;
-                break;
+            case START_OBJECT -> indentLevel++;
+            case END_OBJECT -> indentLevel--;
+            case START_ARRAY -> indentLevel += 10000L;
+            case END_ARRAY -> indentLevel -= 10000L;
         }
     }
 
@@ -133,18 +125,12 @@ public class DocumentParser {
     }
 
     private static DocumentOperationType operationNameToOperationType(String operationName) {
-        switch (operationName) {
-            case PUT:
-            case ID:
-                return DocumentOperationType.PUT;
-            case REMOVE:
-                return DocumentOperationType.REMOVE;
-            case UPDATE:
-                return DocumentOperationType.UPDATE;
-            default:
-                throw new IllegalArgumentException(
-                        "Got " + operationName + " as document operation, only \"put\", " +
-                                "\"remove\" and \"update\" are supported.");
-        }
+        return switch (operationName) {
+            case PUT, ID -> DocumentOperationType.PUT;
+            case REMOVE -> DocumentOperationType.REMOVE;
+            case UPDATE -> DocumentOperationType.UPDATE;
+            default -> throw new IllegalArgumentException("Got " + operationName + " as document operation, only \"put\", " +
+                                                          "\"remove\" and \"update\" are supported.");
+        };
     }
 }
diff --git a/document/src/main/java/com/yahoo/document/json/readers/AddRemoveCreator.java b/document/src/main/java/com/yahoo/document/json/readers/AddRemoveCreator.java
index bc214f18776..35e113fa1d9 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/AddRemoveCreator.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/AddRemoveCreator.java
@@ -41,7 +41,7 @@ public class AddRemoveCreator {
         FieldUpdate singleUpdate;
         int initNesting = buffer.nesting();
 
-        Preconditions.checkState(buffer.currentToken().isStructStart(), "Expected start of composite, got %s", buffer.currentToken());
+        Preconditions.checkState(buffer.current().isStructStart(), "Expected start of composite, got %s", buffer.current());
         if (container instanceof CollectionFieldValue) {
             buffer.next();
             DataType valueType = ((CollectionFieldValue) container).getDataType().getNestedType();
@@ -58,8 +58,8 @@ public class AddRemoveCreator {
             } else {
                 List<FieldValue> arrayContents = new ArrayList<>();
                 ArrayReader.fillArrayUpdate(buffer, initNesting, valueType, arrayContents, ignoreUndefinedFields);
-                if (buffer.currentToken() != JsonToken.END_ARRAY) {
-                    throw new IllegalArgumentException("Expected END_ARRAY. Got '" + buffer.currentToken() + "'.");
+                if (buffer.current() != JsonToken.END_ARRAY) {
+                    throw new IllegalArgumentException("Expected END_ARRAY. Got '" + buffer.current() + "'.");
                 }
                 if (isRemove) {
                     singleUpdate = FieldUpdate.createRemoveAll(field, arrayContents);
@@ -72,7 +72,7 @@ public class AddRemoveCreator {
                     "Trying to add or remove from a field of a type the reader does not know how to handle: "
                             + container.getClass().getName());
         }
-        expectCompositeEnd(buffer.currentToken());
+        expectCompositeEnd(buffer.current());
         update.addAll(singleUpdate);
     }
 }
diff --git a/document/src/main/java/com/yahoo/document/json/readers/ArrayReader.java b/document/src/main/java/com/yahoo/document/json/readers/ArrayReader.java
index e8c4ae85356..c0e2618d6f1 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/ArrayReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/ArrayReader.java
@@ -18,7 +18,7 @@ public class ArrayReader {
     public static void fillArrayUpdate(TokenBuffer buffer, int initNesting, DataType valueType,
                                        List<FieldValue> arrayContents, boolean ignoreUndefinedFields) {
         while (buffer.nesting() >= initNesting) {
-            Preconditions.checkArgument(buffer.currentToken() != JsonToken.VALUE_NULL, "Illegal null value for array entry");
+            Preconditions.checkArgument(buffer.current() != JsonToken.VALUE_NULL, "Illegal null value for array entry");
             arrayContents.add(readSingleValue(buffer, valueType, ignoreUndefinedFields));
             buffer.next();
         }
@@ -27,10 +27,10 @@ public class ArrayReader {
     @SuppressWarnings({ "unchecked", "rawtypes" })
     public static void fillArray(TokenBuffer buffer, CollectionFieldValue parent, DataType valueType, boolean ignoreUndefinedFields) {
         int initNesting = buffer.nesting();
-        expectArrayStart(buffer.currentToken());
+        expectArrayStart(buffer.current());
         buffer.next();
         while (buffer.nesting() >= initNesting) {
-            Preconditions.checkArgument(buffer.currentToken() != JsonToken.VALUE_NULL, "Illegal null value for array entry");
+            Preconditions.checkArgument(buffer.current() != JsonToken.VALUE_NULL, "Illegal null value for array entry");
             parent.add(readSingleValue(buffer, valueType, ignoreUndefinedFields));
             buffer.next();
         }
diff --git a/document/src/main/java/com/yahoo/document/json/readers/CompositeReader.java b/document/src/main/java/com/yahoo/document/json/readers/CompositeReader.java
index 27c18ea2f69..fea6443ee5b 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/CompositeReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/CompositeReader.java
@@ -19,8 +19,8 @@ import static com.yahoo.document.json.readers.WeightedSetReader.fillWeightedSet;
 public class CompositeReader {
 
     public static boolean populateComposite(TokenBuffer buffer, FieldValue fieldValue, boolean ignoreUndefinedFields) {
-        boolean fullyApplied = populateComposite(buffer.currentToken(), buffer, fieldValue, ignoreUndefinedFields);
-        expectCompositeEnd(buffer.currentToken());
+        boolean fullyApplied = populateComposite(buffer.current(), buffer, fieldValue, ignoreUndefinedFields);
+        expectCompositeEnd(buffer.current());
         return fullyApplied;
     }
 
diff --git a/document/src/main/java/com/yahoo/document/json/readers/GeoPositionReader.java b/document/src/main/java/com/yahoo/document/json/readers/GeoPositionReader.java
index eb3919e07d7..e02e2cf05c0 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/GeoPositionReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/GeoPositionReader.java
@@ -1,7 +1,6 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.document.json.readers;
 
-import com.fasterxml.jackson.core.JsonToken;
 import com.yahoo.document.PositionDataType;
 import com.yahoo.document.datatypes.FieldValue;
 import com.yahoo.document.json.TokenBuffer;
@@ -16,7 +15,7 @@ public class GeoPositionReader {
     static void fillGeoPosition(TokenBuffer buffer, FieldValue positionFieldValue) {
         Double latitude = null;
         Double longitude = null;
-        expectObjectStart(buffer.currentToken());
+        expectObjectStart(buffer.current());
         int initNesting = buffer.nesting();
         for (buffer.next(); buffer.nesting() >= initNesting; buffer.next()) {
             String curName = buffer.currentName();
@@ -32,7 +31,7 @@ public class GeoPositionReader {
                 throw new IllegalArgumentException("Unexpected attribute "+curName+" in geo position field");
             }
         }
-        expectObjectEnd(buffer.currentToken());
+        expectObjectEnd(buffer.current());
         if (latitude == null) {
             throw new IllegalArgumentException("Missing 'lat' attribute in geo position field");
         }
diff --git a/document/src/main/java/com/yahoo/document/json/readers/JsonParserHelpers.java b/document/src/main/java/com/yahoo/document/json/readers/JsonParserHelpers.java
index 1723df2bd54..594dfc5ab06 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/JsonParserHelpers.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/JsonParserHelpers.java
@@ -5,6 +5,8 @@ package com.yahoo.document.json.readers;
 import com.fasterxml.jackson.core.JsonToken;
 import com.google.common.base.Preconditions;
 
+import java.util.Arrays;
+
 public class JsonParserHelpers {
 
     public static void expectArrayStart(JsonToken token) {
@@ -61,4 +63,9 @@ public class JsonParserHelpers {
         }
     }
 
+    public static void expectOneOf(JsonToken token, JsonToken ... tokens) {
+        if (Arrays.stream(tokens).noneMatch(t -> t == token))
+            throw new IllegalArgumentException("Expected one of " + tokens + " but got " + token);
+    }
+
 }
diff --git a/document/src/main/java/com/yahoo/document/json/readers/MapReader.java b/document/src/main/java/com/yahoo/document/json/readers/MapReader.java
index 1d4cb85d130..a660613f7f0 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/MapReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/MapReader.java
@@ -32,7 +32,7 @@ public class MapReader {
     public static final String UPDATE_MATCH = "match";
 
     public static void fillMap(TokenBuffer buffer, MapFieldValue parent, boolean ignoreUndefinedFields) {
-        if (buffer.currentToken() == JsonToken.START_ARRAY) {
+        if (buffer.current() == JsonToken.START_ARRAY) {
             MapReader.fillMapFromArray(buffer, parent, ignoreUndefinedFields);
         } else {
             MapReader.fillMapFromObject(buffer, parent, ignoreUndefinedFields);
@@ -41,7 +41,7 @@ public class MapReader {
 
     @SuppressWarnings({ "rawtypes", "cast", "unchecked" })
     public static void fillMapFromArray(TokenBuffer buffer, MapFieldValue parent, boolean ignoreUndefinedFields) {
-        JsonToken token = buffer.currentToken();
+        JsonToken token = buffer.current();
         int initNesting = buffer.nesting();
         expectArrayStart(token);
         token = buffer.next();
@@ -70,7 +70,7 @@ public class MapReader {
 
     @SuppressWarnings({ "rawtypes", "cast", "unchecked" })
     public static void fillMapFromObject(TokenBuffer buffer, MapFieldValue parent, boolean ignoreUndefinedFields) {
-        JsonToken token = buffer.currentToken();
+        JsonToken token = buffer.current();
         int initNesting = buffer.nesting();
         expectObjectStart(token);
         token = buffer.next();
diff --git a/document/src/main/java/com/yahoo/document/json/readers/SingleValueReader.java b/document/src/main/java/com/yahoo/document/json/readers/SingleValueReader.java
index 3ae82676fa8..3a923c8d4bd 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/SingleValueReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/SingleValueReader.java
@@ -10,7 +10,6 @@ import com.yahoo.document.datatypes.FieldValue;
 import com.yahoo.document.json.TokenBuffer;
 import com.yahoo.document.update.ValueUpdate;
 
-import java.util.Base64;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.regex.Matcher;
@@ -42,7 +41,7 @@ public class SingleValueReader {
     }
 
     public static FieldValue readSingleValue(TokenBuffer buffer, DataType expectedType, boolean ignoreUndefinedFields) {
-        if (buffer.currentToken().isScalarValue()) {
+        if (buffer.current().isScalarValue()) {
             return readAtomic(buffer.currentText(), expectedType);
         } else {
             FieldValue fieldValue = expectedType.createFieldValue();
@@ -54,7 +53,7 @@ public class SingleValueReader {
     @SuppressWarnings("rawtypes")
     public static ValueUpdate readSingleUpdate(TokenBuffer buffer, DataType expectedType, String action, boolean ignoreUndefinedFields) {
         return switch (action) {
-            case UPDATE_ASSIGN -> (buffer.currentToken() == JsonToken.VALUE_NULL)
+            case UPDATE_ASSIGN -> (buffer.current() == JsonToken.VALUE_NULL)
                                   ? ValueUpdate.createClear()
                                   : ValueUpdate.createAssign(readSingleValue(buffer, expectedType, ignoreUndefinedFields));
             // double is silly, but it's what is used internally anyway
diff --git a/document/src/main/java/com/yahoo/document/json/readers/StructReader.java b/document/src/main/java/com/yahoo/document/json/readers/StructReader.java
index d489c03eb9d..41b197463b6 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/StructReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/StructReader.java
@@ -36,7 +36,7 @@ public class StructReader {
             }
 
             try {
-                if (buffer.currentToken() != JsonToken.VALUE_NULL) {
+                if (buffer.current() != JsonToken.VALUE_NULL) {
                     FieldValue v = readSingleValue(buffer, field.getDataType(), ignoreUndefinedFields);
                     parent.setFieldValue(field, v);
                 }
diff --git a/document/src/main/java/com/yahoo/document/json/readers/TensorAddUpdateReader.java b/document/src/main/java/com/yahoo/document/json/readers/TensorAddUpdateReader.java
index 0587825da68..02ef46aaafe 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/TensorAddUpdateReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/TensorAddUpdateReader.java
@@ -22,7 +22,7 @@ public class TensorAddUpdateReader {
     }
 
     public static TensorAddUpdate createTensorAddUpdate(TokenBuffer buffer, Field field) {
-        expectObjectStart(buffer.currentToken());
+        expectObjectStart(buffer.current());
         expectTensorTypeHasSparseDimensions(field);
 
         TensorDataType tensorDataType = (TensorDataType)field.getDataType();
diff --git a/document/src/main/java/com/yahoo/document/json/readers/TensorModifyUpdateReader.java b/document/src/main/java/com/yahoo/document/json/readers/TensorModifyUpdateReader.java
index 748b5c046bb..21fa51d5b88 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/TensorModifyUpdateReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/TensorModifyUpdateReader.java
@@ -33,7 +33,7 @@ public class TensorModifyUpdateReader {
     public static TensorModifyUpdate createModifyUpdate(TokenBuffer buffer, Field field) {
         expectFieldIsOfTypeTensor(field);
         expectTensorTypeHasNoIndexedUnboundDimensions(field);
-        expectObjectStart(buffer.currentToken());
+        expectObjectStart(buffer.current());
 
         ModifyUpdateResult result = createModifyUpdateResult(buffer, field);
         expectOperationSpecified(result.operation, field.getName());
diff --git a/document/src/main/java/com/yahoo/document/json/readers/TensorReader.java b/document/src/main/java/com/yahoo/document/json/readers/TensorReader.java
index 193c9491e86..0d971859550 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/TensorReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/TensorReader.java
@@ -23,7 +23,6 @@ import static com.yahoo.tensor.serialization.JsonFormat.decodeHexString;
 public class TensorReader {
 
     public static final String TENSOR_ADDRESS = "address";
-    public static final String TENSOR_DIMENSIONS = "dimensions";
     public static final String TENSOR_CELLS = "cells";
     public static final String TENSOR_VALUES = "values";
     public static final String TENSOR_BLOCKS = "blocks";
@@ -32,41 +31,58 @@ public class TensorReader {
     // MUST be kept in sync with com.yahoo.tensor.serialization.JsonFormat.decode in vespajlib
     static void fillTensor(TokenBuffer buffer, TensorFieldValue tensorFieldValue) {
         Tensor.Builder builder = Tensor.Builder.of(tensorFieldValue.getDataType().getTensorType());
-        expectObjectStart(buffer.currentToken());
+        expectOneOf(buffer.current(), JsonToken.START_OBJECT, JsonToken.START_ARRAY);
         int initNesting = buffer.nesting();
         for (buffer.next(); buffer.nesting() >= initNesting; buffer.next()) {
-            if (TENSOR_CELLS.equals(buffer.currentName()))
+            if (TENSOR_CELLS.equals(buffer.currentName()) && ! primitiveContent(buffer)) {
                 readTensorCells(buffer, builder);
-            else if (TENSOR_VALUES.equals(buffer.currentName()))
+            }
+            else if (TENSOR_VALUES.equals(buffer.currentName()) && builder.type().dimensions().stream().allMatch(d -> d.isIndexed())) {
                 readTensorValues(buffer, builder);
-            else if (TENSOR_BLOCKS.equals(buffer.currentName()))
+            }
+            else if (TENSOR_BLOCKS.equals(buffer.currentName())) {
                 readTensorBlocks(buffer, builder);
-            else if (builder.type().dimensions().stream().anyMatch(d -> d.isIndexed())) // sparse can be empty
-                throw new IllegalArgumentException("Expected a tensor value to contain either 'cells', 'values' or 'blocks', but got: "+buffer.currentName());
+            }
+            else {
+                buffer.previous(); // Back up to the start of the enclosing block
+                readDirectTensorValue(buffer, builder);
+                buffer.previous(); // ... and back up to the end of the enclosing block
+            }
         }
-        expectObjectEnd(buffer.currentToken());
+        expectOneOf(buffer.current(), JsonToken.END_OBJECT, JsonToken.END_ARRAY);
         tensorFieldValue.assign(builder.build());
     }
 
+    static boolean primitiveContent(TokenBuffer buffer) {
+        JsonToken cellsValue = buffer.current();
+        if (cellsValue.isScalarValue()) return true;
+        if (cellsValue == JsonToken.START_ARRAY) {
+            JsonToken firstArrayValue = buffer.peek(1);
+            if (firstArrayValue == JsonToken.END_ARRAY) return false;
+            if (firstArrayValue.isScalarValue()) return true;
+        }
+        return false;
+    }
+
     static void readTensorCells(TokenBuffer buffer, Tensor.Builder builder) {
-        if (buffer.currentToken() == JsonToken.START_ARRAY) {
+        if (buffer.current() == JsonToken.START_ARRAY) {
             int initNesting = buffer.nesting();
             for (buffer.next(); buffer.nesting() >= initNesting; buffer.next())
                 readTensorCell(buffer, builder);
         }
-        else if (buffer.currentToken() == JsonToken.START_OBJECT) { // single dimension short form
+        else if (buffer.current() == JsonToken.START_OBJECT) { // single dimension short form
             int initNesting = buffer.nesting();
             for (buffer.next(); buffer.nesting() >= initNesting; buffer.next())
                 builder.cell(asAddress(buffer.currentName(), builder.type()), readDouble(buffer));
         }
         else {
-            throw new IllegalArgumentException("Expected 'cells' to contain an array or an object, but got " + buffer.currentToken());
+            throw new IllegalArgumentException("Expected 'cells' to contain an array or an object, but got " + buffer.current());
         }
-        expectCompositeEnd(buffer.currentToken());
+        expectCompositeEnd(buffer.current());
     }
 
     private static void readTensorCell(TokenBuffer buffer, Tensor.Builder builder) {
-        expectObjectStart(buffer.currentToken());
+        expectObjectStart(buffer.current());
 
         TensorAddress address = null;
         Double value = null;
@@ -79,7 +95,7 @@ public class TensorReader {
                 value = readDouble(buffer);
             }
         }
-        expectObjectEnd(buffer.currentToken());
+        expectObjectEnd(buffer.current());
         if (address == null)
             throw new IllegalArgumentException("Expected an object in a tensor 'cells' array to contain an 'address' field");
         if (value == null)
@@ -88,11 +104,10 @@ public class TensorReader {
     }
 
     private static void readTensorValues(TokenBuffer buffer, Tensor.Builder builder) {
-        if ( ! (builder instanceof IndexedTensor.BoundBuilder))
+        if ( ! (builder instanceof IndexedTensor.BoundBuilder indexedBuilder))
             throw new IllegalArgumentException("The 'values' field can only be used with dense tensors. " +
                                                "Use 'cells' or 'blocks' instead");
-        IndexedTensor.BoundBuilder indexedBuilder = (IndexedTensor.BoundBuilder)builder;
-        if (buffer.currentToken() == JsonToken.VALUE_STRING) {
+        if (buffer.current() == JsonToken.VALUE_STRING) {
             double[] decoded = decodeHexString(buffer.currentText(), builder.type().valueType());
             if (decoded.length == 0)
                 throw new IllegalArgumentException("The 'values' string does not contain any values");
@@ -108,21 +123,19 @@ public class TensorReader {
         }
         if (index == 0)
             throw new IllegalArgumentException("The 'values' array does not contain any values");
-        expectCompositeEnd(buffer.currentToken());
+        expectCompositeEnd(buffer.current());
     }
 
     static void readTensorBlocks(TokenBuffer buffer, Tensor.Builder builder) {
-        if ( ! (builder instanceof MixedTensor.BoundBuilder))
+        if ( ! (builder instanceof MixedTensor.BoundBuilder mixedBuilder))
             throw new IllegalArgumentException("The 'blocks' field can only be used with mixed tensors with bound dimensions. " +
                                                "Use 'cells' or 'values' instead");
-
-        MixedTensor.BoundBuilder mixedBuilder = (MixedTensor.BoundBuilder) builder;
-        if (buffer.currentToken() == JsonToken.START_ARRAY) {
+        if (buffer.current() == JsonToken.START_ARRAY) {
             int initNesting = buffer.nesting();
             for (buffer.next(); buffer.nesting() >= initNesting; buffer.next())
                 readTensorBlock(buffer, mixedBuilder);
         }
-        else if (buffer.currentToken() == JsonToken.START_OBJECT) {
+        else if (buffer.current() == JsonToken.START_OBJECT) {
             int initNesting = buffer.nesting();
             for (buffer.next(); buffer.nesting() >= initNesting; buffer.next()) {
                 TensorAddress mappedAddress = asAddress(buffer.currentName(), builder.type().mappedSubtype());
@@ -132,14 +145,14 @@ public class TensorReader {
         }
         else {
             throw new IllegalArgumentException("Expected 'blocks' to contain an array or an object, but got " +
-                                               buffer.currentToken());
+                                               buffer.current());
         }
 
-        expectCompositeEnd(buffer.currentToken());
+        expectCompositeEnd(buffer.current());
     }
 
     private static void readTensorBlock(TokenBuffer buffer, MixedTensor.BoundBuilder mixedBuilder) {
-        expectObjectStart(buffer.currentToken());
+        expectObjectStart(buffer.current());
 
         TensorAddress address = null;
         double[] values = null;
@@ -152,7 +165,7 @@ public class TensorReader {
             else if (TensorReader.TENSOR_VALUES.equals(currentName))
                 values = readValues(buffer, (int)mixedBuilder.denseSubspaceSize(), address, mixedBuilder.type());
         }
-        expectObjectEnd(buffer.currentToken());
+        expectObjectEnd(buffer.current());
         if (address == null)
             throw new IllegalArgumentException("Expected a 'blocks' array object to contain an object 'address'");
         if (values == null)
@@ -160,13 +173,26 @@ public class TensorReader {
         mixedBuilder.block(address, values);
     }
 
+    /** Reads a tensor value directly at the root, where the format is decided by the tensor type. */
+    private static void readDirectTensorValue(TokenBuffer buffer, Tensor.Builder builder) {
+        boolean hasIndexed = builder.type().dimensions().stream().anyMatch(TensorType.Dimension::isIndexed);
+        boolean hasMapped = builder.type().dimensions().stream().anyMatch(TensorType.Dimension::isMapped);
+
+        if ( ! hasMapped)
+            readTensorValues(buffer, builder);
+        else if (hasMapped && hasIndexed)
+            readTensorBlocks(buffer, builder);
+        else
+            readTensorCells(buffer, builder);
+    }
+
     private static TensorAddress readAddress(TokenBuffer buffer, TensorType type) {
-        expectObjectStart(buffer.currentToken());
+        expectObjectStart(buffer.current());
         TensorAddress.Builder builder = new TensorAddress.Builder(type);
         int initNesting = buffer.nesting();
         for (buffer.next(); buffer.nesting() >= initNesting; buffer.next())
             builder.add(buffer.currentName(), buffer.currentText());
-        expectObjectEnd(buffer.currentToken());
+        expectObjectEnd(buffer.current());
         return builder.build();
     }
 
@@ -182,15 +208,15 @@ public class TensorReader {
     private static double[] readValues(TokenBuffer buffer, int size, TensorAddress address, TensorType type) {
         int index = 0;
         double[] values = new double[size];
-        if (buffer.currentToken() == JsonToken.VALUE_STRING) {
+        if (buffer.current() == JsonToken.VALUE_STRING) {
             values = decodeHexString(buffer.currentText(), type.valueType());
             index = values.length;
         } else {
-            expectArrayStart(buffer.currentToken());
+            expectArrayStart(buffer.current());
             int initNesting = buffer.nesting();
             for (buffer.next(); buffer.nesting() >= initNesting; buffer.next())
                 values[index++] = readDouble(buffer);
-            expectCompositeEnd(buffer.currentToken());
+            expectCompositeEnd(buffer.current());
         }
         if (index != size)
             throw new IllegalArgumentException((address != null ? "At " + address.toString(type) + ": " : "") +
diff --git a/document/src/main/java/com/yahoo/document/json/readers/TensorRemoveUpdateReader.java b/document/src/main/java/com/yahoo/document/json/readers/TensorRemoveUpdateReader.java
index 6db55f8fb26..d173aade4eb 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/TensorRemoveUpdateReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/TensorRemoveUpdateReader.java
@@ -26,7 +26,7 @@ public class TensorRemoveUpdateReader {
     public static final String TENSOR_ADDRESSES = "addresses";
 
     static TensorRemoveUpdate createTensorRemoveUpdate(TokenBuffer buffer, Field field) {
-        expectObjectStart(buffer.currentToken());
+        expectObjectStart(buffer.current());
         expectTensorTypeHasSparseDimensions(field);
 
         TensorDataType tensorDataType = (TensorDataType)field.getDataType();
@@ -59,11 +59,11 @@ public class TensorRemoveUpdateReader {
      */
     private static Tensor readRemoveUpdateTensor(TokenBuffer buffer, TensorType sparseType, TensorType originalType) {
         Tensor.Builder builder = null;
-        expectObjectStart(buffer.currentToken());
+        expectObjectStart(buffer.current());
         int initNesting = buffer.nesting();
         for (buffer.next(); buffer.nesting() >= initNesting; buffer.next()) {
             if (TENSOR_ADDRESSES.equals(buffer.currentName())) {
-                expectArrayStart(buffer.currentToken());
+                expectArrayStart(buffer.current());
                 int nesting = buffer.nesting();
                 for (buffer.next(); buffer.nesting() >= nesting; buffer.next()) {
                     if (builder == null) {
@@ -74,10 +74,10 @@ public class TensorRemoveUpdateReader {
                         builder.cell(readTensorAddress(buffer, builder.type(), originalType), 1.0);
                     }
                 }
-                expectCompositeEnd(buffer.currentToken());
+                expectCompositeEnd(buffer.current());
             }
         }
-        expectObjectEnd(buffer.currentToken());
+        expectObjectEnd(buffer.current());
         return (builder != null) ? builder.build() : Tensor.Builder.of(sparseType).build();
     }
 
@@ -88,7 +88,7 @@ public class TensorRemoveUpdateReader {
     private static Pair<TensorType, TensorAddress> readFirstTensorAddress(TokenBuffer buffer, TensorType sparseType, TensorType originalType) {
         var typeBuilder = new TensorType.Builder(sparseType.valueType());
         var rawAddress = new HashMap<String, String>();
-        expectObjectStart(buffer.currentToken());
+        expectObjectStart(buffer.current());
         int initNesting = buffer.nesting();
         for (buffer.next(); buffer.nesting() >= initNesting; buffer.next()) {
             var elem = readRawElement(buffer, sparseType, originalType);
@@ -100,7 +100,7 @@ public class TensorRemoveUpdateReader {
                 throw new IllegalArgumentException(originalType + " does not contain dimension '" + elem.getFirst() + "'");
             }
         }
-        expectObjectEnd(buffer.currentToken());
+        expectObjectEnd(buffer.current());
         var type = typeBuilder.build();
         var builder = new TensorAddress.Builder(type);
         rawAddress.forEach((dimension, label) -> builder.add(dimension, label));
@@ -119,13 +119,13 @@ public class TensorRemoveUpdateReader {
 
     private static TensorAddress readTensorAddress(TokenBuffer buffer, TensorType type, TensorType originalType) {
         TensorAddress.Builder builder = new TensorAddress.Builder(type);
-        expectObjectStart(buffer.currentToken());
+        expectObjectStart(buffer.current());
         int initNesting = buffer.nesting();
         for (buffer.next(); buffer.nesting() >= initNesting; buffer.next()) {
             var elem = readRawElement(buffer, type, originalType);
             builder.add(elem.getFirst(), elem.getSecond());
         }
-        expectObjectEnd(buffer.currentToken());
+        expectObjectEnd(buffer.current());
         return builder.build();
     }
 
diff --git a/document/src/main/java/com/yahoo/document/json/readers/VespaJsonDocumentReader.java b/document/src/main/java/com/yahoo/document/json/readers/VespaJsonDocumentReader.java
index 22a9e7a1119..c406bcdb2b3 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/VespaJsonDocumentReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/VespaJsonDocumentReader.java
@@ -96,13 +96,13 @@ public class VespaJsonDocumentReader {
     public boolean readUpdate(TokenBuffer buffer, DocumentUpdate update) {
         if (buffer.isEmpty())
             throw new IllegalArgumentException("Update of document " + update.getId() + " is missing a 'fields' map");
-        expectObjectStart(buffer.currentToken());
+        expectObjectStart(buffer.current());
         int localNesting = buffer.nesting();
 
         buffer.next();
         boolean fullyApplied = true;
         while (localNesting <= buffer.nesting()) {
-            expectObjectStart(buffer.currentToken());
+            expectObjectStart(buffer.current());
 
             String fieldName = buffer.currentName();
             try {
@@ -111,7 +111,7 @@ public class VespaJsonDocumentReader {
                 } else {
                     fullyApplied &= addFieldUpdates(update, buffer, fieldName);
                 }
-                expectObjectEnd(buffer.currentToken());
+                expectObjectEnd(buffer.current());
             }
             catch (IllegalArgumentException | IndexOutOfBoundsException e) {
                 throw new IllegalArgumentException("Error in '" + fieldName + "'", e);
@@ -211,7 +211,7 @@ public class VespaJsonDocumentReader {
     }
 
     private RemoveFieldPathUpdate readRemoveFieldPathUpdate(DocumentType documentType, String fieldPath, TokenBuffer buffer) {
-        expectScalarValue(buffer.currentToken());
+        expectScalarValue(buffer.current());
         return new RemoveFieldPathUpdate(documentType, fieldPath);
     }
 
@@ -219,7 +219,7 @@ public class VespaJsonDocumentReader {
                                                                 TokenBuffer buffer, String fieldPathOperation) {
         AssignFieldPathUpdate fieldPathUpdate = new AssignFieldPathUpdate(documentType, fieldPath);
         String arithmeticSign = SingleValueReader.UPDATE_OPERATION_TO_ARITHMETIC_SIGN.get(fieldPathOperation);
-        double value = Double.valueOf(buffer.currentText());
+        double value = Double.parseDouble(buffer.currentText());
         String expression = String.format("$value %s %s", arithmeticSign, value);
         fieldPathUpdate.setExpression(expression);
         return fieldPathUpdate;
@@ -231,11 +231,11 @@ public class VespaJsonDocumentReader {
     }
 
     private static void verifyEndState(TokenBuffer buffer, JsonToken expectedFinalToken) {
-        Preconditions.checkState(buffer.currentToken() == expectedFinalToken,
-                                 "Expected end of JSON struct (%s), got %s", expectedFinalToken, buffer.currentToken());
+        Preconditions.checkState(buffer.current() == expectedFinalToken,
+                                 "Expected end of JSON struct (%s), got %s", expectedFinalToken, buffer.current());
         Preconditions.checkState(buffer.nesting() == 0, "Nesting not zero at end of operation");
         Preconditions.checkState(buffer.next() == null, "Dangling data at end of operation");
-        Preconditions.checkState(buffer.size() == 0, "Dangling data at end of operation");
+        Preconditions.checkState(buffer.remaining() == 0, "Dangling data at end of operation");
     }
 
 }
diff --git a/document/src/main/java/com/yahoo/document/json/readers/WeightedSetReader.java b/document/src/main/java/com/yahoo/document/json/readers/WeightedSetReader.java
index 7a9921498ef..2bdb5ee4997 100644
--- a/document/src/main/java/com/yahoo/document/json/readers/WeightedSetReader.java
+++ b/document/src/main/java/com/yahoo/document/json/readers/WeightedSetReader.java
@@ -13,7 +13,7 @@ public class WeightedSetReader {
 
     public static void fillWeightedSet(TokenBuffer buffer, DataType valueType, @SuppressWarnings("rawtypes") WeightedSet weightedSet) {
         int initNesting = buffer.nesting();
-        expectObjectStart(buffer.currentToken());
+        expectObjectStart(buffer.current());
         buffer.next();
         iterateThroughWeightedSet(buffer, initNesting, valueType, weightedSet);
     }
diff --git a/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java b/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java
index 41d607b0d8e..c19094ff231 100644
--- a/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java
+++ b/document/src/test/java/com/yahoo/document/json/JsonReaderTestCase.java
@@ -1403,12 +1403,6 @@ public class JsonReaderTestCase {
     }
 
     @Test
-    public void testParsingOfTensorWithEmptyDimensions() {
-        assertSparseTensorField("tensor(x{},y{}):{}",
-                                createPutWithSparseTensor(inputJson("{ 'dimensions': [] }")));
-    }
-
-    @Test
     public void testParsingOfTensorWithEmptyCells() {
         assertSparseTensorField("tensor(x{},y{}):{}",
                                 createPutWithSparseTensor(inputJson("{ 'cells': [] }")));
@@ -1511,6 +1505,32 @@ public class JsonReaderTestCase {
         Tensor tensor = assertTensorField(expected, put, "mixed_bfloat16_tensor");
     }
 
+    /** Tests parsing of various tensor values set at the root, i.e. no 'cells', 'blocks' or 'values' */
+    @Test
+    public void testDirectValue() {
+        assertTensorField("tensor(x{}):{a:2, b:3}", "sparse_single_dimension_tensor", "{'a':2.0, 'b':3.0}");
+        assertTensorField("tensor(x[2],y[3]):[2, 3, 4, 5, 6, 7]]", "dense_tensor", "[2, 3, 4, 5, 6, 7]");
+        assertTensorField("tensor(x{},y[3]):{a:[2, 3, 4], b:[4, 5, 6]}", "mixed_tensor", "{'a':[2, 3, 4], 'b':[4, 5, 6]}");
+        assertTensorField("tensor(x{},y{}):{{x:a,y:0}:2, {x:b,y:1}:3}", "sparse_tensor",
+                          "[{'address':{'x':'a','y':'0'},'value':2}, {'address':{'x':'b','y':'1'},'value':3}]");
+    }
+
+    @Test
+    public void testDirectValueReservedNameKeys() {
+        // Single-valued
+        assertTensorField("tensor(x{}):{cells:2, b:3}", "sparse_single_dimension_tensor", "{'cells':2.0, 'b':3.0}");
+        assertTensorField("tensor(x{}):{values:2, b:3}", "sparse_single_dimension_tensor", "{'values':2.0, 'b':3.0}");
+        assertTensorField("tensor(x{}):{block:2, b:3}", "sparse_single_dimension_tensor", "{'block':2.0, 'b':3.0}");
+
+        // Multi-valued
+        assertTensorField("tensor(x{},y[3]):{cells:[2, 3, 4], b:[4, 5, 6]}", "mixed_tensor",
+                          "{'cells':[2, 3, 4], 'b':[4, 5, 6]}");
+        assertTensorField("tensor(x{},y[3]):{values:[2, 3, 4], b:[4, 5, 6]}", "mixed_tensor",
+                          "{'values':[2, 3, 4], 'b':[4, 5, 6]}");
+        assertTensorField("tensor(x{},y[3]):{block:[2, 3, 4], b:[4, 5, 6]}", "mixed_tensor",
+                          "{'block':[2, 3, 4], 'b':[4, 5, 6]}");
+    }
+
     @Test
     public void testParsingOfMixedTensorOnMixedForm() {
         Tensor.Builder builder = Tensor.Builder.of(TensorType.fromSpec("tensor(x{},y[3])"));
@@ -2070,6 +2090,9 @@ public class JsonReaderTestCase {
     private static Tensor assertSparseTensorField(String expectedTensor, DocumentPut put) {
         return assertTensorField(expectedTensor, put, "sparse_tensor");
     }
+    private Tensor assertTensorField(String expectedTensor, String fieldName, String inputJson) {
+        return assertTensorField(expectedTensor, createPutWithTensor(inputJson, fieldName), fieldName);
+    }
     private static Tensor assertTensorField(String expectedTensor, DocumentPut put, String tensorFieldName) {
         return assertTensorField(Tensor.from(expectedTensor), put, tensorFieldName);
     }