Tensor implementation independent builder API

1 files changed, 37 insertions, 3 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java b/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java
index e946f80eb97..faa1ca9a583 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java
@@ -218,7 +218,7 @@ public class IndexedTensor implements Tensor {
 
     }
 
-    public static class Builder {
+    public static class Builder implements Tensor.Builder {
     
         private final TensorType type;
         
@@ -229,6 +229,7 @@ public class IndexedTensor implements Tensor {
             this.type = type;
         }
     
+        @Override
         public IndexedTensor build() {
             // TODO: Enforce that all values in all dimensions are equally large
             if (firstDimension == null) // empty
@@ -267,7 +268,7 @@ public class IndexedTensor implements Tensor {
         }
 
         /**
-         * Set a value. The number of indexes must be the same as the dimensions in the type of this.
+         * Set a value using an index API. The number of indexes must be the same as the dimensions in the type of this.
          * Values can be written in any order but all values needed to make this dense must be provided
          * before building this.
          * 
@@ -300,11 +301,44 @@ public class IndexedTensor implements Tensor {
             return this;
         }
         
+        @Override
+        public IndexedCellBuilder cell() {
+            return new IndexedCellBuilder();
+        }
+
         /** Fill the given list with nulls if necessary to make sure it has a (possibly null) value at the given index */
         private void ensureCapacity(int index, List<Object> list) {
             while (list.size() <= index)
                 list.add(list.size(), null);
         }
-    
+
+        public class IndexedCellBuilder implements Tensor.Builder.CellBuilder {
+
+            private final TensorAddress.Builder addressBuilder = new TensorAddress.Builder(IndexedTensor.Builder.this.type);
+
+            @Override
+            public IndexedTensor.Builder.IndexedCellBuilder label(String dimension, String label) {
+                addressBuilder.add(dimension, label);
+                return this;
+            }
+
+            @Override
+            public IndexedTensor.Builder value(double cellValue) {
+                TensorAddress address = addressBuilder.build();
+                int[] indexes = new int[address.labels().size()];
+                for (int i = 0; i < address.labels().size(); i++) {
+                    try {
+                        indexes[i] = Integer.parseInt(address.labels().get(i));
+                    } catch (NumberFormatException e) {
+                        throw new IllegalArgumentException("Labels in an indexed tensor must be integers, not '" + 
+                                                           address.labels().get(i) + "'");
+                    }
+                }
+                IndexedTensor.Builder.this.set(cellValue, indexes);
+                return IndexedTensor.Builder.this;
+            }
+
+        }
+
     }
 }