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

public static Field getFieldForColumn(String name, ColumnType columnType) {
    switch (columnType) {
        case Integer:
            return field(name, new ArrowType.Int(32, false));
        case Float:
            return field(name, new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE));
        case Double:
            return field(name, new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE));
        case Long:
            return field(name, new ArrowType.Int(64, false));
        case NDArray:
            return field(name, new ArrowType.Binary());
        case Boolean:
            return field(name, new ArrowType.Bool());
        case Categorical:
            return field(name, new ArrowType.Utf8());
        case Time:
            return field(name, new ArrowType.Date(DateUnit.MILLISECOND));
        case Bytes:
            return field(name, new ArrowType.Binary());
        case String:
            return field(name, new ArrowType.Utf8());
        default:
            throw new IllegalArgumentException("Column type invalid " + columnType);
    }
}
 

/**
 * Create a field given the input {@link ColumnType}
 * and name
 * @param name the name of the field
 * @param columnType the column type to add
 * @return
 */
public static Field getFieldForColumn(String name,ColumnType columnType) {
    switch(columnType) {
        case Long: return field(name,new ArrowType.Int(64,false));
        case Integer: return field(name,new ArrowType.Int(32,false));
        case Double: return field(name,new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE));
        case Float: return field(name,new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE));
        case Boolean: return field(name, new ArrowType.Bool());
        case Categorical: return field(name,new ArrowType.Utf8());
        case Time: return field(name,new ArrowType.Date(DateUnit.MILLISECOND));
        case Bytes: return field(name,new ArrowType.Binary());
        case NDArray: return field(name,new ArrowType.Binary());
        case String: return field(name,new ArrowType.Utf8());

        default: throw new IllegalArgumentException("Column type invalid " + columnType);
    }
}
 

/**
 * Create a field given the input {@link ColumnType}
 * and name
 * @param name the name of the field
 * @param columnType the column type to add
 * @return
 */
public static Field getFieldForColumn(String name,ColumnType columnType) {
    switch(columnType) {
        case Long: return field(name,new ArrowType.Int(64,false));
        case Integer: return field(name,new ArrowType.Int(32,false));
        case Double: return field(name,new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE));
        case Float: return field(name,new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE));
        case Boolean: return field(name, new ArrowType.Bool());
        case Categorical: return field(name,new ArrowType.Utf8());
        case Time: return field(name,new ArrowType.Date(DateUnit.MILLISECOND));
        case Bytes: return field(name,new ArrowType.Binary());
        case NDArray: return field(name,new ArrowType.Binary());
        case String: return field(name,new ArrowType.Utf8());

        default: throw new IllegalArgumentException("Column type invalid " + columnType);
    }
}
 

@Override
protected ArrowType doTypedDeserialize(JsonParser jparser, DeserializationContext ctxt)
        throws IOException
{
    DateUnit unit = DateUnit.valueOf(getNextStringField(jparser, UNIT_FIELD));
    return new ArrowType.Date(unit);
}
 

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

private GandivaRegistryWrapper() throws GandivaException {
  this.supportedTypes = ExpressionRegistry.getInstance().getSupportedTypes();
  ArrowType.Date dateDay = new ArrowType.Date(DateUnit.DAY);

  Set<FunctionSignature> signatures = ExpressionRegistry.getInstance().getSupportedFunctions();
  Set<FunctionSignature> updatedSignatures = new HashSet<>();
  for (FunctionSignature signature : signatures) {
    FunctionSignature updated = signature;
    List<ArrowType> dateDayArgs =
      signature.getParamTypes().stream().filter(type-> {
        return type.equals(dateDay);
      }).collect(Collectors.toList());

    // suppress all date32 functions. date32 is not a supported dremio type;
    if (!dateDayArgs.isEmpty() || signature.getReturnType().equals(dateDay)) {
      continue;
    }
    // To make this fit in dremio model of type inference, add dummy args for precision and
    // scale.
    if (signature.getName().equals("castDECIMAL") || signature.getName().equals("castDECIMALNullOnOverflow")) {
      List<ArrowType> args = new ArrayList<>(signature.getParamTypes());
      args.add(new ArrowType.Int(64, true)); // precision
      args.add(new ArrowType.Int(64, true)); // scale

      updated = new FunctionSignature(signature.getName(), signature.getReturnType(), args);
    }
    updatedSignatures.add(updated);
    addNonDecimalMethods(signature, updated);
  }
  this.supportedFunctionsDecimal = updatedSignatures;
}
 

public static Schema generateTestSchema()
{
    /**
     * Generate and write the schema
     */
    SchemaBuilder schemaBuilder = new SchemaBuilder();
    schemaBuilder.addMetadata("meta1", "meta-value-1");
    schemaBuilder.addMetadata("meta2", "meta-value-2");
    schemaBuilder.addField("intfield1", new ArrowType.Int(32, true));
    schemaBuilder.addField("doublefield2", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE));
    schemaBuilder.addField("varcharfield3", new ArrowType.Utf8());

    schemaBuilder.addField("datemillifield4", new ArrowType.Date(DateUnit.MILLISECOND));
    schemaBuilder.addField("tinyintfield5", new ArrowType.Int(8, true));
    schemaBuilder.addField("uint1field6", new ArrowType.Int(8, false));
    schemaBuilder.addField("smallintfield7", new ArrowType.Int(16, true));
    schemaBuilder.addField("uint2field8", new ArrowType.Int(16, false));
    schemaBuilder.addField("datedayfield9", new ArrowType.Date(DateUnit.DAY));
    schemaBuilder.addField("uint4field10", new ArrowType.Int(32, false));
    schemaBuilder.addField("bigintfield11", new ArrowType.Int(64, true));
    schemaBuilder.addField("decimalfield12", new ArrowType.Decimal(10, 2));
    schemaBuilder.addField("floatfield13", new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE));
    schemaBuilder.addField("varbinaryfield14", new ArrowType.Binary());
    schemaBuilder.addField("bitfield15", new ArrowType.Bool());

    schemaBuilder.addListField("varcharlist16", Types.MinorType.VARCHAR.getType());
    schemaBuilder.addListField("intlist17", Types.MinorType.INT.getType());
    schemaBuilder.addListField("bigintlist18", Types.MinorType.BIGINT.getType());
    schemaBuilder.addListField("tinyintlist19", Types.MinorType.TINYINT.getType());
    schemaBuilder.addListField("smallintlist20", Types.MinorType.SMALLINT.getType());
    schemaBuilder.addListField("float4list21", Types.MinorType.FLOAT4.getType());
    schemaBuilder.addListField("float8list22", Types.MinorType.FLOAT8.getType());
    schemaBuilder.addListField("shortdeclist23", new ArrowType.Decimal(10, 2));
    schemaBuilder.addListField("londdeclist24", new ArrowType.Decimal(21, 2));
    schemaBuilder.addListField("varbinarylist25", Types.MinorType.VARBINARY.getType());
    schemaBuilder.addListField("bitlist26", Types.MinorType.BIT.getType());

    schemaBuilder.addStructField("structField27");
    schemaBuilder.addChildField("structField27", "nestedBigInt", Types.MinorType.BIGINT.getType());
    schemaBuilder.addChildField("structField27", "nestedString", Types.MinorType.VARCHAR.getType());
    schemaBuilder.addChildField("structField27", "tinyintcol", Types.MinorType.TINYINT.getType());
    schemaBuilder.addChildField("structField27", "smallintcol", Types.MinorType.SMALLINT.getType());
    schemaBuilder.addChildField("structField27", "float4Col", Types.MinorType.FLOAT4.getType());
    schemaBuilder.addChildField("structField27", "float8Col", Types.MinorType.FLOAT8.getType());
    schemaBuilder.addChildField("structField27", "shortDecCol", new ArrowType.Decimal(10, 2));
    schemaBuilder.addChildField("structField27", "longDecCol", new ArrowType.Decimal(21, 2));
    schemaBuilder.addChildField("structField27", "binaryCol", Types.MinorType.VARBINARY.getType());
    schemaBuilder.addChildField("structField27", "bitCol", Types.MinorType.BIT.getType());
    schemaBuilder.addStructField("structFieldNested28");
    schemaBuilder.addChildField("structFieldNested28", "bitCol", Types.MinorType.BIT.getType());
    schemaBuilder.addChildField("structFieldNested28",
            FieldBuilder.newBuilder("nestedStruct", new ArrowType.Struct())
                    .addField("nestedString", Types.MinorType.VARCHAR.getType(), null)
                    .addField("nestedBigInt", Types.MinorType.BIGINT.getType(), null)
                    .addListField("nestedList", Types.MinorType.VARCHAR.getType())
                    .addListField("nestedListDec", new ArrowType.Decimal(10, 2))
                    .build());
    return schemaBuilder.build();
}
 

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);
  }
}
 

@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 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;
}
 

@Override
public ArrowType visit(DateType dateType) {
	return new ArrowType.Date(DateUnit.DAY);
}
 

@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);
}