Examples with ChoiceDataType org.apache.nifi.serialization.record.type.ChoiceDataType used on opensource projects

Example 1 with ChoiceDataType

use of org.apache.nifi.serialization.record.type.ChoiceDataType in project nifi by apache.

the class PutElasticsearchHttpRecord method writeValue.

@SuppressWarnings("unchecked")
private void writeValue(final JsonGenerator generator, final Object value, final String fieldName, final DataType dataType) throws IOException {
    if (value == null) {
        if (nullSuppression.equals(NEVER_SUPPRESS.getValue()) || ((nullSuppression.equals(SUPPRESS_MISSING.getValue())) && fieldName != null && !fieldName.equals(""))) {
            generator.writeNullField(fieldName);
        }
        return;
    }
    final DataType chosenDataType = dataType.getFieldType() == RecordFieldType.CHOICE ? DataTypeUtils.chooseDataType(value, (ChoiceDataType) dataType) : dataType;
    final Object coercedValue = DataTypeUtils.convertType(value, chosenDataType, fieldName);
    if (coercedValue == null) {
        generator.writeNull();
        return;
    }
    switch(chosenDataType.getFieldType()) {
        case DATE:
            {
                final String stringValue = DataTypeUtils.toString(coercedValue, () -> DataTypeUtils.getDateFormat(RecordFieldType.DATE.getDefaultFormat()));
                if (DataTypeUtils.isLongTypeCompatible(stringValue)) {
                    generator.writeNumber(DataTypeUtils.toLong(coercedValue, fieldName));
                } else {
                    generator.writeString(stringValue);
                }
                break;
            }
        case TIME:
            {
                final String stringValue = DataTypeUtils.toString(coercedValue, () -> DataTypeUtils.getDateFormat(RecordFieldType.TIME.getDefaultFormat()));
                if (DataTypeUtils.isLongTypeCompatible(stringValue)) {
                    generator.writeNumber(DataTypeUtils.toLong(coercedValue, fieldName));
                } else {
                    generator.writeString(stringValue);
                }
                break;
            }
        case TIMESTAMP:
            {
                final String stringValue = DataTypeUtils.toString(coercedValue, () -> DataTypeUtils.getDateFormat(RecordFieldType.TIMESTAMP.getDefaultFormat()));
                if (DataTypeUtils.isLongTypeCompatible(stringValue)) {
                    generator.writeNumber(DataTypeUtils.toLong(coercedValue, fieldName));
                } else {
                    generator.writeString(stringValue);
                }
                break;
            }
        case DOUBLE:
            generator.writeNumber(DataTypeUtils.toDouble(coercedValue, fieldName));
            break;
        case FLOAT:
            generator.writeNumber(DataTypeUtils.toFloat(coercedValue, fieldName));
            break;
        case LONG:
            generator.writeNumber(DataTypeUtils.toLong(coercedValue, fieldName));
            break;
        case INT:
        case BYTE:
        case SHORT:
            generator.writeNumber(DataTypeUtils.toInteger(coercedValue, fieldName));
            break;
        case CHAR:
        case STRING:
            generator.writeString(coercedValue.toString());
            break;
        case BIGINT:
            if (coercedValue instanceof Long) {
                generator.writeNumber((Long) coercedValue);
            } else {
                generator.writeNumber((BigInteger) coercedValue);
            }
            break;
        case BOOLEAN:
            final String stringValue = coercedValue.toString();
            if ("true".equalsIgnoreCase(stringValue)) {
                generator.writeBoolean(true);
            } else if ("false".equalsIgnoreCase(stringValue)) {
                generator.writeBoolean(false);
            } else {
                generator.writeString(stringValue);
            }
            break;
        case RECORD:
            {
                final Record record = (Record) coercedValue;
                final RecordDataType recordDataType = (RecordDataType) chosenDataType;
                final RecordSchema childSchema = recordDataType.getChildSchema();
                writeRecord(record, childSchema, generator);
                break;
            }
        case MAP:
            {
                final MapDataType mapDataType = (MapDataType) chosenDataType;
                final DataType valueDataType = mapDataType.getValueType();
                final Map<String, ?> map = (Map<String, ?>) coercedValue;
                generator.writeStartObject();
                for (final Map.Entry<String, ?> entry : map.entrySet()) {
                    final String mapKey = entry.getKey();
                    final Object mapValue = entry.getValue();
                    generator.writeFieldName(mapKey);
                    writeValue(generator, mapValue, fieldName + "." + mapKey, valueDataType);
                }
                generator.writeEndObject();
                break;
            }
        case ARRAY:
        default:
            if (coercedValue instanceof Object[]) {
                final Object[] values = (Object[]) coercedValue;
                final ArrayDataType arrayDataType = (ArrayDataType) dataType;
                final DataType elementType = arrayDataType.getElementType();
                writeArray(values, fieldName, generator, elementType);
            } else {
                generator.writeString(coercedValue.toString());
            }
            break;
    }
}

Example 2 with ChoiceDataType

use of org.apache.nifi.serialization.record.type.ChoiceDataType in project nifi by apache.

the class WriteJsonResult method writeValue.

@SuppressWarnings("unchecked")
private void writeValue(final JsonGenerator generator, final Object value, final String fieldName, final DataType dataType) throws JsonGenerationException, IOException {
    if (value == null) {
        generator.writeNull();
        return;
    }
    final DataType chosenDataType = dataType.getFieldType() == RecordFieldType.CHOICE ? DataTypeUtils.chooseDataType(value, (ChoiceDataType) dataType) : dataType;
    final Object coercedValue = DataTypeUtils.convertType(value, chosenDataType, LAZY_DATE_FORMAT, LAZY_TIME_FORMAT, LAZY_TIMESTAMP_FORMAT, fieldName);
    if (coercedValue == null) {
        generator.writeNull();
        return;
    }
    switch(chosenDataType.getFieldType()) {
        case DATE:
            {
                final String stringValue = DataTypeUtils.toString(coercedValue, LAZY_DATE_FORMAT);
                if (DataTypeUtils.isLongTypeCompatible(stringValue)) {
                    generator.writeNumber(DataTypeUtils.toLong(coercedValue, fieldName));
                } else {
                    generator.writeString(stringValue);
                }
                break;
            }
        case TIME:
            {
                final String stringValue = DataTypeUtils.toString(coercedValue, LAZY_TIME_FORMAT);
                if (DataTypeUtils.isLongTypeCompatible(stringValue)) {
                    generator.writeNumber(DataTypeUtils.toLong(coercedValue, fieldName));
                } else {
                    generator.writeString(stringValue);
                }
                break;
            }
        case TIMESTAMP:
            {
                final String stringValue = DataTypeUtils.toString(coercedValue, LAZY_TIMESTAMP_FORMAT);
                if (DataTypeUtils.isLongTypeCompatible(stringValue)) {
                    generator.writeNumber(DataTypeUtils.toLong(coercedValue, fieldName));
                } else {
                    generator.writeString(stringValue);
                }
                break;
            }
        case DOUBLE:
            generator.writeNumber(DataTypeUtils.toDouble(coercedValue, fieldName));
            break;
        case FLOAT:
            generator.writeNumber(DataTypeUtils.toFloat(coercedValue, fieldName));
            break;
        case LONG:
            generator.writeNumber(DataTypeUtils.toLong(coercedValue, fieldName));
            break;
        case INT:
        case BYTE:
        case SHORT:
            generator.writeNumber(DataTypeUtils.toInteger(coercedValue, fieldName));
            break;
        case CHAR:
        case STRING:
            generator.writeString(coercedValue.toString());
            break;
        case BIGINT:
            if (coercedValue instanceof Long) {
                generator.writeNumber(((Long) coercedValue).longValue());
            } else {
                generator.writeNumber((BigInteger) coercedValue);
            }
            break;
        case BOOLEAN:
            final String stringValue = coercedValue.toString();
            if ("true".equalsIgnoreCase(stringValue)) {
                generator.writeBoolean(true);
            } else if ("false".equalsIgnoreCase(stringValue)) {
                generator.writeBoolean(false);
            } else {
                generator.writeString(stringValue);
            }
            break;
        case RECORD:
            {
                final Record record = (Record) coercedValue;
                final RecordDataType recordDataType = (RecordDataType) chosenDataType;
                final RecordSchema childSchema = recordDataType.getChildSchema();
                writeRecord(record, childSchema, generator, gen -> gen.writeStartObject(), gen -> gen.writeEndObject(), true);
                break;
            }
        case MAP:
            {
                final MapDataType mapDataType = (MapDataType) chosenDataType;
                final DataType valueDataType = mapDataType.getValueType();
                final Map<String, ?> map = (Map<String, ?>) coercedValue;
                generator.writeStartObject();
                for (final Map.Entry<String, ?> entry : map.entrySet()) {
                    final String mapKey = entry.getKey();
                    final Object mapValue = entry.getValue();
                    generator.writeFieldName(mapKey);
                    writeValue(generator, mapValue, fieldName + "." + mapKey, valueDataType);
                }
                generator.writeEndObject();
                break;
            }
        case ARRAY:
        default:
            if (coercedValue instanceof Object[]) {
                final Object[] values = (Object[]) coercedValue;
                final ArrayDataType arrayDataType = (ArrayDataType) dataType;
                final DataType elementType = arrayDataType.getElementType();
                writeArray(values, fieldName, generator, elementType);
            } else {
                generator.writeString(coercedValue.toString());
            }
            break;
    }
}

Example 3 with ChoiceDataType

use of org.apache.nifi.serialization.record.type.ChoiceDataType in project nifi by apache.

the class TestJsonTreeRowRecordReader method testSingleJsonElementWithChoiceFields.

@Test
public void testSingleJsonElementWithChoiceFields() throws IOException, MalformedRecordException {
    // Wraps default fields by Choice data type to test mapping to a Choice type.
    final List<RecordField> choiceFields = getDefaultFields().stream().map(f -> new RecordField(f.getFieldName(), RecordFieldType.CHOICE.getChoiceDataType(f.getDataType()))).collect(Collectors.toList());
    final RecordSchema schema = new SimpleRecordSchema(choiceFields);
    try (final InputStream in = new FileInputStream(new File("src/test/resources/json/single-bank-account.json"));
        final JsonTreeRowRecordReader reader = new JsonTreeRowRecordReader(in, Mockito.mock(ComponentLog.class), schema, dateFormat, timeFormat, timestampFormat)) {
        final List<String> fieldNames = schema.getFieldNames();
        final List<String> expectedFieldNames = Arrays.asList(new String[] { "id", "name", "balance", "address", "city", "state", "zipCode", "country" });
        assertEquals(expectedFieldNames, fieldNames);
        final List<RecordFieldType> expectedTypes = Arrays.asList(new RecordFieldType[] { RecordFieldType.INT, RecordFieldType.STRING, RecordFieldType.DOUBLE, RecordFieldType.STRING, RecordFieldType.STRING, RecordFieldType.STRING, RecordFieldType.STRING, RecordFieldType.STRING });
        final List<RecordField> fields = schema.getFields();
        for (int i = 0; i < schema.getFields().size(); i++) {
            assertTrue(fields.get(i).getDataType() instanceof ChoiceDataType);
            final ChoiceDataType choiceDataType = (ChoiceDataType) fields.get(i).getDataType();
            assertEquals(expectedTypes.get(i), choiceDataType.getPossibleSubTypes().get(0).getFieldType());
        }
        final Object[] firstRecordValues = reader.nextRecord().getValues();
        Assert.assertArrayEquals(new Object[] { 1, "John Doe", 4750.89, "123 My Street", "My City", "MS", "11111", "USA" }, firstRecordValues);
        assertNull(reader.nextRecord());
    }
}

Example 4 with ChoiceDataType

use of org.apache.nifi.serialization.record.type.ChoiceDataType in project nifi by apache.

the class DataTypeUtils method isScalarValue.

public static boolean isScalarValue(final DataType dataType, final Object value) {
    final RecordFieldType fieldType = dataType.getFieldType();
    final RecordFieldType chosenType;
    if (fieldType == RecordFieldType.CHOICE) {
        final ChoiceDataType choiceDataType = (ChoiceDataType) dataType;
        final DataType chosenDataType = chooseDataType(value, choiceDataType);
        if (chosenDataType == null) {
            return false;
        }
        chosenType = chosenDataType.getFieldType();
    } else {
        chosenType = fieldType;
    }
    switch(chosenType) {
        case ARRAY:
        case MAP:
        case RECORD:
            return false;
    }
    return true;
}

Example 5 with ChoiceDataType

use of org.apache.nifi.serialization.record.type.ChoiceDataType in project nifi by apache.

the class AvroTypeUtil method buildAvroSchema.

private static Schema buildAvroSchema(final DataType dataType, final String fieldName, final boolean nullable) {
    final Schema schema;
    switch(dataType.getFieldType()) {
        case ARRAY:
            final ArrayDataType arrayDataType = (ArrayDataType) dataType;
            final DataType elementDataType = arrayDataType.getElementType();
            if (RecordFieldType.BYTE.equals(elementDataType.getFieldType())) {
                schema = Schema.create(Type.BYTES);
            } else {
                final Schema elementType = buildAvroSchema(elementDataType, fieldName, false);
                schema = Schema.createArray(elementType);
            }
            break;
        case BIGINT:
            schema = Schema.create(Type.STRING);
            break;
        case BOOLEAN:
            schema = Schema.create(Type.BOOLEAN);
            break;
        case BYTE:
            schema = Schema.create(Type.INT);
            break;
        case CHAR:
            schema = Schema.create(Type.STRING);
            break;
        case CHOICE:
            final ChoiceDataType choiceDataType = (ChoiceDataType) dataType;
            final List<DataType> options = choiceDataType.getPossibleSubTypes();
            // We need to keep track of which types have been added to the union, because if we have
            // two elements in the UNION with the same type, it will fail - even if the logical type is
            // different. So if we have an int and a logical type date (which also has a 'concrete type' of int)
            // then an Exception will be thrown when we try to create the union. To avoid this, we just keep track
            // of the Types and avoid adding it in such a case.
            final List<Schema> unionTypes = new ArrayList<>(options.size());
            final Set<Type> typesAdded = new HashSet<>();
            for (final DataType option : options) {
                final Schema optionSchema = buildAvroSchema(option, fieldName, false);
                if (!typesAdded.contains(optionSchema.getType())) {
                    unionTypes.add(optionSchema);
                    typesAdded.add(optionSchema.getType());
                }
            }
            schema = Schema.createUnion(unionTypes);
            break;
        case DATE:
            schema = Schema.create(Type.INT);
            LogicalTypes.date().addToSchema(schema);
            break;
        case DOUBLE:
            schema = Schema.create(Type.DOUBLE);
            break;
        case FLOAT:
            schema = Schema.create(Type.FLOAT);
            break;
        case INT:
            schema = Schema.create(Type.INT);
            break;
        case LONG:
            schema = Schema.create(Type.LONG);
            break;
        case MAP:
            schema = Schema.createMap(buildAvroSchema(((MapDataType) dataType).getValueType(), fieldName, false));
            break;
        case RECORD:
            final RecordDataType recordDataType = (RecordDataType) dataType;
            final RecordSchema childSchema = recordDataType.getChildSchema();
            final List<Field> childFields = new ArrayList<>(childSchema.getFieldCount());
            for (final RecordField field : childSchema.getFields()) {
                childFields.add(buildAvroField(field));
            }
            schema = Schema.createRecord(fieldName + "Type", null, "org.apache.nifi", false, childFields);
            break;
        case SHORT:
            schema = Schema.create(Type.INT);
            break;
        case STRING:
            schema = Schema.create(Type.STRING);
            break;
        case TIME:
            schema = Schema.create(Type.INT);
            LogicalTypes.timeMillis().addToSchema(schema);
            break;
        case TIMESTAMP:
            schema = Schema.create(Type.LONG);
            LogicalTypes.timestampMillis().addToSchema(schema);
            break;
        default:
            return null;
    }
    if (nullable) {
        return nullable(schema);
    } else {
        return schema;
    }
}