Require equal sizes in join

1 files changed, 17 insertions, 9 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java b/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java
index 58cb151875e..a499807105a 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java
@@ -319,7 +319,7 @@ public class TensorType {
          * [N] + [] = []
          * [] + {} = {}
          */
-        Dimension combineWith(Optional<Dimension> other) {
+        Dimension combineWith(Optional<Dimension> other, boolean allowDifferentSizes) {
             if ( ! other.isPresent()) return this;
             if (this instanceof MappedDimension) return this;
             if (other.get() instanceof MappedDimension) return other.get();
@@ -329,7 +329,11 @@ public class TensorType {
             // both are indexed bound
             IndexedBoundDimension thisIb = (IndexedBoundDimension)this;
             IndexedBoundDimension otherIb = (IndexedBoundDimension)other.get();
-            return thisIb.size().get() < otherIb.size().get() ? thisIb : otherIb;
+            if (allowDifferentSizes)
+                return thisIb.size().get() < otherIb.size().get() ? thisIb : otherIb;
+            if (  ! thisIb.size().equals(otherIb.size()))
+                throw new IllegalArgumentException("Unequal dimension sizes in " + thisIb + " and " + otherIb);
+            return thisIb;
         }
 
         @Override
@@ -483,16 +487,20 @@ public class TensorType {
         /**
          * Creates a builder containing a combination of the dimensions of the given types
          *
-         * If the same dimension is indexed with different size restrictions the largest size will be used.
+         * If the same dimension is indexed with different size restrictions the smallest size will be used.
          * If it is size restricted in one argument but not the other it will not be size restricted.
          * If it is indexed in one and mapped in the other it will become mapped.
          *
          * The value type will be the largest of the value types of the input types
          */
         public Builder(TensorType ... types) {
+            this(true, types);
+        }
+
+        public Builder(boolean allowDifferentSizes, TensorType ... types) {
             this.valueType = TensorType.combinedValueType(types);
             for (TensorType type : types)
-                addDimensionsOf(type);
+                addDimensionsOf(type, allowDifferentSizes);
         }
 
         /** Creates a builder from the given dimensions, having double as the value type */
@@ -510,17 +518,17 @@ public class TensorType {
 
         private static final boolean supportsMixedTypes = false;
 
-        private void addDimensionsOf(TensorType type) {
+        private void addDimensionsOf(TensorType type, boolean allowDifferentSizes) {
             if ( ! supportsMixedTypes) {  // TODO: Support it
-                addDimensionsOfAndDisallowMixedDimensions(type);
+                addDimensionsOfAndDisallowMixedDimensions(type, allowDifferentSizes);
             }
             else {
                 for (Dimension dimension : type.dimensions)
-                    set(dimension.combineWith(Optional.ofNullable(dimensions.get(dimension.name()))));
+                    set(dimension.combineWith(Optional.ofNullable(dimensions.get(dimension.name())), allowDifferentSizes));
             }
         }
 
-        private void addDimensionsOfAndDisallowMixedDimensions(TensorType type) {
+        private void addDimensionsOfAndDisallowMixedDimensions(TensorType type, boolean allowDifferentSizes) {
             boolean containsMapped = dimensions.values().stream().anyMatch(d -> ! d.isIndexed());
             containsMapped = containsMapped || type.dimensions().stream().anyMatch(d -> ! d.isIndexed());
 
@@ -528,7 +536,7 @@ public class TensorType {
                 if (containsMapped)
                     dimension = new MappedDimension(dimension.name());
                 Dimension existing = dimensions.get(dimension.name());
-                set(dimension.combineWith(Optional.ofNullable(existing)));
+                set(dimension.combineWith(Optional.ofNullable(existing), allowDifferentSizes));
             }
         }