Avro union support

docs/development/extensions-core/avro.md

@@ -39,13 +39,30 @@ Make sure to [include](../../development/extensions.md#loading-extensions) `drui
 
 Druid supports most Avro types natively, there are however some exceptions which are detailed here.
 
-`union` types which aren't of the form `[null, otherType]` aren't supported at this time.
+#### Unions
+Druid has two modes for supporting `union` types.
 
+The default mode will treat unions as a single value regardless of the type it is populated with.
+
+If you wish to operate on each different member of a union however you can set `extractUnionsByType` on the Avro parser in which case unions will be expanded into nested objects according to the following rules:
+* Primitive types and unnamed complex types are keyed their type name. i.e `int`, `string`
+* Complex named types are keyed by their names, this includes `record`, `fixed` and `enum`.
+* The Avro null type is elided as its value can only ever be null
+
+This is safe because an Avro union can only contain a single member of each unnamed type and duplicates of the same named type are not allowed.
+i.e only a single array is allowed, multiple records (or other named types) are allowed as long as each has a unique name.
+
+The members can then be accessed using a [flattenSpec](../../ingestion/data-formats.md#flattenspec) similar other nested types.
+
+#### Binary types
 `bytes` and `fixed` Avro types will be returned by default as base64 encoded strings unless the `binaryAsString` option is enabled on the Avro parser.
 This setting will decode these types as UTF-8 strings.
 
+#### Enums
 `enum` types will be returned as `string` of the enum symbol.
 
+#### Complex types
 `record` and `map` types representing nested data can be ingested using [flattenSpec](../../ingestion/data-formats.md#flattenspec) on the parser.
 
-Druid doesn't currently support Avro logical types, they will be ignored and fields will be handled according to the underlying primitive type.
+#### Logical types
+Druid doesn't currently support Avro logical types, they will be ignored and fields will be handled according to the underlying primitive type.

extensions-core/avro-extensions/src/main/java/org/apache/druid/data/input/AvroHadoopInputRowParser.java

@@ -34,18 +34,24 @@ public class AvroHadoopInputRowParser implements InputRowParser<GenericRecord>
 {
   private final ParseSpec parseSpec;
   private final boolean fromPigAvroStorage;
+  private final boolean binaryAsString;
+  private final boolean extractUnionsByType;
   private final ObjectFlattener<GenericRecord> avroFlattener;
   private final MapInputRowParser mapParser;
 
   @JsonCreator
   public AvroHadoopInputRowParser(
       @JsonProperty("parseSpec") ParseSpec parseSpec,
-      @JsonProperty("fromPigAvroStorage") Boolean fromPigAvroStorage
+      @JsonProperty("fromPigAvroStorage") Boolean fromPigAvroStorage,
+      @JsonProperty("binaryAsString") Boolean binaryAsString,
+      @JsonProperty("extractUnionsByType") Boolean extractUnionsByType
   )
   {
     this.parseSpec = parseSpec;
-    this.fromPigAvroStorage = fromPigAvroStorage == null ? false : fromPigAvroStorage;
-    this.avroFlattener = AvroParsers.makeFlattener(parseSpec, this.fromPigAvroStorage, false);
+    this.fromPigAvroStorage = fromPigAvroStorage != null && fromPigAvroStorage;
+    this.binaryAsString = binaryAsString != null && binaryAsString;
+    this.extractUnionsByType = extractUnionsByType != null && extractUnionsByType;
+    this.avroFlattener = AvroParsers.makeFlattener(parseSpec, this.fromPigAvroStorage, this.binaryAsString, this.extractUnionsByType);
     this.mapParser = new MapInputRowParser(parseSpec);
   }
 
@@ -71,6 +77,6 @@ public boolean isFromPigAvroStorage()
   @Override
   public InputRowParser withParseSpec(ParseSpec parseSpec)
   {
-    return new AvroHadoopInputRowParser(parseSpec, fromPigAvroStorage);
+    return new AvroHadoopInputRowParser(parseSpec, fromPigAvroStorage, binaryAsString, extractUnionsByType);
   }
 }

extensions-core/avro-extensions/src/main/java/org/apache/druid/data/input/AvroStreamInputRowParser.java

@@ -36,19 +36,25 @@
 public class AvroStreamInputRowParser implements ByteBufferInputRowParser
 {
   private final ParseSpec parseSpec;
+  private final Boolean binaryAsString;
+  private final Boolean extractUnionsByType;
   private final AvroBytesDecoder avroBytesDecoder;
   private final ObjectFlattener<GenericRecord> avroFlattener;
   private final MapInputRowParser mapParser;
 
   @JsonCreator
   public AvroStreamInputRowParser(
       @JsonProperty("parseSpec") ParseSpec parseSpec,
-      @JsonProperty("avroBytesDecoder") AvroBytesDecoder avroBytesDecoder
+      @JsonProperty("avroBytesDecoder") AvroBytesDecoder avroBytesDecoder,
+      @JsonProperty("binaryAsString") Boolean binaryAsString,
+      @JsonProperty("extractUnionsByType") Boolean extractUnionsByType
   )
   {
     this.parseSpec = Preconditions.checkNotNull(parseSpec, "parseSpec");
     this.avroBytesDecoder = Preconditions.checkNotNull(avroBytesDecoder, "avroBytesDecoder");
-    this.avroFlattener = AvroParsers.makeFlattener(parseSpec, false, false);
+    this.binaryAsString = binaryAsString != null && binaryAsString;
+    this.extractUnionsByType = extractUnionsByType != null && extractUnionsByType;
+    this.avroFlattener = AvroParsers.makeFlattener(parseSpec, false, this.binaryAsString, this.extractUnionsByType);
     this.mapParser = new MapInputRowParser(parseSpec);
   }
 
@@ -76,7 +82,9 @@ public ByteBufferInputRowParser withParseSpec(ParseSpec parseSpec)
   {
     return new AvroStreamInputRowParser(
         parseSpec,
-        avroBytesDecoder
+        avroBytesDecoder,
+        binaryAsString,
+        extractUnionsByType
     );
   }
 

extensions-core/avro-extensions/src/main/java/org/apache/druid/data/input/avro/AvroFlattenerMaker.java

@@ -44,13 +44,8 @@
 
 public class AvroFlattenerMaker implements ObjectFlatteners.FlattenerMaker<GenericRecord>
 {
-  private static final JsonProvider AVRO_JSON_PROVIDER = new GenericAvroJsonProvider();
-  private static final Configuration JSONPATH_CONFIGURATION =
-      Configuration.builder()
-                   .jsonProvider(AVRO_JSON_PROVIDER)
-                   .mappingProvider(new NotImplementedMappingProvider())
-                   .options(EnumSet.of(Option.SUPPRESS_EXCEPTIONS))
-                   .build();
+  private final JsonProvider avroJsonProvider;
+  private final Configuration jsonPathConfiguration;
 
   private static final EnumSet<Schema.Type> ROOT_TYPES = EnumSet.of(
       Schema.Type.STRING,
@@ -99,10 +94,18 @@ private static boolean isFieldPrimitive(Schema.Field field)
    * @param fromPigAvroStorage boolean to specify the data file is stored using AvroStorage
    * @param binaryAsString boolean to encode the byte[] as a string.
    */
-  public AvroFlattenerMaker(final boolean fromPigAvroStorage, final boolean binaryAsString)
+  public AvroFlattenerMaker(final boolean fromPigAvroStorage, final boolean binaryAsString, final boolean extractUnionsByType)
   {
     this.fromPigAvroStorage = fromPigAvroStorage;
     this.binaryAsString = binaryAsString;
+
+    this.avroJsonProvider = new GenericAvroJsonProvider(extractUnionsByType);
+    this.jsonPathConfiguration =
+        Configuration.builder()
+                     .jsonProvider(avroJsonProvider)
+                     .mappingProvider(new NotImplementedMappingProvider())
+                     .options(EnumSet.of(Option.SUPPRESS_EXCEPTIONS))
+                     .build();
   }
 
   @Override
@@ -126,7 +129,7 @@ public Object getRootField(final GenericRecord record, final String key)
   public Function<GenericRecord, Object> makeJsonPathExtractor(final String expr)
   {
     final JsonPath jsonPath = JsonPath.compile(expr);
-    return record -> transformValue(jsonPath.read(record, JSONPATH_CONFIGURATION));
+    return record -> transformValue(jsonPath.read(record, jsonPathConfiguration));
   }
 
   @Override
@@ -138,7 +141,7 @@ public Function<GenericRecord, Object> makeJsonQueryExtractor(final String expr)
   @Override
   public JsonProvider getJsonProvider()
   {
-    return AVRO_JSON_PROVIDER;
+    return avroJsonProvider;
   }
 
   @Override

extensions-core/avro-extensions/src/main/java/org/apache/druid/data/input/avro/AvroOCFInputFormat.java

@@ -42,6 +42,7 @@ public class AvroOCFInputFormat extends NestedInputFormat
   private static final Logger LOGGER = new Logger(AvroOCFInputFormat.class);
 
   private final boolean binaryAsString;
+  private final boolean extractUnionsByType;
   @Nullable
   private final Schema readerSchema;
 
@@ -50,7 +51,8 @@ public AvroOCFInputFormat(
       @JacksonInject @Json ObjectMapper mapper,
       @JsonProperty("flattenSpec") @Nullable JSONPathSpec flattenSpec,
       @JsonProperty("schema") @Nullable Map<String, Object> schema,
-      @JsonProperty("binaryAsString") @Nullable Boolean binaryAsString
+      @JsonProperty("binaryAsString") @Nullable Boolean binaryAsString,
+      @JsonProperty("extractUnionsByType") @Nullable Boolean extractUnionsByType
   ) throws Exception
   {
     super(flattenSpec);
@@ -62,7 +64,8 @@ public AvroOCFInputFormat(
     } else {
       this.readerSchema = null;
     }
-    this.binaryAsString = binaryAsString == null ? false : binaryAsString;
+    this.binaryAsString = binaryAsString != null && binaryAsString;
+    this.extractUnionsByType = extractUnionsByType != null && extractUnionsByType;
   }
 
   @Override
@@ -82,7 +85,8 @@ public InputEntityReader createReader(InputRowSchema inputRowSchema, InputEntity
         temporaryDirectory,
         readerSchema,
         getFlattenSpec(),
-        binaryAsString
+        binaryAsString,
+        extractUnionsByType
     );
   }
 

extensions-core/avro-extensions/src/main/java/org/apache/druid/data/input/avro/AvroOCFReader.java

@@ -56,14 +56,15 @@ public class AvroOCFReader extends IntermediateRowParsingReader<GenericRecord>
       File temporaryDirectory,
       @Nullable Schema readerSchema,
       JSONPathSpec flattenSpec,
-      boolean binaryAsString
+      boolean binaryAsString,
+      boolean extractUnionsByType
   )
   {
     this.inputRowSchema = inputRowSchema;
     this.source = source;
     this.temporaryDirectory = temporaryDirectory;
     this.readerSchema = readerSchema;
-    this.recordFlattener = ObjectFlatteners.create(flattenSpec, new AvroFlattenerMaker(false, binaryAsString));
+    this.recordFlattener = ObjectFlatteners.create(flattenSpec, new AvroFlattenerMaker(false, binaryAsString, extractUnionsByType));
   }
 
   @Override

extensions-core/avro-extensions/src/main/java/org/apache/druid/data/input/avro/AvroParsers.java

@@ -39,7 +39,8 @@ private AvroParsers()
   public static ObjectFlattener<GenericRecord> makeFlattener(
       final ParseSpec parseSpec,
       final boolean fromPigAvroStorage,
-      final boolean binaryAsString
+      final boolean binaryAsString,
+      final boolean extractUnionsByType
   )
   {
     final JSONPathSpec flattenSpec;
@@ -49,7 +50,7 @@ public static ObjectFlattener<GenericRecord> makeFlattener(
       flattenSpec = JSONPathSpec.DEFAULT;
     }
 
-    return ObjectFlatteners.create(flattenSpec, new AvroFlattenerMaker(fromPigAvroStorage, binaryAsString));
+    return ObjectFlatteners.create(flattenSpec, new AvroFlattenerMaker(fromPigAvroStorage, binaryAsString, extractUnionsByType));
   }
 
   public static List<InputRow> parseGenericRecord(

extensions-core/avro-extensions/src/main/java/org/apache/druid/data/input/avro/AvroStreamInputFormat.java

@@ -35,19 +35,22 @@
 public class AvroStreamInputFormat extends NestedInputFormat
 {
   private final boolean binaryAsString;
+  private final boolean extractUnionsByType;
 
   private final AvroBytesDecoder avroBytesDecoder;
 
   @JsonCreator
   public AvroStreamInputFormat(
       @JsonProperty("flattenSpec") @Nullable JSONPathSpec flattenSpec,
       @JsonProperty("avroBytesDecoder") AvroBytesDecoder avroBytesDecoder,
-      @JsonProperty("binaryAsString") @Nullable Boolean binaryAsString
+      @JsonProperty("binaryAsString") @Nullable Boolean binaryAsString,
+      @JsonProperty("extractUnionsByType") @Nullable Boolean extractUnionsByType
   )
   {
     super(flattenSpec);
     this.avroBytesDecoder = avroBytesDecoder;
-    this.binaryAsString = binaryAsString == null ? false : binaryAsString;
+    this.binaryAsString = binaryAsString != null && binaryAsString;
+    this.extractUnionsByType = extractUnionsByType != null && extractUnionsByType;
   }
 
   @Override
@@ -76,7 +79,8 @@ public InputEntityReader createReader(InputRowSchema inputRowSchema, InputEntity
         source,
         avroBytesDecoder,
         getFlattenSpec(),
-        binaryAsString
+        binaryAsString,
+        extractUnionsByType
     );
   }
 

extensions-core/avro-extensions/src/main/java/org/apache/druid/data/input/avro/AvroStreamReader.java

@@ -52,13 +52,14 @@ public class AvroStreamReader extends IntermediateRowParsingReader<GenericRecord
       InputEntity source,
       AvroBytesDecoder avroBytesDecoder,
       JSONPathSpec flattenSpec,
-      boolean binaryAsString
+      boolean binaryAsString,
+      boolean extractUnionsByType
   )
   {
     this.inputRowSchema = inputRowSchema;
     this.source = source;
     this.avroBytesDecoder = avroBytesDecoder;
-    this.recordFlattener = ObjectFlatteners.create(flattenSpec, new AvroFlattenerMaker(false, binaryAsString));
+    this.recordFlattener = ObjectFlatteners.create(flattenSpec, new AvroFlattenerMaker(false, binaryAsString, extractUnionsByType));
   }
 
   @Override

extensions-core/avro-extensions/src/main/java/org/apache/druid/data/input/avro/GenericAvroJsonProvider.java

@@ -19,15 +19,19 @@
 
 package org.apache.druid.data.input.avro;
 
+import com.google.common.collect.ImmutableMap;
 import com.jayway.jsonpath.InvalidJsonException;
 import com.jayway.jsonpath.spi.json.JsonProvider;
 import org.apache.avro.Schema;
+import org.apache.avro.generic.GenericEnumSymbol;
+import org.apache.avro.generic.GenericFixed;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.avro.util.Utf8;
 
 import javax.annotation.Nullable;
 
 import java.io.InputStream;
+import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -41,6 +45,13 @@
  */
 public class GenericAvroJsonProvider implements JsonProvider
 {
+  private final boolean extractUnionsByType;
+
+  GenericAvroJsonProvider(final boolean extractUnionsByType)
+  {
+    this.extractUnionsByType = extractUnionsByType;
+  }
+
   @Override
   public Object parse(final String s) throws InvalidJsonException
   {
@@ -148,7 +159,11 @@ public Object getMapValue(final Object o, final String s)
     if (o == null) {
       return null;
     } else if (o instanceof GenericRecord) {
-      return ((GenericRecord) o).get(s);
+      final GenericRecord record = (GenericRecord) o;
+      if (extractUnionsByType && isExtractableUnion(record.getSchema().getField(s))) {
+        return extractUnionTypes(record.get(s));
+      }
+      return record.get(s);
     } else if (o instanceof Map) {
       final Map theMap = (Map) o;
       if (theMap.containsKey(s)) {
@@ -195,4 +210,46 @@ public Object unwrap(final Object o)
   {
     return o;
   }
+
+  private boolean isExtractableUnion(final Schema.Field field)
+  {
+    return field.schema().isUnion() &&
+           field.schema().getTypes().stream().filter(type -> type.getType() != Schema.Type.NULL).count() > 1;
+  }
+
+  private Map<String, Object> extractUnionTypes(final Object o)
+  {
+    // Primitive types and unnamped complex types are keyed their type name.
+    // Complex named types are keyed by their names.
+    // This is safe because an Avro union can only contain a single member of each unnamed type and duplicates
+    // of the same named type are not allowed. i.e only a single array is allowed, multiple records are allowed as
+    // long as each has a unique name.
+    // The Avro null type is elided as it's value can only ever be null
+    if (o instanceof Integer) {
+      return ImmutableMap.of("int", o);
+    } else if (o instanceof Long) {
+      return ImmutableMap.of("long", o);
+    } else if (o instanceof Float) {
+      return ImmutableMap.of("float", o);
+    } else if (o instanceof Double) {
+      return ImmutableMap.of("double", o);
+    } else if (o instanceof Boolean) {
+      return ImmutableMap.of("boolean", o);
+    } else if (o instanceof Utf8) {
+      return ImmutableMap.of("string", o);
+    } else if (o instanceof ByteBuffer) {
+      return ImmutableMap.of("bytes", o);
+    } else if (o instanceof Map) {
+      return ImmutableMap.of("map", o);
+    } else if (o instanceof List) {
+      return ImmutableMap.of("array", o);
+    } else if (o instanceof GenericRecord) {
+      return ImmutableMap.of(((GenericRecord) o).getSchema().getName(), o);
+    } else if (o instanceof GenericFixed) {
+      return ImmutableMap.of(((GenericFixed) o).getSchema().getName(), o);
+    } else if (o instanceof GenericEnumSymbol) {
+      return ImmutableMap.of(((GenericEnumSymbol<?>) o).getSchema().getName(), o);
+    }
+    return ImmutableMap.of();
+  }
 }