Search in sources :

Example 1 with TypeName

use of org.apache.beam.sdk.schemas.Schema.TypeName in project beam by apache.

the class BigQueryUtils method convertAvroFormat.

/**
 * Tries to convert an Avro decoded value to a Beam field value based on the target type of the
 * Beam field.
 *
 * <p>For the Avro formats of BigQuery types, see
 * https://cloud.google.com/bigquery/docs/exporting-data#avro_export_details and
 * https://cloud.google.com/bigquery/docs/loading-data-cloud-storage-avro#avro_conversions
 */
public static Object convertAvroFormat(FieldType beamFieldType, Object avroValue, BigQueryUtils.ConversionOptions options) {
    TypeName beamFieldTypeName = beamFieldType.getTypeName();
    if (avroValue == null) {
        if (beamFieldType.getNullable()) {
            return null;
        } else {
            throw new IllegalArgumentException(String.format("Field %s not nullable", beamFieldType));
        }
    }
    switch(beamFieldTypeName) {
        case BYTE:
        case INT16:
        case INT32:
        case INT64:
        case FLOAT:
        case DOUBLE:
        case STRING:
        case BYTES:
        case BOOLEAN:
            return convertAvroPrimitiveTypes(beamFieldTypeName, avroValue);
        case DATETIME:
            // Expecting value in microseconds.
            switch(options.getTruncateTimestamps()) {
                case TRUNCATE:
                    return truncateToMillis(avroValue);
                case REJECT:
                    return safeToMillis(avroValue);
                default:
                    throw new IllegalArgumentException(String.format("Unknown timestamp truncation option: %s", options.getTruncateTimestamps()));
            }
        case DECIMAL:
            return convertAvroNumeric(avroValue);
        case ARRAY:
            return convertAvroArray(beamFieldType, avroValue, options);
        case LOGICAL_TYPE:
            String identifier = beamFieldType.getLogicalType().getIdentifier();
            if (SqlTypes.DATE.getIdentifier().equals(identifier)) {
                return convertAvroDate(avroValue);
            } else if (SqlTypes.TIME.getIdentifier().equals(identifier)) {
                return convertAvroTime(avroValue);
            } else if (SqlTypes.DATETIME.getIdentifier().equals(identifier)) {
                return convertAvroDateTime(avroValue);
            } else if (SQL_DATE_TIME_TYPES.contains(identifier)) {
                switch(options.getTruncateTimestamps()) {
                    case TRUNCATE:
                        return truncateToMillis(avroValue);
                    case REJECT:
                        return safeToMillis(avroValue);
                    default:
                        throw new IllegalArgumentException(String.format("Unknown timestamp truncation option: %s", options.getTruncateTimestamps()));
                }
            } else if (SQL_STRING_TYPES.contains(identifier)) {
                return convertAvroPrimitiveTypes(TypeName.STRING, avroValue);
            } else {
                throw new RuntimeException("Unknown logical type " + identifier);
            }
        case ROW:
            Schema rowSchema = beamFieldType.getRowSchema();
            if (rowSchema == null) {
                throw new IllegalArgumentException("Nested ROW missing row schema");
            }
            GenericData.Record record = (GenericData.Record) avroValue;
            return toBeamRow(record, rowSchema, options);
        case MAP:
            return convertAvroRecordToMap(beamFieldType, avroValue, options);
        default:
            throw new RuntimeException("Does not support converting unknown type value: " + beamFieldTypeName);
    }
}
Also used : TypeName(org.apache.beam.sdk.schemas.Schema.TypeName) TableSchema(com.google.api.services.bigquery.model.TableSchema) TableFieldSchema(com.google.api.services.bigquery.model.TableFieldSchema) Schema(org.apache.beam.sdk.schemas.Schema) GenericRecord(org.apache.avro.generic.GenericRecord) GenericData(org.apache.avro.generic.GenericData)

Example 2 with TypeName

use of org.apache.beam.sdk.schemas.Schema.TypeName in project beam by apache.

the class Cast method castValue.

@SuppressWarnings("unchecked")
public static Object castValue(Object inputValue, FieldType input, FieldType output) {
    TypeName inputType = input.getTypeName();
    TypeName outputType = output.getTypeName();
    if (inputValue == null) {
        return null;
    }
    switch(inputType) {
        case ROW:
            return castRow((Row) inputValue, input.getRowSchema(), output.getRowSchema());
        case ARRAY:
        case ITERABLE:
            ;
            Iterable<Object> inputValues = (Iterable<Object>) inputValue;
            List<Object> outputValues = new ArrayList<>(Iterables.size(inputValues));
            for (Object elem : inputValues) {
                outputValues.add(castValue(elem, input.getCollectionElementType(), output.getCollectionElementType()));
            }
            return outputValues;
        case MAP:
            Map<Object, Object> inputMap = (Map<Object, Object>) inputValue;
            Map<Object, Object> outputMap = Maps.newHashMapWithExpectedSize(inputMap.size());
            for (Map.Entry<Object, Object> entry : inputMap.entrySet()) {
                Object outputKey = castValue(entry.getKey(), input.getMapKeyType(), output.getMapKeyType());
                Object outputValue = castValue(entry.getValue(), input.getMapValueType(), output.getMapValueType());
                outputMap.put(outputKey, outputValue);
            }
            return outputMap;
        default:
            if (inputType.equals(outputType)) {
                return inputValue;
            }
            if (inputType.isNumericType()) {
                return castNumber((Number) inputValue, inputType, outputType);
            } else {
                throw new IllegalArgumentException("input should be array, map, numeric or row");
            }
    }
}
Also used : TypeName(org.apache.beam.sdk.schemas.Schema.TypeName) ArrayList(java.util.ArrayList) Map(java.util.Map)

Example 3 with TypeName

use of org.apache.beam.sdk.schemas.Schema.TypeName in project beam by apache.

the class BigDecimalConverterTest method testReturnsConverterForNumericTypes.

@Test
public void testReturnsConverterForNumericTypes() {
    for (TypeName numericType : TypeName.NUMERIC_TYPES) {
        SerializableFunction<BigDecimal, ? extends Number> converter = BigDecimalConverter.forSqlType(numericType);
        assertNotNull(converter);
        assertNotNull(converter.apply(BigDecimal.TEN));
    }
}
Also used : TypeName(org.apache.beam.sdk.schemas.Schema.TypeName) BigDecimal(java.math.BigDecimal) Test(org.junit.Test)

Aggregations

TypeName (org.apache.beam.sdk.schemas.Schema.TypeName)3 TableFieldSchema (com.google.api.services.bigquery.model.TableFieldSchema)1 TableSchema (com.google.api.services.bigquery.model.TableSchema)1 BigDecimal (java.math.BigDecimal)1 ArrayList (java.util.ArrayList)1 Map (java.util.Map)1 GenericData (org.apache.avro.generic.GenericData)1 GenericRecord (org.apache.avro.generic.GenericRecord)1 Schema (org.apache.beam.sdk.schemas.Schema)1 Test (org.junit.Test)1