Fixing bug with type compatibility on Timestamp with TZ type

awslabs · May 6, 2020 · 521683b · 521683b
1 parent 9ab42e8
commit 521683b
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 3 deletions.
diff --git a/.../src/main/java/com/amazonaws/athena/connector/lambda/domain/predicate/SortedRangeSet.java b/.../src/main/java/com/amazonaws/athena/connector/lambda/domain/predicate/SortedRangeSet.java
@@ -467,12 +467,31 @@ private SortedRangeSet checkCompatibility(ValueSet other)
 
     private void checkTypeCompatibility(Marker marker)
     {
-        if (!getType().equals(marker.getType())) {
+        if (!getType().equals(marker.getType())
+                && !checkTypeCompatibilityForTimeStamp(marker)) {
             throw new IllegalStateException(String.format("Marker of %s does not match SortedRangeSet of %s",
                     marker.getType(), getType()));
         }
     }
 
+    /**
+     * Since we are mapping MinorType.TIMESTAMPMILLITZ to ArrowType.Timestamp(MILLI, ZoneId.systemDefault().getId())
+     * with UTC being a place holder,
+     * We cannot check the type compatibility of such types, as UTC (place holder) can be/will be overwritten by
+     * the TZ value coming from the raw data.
+     *
+     * Comparison can still be done for such types with different TZ as we convert to data to ZonedDateType to compare.
+     *
+     * @param marker
+     * @return if both types are ArrowType.Timestamp, returns the equality of unit
+     *         else false
+     */
+    private boolean checkTypeCompatibilityForTimeStamp(Marker marker)
+    {
+        return getType() instanceof ArrowType.Timestamp &&
+                ((ArrowType.Timestamp) getType()).getUnit().equals(((ArrowType.Timestamp) marker.getType()).getUnit());
+    }
+
     @Override
     public int hashCode()
     {

diff --git a/...tion-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/v2/ArrowSerDeTest.java b/...tion-sdk/src/test/java/com/amazonaws/athena/connector/lambda/serde/v2/ArrowSerDeTest.java
@@ -22,6 +22,7 @@
 import com.amazonaws.athena.connector.lambda.data.SupportedTypes;
 import com.amazonaws.athena.connector.lambda.serde.TypedDeserializer;
 import com.amazonaws.athena.connector.lambda.serde.TypedSerializer;
+import com.google.common.collect.ImmutableMap;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.junit.Test;
 
@@ -32,6 +33,13 @@
 
 public class ArrowSerDeTest
 {
+    // we need to have a specific fall back for MinorType.TIMESTAMPMILLITZ since it doesn't have default mapping to ArrowType
+    // AND the simple names are different (i.e. TIMESTAMPMILLITZ != Timestamp).
+    // MinorType.DECIMAL which is a similar case of no default mapping to ArrowType does not need this as
+    // MinorType name and ArrowType name are the same so simple fall back of enum name works.
+    private static final ImmutableMap<SupportedTypes, String> FALL_BACK_ARROW_TYPE_CLASS = ImmutableMap.of(
+            SupportedTypes.TIMESTAMPMILLITZ, ArrowType.Timestamp.class.getSimpleName());
+
     @Test
     public void testSupportedTypesHaveSerializers()
     {
@@ -46,7 +54,7 @@ public void testSupportedTypesHaveSerializers()
             }
             catch (UnsupportedOperationException e) {
                 // fall back to enum name
-                arrowTypeClass = supportedType.name();
+                arrowTypeClass = FALL_BACK_ARROW_TYPE_CLASS.getOrDefault(supportedType, supportedType.name());
             }
             assertTrue("No serializer for supported type " + supportedType + " with ArrowType " + arrowTypeClass, delegateSerDeMap.containsKey(arrowTypeClass));
         }
@@ -66,7 +74,7 @@ public void testSupportedTypesHaveDeserializers()
             }
             catch (UnsupportedOperationException e) {
                 // fall back to enum name
-                arrowTypeClass = supportedType.name();
+                arrowTypeClass = FALL_BACK_ARROW_TYPE_CLASS.getOrDefault(supportedType, supportedType.name());
             }
             assertTrue("No deserializer for supported type " + supportedType + " with ArrowType " + arrowTypeClass, delegateSerDeMap.containsKey(arrowTypeClass));
         }