Search in sources :

Example 11 with EnumerationType

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

the class AvroUtils method genericFromBeamField.

@Nullable
private static Object genericFromBeamField(Schema.FieldType fieldType, org.apache.avro.Schema avroSchema, @Nullable Object value) {
    TypeWithNullability typeWithNullability = new TypeWithNullability(avroSchema);
    if (!fieldType.getNullable().equals(typeWithNullability.nullable)) {
        throw new IllegalArgumentException("FieldType " + fieldType + " and AVRO schema " + avroSchema + " don't have matching nullability");
    }
    if (value == null) {
        return value;
    }
    switch(fieldType.getTypeName()) {
        case BYTE:
        case INT16:
        case INT32:
        case INT64:
        case FLOAT:
        case DOUBLE:
        case BOOLEAN:
            return value;
        case STRING:
            return new Utf8((String) value);
        case DECIMAL:
            BigDecimal decimal = (BigDecimal) value;
            LogicalType logicalType = typeWithNullability.type.getLogicalType();
            return new Conversions.DecimalConversion().toBytes(decimal, null, logicalType);
        case DATETIME:
            if (typeWithNullability.type.getType() == Type.INT) {
                ReadableInstant instant = (ReadableInstant) value;
                return (int) Days.daysBetween(Instant.EPOCH, instant).getDays();
            } else if (typeWithNullability.type.getType() == Type.LONG) {
                ReadableInstant instant = (ReadableInstant) value;
                return (long) instant.getMillis();
            } else {
                throw new IllegalArgumentException("Can't represent " + fieldType + " as " + typeWithNullability.type.getType());
            }
        case BYTES:
            return ByteBuffer.wrap((byte[]) value);
        case LOGICAL_TYPE:
            switch(fieldType.getLogicalType().getIdentifier()) {
                case FixedBytes.IDENTIFIER:
                    FixedBytesField fixedBytesField = checkNotNull(FixedBytesField.fromBeamFieldType(fieldType));
                    byte[] byteArray = (byte[]) value;
                    if (byteArray.length != fixedBytesField.getSize()) {
                        throw new IllegalArgumentException("Incorrectly sized byte array.");
                    }
                    return GenericData.get().createFixed(null, (byte[]) value, typeWithNullability.type);
                case EnumerationType.IDENTIFIER:
                    EnumerationType enumerationType = fieldType.getLogicalType(EnumerationType.class);
                    return GenericData.get().createEnum(enumerationType.toString((EnumerationType.Value) value), typeWithNullability.type);
                case OneOfType.IDENTIFIER:
                    OneOfType oneOfType = fieldType.getLogicalType(OneOfType.class);
                    OneOfType.Value oneOfValue = (OneOfType.Value) value;
                    FieldType innerFieldType = oneOfType.getFieldType(oneOfValue);
                    if (typeWithNullability.nullable && oneOfValue.getValue() == null) {
                        return null;
                    } else {
                        return genericFromBeamField(innerFieldType.withNullable(false), typeWithNullability.type.getTypes().get(oneOfValue.getCaseType().getValue()), oneOfValue.getValue());
                    }
                case "NVARCHAR":
                case "VARCHAR":
                case "LONGNVARCHAR":
                case "LONGVARCHAR":
                    return new Utf8((String) value);
                case "DATE":
                    return Days.daysBetween(Instant.EPOCH, (Instant) value).getDays();
                case "TIME":
                    return (int) ((Instant) value).getMillis();
                default:
                    throw new RuntimeException("Unhandled logical type " + fieldType.getLogicalType().getIdentifier());
            }
        case ARRAY:
        case ITERABLE:
            Iterable iterable = (Iterable) value;
            List<Object> translatedArray = Lists.newArrayListWithExpectedSize(Iterables.size(iterable));
            for (Object arrayElement : iterable) {
                translatedArray.add(genericFromBeamField(fieldType.getCollectionElementType(), typeWithNullability.type.getElementType(), arrayElement));
            }
            return translatedArray;
        case MAP:
            Map map = Maps.newHashMap();
            Map<Object, Object> valueMap = (Map<Object, Object>) value;
            for (Map.Entry entry : valueMap.entrySet()) {
                Utf8 key = new Utf8((String) entry.getKey());
                map.put(key, genericFromBeamField(fieldType.getMapValueType(), typeWithNullability.type.getValueType(), entry.getValue()));
            }
            return map;
        case ROW:
            return toGenericRecord((Row) value, typeWithNullability.type);
        default:
            throw new IllegalArgumentException("Unsupported type " + fieldType);
    }
}
Also used : ReadableInstant(org.joda.time.ReadableInstant) ReadableInstant(org.joda.time.ReadableInstant) Instant(org.joda.time.Instant) EnumerationType(org.apache.beam.sdk.schemas.logicaltypes.EnumerationType) LogicalType(org.apache.avro.LogicalType) BigDecimal(java.math.BigDecimal) FieldType(org.apache.beam.sdk.schemas.Schema.FieldType) AvroRuntimeException(org.apache.avro.AvroRuntimeException) Conversions(org.apache.avro.Conversions) Utf8(org.apache.avro.util.Utf8) Map(java.util.Map) HashMap(java.util.HashMap) OneOfType(org.apache.beam.sdk.schemas.logicaltypes.OneOfType) Nullable(org.checkerframework.checker.nullness.qual.Nullable)

Example 12 with EnumerationType

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

the class RowCoderTest method testLogicalType.

@Test
public void testLogicalType() throws Exception {
    EnumerationType enumeration = EnumerationType.create("one", "two", "three");
    Schema schema = Schema.builder().addLogicalTypeField("f_enum", enumeration).build();
    Row row = Row.withSchema(schema).addValue(enumeration.valueOf("two")).build();
    CoderProperties.coderDecodeEncodeEqual(RowCoder.of(schema), row);
}
Also used : Schema(org.apache.beam.sdk.schemas.Schema) EnumerationType(org.apache.beam.sdk.schemas.logicaltypes.EnumerationType) Row(org.apache.beam.sdk.values.Row) Test(org.junit.Test)

Example 13 with EnumerationType

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

the class JavaFieldSchemaTest method testEnumFieldToRow.

@Test
public void testEnumFieldToRow() 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();
    List<Color> allColorsJava = Lists.newArrayList(Color.RED, Color.GREEN, Color.BLUE);
    SerializableFunction<PojoWithEnum, Row> toRow = registry.getToRowFunction(PojoWithEnum.class);
    assertEquals(redRow, toRow.apply(new PojoWithEnum(Color.RED, allColorsJava)));
    assertEquals(greenRow, toRow.apply(new PojoWithEnum(Color.GREEN, allColorsJava)));
    assertEquals(blueRow, toRow.apply(new PojoWithEnum(Color.BLUE, allColorsJava)));
}
Also used : Color(org.apache.beam.sdk.schemas.utils.TestPOJOs.PojoWithEnum.Color) EnumerationType(org.apache.beam.sdk.schemas.logicaltypes.EnumerationType) PojoWithEnum(org.apache.beam.sdk.schemas.utils.TestPOJOs.PojoWithEnum) Row(org.apache.beam.sdk.values.Row) Test(org.junit.Test)

Aggregations

EnumerationType (org.apache.beam.sdk.schemas.logicaltypes.EnumerationType)13 Test (org.junit.Test)9 Schema (org.apache.beam.sdk.schemas.Schema)8 Row (org.apache.beam.sdk.values.Row)7 Schema.toSchema (org.apache.beam.sdk.schemas.Schema.toSchema)4 Row.toRow (org.apache.beam.sdk.values.Row.toRow)4 Map (java.util.Map)3 OneOfType (org.apache.beam.sdk.schemas.logicaltypes.OneOfType)3 Nullable (org.checkerframework.checker.nullness.qual.Nullable)3 BigDecimal (java.math.BigDecimal)2 HashMap (java.util.HashMap)2 AvroRuntimeException (org.apache.avro.AvroRuntimeException)2 Conversions (org.apache.avro.Conversions)2 LogicalType (org.apache.avro.LogicalType)2 Utf8 (org.apache.avro.util.Utf8)2 ByteArrayInputStream (java.io.ByteArrayInputStream)1 ByteArrayOutputStream (java.io.ByteArrayOutputStream)1 IOException (java.io.IOException)1 ObjectInputStream (java.io.ObjectInputStream)1 ObjectOutputStream (java.io.ObjectOutputStream)1