Allow dimension replacing

2 files changed, 14 insertions, 1 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java b/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java
index 5ab0bd75a5a..4ad0835cfb4 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java
@@ -61,7 +61,7 @@ public class TensorType {
             b.add(thisDimension);
         for (Dimension otherDimension : other.dimensions) {
             Dimension thisDimension = b.dimensions.get(otherDimension.name());
-            b.add(otherDimension.combineWith(Optional.ofNullable(thisDimension)));
+            b.addOrReplace(otherDimension.combineWith(Optional.ofNullable(thisDimension)));
         }
         return b.build();
     }
@@ -240,6 +240,18 @@ public class TensorType {
             return this;
         }
 
+        private Builder addOrReplace(Dimension dimension) { // TODO: Not quite sure I like this solution
+            Objects.requireNonNull(dimension, "A dimension cannot be null");
+
+            if ( ! dimensions.isEmpty()) {
+                validateDimensionType(dimension);
+            }
+
+            dimensions.put(dimension.name(), dimension);
+            prevDimension = dimension;
+            return this;
+        }
+
         private void validateDimensionName(Dimension newDimension) {
             Dimension prevDimension = dimensions.get(newDimension.name());
             if (prevDimension != null) {
diff --git a/vespajlib/src/test/java/com/yahoo/tensor/TensorFunctionBenchmark.java b/vespajlib/src/test/java/com/yahoo/tensor/TensorFunctionBenchmark.java
index 3ba38205f59..c0d2e37cbdc 100644
--- a/vespajlib/src/test/java/com/yahoo/tensor/TensorFunctionBenchmark.java
+++ b/vespajlib/src/test/java/com/yahoo/tensor/TensorFunctionBenchmark.java
@@ -58,6 +58,7 @@ public class TensorFunctionBenchmark {
     
     public static void main(String[] args) {
         // Was: 144 ms
+        // After adding type: 290 ms
         new TensorFunctionBenchmark().benchmark(100);
     }