Add an option to convert Int96 to Arrow Timestamp

Author: yongyanw

parquet-arrow/src/main/java/org/apache/parquet/arrow/schema/SchemaConverter.java

@@ -86,10 +86,18 @@
  */
 public class SchemaConverter {
 
+  // Indicates if Int96 should be converted to timestamp arrow type.
+  private final boolean convertInt96ToArrowTimestamp;
+
   /**
    * For when we'll need this to be configurable
    */
   public SchemaConverter() {
+    this(false);
+  }
+
+  public SchemaConverter(final boolean convertInt96ToArrowTimestamp) {
+    this.convertInt96ToArrowTimestamp = convertInt96ToArrowTimestamp;
   }
 
   /**
@@ -492,8 +500,11 @@ private String getTimeZone(LogicalTypeAnnotation.TimestampLogicalTypeAnnotation
 
       @Override
       public TypeMapping convertINT96(PrimitiveTypeName primitiveTypeName) throws RuntimeException {
-        // Possibly timestamp
-        return field(new ArrowType.Binary());
+        if (convertInt96ToArrowTimestamp) {
+          return field(new ArrowType.Timestamp(TimeUnit.NANOSECOND, null));
+        } else {
+          return field(new ArrowType.Binary());
+        }
       }
 
       @Override

parquet-arrow/src/test/java/org/apache/parquet/arrow/schema/TestSchemaConverter.java

@@ -47,6 +47,7 @@
 import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.FLOAT;
 import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT32;
 import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT64;
+import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT96;
 
 import java.io.IOException;
 import java.util.List;
@@ -419,6 +420,27 @@ public void testParquetInt64TimeMicrosToArrow() {
     Assert.assertEquals(expected, converter.fromParquet(parquet).getArrowSchema());
   }
 
+  @Test
+  public void testParquetInt96ToArrowBinary() {
+    MessageType parquet = Types.buildMessage()
+      .addField(Types.optional(INT96).named("a")).named("root");
+    Schema expected = new Schema(asList(
+      field("a", new ArrowType.Binary())
+    ));
+    Assert.assertEquals(expected, converter.fromParquet(parquet).getArrowSchema());
+  }
+
+  @Test
+  public void testParquetInt96ToArrowTimestamp() {
+    final SchemaConverter converterInt96ToTimestamp = new SchemaConverter(true);
+    MessageType parquet = Types.buildMessage()
+      .addField(Types.optional(INT96).named("a")).named("root");
+    Schema expected = new Schema(asList(
+      field("a", new ArrowType.Timestamp(TimeUnit.NANOSECOND, null))
+    ));
+    Assert.assertEquals(expected, converterInt96ToTimestamp.fromParquet(parquet).getArrowSchema());
+  }
+
   @Test(expected = IllegalStateException.class)
   public void testParquetInt64TimeMillisToArrow() {
     converter.fromParquet(Types.buildMessage()