Merge pull request #2427 from yahoo/lesters/add-tensors-to-documentgen

Add tensor support to vespa-documentgen-plugin

4 files changed, 27 insertions, 1 deletions

diff --git a/document/src/main/java/com/yahoo/document/datatypes/TensorFieldValue.java b/document/src/main/java/com/yahoo/document/datatypes/TensorFieldValue.java
index df09b312738..4d18c5ca51a 100644
--- a/document/src/main/java/com/yahoo/document/datatypes/TensorFieldValue.java
+++ b/document/src/main/java/com/yahoo/document/datatypes/TensorFieldValue.java
@@ -95,5 +95,10 @@ public class TensorFieldValue extends FieldValue {
         return true;
     }
 
+    @Override
+    public Object getWrappedValue() {
+        return tensor.orElse(null);
+    }
+
 }
 
diff --git a/documentgen-test/etc/complex/book.sd b/documentgen-test/etc/complex/book.sd
index d25995044c2..bccf72e057e 100644
--- a/documentgen-test/etc/complex/book.sd
+++ b/documentgen-test/etc/complex/book.sd
@@ -69,6 +69,10 @@ search book {
         field ref type reference<parent> {
             indexing: attribute
         }
+        field vector type tensor(x{}) {
+            indexing: attribute | summary
+            attribute: tensor(x{})
+        }
     }
     field sw1 type float {
     }
diff --git a/documentgen-test/src/test/java/com/yahoo/vespa/config/DocumentGenPluginTest.java b/documentgen-test/src/test/java/com/yahoo/vespa/config/DocumentGenPluginTest.java
index 7e9365a0e6b..b113dc365ba 100644
--- a/documentgen-test/src/test/java/com/yahoo/vespa/config/DocumentGenPluginTest.java
+++ b/documentgen-test/src/test/java/com/yahoo/vespa/config/DocumentGenPluginTest.java
@@ -14,6 +14,7 @@ import com.yahoo.document.datatypes.*;
 import com.yahoo.document.serialization.*;
 import com.yahoo.io.GrowableByteBuffer;
 import com.yahoo.searchdefinition.derived.Deriver;
+import com.yahoo.tensor.Tensor;
 import com.yahoo.vespa.document.NodeImpl;
 import com.yahoo.vespa.document.dom.DocumentImpl;
 import com.yahoo.vespa.documentgen.test.*;
@@ -918,5 +919,13 @@ public class DocumentGenPluginTest {
         System.out.println(unmasked);
         assertEquals(unmasked.size(), 0); // probably not needed
     }
+
+    @Test
+    public void testTensorType() {
+        Book book = new Book(new DocumentId("doc:book:0"));
+        assertNull(book.getVector());
+        book.setVector(Tensor.from("{{x:0}:1.0, {x:1}:2.0, {x:2}:3.0}"));
+        assertEquals("{{x:0}:1.0,{x:1}:2.0,{x:2}:3.0}", book.getVector().toString());
+    }
     
 }
diff --git a/vespa-documentgen-plugin/src/main/java/com/yahoo/vespa/DocumentGenMojo.java b/vespa-documentgen-plugin/src/main/java/com/yahoo/vespa/DocumentGenMojo.java
index 67627ca9410..fff93ec4525 100644
--- a/vespa-documentgen-plugin/src/main/java/com/yahoo/vespa/DocumentGenMojo.java
+++ b/vespa-documentgen-plugin/src/main/java/com/yahoo/vespa/DocumentGenMojo.java
@@ -10,6 +10,7 @@ import com.yahoo.searchdefinition.Search;
 import com.yahoo.searchdefinition.SearchBuilder;
 import com.yahoo.searchdefinition.UnprocessingSearchBuilder;
 import com.yahoo.searchdefinition.parser.ParseException;
+import com.yahoo.tensor.TensorType;
 import org.apache.maven.plugin.AbstractMojo;
 import org.apache.maven.plugin.MojoExecutionException;
 import org.apache.maven.plugin.MojoFailureException;
@@ -841,6 +842,9 @@ public class DocumentGenMojo extends AbstractMojo {
         if (dt instanceof ReferenceDataType) {
             return "com.yahoo.document.DocumentId";
         }
+        if (dt instanceof TensorDataType) {
+            return "com.yahoo.tensor.Tensor";
+        }
         return "byte[]";
     }
 
@@ -871,7 +875,11 @@ public class DocumentGenMojo extends AbstractMojo {
             return String.format("new com.yahoo.document.ReferenceDataType(%s.type, %d)",
                     className(((ReferenceDataType) dt).getTargetType().getName()), dt.getId());
         }
-        return "DataType.RAW";
+        if (dt instanceof TensorDataType) {
+            return String.format("new com.yahoo.document.TensorDataType(com.yahoo.tensor.TensorType.fromSpec(\"%s\"))",
+                    ((TensorDataType)dt).getTensorType().toString());
+        }
+        return "com.yahoo.document.DataType.RAW";
     }
 
     @Override