apacheGH-38614: [Java] Add VarBinary and VarCharWriter helper methods…

… to more writers (apache#38631) ### Rationale for this change Add the overrides for new convenience Writer methods added to VarCharWriter and VarBinaryWriter so that classes that use composition such as UnionWriter and PromotableWriter can invoke them properly. ### What changes are included in this PR? - Rename from writeTo$type to write$type for consistency with other methods - Add new helper methods to PromotableWriter - Add new helper methods to complex writers such as list and union ### Are these changes tested? Yes. New unit tests added for several Writer classes. **This PR includes breaking changes to public APIs.** The writeTo<Type>() and similar methods in Writers have been renamed to just write<Type>() * Closes: apache#38614 Authored-by: James Duong <james.duong@improving.com> Signed-off-by: David Li <li.davidm96@gmail.com>
Bit-Quill · Nov 9, 2023 · 0acf4c8 · 0acf4c8
1 parent 79e328b
commit 0acf4c8
Show file tree

Hide file tree

Showing 11 changed files with 741 additions and 26 deletions.
diff --git a/java/vector/src/main/codegen/templates/AbstractFieldWriter.java b/java/vector/src/main/codegen/templates/AbstractFieldWriter.java
@@ -27,6 +27,9 @@
 
 /*
  * This class is generated using freemarker and the ${.template_name} template.
+ * Note that changes to the AbstractFieldWriter template should also get reflected in the
+ * AbstractPromotableFieldWriter, ComplexWriters, UnionFixedSizeListWriter, UnionListWriter
+ * and UnionWriter templates and the PromotableWriter concrete code.
  */
 @SuppressWarnings("unused")
 abstract class AbstractFieldWriter extends AbstractBaseWriter implements FieldWriter {
@@ -125,19 +128,19 @@ public void write(${name}Holder holder) {
   </#if>
 
   <#if minor.class?ends_with("VarBinary")>
-  public void writeTo${minor.class}(byte[] value) {
+  public void write${minor.class}(byte[] value) {
     fail("${name}");
   }
 
-  public void writeTo${minor.class}(byte[] value, int offset, int length) {
+  public void write${minor.class}(byte[] value, int offset, int length) {
     fail("${name}");
   }
 
-  public void writeTo${minor.class}(ByteBuffer value) {
+  public void write${minor.class}(ByteBuffer value) {
     fail("${name}");
   }
 
-  public void writeTo${minor.class}(ByteBuffer value, int offset, int length) {
+  public void write${minor.class}(ByteBuffer value, int offset, int length) {
     fail("${name}");
   }
   </#if>

diff --git a/java/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java b/java/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java
@@ -221,6 +221,38 @@ public void write(${name}Holder holder) {
   }
   </#if>
 
+  <#if minor.class?ends_with("VarBinary")>
+  @Override
+  public void write${minor.class}(byte[] value) {
+    getWriter(MinorType.${name?upper_case}).write${minor.class}(value);
+  }
+
+  @Override
+  public void write${minor.class}(byte[] value, int offset, int length) {
+    getWriter(MinorType.${name?upper_case}).write${minor.class}(value, offset, length);
+  }
+
+  @Override
+  public void write${minor.class}(ByteBuffer value) {
+    getWriter(MinorType.${name?upper_case}).write${minor.class}(value);
+  }
+
+  @Override
+  public void write${minor.class}(ByteBuffer value, int offset, int length) {
+    getWriter(MinorType.${name?upper_case}).write${minor.class}(value, offset, length);
+  }
+  <#elseif minor.class?ends_with("VarChar")>
+  @Override
+  public void write${minor.class}(Text value) {
+    getWriter(MinorType.${name?upper_case}).write${minor.class}(value);
+  }
+
+  @Override
+  public void write${minor.class}(String value) {
+    getWriter(MinorType.${name?upper_case}).write${minor.class}(value);
+  }
+  </#if>
+
   </#list></#list>
   public void writeNull() {
   }

diff --git a/java/vector/src/main/codegen/templates/ComplexWriters.java b/java/vector/src/main/codegen/templates/ComplexWriters.java
@@ -194,22 +194,22 @@ public void writeNull() {
   </#if>
 
   <#if minor.class?ends_with("VarBinary")>
-  public void writeTo${minor.class}(byte[] value) {
+  public void write${minor.class}(byte[] value) {
     vector.setSafe(idx(), value);
     vector.setValueCount(idx() + 1);
   }
 
-  public void writeTo${minor.class}(byte[] value, int offset, int length) {
+  public void write${minor.class}(byte[] value, int offset, int length) {
     vector.setSafe(idx(), value, offset, length);
     vector.setValueCount(idx() + 1);
   }
 
-  public void writeTo${minor.class}(ByteBuffer value) {
+  public void write${minor.class}(ByteBuffer value) {
     vector.setSafe(idx(), value, 0, value.remaining());
     vector.setValueCount(idx() + 1);
   }
 
-  public void writeTo${minor.class}(ByteBuffer value, int offset, int length) {
+  public void write${minor.class}(ByteBuffer value, int offset, int length) {
     vector.setSafe(idx(), value, offset, length);
     vector.setValueCount(idx() + 1);
   }
@@ -259,13 +259,13 @@ public interface ${eName}Writer extends BaseWriter {
 </#if>
 
 <#if minor.class?ends_with("VarBinary")>
-  public void writeTo${minor.class}(byte[] value);
+  public void write${minor.class}(byte[] value);
 
-  public void writeTo${minor.class}(byte[] value, int offset, int length);
+  public void write${minor.class}(byte[] value, int offset, int length);
 
-  public void writeTo${minor.class}(ByteBuffer value);
+  public void write${minor.class}(ByteBuffer value);
 
-  public void writeTo${minor.class}(ByteBuffer value, int offset, int length);
+  public void write${minor.class}(ByteBuffer value, int offset, int length);
 </#if>
 
 <#if minor.class?ends_with("VarChar")>

diff --git a/java/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java b/java/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java
@@ -295,6 +295,62 @@ public void writeBigEndianBytesToDecimal256(byte[] value, ArrowType arrowType) {
       <#assign name = minor.class?cap_first />
       <#assign fields = minor.fields!type.fields />
       <#assign uncappedName = name?uncap_first/>
+      <#if minor.class?ends_with("VarBinary")>
+  @Override
+  public void write${minor.class}(byte[] value) {
+    if (writer.idx() >= (idx() + 1) * listSize) {
+      throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
+    }
+    writer.write${minor.class}(value);
+    writer.setPosition(writer.idx() + 1);
+  }
+
+  @Override
+  public void write${minor.class}(byte[] value, int offset, int length) {
+    if (writer.idx() >= (idx() + 1) * listSize) {
+      throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
+    }
+    writer.write${minor.class}(value, offset, length);
+    writer.setPosition(writer.idx() + 1);
+  }
+
+  @Override
+  public void write${minor.class}(ByteBuffer value) {
+    if (writer.idx() >= (idx() + 1) * listSize) {
+      throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
+    }
+    writer.write${minor.class}(value);
+    writer.setPosition(writer.idx() + 1);
+  }
+
+  @Override
+  public void write${minor.class}(ByteBuffer value, int offset, int length) {
+    if (writer.idx() >= (idx() + 1) * listSize) {
+      throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
+    }
+    writer.write${minor.class}(value, offset, length);
+    writer.setPosition(writer.idx() + 1);
+  }
+  <#elseif minor.class?ends_with("VarChar")>
+  @Override
+  public void write${minor.class}(Text value) {
+    if (writer.idx() >= (idx() + 1) * listSize) {
+      throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
+    }
+    writer.write${minor.class}(value);
+    writer.setPosition(writer.idx() + 1);
+  }
+
+  @Override
+  public void write${minor.class}(String value) {
+    if (writer.idx() >= (idx() + 1) * listSize) {
+      throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
+    }
+    writer.write${minor.class}(value);
+    writer.setPosition(writer.idx() + 1);
+  }
+      </#if>
+
       <#if !minor.typeParams?? >
   @Override
   public void write${name}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, </#if></#list>) {

diff --git a/java/vector/src/main/codegen/templates/UnionListWriter.java b/java/vector/src/main/codegen/templates/UnionListWriter.java
@@ -274,6 +274,43 @@ public void write(${name}Holder holder) {
     writer.write${name}(<#list fields as field>holder.${field.name}<#if field_has_next>, </#if></#list>);
     writer.setPosition(writer.idx()+1);
   }
+  </#if>
+
+  <#if minor.class?ends_with("VarBinary")>
+  @Override
+  public void write${minor.class}(byte[] value) {
+    writer.write${minor.class}(value);
+    writer.setPosition(writer.idx() + 1);
+  }
+
+  @Override
+  public void write${minor.class}(byte[] value, int offset, int length) {
+    writer.write${minor.class}(value, offset, length);
+    writer.setPosition(writer.idx() + 1);
+  }
+
+  @Override
+  public void write${minor.class}(ByteBuffer value) {
+    writer.write${minor.class}(value);
+    writer.setPosition(writer.idx() + 1);
+  }
+
+  @Override
+  public void write${minor.class}(ByteBuffer value, int offset, int length) {
+    writer.write${minor.class}(value, offset, length);
+    writer.setPosition(writer.idx() + 1);
+  }
+  <#elseif minor.class?ends_with("VarChar")>
+  @Override
+  public void write${minor.class}(Text value) {
+    writer.write${minor.class}(value);
+    writer.setPosition(writer.idx() + 1);
+  }
+
+  public void write${minor.class}(String value) {
+    writer.write${minor.class}(value);
+    writer.setPosition(writer.idx() + 1);
+  }
   </#if>
 
     </#list>

diff --git a/java/vector/src/main/codegen/templates/UnionWriter.java b/java/vector/src/main/codegen/templates/UnionWriter.java
@@ -302,6 +302,42 @@ public void write(${name}Holder holder) {
     get${name}Writer(arrowType).setPosition(idx());
     get${name}Writer(arrowType).writeBigEndianBytesTo${name}(value, arrowType);
   }
+  <#elseif minor.class?ends_with("VarBinary")>
+  @Override
+  public void write${minor.class}(byte[] value) {
+    get${name}Writer().setPosition(idx());
+    get${name}Writer().write${minor.class}(value);
+  }
+
+  @Override
+  public void write${minor.class}(byte[] value, int offset, int length) {
+    get${name}Writer().setPosition(idx());
+    get${name}Writer().write${minor.class}(value, offset, length);
+  }
+
+  @Override
+  public void write${minor.class}(ByteBuffer value) {
+    get${name}Writer().setPosition(idx());
+    get${name}Writer().write${minor.class}(value);
+  }
+
+  @Override
+  public void write${minor.class}(ByteBuffer value, int offset, int length) {
+    get${name}Writer().setPosition(idx());
+    get${name}Writer().write${minor.class}(value, offset, length);
+  }
+  <#elseif minor.class?ends_with("VarChar")>
+  @Override
+  public void write${minor.class}(${friendlyType} value) {
+    get${name}Writer().setPosition(idx());
+    get${name}Writer().write${minor.class}(value);
+  }
+
+  @Override
+  public void write${minor.class}(String value) {
+    get${name}Writer().setPosition(idx());
+    get${name}Writer().write${minor.class}(value);
+  }
   </#if>
       </#if>
     </#list>

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/PromotableWriter.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/PromotableWriter.java
@@ -18,6 +18,7 @@
 package org.apache.arrow.vector.complex.impl;
 
 import java.math.BigDecimal;
+import java.nio.ByteBuffer;
 
 import org.apache.arrow.memory.ArrowBuf;
 import org.apache.arrow.vector.FieldVector;
@@ -37,6 +38,7 @@
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.FieldType;
+import org.apache.arrow.vector.util.Text;
 import org.apache.arrow.vector.util.TransferPair;
 
 /**
@@ -378,7 +380,66 @@ public void writeBigEndianBytesToDecimal256(byte[] value, ArrowType arrowType) {
           /*bitWidth=*/256)).writeBigEndianBytesToDecimal256(value, arrowType);
   }
 
-
+  @Override
+  public void writeVarBinary(byte[] value) {
+    getWriter(MinorType.VARBINARY).writeVarBinary(value);
+  }
+
+  @Override
+  public void writeVarBinary(byte[] value, int offset, int length) {
+    getWriter(MinorType.VARBINARY).writeVarBinary(value, offset, length);
+  }
+
+  @Override
+  public void writeVarBinary(ByteBuffer value) {
+    getWriter(MinorType.VARBINARY).writeVarBinary(value);
+  }
+
+  @Override
+  public void writeVarBinary(ByteBuffer value, int offset, int length) {
+    getWriter(MinorType.VARBINARY).writeVarBinary(value, offset, length);
+  }
+
+  @Override
+  public void writeLargeVarBinary(byte[] value) {
+    getWriter(MinorType.LARGEVARBINARY).writeLargeVarBinary(value);
+  }
+
+  @Override
+  public void writeLargeVarBinary(byte[] value, int offset, int length) {
+    getWriter(MinorType.LARGEVARBINARY).writeLargeVarBinary(value, offset, length);
+  }
+
+  @Override
+  public void writeLargeVarBinary(ByteBuffer value) {
+    getWriter(MinorType.LARGEVARBINARY).writeLargeVarBinary(value);
+  }
+
+  @Override
+  public void writeLargeVarBinary(ByteBuffer value, int offset, int length) {
+    getWriter(MinorType.LARGEVARBINARY).writeLargeVarBinary(value, offset, length);
+  }
+
+  @Override
+  public void writeVarChar(Text value) {
+    getWriter(MinorType.VARCHAR).writeVarChar(value);
+  }
+
+  @Override
+  public void writeVarChar(String value) {
+    getWriter(MinorType.VARCHAR).writeVarChar(value);
+  }
+
+  @Override
+  public void writeLargeVarChar(Text value) {
+    getWriter(MinorType.LARGEVARCHAR).writeLargeVarChar(value);
+  }
+
+  @Override
+  public void writeLargeVarChar(String value) {
+    getWriter(MinorType.LARGEVARCHAR).writeLargeVarChar(value);
+  }
+
   @Override
   public void allocate() {
     getWriter().allocate();