GH-87: [Vector] Add ExtensionWriter (#697)

Based on changes from apache/arrow#41731.

## What's Changed

Added writer ExtensionWriter with 3 methods:
- write method  for writing values from Extension holders;
- writeExtensionType method for writing values (arguments is Object
because we don't know exact type);
- addExtensionTypeFactory method - because the exact vector and value
type are unknown, the user should create their own extension type
vector, write for it, and ExtensionTypeFactory, which should map the
vector and writer.

Closes #87.

Co-authored-by: Finn Völkel <[email protected]>

Author: xxlaykxx

vector/src/main/codegen/templates/AbstractFieldWriter.java

@@ -107,6 +107,16 @@ public void endEntry() {
     throw new IllegalStateException(String.format("You tried to end a map entry when you are using a ValueWriter of type %s.", this.getClass().getSimpleName()));
   }
 
+  public void write(ExtensionHolder var1)  {
+    this.fail("ExtensionType");
+  }
+  public void writeExtension(Object var1)  {
+    this.fail("ExtensionType");
+  }
+  public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory var1) {
+    this.fail("ExtensionType");
+  }
+
   <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
   <#assign fields = minor.fields!type.fields />
   <#assign friendlyType = (minor.friendlyType!minor.boxedType!type.boxedType) />
@@ -241,6 +251,18 @@ public MapWriter map(String name, boolean keysSorted) {
     fail("Map");
     return null;
   }
+
+  @Override
+  public ExtensionWriter extension(String name, ArrowType arrowType) {
+    fail("Extension");
+    return null;
+  }
+
+  @Override
+  public ExtensionWriter extension(ArrowType arrowType) {
+    fail("Extension");
+    return null;
+  }
   <#list vv.types as type><#list type.minor as minor>
   <#assign lowerName = minor.class?uncap_first />
   <#if lowerName == "int" ><#assign lowerName = "integer" /></#if>

vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java

@@ -293,6 +293,11 @@ public MapWriter map(boolean keysSorted) {
     return getWriter(MinorType.MAP, new ArrowType.Map(keysSorted));
   }
 
+  @Override
+  public ExtensionWriter extension(ArrowType arrowType) {
+    return getWriter(MinorType.EXTENSIONTYPE).extension(arrowType);
+  }
+
   @Override
   public StructWriter struct(String name) {
     return getWriter(MinorType.STRUCT).struct(name);
@@ -318,6 +323,11 @@ public MapWriter map(String name, boolean keysSorted) {
     return getWriter(MinorType.STRUCT).map(name, keysSorted);
   }
 
+  @Override
+  public ExtensionWriter extension(String name, ArrowType arrowType) {
+    return getWriter(MinorType.EXTENSIONTYPE).extension(name, arrowType);
+  }
+
   <#list vv.types as type><#list type.minor as minor>
   <#assign lowerName = minor.class?uncap_first />
   <#if lowerName == "int" ><#assign lowerName = "integer" /></#if>

vector/src/main/codegen/templates/BaseWriter.java

@@ -61,6 +61,7 @@ public interface StructWriter extends BaseWriter {
 
     void copyReaderToField(String name, FieldReader reader);
     StructWriter struct(String name);
+    ExtensionWriter extension(String name, ArrowType arrowType);
     ListWriter list(String name);
     ListWriter listView(String name);
     MapWriter map(String name);
@@ -79,6 +80,7 @@ public interface ListWriter extends BaseWriter {
     ListWriter listView();
     MapWriter map();
     MapWriter map(boolean keysSorted);
+    ExtensionWriter extension(ArrowType arrowType);
     void copyReader(FieldReader reader);
 
     <#list vv.types as type><#list type.minor as minor>
@@ -101,6 +103,35 @@ public interface MapWriter extends ListWriter {
     MapWriter value();
   }
 
+  public interface ExtensionWriter extends BaseWriter {
+
+    /**
+     * Writes a null value.
+     */
+    void writeNull();
+
+    /**
+     * Writes value from the given extension holder.
+     *
+     * @param holder the extension holder to write
+     */
+    void write(ExtensionHolder holder);
+
+    /**
+     * Writes the given extension type value.
+     *
+     * @param value the extension type value to write
+     */
+    void writeExtension(Object value);
+
+    /**
+     * Adds the given extension type factory. This factory allows configuring writer implementations for specific ExtensionTypeVector.
+     *
+     * @param factory the extension type factory to add
+     */
+    void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory);
+  }
+
   public interface ScalarWriter extends
   <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> ${name}Writer, </#list></#list> BaseWriter {}
 

vector/src/main/codegen/templates/PromotableWriter.java

@@ -285,6 +285,9 @@ protected void setWriter(ValueVector v) {
       case UNION:
         writer = new UnionWriter((UnionVector) vector, nullableStructWriterFactory);
         break;
+      case EXTENSIONTYPE:
+        writer = new UnionExtensionWriter((ExtensionTypeVector) vector);
+        break;
       default:
         writer = type.getNewFieldWriter(vector);
         break;
@@ -316,6 +319,7 @@ protected boolean requiresArrowType(MinorType type) {
         || type == MinorType.MAP
         || type == MinorType.DURATION
         || type == MinorType.FIXEDSIZEBINARY
+        || type == MinorType.EXTENSIONTYPE
         || (type.name().startsWith("TIMESTAMP") && type.name().endsWith("TZ"));
   }
 
@@ -536,6 +540,16 @@ public void writeLargeVarChar(String value) {
     getWriter(MinorType.LARGEVARCHAR).writeLargeVarChar(value);
   }
 
+  @Override
+  public void writeExtension(Object value) {
+    getWriter(MinorType.EXTENSIONTYPE).writeExtension(value);
+  }
+
+  @Override
+  public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory) {
+    getWriter(MinorType.EXTENSIONTYPE).addExtensionTypeWriterFactory(factory);
+  }
+
   @Override
   public void allocate() {
     getWriter().allocate();

vector/src/main/codegen/templates/StructWriters.java

@@ -83,6 +83,9 @@ public class ${mode}StructWriter extends AbstractFieldWriter {
         fields.put(handleCase(child.getName()), writer);
         break;
       }
+      case EXTENSIONTYPE:
+        extension(child.getName(), child.getType());
+        break;
       case UNION:
         FieldType fieldType = new FieldType(addVectorAsNullable, MinorType.UNION.getType(), null, null);
         UnionWriter writer = new UnionWriter(container.addOrGet(child.getName(), fieldType, UnionVector.class), getNullableStructWriterFactory());
@@ -159,6 +162,29 @@ public StructWriter struct(String name) {
     return writer;
   }
 
+  @Override
+  public ExtensionWriter extension(String name, ArrowType arrowType) {
+    String finalName = handleCase(name);
+    FieldWriter writer = fields.get(finalName);
+    if(writer == null){
+      int vectorCount=container.size();
+      FieldType fieldType = new FieldType(addVectorAsNullable, arrowType, null, null);
+      ExtensionTypeVector vector = container.addOrGet(name, fieldType, ExtensionTypeVector.class);
+      writer = new PromotableWriter(vector, container, getNullableStructWriterFactory());
+      if(vectorCount != container.size()) {
+        writer.allocate();
+      }
+      writer.setPosition(idx());
+      fields.put(finalName, writer);
+    } else {
+      if (writer instanceof PromotableWriter) {
+        // ensure writers are initialized
+        ((PromotableWriter)writer).getWriter(MinorType.EXTENSIONTYPE, arrowType);
+      }
+    }
+    return (ExtensionWriter) writer;
+  }
+
   @Override
   public void close() throws Exception {
     clear();

vector/src/main/codegen/templates/UnionListWriter.java

@@ -201,6 +201,17 @@ public MapWriter map(String name, boolean keysSorted) {
     return mapWriter;
   }
 
+  @Override
+  public ExtensionWriter extension(ArrowType arrowType) {
+    writer.extension(arrowType);
+    return writer;
+  }
+  @Override
+  public ExtensionWriter extension(String name, ArrowType arrowType) {
+    ExtensionWriter extensionWriter = writer.extension(name, arrowType);
+    return extensionWriter;
+  }
+
   <#if listName == "LargeList">
   @Override
   public void startList() {
@@ -323,6 +334,18 @@ public void writeNull() {
     }
   }
 
+  @Override
+  public void writeExtension(Object value) {
+    writer.writeExtension(value);
+  }
+  @Override
+  public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory var1) {
+    writer.addExtensionTypeWriterFactory(var1);
+  }
+  public void write(ExtensionHolder var1) {
+    writer.write(var1);
+  }
+
   <#list vv.types as type>
     <#list type.minor as minor>
       <#assign name = minor.class?cap_first />

vector/src/main/codegen/templates/UnionMapWriter.java

@@ -231,4 +231,16 @@ public MapWriter map() {
         return super.map();
     }
   }
+
+  @Override
+  public ExtensionWriter extension(ArrowType type) {
+    switch (mode) {
+      case KEY:
+        return entryWriter.extension(MapVector.KEY_NAME, type);
+      case VALUE:
+        return entryWriter.extension(MapVector.VALUE_NAME, type);
+      default:
+        return super.extension(type);
+    }
+  }
 }

vector/src/main/codegen/templates/UnionWriter.java

@@ -213,6 +213,10 @@ public MapWriter asMap(ArrowType arrowType) {
     return getMapWriter(arrowType);
   }
 
+  private ExtensionWriter getExtensionWriter(ArrowType arrowType) {
+    throw new UnsupportedOperationException("ExtensionTypes are not supported yet.");
+  }
+
   BaseWriter getWriter(MinorType minorType) {
     return getWriter(minorType, null);
   }
@@ -227,6 +231,8 @@ BaseWriter getWriter(MinorType minorType, ArrowType arrowType) {
       return getListViewWriter();
     case MAP:
       return getMapWriter(arrowType);
+    case EXTENSIONTYPE:
+      return getExtensionWriter(arrowType);
     <#list vv.types as type>
       <#list type.minor as minor>
         <#assign name = minor.class?cap_first />
@@ -460,6 +466,20 @@ public MapWriter map(String name, boolean keysSorted) {
     return getStructWriter().map(name, keysSorted);
   }
 
+  @Override
+  public ExtensionWriter extension(ArrowType arrowType) {
+    data.setType(idx(), MinorType.EXTENSIONTYPE);
+    getListWriter().setPosition(idx());
+    return getListWriter().extension(arrowType);
+  }
+
+  @Override
+  public ExtensionWriter extension(String name, ArrowType arrowType) {
+    data.setType(idx(), MinorType.EXTENSIONTYPE);
+    getStructWriter().setPosition(idx());
+    return getStructWriter().extension(name, arrowType);
+  }
+
   <#list vv.types as type><#list type.minor as minor>
   <#assign lowerName = minor.class?uncap_first />
   <#if lowerName == "int" ><#assign lowerName = "integer" /></#if>

vector/src/main/java/org/apache/arrow/vector/complex/impl/AbstractExtensionTypeWriter.java

@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.vector.complex.impl;
+
+import org.apache.arrow.vector.ExtensionTypeVector;
+import org.apache.arrow.vector.types.pojo.Field;
+
+/**
+ * Base {@link AbstractFieldWriter} class for an {@link
+ * org.apache.arrow.vector.ExtensionTypeVector}.
+ *
+ * @param <T> a specific {@link ExtensionTypeVector}.
+ */
+public class AbstractExtensionTypeWriter<T extends ExtensionTypeVector>
+    extends AbstractFieldWriter {
+  protected final T vector;
+
+  public AbstractExtensionTypeWriter(T vector) {
+    this.vector = vector;
+  }
+
+  @Override
+  public Field getField() {
+    return this.vector.getField();
+  }
+
+  @Override
+  public int getValueCapacity() {
+    return this.vector.getValueCapacity();
+  }
+
+  @Override
+  public void allocate() {
+    this.vector.allocateNew();
+  }
+
+  @Override
+  public void close() {
+    this.vector.close();
+  }
+
+  @Override
+  public void clear() {
+    this.vector.clear();
+  }
+
+  @Override
+  public void writeNull() {
+    this.vector.setNull(getPosition());
+    this.vector.setValueCount(getPosition() + 1);
+  }
+}

vector/src/main/java/org/apache/arrow/vector/complex/impl/ExtensionTypeWriterFactory.java

@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.vector.complex.impl;
+
+import org.apache.arrow.vector.ExtensionTypeVector;
+import org.apache.arrow.vector.complex.writer.FieldWriter;
+
+/**
+ * A factory interface for creating instances of {@link ExtensionTypeWriter}. This factory allows
+ * configuring writer implementations for specific {@link ExtensionTypeVector}.
+ *
+ * @param <T> the type of writer implementation for a specific {@link ExtensionTypeVector}.
+ */
+public interface ExtensionTypeWriterFactory<T extends FieldWriter> {
+
+  /**
+   * Returns an instance of the writer implementation for the given {@link ExtensionTypeVector}.
+   *
+   * @param vector the {@link ExtensionTypeVector} for which the writer implementation is to be
+   *     returned.
+   * @return an instance of the writer implementation for the given {@link ExtensionTypeVector}.
+   */
+  T getWriterImpl(ExtensionTypeVector vector);
+}