ARROW-977: [java] Add Timezone aware timestamp vectors

Author: Julien Le Dem <julien@apache.org> Closes #644 from julienledem/TZ and squashes the following commits: 37987b9 [Julien Le Dem] add integration tests a58fcae [Julien Le Dem] add integration test 39966aa [Julien Le Dem] add other vectors and tests bf245ce [Julien Le Dem] add TZ vectors
apache · May 15, 2017 · 681afab · 681afab
1 parent 4381845
commit 681afab
Show file tree

Hide file tree

Showing 17 changed files with 429 additions and 191 deletions.
diff --git a/integration/integration_test.py b/integration/integration_test.py
@@ -620,7 +620,11 @@ def generate_datetime_case():
         TimestampType('f7', 'ms'),
         TimestampType('f8', 'us'),
         TimestampType('f9', 'ns'),
-        TimestampType('f10', 'ms', tz=None)
+        TimestampType('f10', 'ms', tz=None),
+        TimestampType('f11', 's', tz='UTC'),
+        TimestampType('f12', 'ms', tz='US/Eastern'),
+        TimestampType('f13', 'us', tz='Europe/Paris'),
+        TimestampType('f14', 'ns', tz='US/Pacific')
     ]
 
     batch_sizes = [7, 10]

diff --git a/java/vector/src/main/codegen/data/ValueVectorTypes.tdd b/java/vector/src/main/codegen/data/ValueVectorTypes.tdd
@@ -71,12 +71,28 @@
       minor: [
         { class: "BigInt"},
         { class: "UInt8" },
-        { class: "Float8",         javaType: "double", boxedType: "Double",                              fields: [{name: "value", type: "double"}], },
-        { class: "DateMilli",      javaType: "long",                      friendlyType: "LocalDateTime" },
-        { class: "TimeStampSec",   javaType: "long",   boxedType: "Long", friendlyType: "LocalDateTime" },
-        { class: "TimeStampMilli", javaType: "long",   boxedType: "Long", friendlyType: "LocalDateTime" },
-        { class: "TimeStampMicro", javaType: "long",   boxedType: "Long", friendlyType: "LocalDateTime" },
-        { class: "TimeStampNano",  javaType: "long",   boxedType: "Long", friendlyType: "LocalDateTime" },
+        { class: "Float8",           javaType: "double", boxedType: "Double", fields: [{name: "value", type: "double"}] },
+        { class: "DateMilli",        javaType: "long",                      friendlyType: "LocalDateTime" },
+        { class: "TimeStampSec",     javaType: "long",   boxedType: "Long", friendlyType: "LocalDateTime" },
+        { class: "TimeStampMilli",   javaType: "long",   boxedType: "Long", friendlyType: "LocalDateTime" },
+        { class: "TimeStampMicro",   javaType: "long",   boxedType: "Long", friendlyType: "LocalDateTime" },
+        { class: "TimeStampNano",    javaType: "long",   boxedType: "Long", friendlyType: "LocalDateTime" },
+        { class: "TimeStampSecTZ", javaType: "long",   boxedType: "Long",
+                                     typeParams: [ {name: "timezone", type: "String"} ],
+                                     arrowType: "org.apache.arrow.vector.types.pojo.ArrowType.Timestamp",
+                                     arrowTypeConstructorParams: ["org.apache.arrow.vector.types.TimeUnit.SECOND", "timezone"] },
+        { class: "TimeStampMilliTZ", javaType: "long",   boxedType: "Long",
+                                     typeParams: [ {name: "timezone", type: "String"} ],
+                                     arrowType: "org.apache.arrow.vector.types.pojo.ArrowType.Timestamp",
+                                     arrowTypeConstructorParams: ["org.apache.arrow.vector.types.TimeUnit.MILLISECOND", "timezone"] },
+        { class: "TimeStampMicroTZ", javaType: "long",   boxedType: "Long",
+                                     typeParams: [ {name: "timezone", type: "String"} ],
+                                     arrowType: "org.apache.arrow.vector.types.pojo.ArrowType.Timestamp",
+                                     arrowTypeConstructorParams: ["org.apache.arrow.vector.types.TimeUnit.MICROSECOND", "timezone"] },
+        { class: "TimeStampNanoTZ", javaType: "long",   boxedType: "Long",
+                                     typeParams: [ {name: "timezone", type: "String"} ],
+                                     arrowType: "org.apache.arrow.vector.types.pojo.ArrowType.Timestamp",
+                                     arrowTypeConstructorParams: ["org.apache.arrow.vector.types.TimeUnit.NANOSECOND", "timezone"] },
         { class: "TimeMicro" },
         { class: "TimeNano" }
       ]
@@ -97,7 +113,13 @@
       boxedType: "ArrowBuf",
 
       minor: [
-        { class: "Decimal", maxPrecisionDigits: 38, nDecimalDigits: 4, friendlyType: "BigDecimal", fields: [{name: "start", type: "int"}, {name: "buffer", type: "ArrowBuf"}, {name: "scale", type: "int", include: false}, {name: "precision", type: "int", include: false}] }
+        {
+          class: "Decimal",
+          maxPrecisionDigits: 38, nDecimalDigits: 4, friendlyType: "BigDecimal",
+          typeParams: [ {name: "scale", type: "int"}, { name: "precision", type: "int"}],
+          arrowType: "org.apache.arrow.vector.types.pojo.ArrowType.Decimal",
+          fields: [{name: "start", type: "int"}, {name: "buffer", type: "ArrowBuf"}, {name: "scale", type: "int", include: false}, {name: "precision", type: "int", include: false}]
+        }
       ]
     },
     {

diff --git a/java/vector/src/main/codegen/templates/AbstractFieldWriter.java b/java/vector/src/main/codegen/templates/AbstractFieldWriter.java
@@ -59,7 +59,7 @@ public void write(${name}Holder holder) {
   }
 
   <#if minor.class == "Decimal">
-  public void writeDecimal(int start, ArrowBuf buffer) {
+  public void write${minor.class}(int start, ArrowBuf buffer) {
     fail("${name}");
   }
   <#else>
@@ -114,9 +114,11 @@ public ListWriter list(String name) {
   <#if lowerName == "int" ><#assign lowerName = "integer" /></#if>
   <#assign upperName = minor.class?upper_case />
   <#assign capName = minor.class?cap_first />
-  <#if minor.class?starts_with("Decimal") >
-  public ${capName}Writer ${lowerName}(String name, int scale, int precision) {
-    fail("${capName}");
+  <#if minor.typeParams?? >
+
+  @Override
+  public ${capName}Writer ${lowerName}(String name<#list minor.typeParams as typeParam>, ${typeParam.type} ${typeParam.name}</#list>) {
+    fail("${capName}(" + <#list minor.typeParams as typeParam>"${typeParam.name}: " + ${typeParam.name} + ", " + </#list>")");
     return null;
   }
   </#if>

diff --git a/java/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java b/java/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java
@@ -72,31 +72,23 @@ public void endList() {
   }
 
   <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
-  <#assign fields = minor.fields!type.fields />
-  <#if !minor.class?starts_with("Decimal") >
+    <#assign fields = minor.fields!type.fields />
   @Override
   public void write(${name}Holder holder) {
     getWriter(MinorType.${name?upper_case}).write(holder);
   }
 
+    <#if minor.class == "Decimal">
+  public void write${minor.class}(int start, ArrowBuf buffer) {
+    getWriter(MinorType.${name?upper_case}).write${minor.class}(start, buffer);
+  }
+    <#else>
   public void write${minor.class}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, </#if></#list>) {
     getWriter(MinorType.${name?upper_case}).write${minor.class}(<#list fields as field>${field.name}<#if field_has_next>, </#if></#list>);
   }
-
-  <#else>
-  @Override
-  public void write(DecimalHolder holder) {
-    getWriter(MinorType.DECIMAL).write(holder);
-  }
-
-  public void writeDecimal(int start, ArrowBuf buffer) {
-    getWriter(MinorType.DECIMAL).writeDecimal(start, buffer);
-  }
-
-  </#if>
+    </#if>
 
   </#list></#list>
-
   public void writeNull() {
   }
 
@@ -125,10 +117,13 @@ public ListWriter list(String name) {
   <#if lowerName == "int" ><#assign lowerName = "integer" /></#if>
   <#assign upperName = minor.class?upper_case />
   <#assign capName = minor.class?cap_first />
-  <#if minor.class?starts_with("Decimal") >
-  public ${capName}Writer ${lowerName}(String name, int scale, int precision) {
-    return getWriter(MinorType.MAP).${lowerName}(name, scale, precision);
+
+  <#if minor.typeParams?? >
+  @Override
+  public ${capName}Writer ${lowerName}(String name<#list minor.typeParams as typeParam>, ${typeParam.type} ${typeParam.name}</#list>) {
+    return getWriter(MinorType.MAP).${lowerName}(name<#list minor.typeParams as typeParam>, ${typeParam.name}</#list>);
   }
+
   </#if>
   @Override
   public ${capName}Writer ${lowerName}(String name) {

diff --git a/java/vector/src/main/codegen/templates/BaseWriter.java b/java/vector/src/main/codegen/templates/BaseWriter.java
@@ -53,8 +53,8 @@ public interface MapWriter extends BaseWriter {
     <#if lowerName == "int" ><#assign lowerName = "integer" /></#if>
     <#assign upperName = minor.class?upper_case />
     <#assign capName = minor.class?cap_first />
-    <#if minor.class?starts_with("Decimal") >
-    ${capName}Writer ${lowerName}(String name, int scale, int precision);
+    <#if minor.typeParams?? >
+    ${capName}Writer ${lowerName}(String name<#list minor.typeParams as typeParam>, ${typeParam.type} ${typeParam.name}</#list>);
     </#if>
     ${capName}Writer ${lowerName}(String name);
     </#list></#list>

diff --git a/java/vector/src/main/codegen/templates/ComplexCopier.java b/java/vector/src/main/codegen/templates/ComplexCopier.java
@@ -72,7 +72,8 @@ private static void writeValue(FieldReader reader, FieldWriter writer) {
   <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
   <#assign fields = minor.fields!type.fields />
   <#assign uncappedName = name?uncap_first/>
-  <#if !minor.class?starts_with("Decimal")>
+
+  <#if !minor.typeParams?? >
 
       case ${name?upper_case}:
         if (reader.isSet()) {
@@ -94,7 +95,7 @@ private static FieldWriter getMapWriterForReader(FieldReader reader, MapWriter w
     <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
     <#assign fields = minor.fields!type.fields />
     <#assign uncappedName = name?uncap_first/>
-    <#if !minor.class?starts_with("Decimal")>
+    <#if !minor.typeParams?? >
     case ${name?upper_case}:
       return (FieldWriter) writer.<#if name == "Int">integer<#else>${uncappedName}</#if>(name);
     </#if>
@@ -113,7 +114,7 @@ private static FieldWriter getListWriterForReader(FieldReader reader, ListWriter
     <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
     <#assign fields = minor.fields!type.fields />
     <#assign uncappedName = name?uncap_first/>
-    <#if !minor.class?starts_with("Decimal")>
+    <#if !minor.typeParams?? >
     case ${name?upper_case}:
     return (FieldWriter) writer.<#if name == "Int">integer<#else>${uncappedName}</#if>();
     </#if>

diff --git a/java/vector/src/main/codegen/templates/FixedValueVectors.java b/java/vector/src/main/codegen/templates/FixedValueVectors.java
@@ -25,9 +25,10 @@
 <#list vv.types as type>
 <#list type.minor as minor>
 <#assign friendlyType = (minor.friendlyType!minor.boxedType!type.boxedType) />
+<#assign className = "${minor.class}Vector" />
 
 <#if type.major == "Fixed">
-<@pp.changeOutputFile name="/org/apache/arrow/vector/${minor.class}Vector.java" />
+<@pp.changeOutputFile name="/org/apache/arrow/vector/${className}.java" />
 <#include "/@includes/license.ftl" />
 
 package org.apache.arrow.vector;
@@ -43,8 +44,8 @@
  *
  * NB: this class is automatically generated from ${.template_name} and ValueVectorTypes.tdd using FreeMarker.
  */
-public final class ${minor.class}Vector extends BaseDataValueVector implements FixedWidthVector{
-  private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(${minor.class}Vector.class);
+public final class ${className} extends BaseDataValueVector implements FixedWidthVector{
+  private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(${className}.class);
 
   public static final int TYPE_WIDTH = ${type.width};
 
@@ -53,24 +54,25 @@ public final class ${minor.class}Vector extends BaseDataValueVector implements F
 
   private int allocationSizeInBytes = INITIAL_VALUE_ALLOCATION * ${type.width};
   private int allocationMonitor = 0;
+  <#if minor.typeParams??>
 
-  <#if minor.class == "Decimal">
+     <#list minor.typeParams as typeParam>
+  private final ${typeParam.type} ${typeParam.name};
+    </#list>
 
-  private int precision;
-  private int scale;
-
-  public ${minor.class}Vector(String name, BufferAllocator allocator, int precision, int scale) {
+  public ${className}(String name, BufferAllocator allocator<#list minor.typeParams as typeParam>, ${typeParam.type} ${typeParam.name}</#list>) {
     super(name, allocator);
-    this.precision = precision;
-    this.scale = scale;
+    <#list minor.typeParams as typeParam>
+    this.${typeParam.name} = ${typeParam.name};
+    </#list>
   }
   <#else>
-  public ${minor.class}Vector(String name, BufferAllocator allocator) {
+
+  public ${className}(String name, BufferAllocator allocator) {
     super(name, allocator);
   }
   </#if>
 
-
   @Override
   public MinorType getMinorType() {
     return MinorType.${minor.class?upper_case};
@@ -219,17 +221,17 @@ public TransferPair getTransferPair(String ref, BufferAllocator allocator){
 
   @Override
   public TransferPair makeTransferPair(ValueVector to) {
-    return new TransferImpl((${minor.class}Vector) to);
+    return new TransferImpl((${className}) to);
   }
 
-  public void transferTo(${minor.class}Vector target){
+  public void transferTo(${className} target){
     target.clear();
     target.data = data.transferOwnership(target.allocator).buffer;
     target.data.writerIndex(data.writerIndex());
     clear();
   }
 
-  public void splitAndTransferTo(int startIndex, int length, ${minor.class}Vector target) {
+  public void splitAndTransferTo(int startIndex, int length, ${className} target) {
     final int startPoint = startIndex * ${type.width};
     final int sliceLength = length * ${type.width};
     target.clear();
@@ -238,22 +240,18 @@ public void splitAndTransferTo(int startIndex, int length, ${minor.class}Vector
   }
 
   private class TransferImpl implements TransferPair{
-    private ${minor.class}Vector to;
+    private ${className} to;
 
     public TransferImpl(String name, BufferAllocator allocator){
-      <#if minor.class == "Decimal">
-      to = new ${minor.class}Vector(name, allocator, precision, scale);
-      <#else>
-      to = new ${minor.class}Vector(name, allocator);
-      </#if>
+      to = new ${className}(name, allocator<#if minor.typeParams??><#list minor.typeParams as typeParam>,  ${className}.this.${typeParam.name}</#list></#if>);
     }
 
-    public TransferImpl(${minor.class}Vector to) {
+    public TransferImpl(${className} to) {
       this.to = to;
     }
 
     @Override
-    public ${minor.class}Vector getTo(){
+    public ${className} getTo(){
       return to;
     }
 
@@ -269,11 +267,11 @@ public void splitAndTransfer(int startIndex, int length) {
 
     @Override
     public void copyValueSafe(int fromIndex, int toIndex) {
-      to.copyFromSafe(fromIndex, toIndex, ${minor.class}Vector.this);
+      to.copyFromSafe(fromIndex, toIndex, ${className}.this);
     }
   }
 
-  public void copyFrom(int fromIndex, int thisIndex, ${minor.class}Vector from){
+  public void copyFrom(int fromIndex, int thisIndex, ${className} from){
     <#if (type.width > 8 || minor.class == "IntervalDay")>
     from.data.getBytes(fromIndex * ${type.width}, data, thisIndex * ${type.width}, ${type.width});
     <#else> <#-- type.width <= 8 -->
@@ -283,7 +281,7 @@ public void copyFrom(int fromIndex, int thisIndex, ${minor.class}Vector from){
     </#if> <#-- type.width -->
   }
 
-  public void copyFromSafe(int fromIndex, int thisIndex, ${minor.class}Vector from){
+  public void copyFromSafe(int fromIndex, int thisIndex, ${className} from){
     while(thisIndex >= getValueCapacity()) {
         reAlloc();
     }

diff --git a/java/vector/src/main/codegen/templates/MapWriters.java b/java/vector/src/main/codegen/templates/MapWriters.java
@@ -56,7 +56,8 @@ public class ${mode}MapWriter extends AbstractFieldWriter {
     </#if>
     this.container = container;
     for (Field child : container.getField().getChildren()) {
-      switch (Types.getMinorTypeForArrowType(child.getType())) {
+      MinorType minorType = Types.getMinorTypeForArrowType(child.getType());
+      switch (minorType) {
       case MAP:
         map(child.getName());
         break;
@@ -71,15 +72,18 @@ public class ${mode}MapWriter extends AbstractFieldWriter {
 <#assign lowerName = minor.class?uncap_first />
 <#if lowerName == "int" ><#assign lowerName = "integer" /></#if>
 <#assign upperName = minor.class?upper_case />
-      case ${upperName}:
-        <#if lowerName == "decimal" >
-        Decimal decimal = (Decimal)child.getType();
-        decimal(child.getName(), decimal.getScale(), decimal.getPrecision());
+      case ${upperName}: {
+        <#if minor.typeParams?? >
+        ${minor.arrowType} arrowType = (${minor.arrowType})child.getType();
+        ${lowerName}(child.getName()<#list minor.typeParams as typeParam>, arrowType.get${typeParam.name?cap_first}()</#list>);
         <#else>
         ${lowerName}(child.getName());
-       </#if>
+        </#if>
         break;
+      }
 </#list></#list>
+        default:
+          throw new UnsupportedOperationException("Unknown type: " + minorType);
       }
     }
   }
@@ -205,15 +209,17 @@ public void end() {
   <#assign vectName = capName />
   <#assign vectName = "Nullable${capName}" />
 
-  <#if minor.class?starts_with("Decimal") >
+  <#if minor.typeParams?? >
+  @Override
   public ${minor.class}Writer ${lowerName}(String name) {
     // returns existing writer
     final FieldWriter writer = fields.get(handleCase(name));
     assert writer != null;
     return writer;
   }
 
-  public ${minor.class}Writer ${lowerName}(String name, int scale, int precision) {
+  @Override
+  public ${minor.class}Writer ${lowerName}(String name<#list minor.typeParams as typeParam>, ${typeParam.type} ${typeParam.name}</#list>) {
   <#else>
   @Override
   public ${minor.class}Writer ${lowerName}(String name) {
@@ -223,7 +229,21 @@ public void end() {
       ValueVector vector;
       ValueVector currentVector = container.getChild(name);
       ${vectName}Vector v = container.addOrGet(name, 
-          FieldType.nullable(<#if minor.class == "Decimal">new Decimal(precision, scale)<#else>MinorType.${upperName}.getType()</#if>),
+          FieldType.nullable(
+          <#if minor.typeParams??>
+            <#if minor.arrowTypeConstructorParams??>
+              <#assign constructorParams = minor.arrowTypeConstructorParams />
+            <#else>
+              <#assign constructorParams = [] />
+              <#list minor.typeParams as typeParam>
+                <#assign constructorParams = constructorParams + [ typeParam.name ] />
+              </#list>
+            </#if>    
+            new ${minor.arrowType}(${constructorParams?join(", ")})
+          <#else>
+            MinorType.${upperName}.getType()
+          </#if>
+          ),
           ${vectName}Vector.class);
       writer = new PromotableWriter(v, container, getNullableMapWriterFactory());
       vector = v;