Examples with RecordSetWriter org.apache.nifi.serialization.RecordSetWriter used on opensource projects

Example 21 with RecordSetWriter

use of org.apache.nifi.serialization.RecordSetWriter in project nifi by apache.

the class TestWriteAvroResult method testLogicalTypes.

private void testLogicalTypes(Schema schema) throws ParseException, IOException {
    final ByteArrayOutputStream baos = new ByteArrayOutputStream();
    final List<RecordField> fields = new ArrayList<>();
    fields.add(new RecordField("timeMillis", RecordFieldType.TIME.getDataType()));
    fields.add(new RecordField("timeMicros", RecordFieldType.TIME.getDataType()));
    fields.add(new RecordField("timestampMillis", RecordFieldType.TIMESTAMP.getDataType()));
    fields.add(new RecordField("timestampMicros", RecordFieldType.TIMESTAMP.getDataType()));
    fields.add(new RecordField("date", RecordFieldType.DATE.getDataType()));
    // Avro decimal is represented as double in NiFi type system.
    fields.add(new RecordField("decimal", RecordFieldType.DOUBLE.getDataType()));
    final RecordSchema recordSchema = new SimpleRecordSchema(fields);
    final String expectedTime = "2017-04-04 14:20:33.789";
    final DateFormat df = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS");
    df.setTimeZone(TimeZone.getTimeZone("gmt"));
    final long timeLong = df.parse(expectedTime).getTime();
    final Map<String, Object> values = new HashMap<>();
    values.put("timeMillis", new Time(timeLong));
    values.put("timeMicros", new Time(timeLong));
    values.put("timestampMillis", new Timestamp(timeLong));
    values.put("timestampMicros", new Timestamp(timeLong));
    values.put("date", new Date(timeLong));
    // Avro decimal is represented as double in NiFi type system.
    final BigDecimal expectedDecimal = new BigDecimal("123.45");
    values.put("decimal", expectedDecimal.doubleValue());
    final Record record = new MapRecord(recordSchema, values);
    try (final RecordSetWriter writer = createWriter(schema, baos)) {
        writer.write(RecordSet.of(record.getSchema(), record));
    }
    final byte[] data = baos.toByteArray();
    try (final InputStream in = new ByteArrayInputStream(data)) {
        final GenericRecord avroRecord = readRecord(in, schema);
        final long secondsSinceMidnight = 33 + (20 * 60) + (14 * 60 * 60);
        final long millisSinceMidnight = (secondsSinceMidnight * 1000L) + 789;
        assertEquals((int) millisSinceMidnight, avroRecord.get("timeMillis"));
        assertEquals(millisSinceMidnight * 1000L, avroRecord.get("timeMicros"));
        assertEquals(timeLong, avroRecord.get("timestampMillis"));
        assertEquals(timeLong * 1000L, avroRecord.get("timestampMicros"));
        assertEquals(17260, avroRecord.get("date"));
        // Double value will be converted into logical decimal if Avro schema is defined as logical decimal.
        final Schema decimalSchema = schema.getField("decimal").schema();
        final LogicalType logicalType = decimalSchema.getLogicalType() != null ? decimalSchema.getLogicalType() : // Union type doesn't return logical type. Find the first logical type defined within the union.
        decimalSchema.getTypes().stream().map(s -> s.getLogicalType()).filter(Objects::nonNull).findFirst().get();
        final BigDecimal decimal = new Conversions.DecimalConversion().fromBytes((ByteBuffer) avroRecord.get("decimal"), decimalSchema, logicalType);
        assertEquals(expectedDecimal, decimal);
    }
}