Examples with OneOfType org.apache.beam.sdk.schemas.logicaltypes.OneOfType used on opensource projects

Example 1 with OneOfType

use of org.apache.beam.sdk.schemas.logicaltypes.OneOfType in project beam by apache.

the class AvroUtils method getFieldSchema.

private static org.apache.avro.Schema getFieldSchema(Schema.FieldType fieldType, String fieldName, String namespace) {
    org.apache.avro.Schema baseType;
    switch(fieldType.getTypeName()) {
        case BYTE:
        case INT16:
        case INT32:
            baseType = org.apache.avro.Schema.create(Type.INT);
            break;
        case INT64:
            baseType = org.apache.avro.Schema.create(Type.LONG);
            break;
        case DECIMAL:
            baseType = LogicalTypes.decimal(Integer.MAX_VALUE).addToSchema(org.apache.avro.Schema.create(Type.BYTES));
            break;
        case FLOAT:
            baseType = org.apache.avro.Schema.create(Type.FLOAT);
            break;
        case DOUBLE:
            baseType = org.apache.avro.Schema.create(Type.DOUBLE);
            break;
        case STRING:
            baseType = org.apache.avro.Schema.create(Type.STRING);
            break;
        case DATETIME:
            // TODO: There is a desire to move Beam schema DATETIME to a micros representation. When
            // this is done, this logical type needs to be changed.
            baseType = LogicalTypes.timestampMillis().addToSchema(org.apache.avro.Schema.create(Type.LONG));
            break;
        case BOOLEAN:
            baseType = org.apache.avro.Schema.create(Type.BOOLEAN);
            break;
        case BYTES:
            baseType = org.apache.avro.Schema.create(Type.BYTES);
            break;
        case LOGICAL_TYPE:
            switch(fieldType.getLogicalType().getIdentifier()) {
                case FixedBytes.IDENTIFIER:
                    FixedBytesField fixedBytesField = checkNotNull(FixedBytesField.fromBeamFieldType(fieldType));
                    baseType = fixedBytesField.toAvroType("fixed", namespace + "." + fieldName);
                    break;
                case EnumerationType.IDENTIFIER:
                    EnumerationType enumerationType = fieldType.getLogicalType(EnumerationType.class);
                    baseType = org.apache.avro.Schema.createEnum(fieldName, "", "", enumerationType.getValues());
                    break;
                case OneOfType.IDENTIFIER:
                    OneOfType oneOfType = fieldType.getLogicalType(OneOfType.class);
                    baseType = org.apache.avro.Schema.createUnion(oneOfType.getOneOfSchema().getFields().stream().map(x -> getFieldSchema(x.getType(), x.getName(), namespace)).collect(Collectors.toList()));
                    break;
                case "CHAR":
                case "NCHAR":
                    baseType = buildHiveLogicalTypeSchema("char", (int) fieldType.getLogicalType().getArgument());
                    break;
                case "NVARCHAR":
                case "VARCHAR":
                case "LONGNVARCHAR":
                case "LONGVARCHAR":
                    baseType = buildHiveLogicalTypeSchema("varchar", (int) fieldType.getLogicalType().getArgument());
                    break;
                case "DATE":
                    baseType = LogicalTypes.date().addToSchema(org.apache.avro.Schema.create(Type.INT));
                    break;
                case "TIME":
                    baseType = LogicalTypes.timeMillis().addToSchema(org.apache.avro.Schema.create(Type.INT));
                    break;
                default:
                    throw new RuntimeException("Unhandled logical type " + fieldType.getLogicalType().getIdentifier());
            }
            break;
        case ARRAY:
        case ITERABLE:
            baseType = org.apache.avro.Schema.createArray(getFieldSchema(fieldType.getCollectionElementType(), fieldName, namespace));
            break;
        case MAP:
            if (fieldType.getMapKeyType().getTypeName().isStringType()) {
                // Avro only supports string keys in maps.
                baseType = org.apache.avro.Schema.createMap(getFieldSchema(fieldType.getMapValueType(), fieldName, namespace));
            } else {
                throw new IllegalArgumentException("Avro only supports maps with string keys");
            }
            break;
        case ROW:
            baseType = toAvroSchema(fieldType.getRowSchema(), fieldName, namespace);
            break;
        default:
            throw new IllegalArgumentException("Unexpected type " + fieldType);
    }
    return fieldType.getNullable() ? ReflectData.makeNullable(baseType) : baseType;
}

Example 2 with OneOfType

use of org.apache.beam.sdk.schemas.logicaltypes.OneOfType in project beam by apache.

the class ProtoByteBuddyUtils method getProtoFieldValueSetter.

private static <ProtoBuilderT extends MessageLite.Builder> FieldValueSetter<ProtoBuilderT, Object> getProtoFieldValueSetter(Field field, Multimap<String, Method> methods, Class<ProtoBuilderT> builderClass) {
    if (field.getType().isLogicalType(OneOfType.IDENTIFIER)) {
        OneOfType oneOfType = field.getType().getLogicalType(OneOfType.class);
        TreeMap<Integer, FieldValueSetter<ProtoBuilderT, Object>> oneOfSetters = Maps.newTreeMap();
        for (Field oneOfField : oneOfType.getOneOfSchema().getFields()) {
            FieldValueSetter setter = getProtoFieldValueSetter(oneOfField, methods, builderClass);
            oneOfSetters.put(getFieldNumber(oneOfField), setter);
        }
        return createOneOfSetter(field.getName(), oneOfSetters, builderClass);
    } else {
        Method method = getProtoSetter(methods, field.getName(), field.getType());
        return JavaBeanUtils.createSetter(FieldValueTypeInformation.forSetter(method, protoSetterPrefix(field.getType())), new ProtoTypeConversionsFactory());
    }
}

Example 3 with OneOfType

use of org.apache.beam.sdk.schemas.logicaltypes.OneOfType in project beam by apache.

the class AvroUtilsTest method testUnionFieldInAvroSchema.

@Test
public void testUnionFieldInAvroSchema() {
    List<org.apache.avro.Schema.Field> fields = Lists.newArrayList();
    List<org.apache.avro.Schema> unionFields = Lists.newArrayList();
    unionFields.add(org.apache.avro.Schema.create(Type.INT));
    unionFields.add(org.apache.avro.Schema.create(Type.STRING));
    fields.add(new org.apache.avro.Schema.Field("union", org.apache.avro.Schema.createUnion(unionFields), "", null));
    org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("topLevelRecord", null, null, false, fields);
    OneOfType oneOfType = OneOfType.create(Field.of("int", FieldType.INT32), Field.of("string", FieldType.STRING));
    Schema expectedSchema = Schema.builder().addLogicalTypeField("union", oneOfType).build();
    assertEquals(expectedSchema, AvroUtils.toBeamSchema(avroSchema));
    GenericRecord genericRecord = new GenericRecordBuilder(avroSchema).set("union", 23423).build();
    Row expectedRow = Row.withSchema(expectedSchema).addValue(oneOfType.createValue(0, 23423)).build();
    assertEquals(expectedRow, AvroUtils.toBeamRowStrict(genericRecord, expectedSchema));
}

Example 4 with OneOfType

use of org.apache.beam.sdk.schemas.logicaltypes.OneOfType in project beam by apache.

the class AvroUtilsTest method testUnionFieldInBeamSchema.

@Test
public void testUnionFieldInBeamSchema() {
    OneOfType oneOfType = OneOfType.create(Field.of("int", FieldType.INT32), Field.of("string", FieldType.STRING));
    Schema beamSchema = Schema.builder().addLogicalTypeField("union", oneOfType).build();
    List<org.apache.avro.Schema.Field> fields = Lists.newArrayList();
    List<org.apache.avro.Schema> unionFields = Lists.newArrayList();
    unionFields.add(org.apache.avro.Schema.create(Type.INT));
    unionFields.add(org.apache.avro.Schema.create(Type.STRING));
    fields.add(new org.apache.avro.Schema.Field("union", org.apache.avro.Schema.createUnion(unionFields), "", null));
    org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("topLevelRecord", null, null, false, fields);
    GenericRecord expectedGenericRecord = new GenericRecordBuilder(avroSchema).set("union", 23423).build();
    Row row = Row.withSchema(beamSchema).addValue(oneOfType.createValue(0, 23423)).build();
    assertEquals(expectedGenericRecord, AvroUtils.toGenericRecord(row, avroSchema));
}

Example 5 with OneOfType

use of org.apache.beam.sdk.schemas.logicaltypes.OneOfType in project beam by apache.

the class FromRowUsingCreator method fromValue.

@SuppressWarnings("unchecked")
@Nullable
private <ValueT> ValueT fromValue(FieldType type, ValueT value, Type fieldType, FieldValueTypeInformation fieldValueTypeInformation, Factory<List<FieldValueTypeInformation>> typeFactory) {
    FieldValueTypeInformation elementType = fieldValueTypeInformation.getElementType();
    FieldValueTypeInformation keyType = fieldValueTypeInformation.getMapKeyType();
    FieldValueTypeInformation valueType = fieldValueTypeInformation.getMapValueType();
    if (value == null) {
        return null;
    }
    if (TypeName.ROW.equals(type.getTypeName())) {
        return (ValueT) fromRow((Row) value, (Class) fieldType, typeFactory);
    } else if (TypeName.ARRAY.equals(type.getTypeName())) {
        return (ValueT) fromCollectionValue(type.getCollectionElementType(), (Collection) value, elementType, typeFactory);
    } else if (TypeName.ITERABLE.equals(type.getTypeName())) {
        return (ValueT) fromIterableValue(type.getCollectionElementType(), (Iterable) value, elementType, typeFactory);
    }
    if (TypeName.MAP.equals(type.getTypeName())) {
        return (ValueT) fromMapValue(type.getMapKeyType(), type.getMapValueType(), (Map) value, keyType, valueType, typeFactory);
    } else {
        if (type.isLogicalType(OneOfType.IDENTIFIER)) {
            OneOfType oneOfType = type.getLogicalType(OneOfType.class);
            EnumerationType oneOfEnum = oneOfType.getCaseEnumType();
            OneOfType.Value oneOfValue = (OneOfType.Value) value;
            FieldValueTypeInformation oneOfFieldValueTypeInformation = checkNotNull(fieldValueTypeInformation.getOneOfTypes().get(oneOfEnum.toString(oneOfValue.getCaseType())));
            Object fromValue = fromValue(oneOfType.getFieldType(oneOfValue), oneOfValue.getValue(), oneOfFieldValueTypeInformation.getRawType(), oneOfFieldValueTypeInformation, typeFactory);
            return (ValueT) oneOfType.createValue(oneOfValue.getCaseType(), fromValue);
        } else if (type.getTypeName().isLogicalType()) {
            return (ValueT) type.getLogicalType().toBaseType(value);
        }
        return value;
    }
}