use of org.apache.arrow.vector.TimeStampMilliTZVector in project beam by apache.
the class ArrowConversionTest method rowIterator.
@Test
public void rowIterator() {
org.apache.arrow.vector.types.pojo.Schema schema = new org.apache.arrow.vector.types.pojo.Schema(asList(field("int32", new ArrowType.Int(32, true)), field("float64", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), field("string", new ArrowType.Utf8()), field("timestampMicroUTC", new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC")), field("timestampMilliUTC", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), field("int32_list", new ArrowType.List(), field("int32s", new ArrowType.Int(32, true))), field("boolean", new ArrowType.Bool()), field("fixed_size_binary", new ArrowType.FixedSizeBinary(3))));
Schema beamSchema = ArrowConversion.ArrowSchemaTranslator.toBeamSchema(schema);
VectorSchemaRoot expectedSchemaRoot = VectorSchemaRoot.create(schema, allocator);
expectedSchemaRoot.allocateNew();
expectedSchemaRoot.setRowCount(16);
IntVector intVector = (IntVector) expectedSchemaRoot.getFieldVectors().get(0);
Float8Vector floatVector = (Float8Vector) expectedSchemaRoot.getFieldVectors().get(1);
VarCharVector strVector = (VarCharVector) expectedSchemaRoot.getFieldVectors().get(2);
TimeStampMicroTZVector timestampMicroUtcVector = (TimeStampMicroTZVector) expectedSchemaRoot.getFieldVectors().get(3);
TimeStampMilliTZVector timeStampMilliTZVector = (TimeStampMilliTZVector) expectedSchemaRoot.getFieldVectors().get(4);
ListVector int32ListVector = (ListVector) expectedSchemaRoot.getFieldVectors().get(5);
IntVector int32ListElementVector = int32ListVector.<IntVector>addOrGetVector(new org.apache.arrow.vector.types.pojo.FieldType(false, new ArrowType.Int(32, true), null)).getVector();
BitVector boolVector = (BitVector) expectedSchemaRoot.getFieldVectors().get(6);
FixedSizeBinaryVector fixedSizeBinaryVector = (FixedSizeBinaryVector) expectedSchemaRoot.getFieldVectors().get(7);
ArrayList<Row> expectedRows = new ArrayList<>();
for (int i = 0; i < 16; i++) {
DateTime dt = new DateTime(2019, 1, i + 1, i, i, i, DateTimeZone.UTC);
expectedRows.add(Row.withSchema(beamSchema).addValues(i, i + .1 * i, "" + i, dt, dt, ImmutableList.of(i), (i % 2) != 0, new byte[] { (byte) i, (byte) (i + 1), (byte) (i + 2) }).build());
intVector.set(i, i);
floatVector.set(i, i + .1 * i);
strVector.set(i, new Text("" + i));
timestampMicroUtcVector.set(i, dt.getMillis() * 1000);
timeStampMilliTZVector.set(i, dt.getMillis());
int32ListVector.startNewValue(i);
int32ListElementVector.set(i, i);
int32ListVector.endValue(i, 1);
boolVector.set(i, i % 2);
fixedSizeBinaryVector.set(i, new byte[] { (byte) i, (byte) (i + 1), (byte) (i + 2) });
}
assertThat(ImmutableList.copyOf(ArrowConversion.rowsFromRecordBatch(beamSchema, expectedSchemaRoot)), IsIterableContainingInOrder.contains(expectedRows.stream().map((row) -> equalTo(row)).collect(ImmutableList.toImmutableList())));
expectedSchemaRoot.close();
}
Aggregations