Merge pull request #29994 from vespa-engine/balder/fetch-expensive-size-once

Balder/fetch expensive size once

2 files changed, 22 insertions, 30 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/MixedTensor.java b/vespajlib/src/main/java/com/yahoo/tensor/MixedTensor.java
index 30dd1d6dc29..55fb4af551e 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/MixedTensor.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/MixedTensor.java
@@ -3,6 +3,7 @@
 package com.yahoo.tensor;
 
 import com.google.common.collect.ImmutableMap;
+import com.yahoo.tensor.impl.NumericTensorAddress;
 import com.yahoo.tensor.impl.StringTensorAddress;
 
 import java.util.ArrayList;
@@ -326,10 +327,12 @@ public class MixedTensor implements Tensor {
         }
 
         private double[] denseSubspace(TensorAddress sparseAddress) {
-            if (!denseSubspaceMap.containsKey(sparseAddress)) {
-                denseSubspaceMap.put(sparseAddress, new double[(int)denseSubspaceSize()]);
+            double [] values = denseSubspaceMap.get(sparseAddress);
+            if (values == null) {
+                values = new double[(int)denseSubspaceSize()];
+                denseSubspaceMap.put(sparseAddress, values);
             }
-            return denseSubspaceMap.get(sparseAddress);
+            return values;
         }
 
         public IndexedTensor.DirectIndexBuilder denseSubspaceBuilder(TensorAddress sparseAddress) {
@@ -520,7 +523,7 @@ public class MixedTensor implements Tensor {
             return builder.build();
         }
 
-        private TensorAddress denseOffsetToAddress(long denseOffset) {
+        private long[] denseOffsetToAddress(long denseOffset) {
             if (denseOffset < 0 || denseOffset > denseSubspaceSize) {
                 throw new IllegalArgumentException("Offset out of bounds");
             }
@@ -538,17 +541,17 @@ public class MixedTensor implements Tensor {
                 labels[i] = restSize / innerSize;
                 restSize %= innerSize;
             }
-            return TensorAddress.of(labels);
+            return labels;
         }
 
         TensorAddress fullAddressOf(TensorAddress sparsePart, long denseOffset) {
-            TensorAddress densePart = denseOffsetToAddress(denseOffset);
+            long [] densePart = denseOffsetToAddress(denseOffset);
             String[] labels = new String[type.dimensions().size()];
             int mappedIndex = 0;
             int indexedIndex = 0;
             for (TensorType.Dimension d : type.dimensions()) {
                 if (d.isIndexed()) {
-                    labels[mappedIndex + indexedIndex] = densePart.label(indexedIndex);
+                    labels[mappedIndex + indexedIndex] = NumericTensorAddress.asString(densePart[indexedIndex]);
                     indexedIndex++;
                 } else {
                     labels[mappedIndex + indexedIndex] = sparsePart.label(mappedIndex);
@@ -607,8 +610,7 @@ public class MixedTensor implements Tensor {
                     b.append(", ");
 
                 // start brackets
-                for (int i = 0; i < indexes.nextDimensionsAtStart(); i++)
-                    b.append("[");
+                b.append("[".repeat(Math.max(0, indexes.nextDimensionsAtStart())));
 
                 // value
                 switch (type.valueType()) {
@@ -621,8 +623,7 @@ public class MixedTensor implements Tensor {
                 }
 
                 // end bracket
-                for (int i = 0; i < indexes.nextDimensionsAtEnd(); i++)
-                    b.append("]");
+                b.append("]".repeat(Math.max(0, indexes.nextDimensionsAtEnd())));
             }
             return index;
         }
@@ -656,27 +657,16 @@ public class MixedTensor implements Tensor {
         }
     }
 
-    private static class DenseSubspaceBuilder implements IndexedTensor.DirectIndexBuilder {
-
-        private final TensorType type;
-        private final double[] values;
-
-        public DenseSubspaceBuilder(TensorType type, double[] values) {
-            this.type = type;
-            this.values = values;
-        }
-
-        @Override
-        public TensorType type() { return type; }
+    private record DenseSubspaceBuilder(TensorType type, double[] values) implements IndexedTensor.DirectIndexBuilder {
 
         @Override
         public void cellByDirectIndex(long index, double value) {
-            values[(int)index] = value;
+            values[(int) index] = value;
         }
 
         @Override
         public void cellByDirectIndex(long index, float value) {
-            values[(int)index] = value;
+            values[(int) index] = value;
         }
 
     }
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java
index 1ded16636d3..7a336233de0 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java
@@ -114,7 +114,7 @@ public class Join<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAMETYP
     }
 
     private static Tensor indexedVectorJoin(IndexedTensor a, IndexedTensor b, TensorType type, DoubleBinaryOperator combinator) {
-        long joinedRank = Math.min(a.dimensionSizes().size(0), b.dimensionSizes().size(0));
+        int joinedRank = (int)Math.min(a.dimensionSizes().size(0), b.dimensionSizes().size(0));
         Iterator<Double> aIterator = a.valueIterator();
         Iterator<Double> bIterator = b.valueIterator();
         IndexedTensor.Builder builder = IndexedTensor.Builder.of(type, new DimensionSizes.Builder(1).set(0, joinedRank).build());
@@ -170,7 +170,7 @@ public class Join<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAMETYP
                                       Iterator<Tensor.Cell> superspace, long superspaceSize,
                                       boolean reversedArgumentOrder, IndexedTensor.Builder builder,
                                       DoubleBinaryOperator combinator) {
-        long joinedLength = Math.min(subspaceSize, superspaceSize);
+        int joinedLength = (int)Math.min(subspaceSize, superspaceSize);
         if (reversedArgumentOrder) {
             for (int i = 0; i < joinedLength; i++) {
                 Tensor.Cell supercell = superspace.next();
@@ -252,6 +252,7 @@ public class Join<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAMETYP
                                int[] aToIndexes, int[] bToIndexes, Tensor.Builder builder,
                                DoubleBinaryOperator combinator) {
         Set<String> sharedDimensions = Sets.intersection(a.type().dimensionNames(), b.type().dimensionNames());
+        int sharedDimensionSize = sharedDimensions.size();  // Expensive to compute size after intersection
         Set<String> dimensionsOnlyInA = Sets.difference(a.type().dimensionNames(), b.type().dimensionNames());
 
         DimensionSizes aIterateSize = joinedSizeOf(a.type(), joinedType, joinedSize);
@@ -263,7 +264,7 @@ public class Join<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAMETYP
             // for each combination of dimensions in a which is also in b
             while (aSubspace.hasNext()) {
                 Tensor.Cell aCell = aSubspace.next();
-                PartialAddress matchingBCells = partialAddress(a.type(), aSubspace.address(), sharedDimensions);
+                PartialAddress matchingBCells = partialAddress(a.type(), aSubspace.address(), sharedDimensions, sharedDimensionSize);
                 // for each matching combination of dimensions ony in b
                 for (IndexedTensor.SubspaceIterator bSubspace = b.cellIterator(matchingBCells, bIterateSize); bSubspace.hasNext(); ) {
                     Tensor.Cell bCell = bSubspace.next();
@@ -275,8 +276,9 @@ public class Join<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAMETYP
         }
     }
 
-    private static PartialAddress partialAddress(TensorType addressType, TensorAddress address, Set<String> retainDimensions) {
-        PartialAddress.Builder builder = new PartialAddress.Builder(retainDimensions.size());
+    private static PartialAddress partialAddress(TensorType addressType, TensorAddress address,
+                                                 Set<String> retainDimensions, int sharedDimensionSize) {
+        PartialAddress.Builder builder = new PartialAddress.Builder(sharedDimensionSize);
         for (int i = 0; i < addressType.dimensions().size(); i++)
             if (retainDimensions.contains(addressType.dimensions().get(i).name()))
                 builder.add(addressType.dimensions().get(i).name(), address.numericLabel(i));