org.apache.arrow.vector.types.TimeUnit Java Examples

@Override
protected ArrowType doTypedDeserialize(JsonParser jparser, DeserializationContext ctxt)
        throws IOException
{
    TimeUnit unit = TimeUnit.valueOf(getNextStringField(jparser, UNIT_FIELD));
    String timezone = getNextStringField(jparser, TIMEZONE_FIELD);
    return new ArrowType.Timestamp(unit, timezone);
}
 

@Test
public void testParquetInt64TimestampMicrosToArrow() {
  MessageType parquet = Types.buildMessage()
    .addField(Types.optional(INT64).as(TIMESTAMP_MICROS).named("a")).named("root");
  Schema expected = new Schema(asList(
    field("a", new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"))
  ));
  Assert.assertEquals(expected, converter.fromParquet(parquet).getArrowSchema());
}
 

@Test
public void testParquetInt64TimestampMillisToArrow() {
  MessageType parquet = Types.buildMessage()
    .addField(Types.optional(INT64).as(TIMESTAMP_MILLIS).named("a")).named("root");
  Schema expected = new Schema(asList(
    field("a", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"))
  ));
  Assert.assertEquals(expected, converter.fromParquet(parquet).getArrowSchema());
}
 

@Test
public void testArrowTimestampMicrosecondToParquet() {
  MessageType expected = converter.fromArrow(new Schema(asList(
    field("a", new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"))
  ))).getParquetSchema();
  Assert.assertEquals(expected, Types.buildMessage().addField(Types.optional(INT64).as(TIMESTAMP_MICROS).named("a")).named("root"));
}
 

@Test
public void testArrowTimestampMillisecondToParquet() {
  MessageType expected = converter.fromArrow(new Schema(asList(
    field("a", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"))
  ))).getParquetSchema();
  Assert.assertEquals(expected, Types.buildMessage().addField(Types.optional(INT64).as(TIMESTAMP_MILLIS).named("a")).named("root"));
}
 

@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
public void testParquetInt64TimeMicrosToArrow() {
  MessageType parquet = Types.buildMessage()
    .addField(Types.optional(INT64).as(TIME_MICROS).named("a")).named("root");
  Schema expected = new Schema(asList(
    field("a", new ArrowType.Time(TimeUnit.MICROSECOND, 64))
  ));
  Assert.assertEquals(expected, converter.fromParquet(parquet).getArrowSchema());
}
 

@Test
public void testParquetInt32TimeMillisToArrow() {
  MessageType parquet = Types.buildMessage()
    .addField(Types.optional(INT32).as(TIME_MILLIS).named("a")).named("root");
  Schema expected = new Schema(asList(
    field("a", new ArrowType.Time(TimeUnit.MILLISECOND, 32))
  ));
  Assert.assertEquals(expected, converter.fromParquet(parquet).getArrowSchema());
}
 

@Test
public void testArrowTimeMicrosecondToParquet() {
  MessageType expected = converter.fromArrow(new Schema(asList(
    field("a", new ArrowType.Time(TimeUnit.MICROSECOND, 64))
  ))).getParquetSchema();
  Assert.assertEquals(expected,
    Types.buildMessage().addField(Types.optional(INT64).as(timeType(false, MICROS)).named("a")).named("root"));
}
 

@Test
public void testArrowTimeMillisecondToParquet() {
  MessageType expected = converter.fromArrow(new Schema(asList(
    field("a", new ArrowType.Time(TimeUnit.MILLISECOND, 32))
  ))).getParquetSchema();
  Assert.assertEquals(expected,
    Types.buildMessage().addField(Types.optional(INT32).as(timeType(false, MILLIS)).named("a")).named("root"));
}
 

@Override
public ArrowType visit(TimestampType timestampType) {
	if (timestampType.getPrecision() == 0) {
		return new ArrowType.Timestamp(TimeUnit.SECOND, null);
	} else if (timestampType.getPrecision() >= 1 && timestampType.getPrecision() <= 3) {
		return new ArrowType.Timestamp(TimeUnit.MILLISECOND, null);
	} else if (timestampType.getPrecision() >= 4 && timestampType.getPrecision() <= 6) {
		return new ArrowType.Timestamp(TimeUnit.MICROSECOND, null);
	} else {
		return new ArrowType.Timestamp(TimeUnit.NANOSECOND, null);
	}
}
 

@Override
public ArrowType visit(LocalZonedTimestampType localZonedTimestampType) {
	if (localZonedTimestampType.getPrecision() == 0) {
		return new ArrowType.Timestamp(TimeUnit.SECOND, null);
	} else if (localZonedTimestampType.getPrecision() >= 1 && localZonedTimestampType.getPrecision() <= 3) {
		return new ArrowType.Timestamp(TimeUnit.MILLISECOND, null);
	} else if (localZonedTimestampType.getPrecision() >= 4 && localZonedTimestampType.getPrecision() <= 6) {
		return new ArrowType.Timestamp(TimeUnit.MICROSECOND, null);
	} else {
		return new ArrowType.Timestamp(TimeUnit.NANOSECOND, null);
	}
}
 

@Override
public ArrowType visit(TimeType timeType) {
	if (timeType.getPrecision() == 0) {
		return new ArrowType.Time(TimeUnit.SECOND, 32);
	} else if (timeType.getPrecision() >= 1 && timeType.getPrecision() <= 3) {
		return new ArrowType.Time(TimeUnit.MILLISECOND, 32);
	} else if (timeType.getPrecision() >= 4 && timeType.getPrecision() <= 6) {
		return new ArrowType.Time(TimeUnit.MICROSECOND, 64);
	} else {
		return new ArrowType.Time(TimeUnit.NANOSECOND, 64);
	}
}
 

private static Integer getPrecision(TimeUnit unit) {
  switch (unit) {
  case SECOND:
    return 0;
  case MILLISECOND:
    return 3;
  case MICROSECOND:
    return 6;
  case NANOSECOND:
    return 9;
  }
  throw new IllegalArgumentException("unknown unit: " + unit);
}
 

@Override
public String visit(Timestamp type) {
  String name = "timestamp";
  String timezone = type.getTimezone();
  TimeUnit unit = type.getUnit();
  if (timezone == null && unit == TimeUnit.MILLISECOND) {
    return name;
  }
  return String.format("%s(%s,%s)", name, timezone == null ? "?" : timezone, unit.name());
}
 

@Override
public String visit(Time type) {
  String name = "time";
  TimeUnit unit = type.getUnit();
  if (unit == TimeUnit.MILLISECOND) {
    return name;
  }
  return String.format("%s(%s)", name, unit);
}
 

@Override
protected ArrowType doTypedDeserialize(JsonParser jparser, DeserializationContext ctxt)
        throws IOException
{
    TimeUnit unit = TimeUnit.valueOf(getNextStringField(jparser, UNIT_FIELD));
    int bitWidth = getNextIntField(jparser, BIT_WIDTH_FIELD);
    return new ArrowType.Time(unit, bitWidth);
}
 

@BeforeClass
public static void init() {
	testFields = new ArrayList<>();
	testFields.add(Tuple7.of(
		"f1", new TinyIntType(), new ArrowType.Int(8, true), RowTinyIntWriter.class,
		TinyIntWriter.TinyIntWriterForRow.class, TinyIntFieldReader.class, ArrowTinyIntColumnVector.class));

	testFields.add(Tuple7.of("f2", new SmallIntType(), new ArrowType.Int(8 * 2, true),
		RowSmallIntWriter.class, SmallIntWriter.SmallIntWriterForRow.class, SmallIntFieldReader.class, ArrowSmallIntColumnVector.class));

	testFields.add(Tuple7.of("f3", new IntType(), new ArrowType.Int(8 * 4, true),
		RowIntWriter.class, IntWriter.IntWriterForRow.class, IntFieldReader.class, ArrowIntColumnVector.class));

	testFields.add(Tuple7.of("f4", new BigIntType(), new ArrowType.Int(8 * 8, true),
		RowBigIntWriter.class, BigIntWriter.BigIntWriterForRow.class, BigIntFieldReader.class, ArrowBigIntColumnVector.class));

	testFields.add(Tuple7.of("f5", new BooleanType(), new ArrowType.Bool(),
		RowBooleanWriter.class, BooleanWriter.BooleanWriterForRow.class, BooleanFieldReader.class, ArrowBooleanColumnVector.class));

	testFields.add(Tuple7.of("f6", new FloatType(), new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE),
		RowFloatWriter.class, FloatWriter.FloatWriterForRow.class, FloatFieldReader.class, ArrowFloatColumnVector.class));

	testFields.add(Tuple7.of("f7", new DoubleType(), new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE),
		RowDoubleWriter.class, DoubleWriter.DoubleWriterForRow.class, DoubleFieldReader.class, ArrowDoubleColumnVector.class));

	testFields.add(Tuple7.of("f8", new VarCharType(), ArrowType.Utf8.INSTANCE,
		RowVarCharWriter.class, VarCharWriter.VarCharWriterForRow.class, VarCharFieldReader.class, ArrowVarCharColumnVector.class));

	testFields.add(Tuple7.of("f9", new VarBinaryType(), ArrowType.Binary.INSTANCE,
		RowVarBinaryWriter.class, VarBinaryWriter.VarBinaryWriterForRow.class, VarBinaryFieldReader.class, ArrowVarBinaryColumnVector.class));

	testFields.add(Tuple7.of("f10", new DecimalType(10, 3), new ArrowType.Decimal(10, 3),
		RowDecimalWriter.class, DecimalWriter.DecimalWriterForRow.class, DecimalFieldReader.class, ArrowDecimalColumnVector.class));

	testFields.add(Tuple7.of("f11", new DateType(), new ArrowType.Date(DateUnit.DAY),
		RowDateWriter.class, DateWriter.DateWriterForRow.class, DateFieldReader.class, ArrowDateColumnVector.class));

	testFields.add(Tuple7.of("f13", new TimeType(0), new ArrowType.Time(TimeUnit.SECOND, 32),
		RowTimeWriter.class, TimeWriter.TimeWriterForRow.class, TimeFieldReader.class, ArrowTimeColumnVector.class));

	testFields.add(Tuple7.of("f14", new TimeType(2), new ArrowType.Time(TimeUnit.MILLISECOND, 32),
		RowTimeWriter.class, TimeWriter.TimeWriterForRow.class, TimeFieldReader.class, ArrowTimeColumnVector.class));

	testFields.add(Tuple7.of("f15", new TimeType(4), new ArrowType.Time(TimeUnit.MICROSECOND, 64),
		RowTimeWriter.class, TimeWriter.TimeWriterForRow.class, TimeFieldReader.class, ArrowTimeColumnVector.class));

	testFields.add(Tuple7.of("f16", new TimeType(8), new ArrowType.Time(TimeUnit.NANOSECOND, 64),
		RowTimeWriter.class, TimeWriter.TimeWriterForRow.class, TimeFieldReader.class, ArrowTimeColumnVector.class));

	testFields.add(Tuple7.of("f17", new LocalZonedTimestampType(0), new ArrowType.Timestamp(TimeUnit.SECOND, null),
		RowTimestampWriter.class, TimestampWriter.TimestampWriterForRow.class, TimestampFieldReader.class, ArrowTimestampColumnVector.class));

	testFields.add(Tuple7.of("f18", new LocalZonedTimestampType(2), new ArrowType.Timestamp(TimeUnit.MILLISECOND, null),
		RowTimestampWriter.class, TimestampWriter.TimestampWriterForRow.class, TimestampFieldReader.class, ArrowTimestampColumnVector.class));

	testFields.add(Tuple7.of("f19", new LocalZonedTimestampType(4), new ArrowType.Timestamp(TimeUnit.MICROSECOND, null),
		RowTimestampWriter.class, TimestampWriter.TimestampWriterForRow.class, TimestampFieldReader.class, ArrowTimestampColumnVector.class));

	testFields.add(Tuple7.of("f20", new LocalZonedTimestampType(8), new ArrowType.Timestamp(TimeUnit.NANOSECOND, null),
		RowTimestampWriter.class, TimestampWriter.TimestampWriterForRow.class, TimestampFieldReader.class, ArrowTimestampColumnVector.class));

	testFields.add(Tuple7.of("f21", new TimestampType(0), new ArrowType.Timestamp(TimeUnit.SECOND, null),
		RowTimestampWriter.class, TimestampWriter.TimestampWriterForRow.class, TimestampFieldReader.class, ArrowTimestampColumnVector.class));

	testFields.add(Tuple7.of("f22", new TimestampType(2), new ArrowType.Timestamp(TimeUnit.MILLISECOND, null),
		RowTimestampWriter.class, TimestampWriter.TimestampWriterForRow.class, TimestampFieldReader.class, ArrowTimestampColumnVector.class));

	testFields.add(Tuple7.of("f23", new TimestampType(4), new ArrowType.Timestamp(TimeUnit.MICROSECOND, null),
		RowTimestampWriter.class, TimestampWriter.TimestampWriterForRow.class, TimestampFieldReader.class, ArrowTimestampColumnVector.class));

	testFields.add(Tuple7.of("f24", new TimestampType(8), new ArrowType.Timestamp(TimeUnit.NANOSECOND, null),
		RowTimestampWriter.class, TimestampWriter.TimestampWriterForRow.class, TimestampFieldReader.class, ArrowTimestampColumnVector.class));

	testFields.add(Tuple7.of("f25", new ArrayType(new VarCharType()), ArrowType.List.INSTANCE,
		RowArrayWriter.class, ArrayWriter.ArrayWriterForRow.class, ArrayFieldReader.class, ArrowArrayColumnVector.class));

	RowType rowFieldType = new RowType(Arrays.asList(
		new RowType.RowField("a", new IntType()),
		new RowType.RowField("b", new VarCharType()),
		new RowType.RowField("c", new ArrayType(new VarCharType())),
		new RowType.RowField("d", new TimestampType(2)),
		new RowType.RowField("e", new RowType((Arrays.asList(
			new RowType.RowField("e1", new IntType()),
			new RowType.RowField("e2", new VarCharType())))))));
	testFields.add(Tuple7.of("f26", rowFieldType, ArrowType.Struct.INSTANCE,
		RowRowWriter.class, RowWriter.RowWriterForRow.class, RowFieldReader.class, ArrowRowColumnVector.class));

	List<RowType.RowField> rowFields = new ArrayList<>();
	for (Tuple7<String, LogicalType, ArrowType, Class<?>, Class<?>, Class<?>, Class<?>> field : testFields) {
		rowFields.add(new RowType.RowField(field.f0, field.f1));
	}
	rowType = new RowType(rowFields);

	allocator = ArrowUtils.getRootAllocator().newChildAllocator("stdout", 0, Long.MAX_VALUE);
}
 

@Test(expected = UnsupportedOperationException.class)
public void testArrowTimeSecondToParquet() {
  converter.fromArrow(new Schema(asList(
    field("a", new ArrowType.Time(TimeUnit.SECOND, 32))
  ))).getParquetSchema();
}
 

public static Field getArrowFieldFromHivePrimitiveType(String name, TypeInfo typeInfo) {
  switch (typeInfo.getCategory()) {
  case PRIMITIVE:
    PrimitiveTypeInfo pTypeInfo = (PrimitiveTypeInfo) typeInfo;
    switch (pTypeInfo.getPrimitiveCategory()) {
    case BOOLEAN:

      return new Field(name, true, new Bool(), null);
    case BYTE:
      return new Field(name, true, new Int(32, true), null);
    case SHORT:
      return new Field(name, true, new Int(32, true), null);

    case INT:
      return new Field(name, true, new Int(32, true), null);

    case LONG:
      return new Field(name, true, new Int(64, true), null);

    case FLOAT:
      return new Field(name, true, new FloatingPoint(FloatingPointPrecision.SINGLE), null);

    case DOUBLE:
      return new Field(name, true, new FloatingPoint(FloatingPointPrecision.DOUBLE), null);

    case DATE:
      return new Field(name, true, new Date(DateUnit.MILLISECOND), null);

    case TIMESTAMP:
      return new Field(name, true, new Timestamp(TimeUnit.MILLISECOND, null), null);

    case BINARY:
      return new Field(name, true, new Binary(), null);
    case DECIMAL: {
      DecimalTypeInfo decimalTypeInfo = (DecimalTypeInfo) pTypeInfo;
      return new Field(name, true, new Decimal(decimalTypeInfo.getPrecision(), decimalTypeInfo.getScale()), null);
    }

    case STRING:
    case VARCHAR:
    case CHAR: {
      return new Field(name, true, new Utf8(), null);
    }
    case UNKNOWN:
    case VOID:
    default:
      // fall through.
    }
  default:
  }

  return null;
}
 

@Test
public void verifyNativeOutput()
    throws IOException, SAXException, ParserConfigurationException, ParseException {
  when(optionManager.getOption(TableauMessageBodyGenerator.TABLEAU_EXPORT_TYPE))
    .thenReturn(TableauMessageBodyGenerator.TableauExportType.ODBC.toString());
  DatasetConfig datasetConfig = new DatasetConfig();
  datasetConfig.setFullPathList(path.toPathList());

  // create a schema to test the metadata output for native connectors
  datasetConfig.setType(DatasetType.PHYSICAL_DATASET);
  BatchSchema schema = BatchSchema.newBuilder()
    .addField(new Field("string", FieldType.nullable(ArrowType.Utf8.INSTANCE), null))
    .addField(new Field("bool", FieldType.nullable(ArrowType.Bool.INSTANCE), null))
    .addField(new Field("decimal", FieldType.nullable(new ArrowType.Decimal(0, 0)), null))
    .addField(new Field("int", FieldType.nullable(new ArrowType.Int(8, false)), null))
    .addField(new Field("date", FieldType.nullable(new ArrowType.Date(DateUnit.MILLISECOND)), null))
    .addField(new Field("time", FieldType.nullable(new ArrowType.Time(TimeUnit.MILLISECOND, 8)), null))
    .build();
  datasetConfig.setRecordSchema(schema.toByteString());

  TableauMessageBodyGenerator generator = new TableauMessageBodyGenerator(configuration, ENDPOINT, optionManager);
  MultivaluedMap<String, Object> httpHeaders = new MultivaluedHashMap<>();
  ByteArrayOutputStream baos = new ByteArrayOutputStream();
  assertTrue(generator.isWriteable(datasetConfig.getClass(), null, null, WebServer.MediaType.APPLICATION_TDS_DRILL_TYPE));
  generator.writeTo(datasetConfig, DatasetConfig.class, null, new Annotation[] {}, WebServer.MediaType.APPLICATION_TDS_DRILL_TYPE, httpHeaders, baos);

  // Convert the baos into a DOM Tree to verify content
  DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();
  Document document = factory.newDocumentBuilder().parse(new ByteArrayInputStream(baos.toByteArray()));

  NodeList connections = document.getDocumentElement().getElementsByTagName("connection");

  assertEquals(1, connections.getLength());
  Element connection = (Element) connections.item(0);

  assertEquals("drill", connection.getAttribute("class"));
  assertEquals("Direct", connection.getAttribute("connection-type"));
  assertEquals("foo", connection.getAttribute("server"));
  assertEquals("12345", connection.getAttribute("port"));
  assertEquals(path.toParentPath(), connection.getAttribute("schema"));

  NodeList relations = connection.getElementsByTagName("relation");
  assertEquals(1, relations.getLength());
  Element relation = (Element) relations.item(0);
  assertEquals("table", relation.getAttribute("type"));
  assertEquals(tableName, relation.getAttribute("table"));

  // metadata tests
  NodeList metadataRecords = document.getDocumentElement().getElementsByTagName("metadata-record");

  assertEquals(metadataRecords.getLength(), schema.getFieldCount());
  assertEqualsMetadataRecord(metadataRecords.item(0), "[string]", "string");
  assertEqualsMetadataRecord(metadataRecords.item(1), "[bool]", "boolean");
  assertEqualsMetadataRecord(metadataRecords.item(2), "[decimal]", "real");
  assertEqualsMetadataRecord(metadataRecords.item(3), "[int]", "integer");
  assertEqualsMetadataRecord(metadataRecords.item(4), "[date]", "date");
  assertEqualsMetadataRecord(metadataRecords.item(5), "[time]", "datetime");

  // Also check that Content-Disposition header is set with a filename ending by tds
  ContentDisposition contentDisposition = new ContentDisposition((String) httpHeaders.getFirst(HttpHeaders.CONTENT_DISPOSITION));
  assertTrue("filename should end with .tds", contentDisposition.getFileName().endsWith(".tds"));
}
 

public static MajorType getMajorTypeForArrowType(ArrowType arrowType, List<Field> children) {
  MajorType.Builder builder = MajorType.newBuilder()
    .setMinorType(getMinorTypeFromArrowMinorType(getMinorTypeForArrowType(arrowType)))
    .setMode(DataMode.OPTIONAL);
  ArrowTypeID fieldType = arrowType.getTypeID();
  switch(fieldType) {
    case Decimal:
      builder.setPrecision(((Decimal) arrowType).getPrecision()).setScale(((Decimal) arrowType).getScale());
      break;

    case Utf8:
    case Binary:
      builder.setPrecision(CompleteType.DEFAULT_VARCHAR_PRECISION);
      break;

    case Timestamp:
      TimeUnit unit = ((Timestamp) arrowType).getUnit();
      switch(unit) {
        // Only MILLISECONDS is supported, but future-proofing
        case SECOND:
          builder.setPrecision(0);
          break;
        case MILLISECOND:
          builder.setPrecision(3);
          break;
        case MICROSECOND:
          builder.setPrecision(6);
          break;
        case NANOSECOND:
          builder.setPrecision(9);
          break;
        default:
          throw new AssertionError("Arrow TimeUnit " + unit + "not supported");
      }
      break;

    case Union:
      for (Field child : children) {
        builder.addSubType(getMinorTypeFromArrowMinorType(getMinorTypeForArrowType(child.getType())));
      }
      break;

    default:
      // Nothing
  }
  return builder.build();
}
 

@Test(expected = UnsupportedOperationException.class)
public void testArrowTimestampSecondToParquet() {
  converter.fromArrow(new Schema(asList(
    field("a", new ArrowType.Timestamp(TimeUnit.SECOND, "UTC"))
  ))).getParquetSchema();
}
 

public static ArrowType toArrowType(UserProtos.ResultColumnMetadata type) {
  String typeName = type.getDataType();
  switch (typeName) {
    case "NULL":
      return new Null();
    case "MAP":
      return new ArrowType.Map(false); //todo inner type?
    case "ARRAY":
      return new ArrowType.List(); //todo inner type?
    case "UNION":
      throw new UnsupportedOperationException("have not implemented unions");
      //return new Union(); //todo inner type?
    case "TINYINT":
      return new Int(8, true);
    case "SMALLINT":
      return new Int(16, true);
    case "INTEGER":
      return new Int(32, true);
    case "BIGINT":
      return new Int(64, true);
    case "FLOAT":
      return new FloatingPoint(FloatingPointPrecision.SINGLE);
    case "DOUBLE":
      return new FloatingPoint(FloatingPointPrecision.DOUBLE);
    case "CHARACTER VARYING":
      return new Utf8();
    case "BINARY VARYING":
      return new Binary();
    case "BOOLEAN":
      return new Bool();
    case "DECIMAL":
      return new Decimal(type.getPrecision(), type.getScale());
    case "DATE":
      return new Date(DateUnit.MILLISECOND);
    case "TIME":
      return new Time(TimeUnit.MICROSECOND, 64);
    case "TIMESTAMP":
      return new Timestamp(TimeUnit.MICROSECOND, "UTC");
    case "INTERVAL DAY TO SECOND":
      return new Interval(IntervalUnit.DAY_TIME);
    case "INTERVAL YEAR TO MONTH":
      return new Interval(IntervalUnit.YEAR_MONTH);
    case "BINARY":
      return new ArrowType.FixedSizeBinary(50);
    default:
      throw new IllegalStateException("unable to find arrow type for " + typeName);
  }
}