Merge pull request #4 from dremio/extension_writer_dremio_26.1_19.0.1

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

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);
+}