Add value iterator and iterate in vector-vector join

6 files changed, 58 insertions, 16 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java b/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java
index 263c41a6e13..4091fd97a16 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java
@@ -2,7 +2,6 @@
 package com.yahoo.tensor;
 
 import com.google.common.annotations.Beta;
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 
 import java.util.ArrayList;
@@ -41,8 +40,23 @@ public class IndexedTensor implements Tensor {
         return values.length;
     }
 
+    /**
+     * Returns an iterator over the cells of this. 
+     * Cells are returned in order of increasing indexes in each dimension, increasing 
+     * indexes of later dimensions in the dimension type before earlier.
+     */
     @Override
     public Iterator<Map.Entry<TensorAddress, Double>> cellIterator() {
+        return new CellIterator();
+    }
+
+    /**
+     * Returns an iterator over the values of this.
+     * Values are returned in order of increasing indexes in each dimension, increasing 
+     * indexes of later dimensions in the dimension type before earlier.
+     */
+    @Override
+    public Iterator<Double> valueIterator() {
         return new ValueIterator();
     }
 
@@ -383,7 +397,7 @@ public class IndexedTensor implements Tensor {
 
     }
     
-    private class ValueIterator implements Iterator<Map.Entry<TensorAddress, Double>> {
+    private final class CellIterator implements Iterator<Map.Entry<TensorAddress, Double>> {
 
         private int cursor = 0;
         private final int[] tensorIndexes = new int[dimensionSizes.length];
@@ -431,4 +445,25 @@ public class IndexedTensor implements Tensor {
 
     }
 
+    private final class ValueIterator implements Iterator<Double> {
+
+        private int cursor = 0;
+
+        @Override
+        public boolean hasNext() {
+            return cursor < values.length;
+        }
+
+        @Override
+        public Double next() {
+            try {
+                return values[cursor++];
+            }
+            catch (IndexOutOfBoundsException e) {
+                throw new NoSuchElementException("No element at position " + cursor);
+            }
+        }
+
+    }
+
 }
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/MappedTensor.java b/vespajlib/src/main/java/com/yahoo/tensor/MappedTensor.java
index c3dce27c651..f87a97cb44e 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/MappedTensor.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/MappedTensor.java
@@ -38,6 +38,9 @@ public class MappedTensor implements Tensor {
     public Iterator<Map.Entry<TensorAddress, Double>> cellIterator() { return cells.entrySet().iterator(); }
 
     @Override
+    public Iterator<Double> valueIterator() { return cells.values().iterator(); }
+
+    @Override
     public Map<TensorAddress, Double> cells() { return cells; }
 
     @Override
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/Tensor.java b/vespajlib/src/main/java/com/yahoo/tensor/Tensor.java
index cc40f84ccd3..260c48ace7f 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/Tensor.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/Tensor.java
@@ -59,6 +59,8 @@ public interface Tensor {
 
     Iterator<Map.Entry<TensorAddress, Double>> cellIterator();
 
+    Iterator<Double> valueIterator();
+
     /** Returns an immutable map of the cells of this. This may be expensive for some implementations - avoid when possible */
     Map<TensorAddress, Double> cells();
 
@@ -73,7 +75,7 @@ public interface Tensor {
         if (size() == 0) return Double.NaN;
         if (size() > 1)
             throw new IllegalStateException("This tensor does not have a single value, it has " + size());
-        return cellIterator().next().getValue();
+        return valueIterator().next();
     }
     
     // ----------------- Primitive tensor functions
@@ -250,6 +252,7 @@ public interface Tensor {
     interface Builder {
         
         /** Creates a suitable builder for the given type */
+        // TODO: Create version of this which takes size info and use it when possible
         static Builder of(TensorType type) {
             boolean containsIndexed = type.dimensions().stream().anyMatch(d -> d.isIndexed());
             boolean containsMapped = type.dimensions().stream().anyMatch( d ->  ! d.isIndexed());
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 a48b2b5ae4f..19b4ad39af3 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java
@@ -83,9 +83,11 @@ public class Join extends PrimitiveTensorFunction {
     
     private Tensor indexedVectorJoin(IndexedTensor a, IndexedTensor b, TensorType type) {
         int joinedLength = Math.min(a.length(0), b.length(0));
+        Iterator<Double> aIterator = a.valueIterator();
+        Iterator<Double> bIterator = b.valueIterator();
         IndexedTensor.Builder builder = IndexedTensor.Builder.of(type, new int[] { joinedLength});
         for (int i = 0; i < joinedLength; i++)
-            builder.cell(combinator.applyAsDouble(a.get(i), b.get(i)), i);
+            builder.cell(combinator.applyAsDouble(aIterator.next(), bIterator.next()), i);
         return builder.build();
     }
 
@@ -114,8 +116,7 @@ public class Join extends PrimitiveTensorFunction {
                              reversedArgumentOrder ? combinator.applyAsDouble(supercell.getValue(), subspaceValue)
                                                    : combinator.applyAsDouble(subspaceValue, supercell.getValue()));
         }
-        Tensor result = builder.build();
-        return result;
+        return builder.build();
     }
     
     /** Returns the indexes in the superspace type which should be retained to create the subspace type */
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/Reduce.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/Reduce.java
index 96ab8bc6f60..7595d18fb18 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Reduce.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Reduce.java
@@ -142,8 +142,8 @@ public class Reduce extends PrimitiveTensorFunction {
     
     private Tensor reduceAllGeneral(Tensor argument) {
         ValueAggregator valueAggregator = ValueAggregator.ofType(aggregator);
-        for (Iterator<Map.Entry<TensorAddress, Double>> i = argument.cellIterator(); i.hasNext(); )
-            valueAggregator.aggregate(i.next().getValue());
+        for (Iterator<Double> i = argument.valueIterator(); i.hasNext(); )
+            valueAggregator.aggregate(i.next());
         return Tensor.Builder.of(TensorType.empty).cell((valueAggregator.aggregatedValue())).build();
     }
 
diff --git a/vespajlib/src/test/java/com/yahoo/tensor/TensorFunctionBenchmark.java b/vespajlib/src/test/java/com/yahoo/tensor/TensorFunctionBenchmark.java
index 55fa54f088d..bbab92fc16d 100644
--- a/vespajlib/src/test/java/com/yahoo/tensor/TensorFunctionBenchmark.java
+++ b/vespajlib/src/test/java/com/yahoo/tensor/TensorFunctionBenchmark.java
@@ -106,11 +106,11 @@ public class TensorFunctionBenchmark {
 
         // ---------------- Mapped with extra space (sidesteps current special-case optimizations):
         // Initial: 450 ms
-        time = new TensorFunctionBenchmark().benchmark(20, vectors(100, 300, TensorType.Dimension.Type.mapped), TensorType.Dimension.Type.mapped, true);
-        System.out.printf("Mapped vectors, x space  time per join: %1$8.3f ms\n", time);
+        //time = new TensorFunctionBenchmark().benchmark(20, vectors(100, 300, TensorType.Dimension.Type.mapped), TensorType.Dimension.Type.mapped, true);
+        //System.out.printf("Mapped vectors, x space  time per join: %1$8.3f ms\n", time);
         // Initial: 900 ms
-        time = new TensorFunctionBenchmark().benchmark(20, matrix(100, 300, TensorType.Dimension.Type.mapped), TensorType.Dimension.Type.mapped, true);
-        System.out.printf("Mapped matrix, x space   time per join: %1$8.3f ms\n", time);
+        //time = new TensorFunctionBenchmark().benchmark(20, matrix(100, 300, TensorType.Dimension.Type.mapped), TensorType.Dimension.Type.mapped, true);
+        //System.out.printf("Mapped matrix, x space   time per join: %1$8.3f ms\n", time);
 
         // ---------------- Mapped:
         // Initial: 150 ms
@@ -127,12 +127,12 @@ public class TensorFunctionBenchmark {
         // ---------------- Indexed (unbound) with extra space (sidesteps current special-case optimizations):
         // Initial: 1900 ms
         // - After moving to cell iterators: 1100
-        time = new TensorFunctionBenchmark().benchmark(20, vectors(100, 300, TensorType.Dimension.Type.indexedUnbound), TensorType.Dimension.Type.indexedUnbound, true);
-        System.out.printf("Indexed vectors, x space time per join: %1$8.3f ms\n", time);
+        //time = new TensorFunctionBenchmark().benchmark(20, vectors(100, 300, TensorType.Dimension.Type.indexedUnbound), TensorType.Dimension.Type.indexedUnbound, true);
+        //System.out.printf("Indexed vectors, x space time per join: %1$8.3f ms\n", time);
         // Initial: 2200 ms
         // - After moving to cell iterators: 1300
-        time = new TensorFunctionBenchmark().benchmark(20, matrix(100, 300, TensorType.Dimension.Type.indexedUnbound), TensorType.Dimension.Type.indexedUnbound, true);
-        System.out.printf("Indexed matrix, x space  time per join: %1$8.3f ms\n", time);
+        //time = new TensorFunctionBenchmark().benchmark(20, matrix(100, 300, TensorType.Dimension.Type.indexedUnbound), TensorType.Dimension.Type.indexedUnbound, true);
+        //System.out.printf("Indexed matrix, x space  time per join: %1$8.3f ms\n", time);
 
         // ---------------- Indexed unbound:
         // Initial: 718 ms