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

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

Based on changes from apache#41731.

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 <finn.volkel@gmail.com>

* just checking

* Update Brewfile

* revert

* Fix for Cmake 3.xx installation

* Fix for Cmake 3.xx installation

* update thrift

* update thrift url

* update thrift url

* Update github.yml

* fix format

---------

Co-authored-by: Finn Völkel <finn.volkel@gmail.com>

Author: xxlaykxx

cpp/cmake_modules/ThirdpartyToolchain.cmake

@@ -817,6 +817,7 @@ if(DEFINED ENV{ARROW_THRIFT_URL})
   set(THRIFT_SOURCE_URL "$ENV{ARROW_THRIFT_URL}")
 else()
   set(THRIFT_SOURCE_URL
+      "https://www.apache.org/dyn/closer.lua/thrift/${ARROW_THRIFT_BUILD_VERSION}/thrift-${ARROW_THRIFT_BUILD_VERSION}.tar.gz?action=download"
       "https://dlcdn.apache.org/thrift/${ARROW_THRIFT_BUILD_VERSION}/thrift-${ARROW_THRIFT_BUILD_VERSION}.tar.gz"
   )
 endif()

cpp/thirdparty/versions.txt

@@ -164,7 +164,7 @@ DEPENDENCIES=(
   "ARROW_RE2_URL re2-${ARROW_RE2_BUILD_VERSION}.tar.gz https://github.com/google/re2/archive/${ARROW_RE2_BUILD_VERSION}.tar.gz"
   "ARROW_S2N_TLS_URL s2n-${ARROW_S2N_TLS_BUILD_VERSION}.tar.gz https://github.com/aws/s2n-tls/archive/${ARROW_S2N_TLS_BUILD_VERSION}.tar.gz"
   "ARROW_SNAPPY_URL snappy-${ARROW_SNAPPY_BUILD_VERSION}.tar.gz https://github.com/google/snappy/archive/${ARROW_SNAPPY_BUILD_VERSION}.tar.gz"
-  "ARROW_THRIFT_URL thrift-${ARROW_THRIFT_BUILD_VERSION}.tar.gz https://dlcdn.apache.org/thrift/${ARROW_THRIFT_BUILD_VERSION}/thrift-${ARROW_THRIFT_BUILD_VERSION}.tar.gz"
+  "ARROW_THRIFT_URL thrift-${ARROW_THRIFT_BUILD_VERSION}.tar.gz https://www.apache.org/dyn/closer.lua/thrift/${ARROW_THRIFT_BUILD_VERSION}/thrift-${ARROW_THRIFT_BUILD_VERSION}.tar.gz?action=download"
   "ARROW_UCX_URL ucx-${ARROW_UCX_BUILD_VERSION}.tar.gz https://github.com/openucx/ucx/archive/v${ARROW_UCX_BUILD_VERSION}.tar.gz"
   "ARROW_UTF8PROC_URL utf8proc-${ARROW_UTF8PROC_BUILD_VERSION}.tar.gz https://github.com/JuliaStrings/utf8proc/archive/${ARROW_UTF8PROC_BUILD_VERSION}.tar.gz"
   "ARROW_XSIMD_URL xsimd-${ARROW_XSIMD_BUILD_VERSION}.tar.gz https://github.com/xtensor-stack/xsimd/archive/${ARROW_XSIMD_BUILD_VERSION}.tar.gz"

dev/tasks/java-jars/github.yml

@@ -148,6 +148,18 @@ jobs:
           # used on test  We uninstall Homebrew's Protobuf to ensure using
           # bundled Protobuf.
           brew uninstall protobuf
+          # fix cmake and boost versionsAdd commentMore actions
+          brew uninstall -f boost || true
+          brew uninstall -f cmake || true
+          mkdir -p homebrew-custom/Formula
+          curl -o homebrew-custom/Formula/cmake.rb https://raw.githubusercontent.com/Homebrew/homebrew-core/f68532bfe5cb87474093df8a839c3818c6aa44dd/Formula/c/cmake.rb
+          curl -o homebrew-custom/Formula/boost.rb https://raw.githubusercontent.com/Homebrew/homebrew-core/23f9c56c5075dd56b4471e2c93f89f6400b49ddd/Formula/b/boost.rb
+          brew install -v ./homebrew-custom/Formula/cmake.rb
+          brew install -v ./homebrew-custom/Formula/boost.rb
+          brew pin cmake
+          brew pin boost
+          #
+
 
           brew bundle --file=arrow/java/Brewfile
       - name: Build C++ libraries

java/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>

java/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>

java/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 {}
 

java/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();

java/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();

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

@@ -201,6 +201,18 @@ 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 +335,20 @@ 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 />

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