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

Example 1 with EnumerationType

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

the class JavaFieldSchemaTest method testEnumFieldFromRow.

@Test
public void testEnumFieldFromRow() throws NoSuchSchemaException {
    SchemaRegistry registry = SchemaRegistry.createDefault();
    Schema schema = registry.getSchema(PojoWithEnum.class);
    SchemaTestUtils.assertSchemaEquivalent(POJO_WITH_ENUM_SCHEMA, schema);
    EnumerationType enumerationType = ENUMERATION;
    List<EnumerationType.Value> allColors = Lists.newArrayList(enumerationType.valueOf("RED"), enumerationType.valueOf("GREEN"), enumerationType.valueOf("BLUE"));
    Row redRow = Row.withSchema(POJO_WITH_ENUM_SCHEMA).addValues(enumerationType.valueOf("RED"), allColors).build();
    Row greenRow = Row.withSchema(POJO_WITH_ENUM_SCHEMA).addValues(enumerationType.valueOf("GREEN"), allColors).build();
    Row blueRow = Row.withSchema(POJO_WITH_ENUM_SCHEMA).addValues(enumerationType.valueOf("BLUE"), allColors).build();
    SerializableFunction<Row, PojoWithEnum> fromRow = registry.getFromRowFunction(PojoWithEnum.class);
    List<Color> allColorsJava = Lists.newArrayList(Color.RED, Color.GREEN, Color.BLUE);
    assertEquals(new PojoWithEnum(Color.RED, allColorsJava), fromRow.apply(redRow));
    assertEquals(new PojoWithEnum(Color.GREEN, allColorsJava), fromRow.apply(greenRow));
    assertEquals(new PojoWithEnum(Color.BLUE, allColorsJava), fromRow.apply(blueRow));
}

Example 2 with EnumerationType

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

the class RowCoderTest method testLogicalTypeInCollection.

@Test
public void testLogicalTypeInCollection() throws Exception {
    EnumerationType enumeration = EnumerationType.create("one", "two", "three");
    Schema schema = Schema.builder().addArrayField("f_enum_array", FieldType.logicalType(enumeration)).build();
    Row row = Row.withSchema(schema).addArray(enumeration.valueOf("two"), enumeration.valueOf("three")).build();
    CoderProperties.coderDecodeEncodeEqual(RowCoder.of(schema), row);
}

Example 3 with EnumerationType

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

the class RowCoderTest method testNestedLogicalTypes.

@Test
public void testNestedLogicalTypes() throws Exception {
    EnumerationType enumeration = EnumerationType.create("one", "two", "three");
    Schema schema = Schema.builder().addLogicalTypeField("f_nested_logical_type", new NestedLogicalType(enumeration)).build();
    Row row = Row.withSchema(schema).addValue("two").build();
    CoderProperties.coderDecodeEncodeEqual(RowCoder.of(schema), row);
}

Example 4 with EnumerationType

use of org.apache.beam.sdk.schemas.logicaltypes.EnumerationType 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 5 with EnumerationType

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

the class AvroUtils method convertEnumStrict.

private static Object convertEnumStrict(Object value, Schema.FieldType fieldType) {
    checkTypeName(fieldType.getTypeName(), TypeName.LOGICAL_TYPE, "enum");
    checkArgument(fieldType.getLogicalType().getIdentifier().equals(EnumerationType.IDENTIFIER));
    EnumerationType enumerationType = fieldType.getLogicalType(EnumerationType.class);
    return enumerationType.valueOf(value.toString());
}