io.kestra.plugin.serdes.avro.AvroDeserializer Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of plugin-serdes Show documentation
Serialize and deserialize data formats in Kestra workflows.
The newest version!
package io.kestra.plugin.serdes.avro;

import org.apache.avro.Conversions.DecimalConversion;
import org.apache.avro.Conversions.UUIDConversion;
import org.apache.avro.LogicalType;
import org.apache.avro.Schema;
import org.apache.avro.Schema.Type;
import org.apache.avro.data.TimeConversions.*;
import org.apache.avro.generic.GenericFixed;
import org.apache.avro.generic.GenericRecord;

import java.math.BigDecimal;
import java.nio.ByteBuffer;
import java.time.Instant;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.LocalTime;
import java.util.*;
import java.util.stream.Collectors;

public class AvroDeserializer {
    private static final String DECIMAL = "decimal";
    private static final String UUID = "uuid";
    private static final String DATE = "date";
    private static final String TIME_MILLIS = "time-millis";
    private static final String TIME_MICROS = "time-micros";
    private static final String TIMESTAMP_MILLIS = "timestamp-millis";
    private static final String TIMESTAMP_MICROS = "timestamp-micros";
    private static final String LOCAL_TIMESTAMP_MILLIS = "local-timestamp-millis";
    private static final String LOCAL_TIMESTAMP_MICROS = "local-timestamp-micros";

    private static final DecimalConversion DECIMAL_CONVERSION = new DecimalConversion();
    private static final UUIDConversion UUID_CONVERSION = new UUIDConversion();
    private static final DateConversion DATE_CONVERSION = new DateConversion();
    private static final TimeMicrosConversion TIME_MICROS_CONVERSION = new TimeMicrosConversion();
    private static final TimeMillisConversion TIME_MILLIS_CONVERSION = new TimeMillisConversion();
    private static final TimestampMicrosConversion TIMESTAMP_MICROS_CONVERSION = new TimestampMicrosConversion();
    private static final TimestampMillisConversion TIMESTAMP_MILLIS_CONVERSION = new TimestampMillisConversion();
    private static final LocalTimestampMicrosConversion LOCAL_TIMESTAMP_MICROS_CONVERSION = new LocalTimestampMicrosConversion();
    private static final LocalTimestampMillisConversion LOCAL_TIMESTAMP_MILLIS_CONVERSION = new LocalTimestampMillisConversion();

    public static Map recordDeserializer(GenericRecord record) {
        return record
            .getSchema()
            .getFields()
            .stream()
            .collect(
                LinkedHashMap::new, // preserve schema field order
                (m, v) -> m.put(
                    v.name(),
                    AvroDeserializer.objectDeserializer(record.get(v.name()), v.schema())
                ),
                HashMap::putAll
            );
    }

    @SuppressWarnings("unchecked")
    private static Object objectDeserializer(Object value, Schema schema) {
        LogicalType logicalType = schema.getLogicalType();
        Type primitiveType = schema.getType();
        if (logicalType != null) {
            switch (logicalType.getName()) {
                case DATE:
                    return AvroDeserializer.dateDeserializer(value, schema, primitiveType, logicalType);
                case DECIMAL:
                    return AvroDeserializer.decimalDeserializer(value, schema, primitiveType, logicalType);
                case TIME_MICROS:
                    return AvroDeserializer.timeMicrosDeserializer(value, schema, primitiveType, logicalType);
                case TIME_MILLIS:
                    return AvroDeserializer.timeMillisDeserializer(value, schema, primitiveType, logicalType);
                case TIMESTAMP_MICROS:
                    return AvroDeserializer.timestampMicrosDeserializer(value, schema, primitiveType, logicalType);
                case TIMESTAMP_MILLIS:
                    return AvroDeserializer.timestampMillisDeserializer(value, schema, primitiveType, logicalType);
                case LOCAL_TIMESTAMP_MICROS:
                    return AvroDeserializer.localTimestampMicrosDeserializer(value, schema, primitiveType, logicalType);
                case LOCAL_TIMESTAMP_MILLIS:
                    return AvroDeserializer.localTimestampMillisDeserializer(value, schema, primitiveType, logicalType);
                case UUID:
                    return AvroDeserializer.uuidDeserializer(value, schema, primitiveType, logicalType);
                default:
                    throw new IllegalStateException("Unexpected value: " + logicalType);
            }
        } else {
            switch (primitiveType) {
                case UNION:
                    return AvroDeserializer.unionDeserializer(value, schema);
                case MAP:
                    return AvroDeserializer.mapDeserializer((Map) value, schema);
                case RECORD:
                    return AvroDeserializer.recordDeserializer((GenericRecord) value);
                case ENUM:
                    return value.toString();
                case ARRAY:
                    return arrayDeserializer((Collection) value, schema);
                case FIXED:
                    return ((GenericFixed) value).bytes();
                case STRING:
                    return ((CharSequence) value).toString();
                case BYTES:
                    return ((ByteBuffer) value).array();
                case INT:
                case LONG:
                case FLOAT:
                case DOUBLE:
                case BOOLEAN:
                case NULL:
                    return value;
                default:
                    throw new IllegalStateException("Unexpected value: " + primitiveType);
            }
        }
    }

    private static Object unionDeserializer(Object value, Schema schema) {
        // first, if value is null, check if the null type exist to avoid generating a NPE
        if(value == null) {
            if (schema.getTypes().stream().anyMatch(t -> t.getType() == Type.NULL)) {
                return null;
            }
            else {
                throw new NullPointerException("value is null but the schema type is not nullable: " + schema);
            }
        }

        // then, evaluate each type and return the first that didn't generate an exception
        for(var type  : schema.getTypes()) {
            // try to deserialized by each type and return the first that works
            try {
                return AvroDeserializer.objectDeserializer(value, type);
            }
            catch(Exception e) {
                // do nothing : try the next one
            }
        }
        throw new IllegalArgumentException("Unable to deserialize objet " + value + " with schema " + schema);
    }

    private static Map mapDeserializer(Map value, Schema schema) {
        return value
            .entrySet()
            .stream()
            .collect(Collectors.toMap(
                Map.Entry::getKey,
                e -> AvroDeserializer.objectDeserializer(e.getValue(), schema.getValueType()))
            );
    }

    private static Collection arrayDeserializer(Collection value, Schema schema) {
        return value
            .stream()
            .map(e -> AvroDeserializer.objectDeserializer(e, schema.getElementType()))
            .collect(Collectors.toList());
    }

    private static Instant timestampMicrosDeserializer(Object value, Schema schema, Type primitiveType, LogicalType logicalType) {
        if (value instanceof Instant) {
            return (Instant) value;
        } else if (primitiveType == Type.LONG) {
            return AvroDeserializer.TIMESTAMP_MICROS_CONVERSION.fromLong((Long) value, schema, logicalType);
        }

        throw new IllegalStateException("Unexpected value: " + primitiveType);
    }

    private static Instant timestampMillisDeserializer(Object value, Schema schema, Type primitiveType, LogicalType logicalType) {
        if (value instanceof Instant) {
            return (Instant) value;
        } else if (primitiveType == Type.LONG) {
            return AvroDeserializer.TIMESTAMP_MILLIS_CONVERSION.fromLong((Long) value, schema, logicalType);
        }

        throw new IllegalStateException("Unexpected value: " + primitiveType);
    }

    private static LocalDateTime localTimestampMicrosDeserializer(Object value, Schema schema, Type primitiveType, LogicalType logicalType) {
        if (value instanceof LocalDateTime) {
            return (LocalDateTime) value;
        } else if (primitiveType == Type.LONG) {
            return AvroDeserializer.LOCAL_TIMESTAMP_MICROS_CONVERSION.fromLong((Long) value, schema, logicalType);
        }

        throw new IllegalStateException("Unexpected value: " + primitiveType);
    }

    private static LocalDateTime localTimestampMillisDeserializer(Object value, Schema schema, Type primitiveType, LogicalType logicalType) {
        if (value instanceof LocalDateTime) {
            return (LocalDateTime) value;
        } else if (primitiveType == Type.LONG) {
            return AvroDeserializer.LOCAL_TIMESTAMP_MILLIS_CONVERSION.fromLong((Long) value, schema, logicalType);
        }

        throw new IllegalStateException("Unexpected value: " + primitiveType);
    }

    private static LocalTime timeMicrosDeserializer(Object value, Schema schema, Type primitiveType, LogicalType logicalType) {
        if (value instanceof LocalTime) {
            return (LocalTime) value;
        } else if (primitiveType == Type.LONG) {
            return AvroDeserializer.TIME_MICROS_CONVERSION.fromLong((Long) value, schema, logicalType);
        }

        throw new IllegalStateException("Unexpected value: " + primitiveType);
    }

    private static LocalTime timeMillisDeserializer(Object value, Schema schema, Type primitiveType, LogicalType logicalType) {
        if (value instanceof LocalTime) {
            return (LocalTime) value;
        } else if (primitiveType == Type.INT) {
            return AvroDeserializer.TIME_MILLIS_CONVERSION.fromInt((Integer) value, schema, logicalType);
        }

        throw new IllegalStateException("Unexpected value: " + primitiveType);
    }

    private static LocalDate dateDeserializer(Object value, Schema schema, Type primitiveType, LogicalType logicalType) {
        if (value instanceof LocalDate) {
            return (LocalDate) value;
        } else if (primitiveType == Type.INT) {
            return AvroDeserializer.DATE_CONVERSION.fromInt((Integer) value, schema, logicalType);
        }

        throw new IllegalStateException("Unexpected value: " + primitiveType);
    }

    private static UUID uuidDeserializer(Object value, Schema schema, Type primitiveType, LogicalType logicalType) {
        if (value instanceof UUID) {
            return (UUID) value;
        } else if (primitiveType == Type.STRING) {
            return AvroDeserializer.UUID_CONVERSION.fromCharSequence((CharSequence) value, schema, logicalType);
        }

        throw new IllegalStateException("Unexpected value: " + primitiveType);
    }

    private static BigDecimal decimalDeserializer(Object value, Schema schema, Type primitiveType, LogicalType logicalType) {
        if (value instanceof BigDecimal) {
            return (BigDecimal) value;
        } else if (primitiveType == Type.BYTES) {
            return AvroDeserializer.DECIMAL_CONVERSION.fromBytes((ByteBuffer) value, schema, logicalType);
        } else if (primitiveType == Type.FIXED) {
            return AvroDeserializer.DECIMAL_CONVERSION.fromFixed((GenericFixed) value, schema, logicalType);
        }

        throw new IllegalStateException("Unexpected value: " + primitiveType);
    }
}