Support reading Avro complex union types #65

rzhang10 · 2021-05-14T20:47:19Z

Could there be a corner case where the schema is a union that only contains null, like this:

"type": [ "null" ]

If so, it seems this will make the related code in SchemaToType.java return an empty struct when converting an union.

good catch @rzhang10 ! While it is meaningless to have a single-type in a union, the avro Schema.UnionSchema constructor does not perform any checks on the number of types, so we definitely need to guard against such cases (throw if size < 2 ?).

Yeah, it can happen. For example, [null] and [int] will pass the check. However, those are valid Avro unions and Avro does not prevent users from creating those schemas. In addition, our reader will works with those single unions. Maybe it is better to support those than guard it? We can let the user to decide whether it is useful or not.

Should those be converted to structs or simple types?

I'd prefer to convert it to struct for consistency as it is valid union in avro.

shenodaguirguis · 2021-06-03T17:09:23Z

nit: we now maintain index outside the loop, and check if union is option (line 88) every time in the loop. May be use one loop for both branches and fields, and re-org to hoist up the union check (if that is feasible)?

How about re-assigning the type before the loop?

if (AvroSchemaUtil.isOptionSchema(union)) branches = type; else branches = type.asStructType();

rzhang10 · 2021-05-14T20:51:14Z

In this else branch, if it's the corner case (one null union) I mentioned above, it will construct an empty struct, which will prob cause an exception later.

Yes, it will create an empty struct. However, it is a valid Avro schema and our reader can read it without issue. Should we support it instead of guard against it?

thanks for validating this Wenye. I am fine with not guarding against it, given that avro schema also does not guard against this corner case.

shenodaguirguis · 2021-05-04T22:04:48Z

not sure I follow the reasoning why struct is the appropriate way to represent non-optional union. Should we introduce a new Types.unionType (that extends struct may be)? Because using struct per se (with optional fields) is somewhat confusing

Discussed offline with @shenodaguirguis . The reasoning is documented in design doc. The non-optional union can be represented with struct type as the following example
[ “null”, “int”, “string” ] => optional struct ( 1 tag_0: optional int, 2 tag_1: optional string)

-Original file line number
+Diff line change
@@ Expand Up / @@ -130,6 +130,27 @@ public static boolean isOptionSchema(Schema schema) { @@
         return false;
       }
+      /**
+       * This method decides whether a schema is of type union and is complex union and is optional
+       *
+       * Complex union: the number of options in union not equals to 2
+       * Optional: null is present in union
+       *
+       * @param schema input schema
+       * @return true if schema is complex union and it is optional
+       */
+      public static boolean isOptionalComplexUnion(Schema schema) {
+        if (schema.getType() == UNION && schema.getTypes().size() != 2) {
+          for (Schema type : schema.getTypes()) {
+            if (type.getType() == Schema.Type.NULL) {
+              return true;
+            }
+          }
+        }
+        return false;
+      }
       public static Schema toOption(Schema schema) {
         if (schema.getType() == UNION) {
           Preconditions.checkArgument(isOptionSchema(schema),
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up @@
       private static <T> T visitUnion(Type type, Schema union, AvroSchemaWithTypeVisitor<T> visitor) {
         List<Schema> types = union.getTypes();
         List<T> options = Lists.newArrayListWithExpectedSize(types.size());
+        int index = 0;
         for (Schema branch : types) {
           if (branch.getType() == Schema.Type.NULL) {
             options.add(visit((Type) null, branch, visitor));
           } else {
-            options.add(visit(type, branch, visitor));
+            if (AvroSchemaUtil.isOptionSchema(union)) {
+              options.add(visit(type, branch, visitor));
+            } else {
+              options.add(visit(type.asStructType().fields().get(index).type(), branch, visitor));
+            }
+            index++;
           }
         }
         return visitor.union(type, union, options);
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up @@
       @Override
       public Schema union(Schema union, Iterable<Schema> options) {
-        Preconditions.checkState(AvroSchemaUtil.isOptionSchema(union),
-            "Invalid schema: non-option unions are not supported: %s", union);
-        Schema nonNullOriginal = AvroSchemaUtil.fromOption(union);
-        Schema nonNullResult = AvroSchemaUtil.fromOptions(Lists.newArrayList(options));
+        if (AvroSchemaUtil.isOptionSchema(union)) {
+          Schema nonNullOriginal = AvroSchemaUtil.fromOption(union);
+          Schema nonNullResult = AvroSchemaUtil.fromOptions(Lists.newArrayList(options));
-        if (nonNullOriginal != nonNullResult) {
-          return AvroSchemaUtil.toOption(nonNullResult);
+          if (nonNullOriginal != nonNullResult) {
+            return AvroSchemaUtil.toOption(nonNullResult);
+          }
         }
         return union;
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up @@
       @Override
       public Schema union(Schema union, List<Schema> options) {
-        Preconditions.checkState(AvroSchemaUtil.isOptionSchema(union),
-            "Invalid schema: non-option unions are not supported: %s", union);
-        // only unions with null are allowed, and a null schema results in null
-        Schema pruned = null;
-        if (options.get(0) != null) {
-          pruned = options.get(0);
-        } else if (options.get(1) != null) {
-          pruned = options.get(1);
-        }
+        if (AvroSchemaUtil.isOptionSchema(union)) {
+          // case option union
+          Schema pruned = null;
+          if (options.get(0) != null) {
+            pruned = options.get(0);
+          } else if (options.get(1) != null) {
+            pruned = options.get(1);
+          }
-        if (pruned != null) {
-          if (pruned != AvroSchemaUtil.fromOption(union)) {
-            return AvroSchemaUtil.toOption(pruned);
+          if (pruned != null) {
+            if (pruned != AvroSchemaUtil.fromOption(union)) {
+              return AvroSchemaUtil.toOption(pruned);
+            }
+            return union;
           }
+          return null;
+        } else {
+          // Complex union case
           return union;
         }
-        return null;
       }
       @Override
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up @@
           Object defaultValue = field.hasDefaultValue() && !(field.defaultVal() instanceof JsonProperties.Null) ?
               field.defaultVal() : null;
-          if (AvroSchemaUtil.isOptionSchema(field.schema())) {
+          if (AvroSchemaUtil.isOptionSchema(field.schema()) || AvroSchemaUtil.isOptionalComplexUnion(field.schema())) {
             newFields.add(Types.NestedField.optional(fieldId, field.name(), fieldType, defaultValue, null));
           } else if (defaultValue != null) {
             newFields.add(Types.NestedField.required(fieldId, field.name(), fieldType, defaultValue, null));
@@ Expand All @@
       @Override
       public Type union(Schema union, List<Type> options) {
-        Preconditions.checkArgument(AvroSchemaUtil.isOptionSchema(union),
-            "Unsupported type: non-option union: %s", union);
-        // records, arrays, and maps will check nullability later
-        if (options.get(0) == null) {
-          return options.get(1);
+        if (AvroSchemaUtil.isOptionSchema(union)) {
+          // Optional simple union
+          // records, arrays, and maps will check nullability later
+          if (options.get(0) == null) {
+            return options.get(1);
+          } else {
+            return options.get(0);
+          }
         } else {
-          return options.get(0);
+          // Complex union
+          List<Types.NestedField> newFields = Lists.newArrayListWithExpectedSize(options.size());
+          int tagIndex = 0;
+          for (Type type : options) {
+            if (type != null) {
+              newFields.add(Types.NestedField.optional(allocateId(), "tag_" + tagIndex++, type));
+            }
+          }
+          return Types.StructType.of(newFields);
         }
       }
@@ Expand Down @@

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support reading Avro complex union types #65

Uh oh!

Diff view

Diff view

There are no files selected for viewing

rzhang10 May 14, 2021

Uh oh!

shenodaguirguis Jun 3, 2021

Uh oh!

funcheetah Jun 3, 2021

Uh oh!

wmoustafa Jun 3, 2021

Uh oh!

funcheetah Jun 4, 2021

Uh oh!

shenodaguirguis Jun 3, 2021

Uh oh!

wmoustafa Jun 4, 2021

Uh oh!

rzhang10 May 14, 2021

Uh oh!

shenodaguirguis Jun 3, 2021

Uh oh!

funcheetah Jun 3, 2021

Uh oh!

shenodaguirguis Jun 3, 2021

Uh oh!

shenodaguirguis May 4, 2021

Uh oh!

funcheetah May 18, 2021

Uh oh!

Uh oh!

-Original file line number
+Diff line change
@@ -0,0 +1,129 @@
+    /*
+     * Licensed to the Apache Software Foundation (ASF) under one
+     * or more contributor license agreements.  See the NOTICE file
+     * distributed with this work for additional information
+     * regarding copyright ownership.  The ASF licenses this file
+     * to you under the Apache License, Version 2.0 (the
+     * "License"); you may not use this file except in compliance
+     * with the License.  You may obtain a copy of the License at
+     *
+     *   http://www.apache.org/licenses/LICENSE-2.0
+     *
+     * Unless required by applicable law or agreed to in writing,
+     * software distributed under the License is distributed on an
+     * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+     * KIND, either express or implied.  See the License for the
+     * specific language governing permissions and limitations
+     * under the License.
+     */
+    package org.apache.iceberg.avro;
+    import org.apache.avro.Schema;
+    import org.apache.avro.SchemaBuilder;
+    import org.junit.Assert;
+    import org.junit.Test;
+    public class TestAvroComplexUnion {
+      @Test
+      public void testRequiredComplexUnion() {
+        Schema avroSchema = SchemaBuilder.record("root")
+            .fields()
+            .name("unionCol")
+            .type()
+            .unionOf()
+            .intType()
+            .and()
+            .stringType()
+            .endUnion()
+            .noDefault()
+            .endRecord();
+        org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema);
+        String expectedIcebergSchema = "table {\n" +
+            "  0: unionCol: required struct<1: tag_0: optional int, 2: tag_1: optional string>\n" + "}";
+        Assert.assertEquals(expectedIcebergSchema, icebergSchema.toString());
+      }
+      @Test
+      public void testOptionalComplexUnion() {
+        Schema avroSchema = SchemaBuilder.record("root")
+            .fields()
+            .name("unionCol")
+            .type()
+            .unionOf()
+            .nullType()
+            .and()
+            .intType()
+            .and()
+            .stringType()
+            .endUnion()
+            .noDefault()
+            .endRecord();
+        org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema);
+        String expectedIcebergSchema =
+            "table {\n" + "  0: unionCol: optional struct<1: tag_0: optional int, 2: tag_1: optional string>\n" + "}";
+        Assert.assertEquals(expectedIcebergSchema, icebergSchema.toString());
+      }
+      @Test
+      public void testSingleComponentUnion() {
+        Schema avroSchema = SchemaBuilder.record("root")
+            .fields()
+            .name("unionCol")
+            .type()
+            .unionOf()
+            .intType()
+            .endUnion()
+            .noDefault()
+            .endRecord();
+        org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema);
+        String expectedIcebergSchema = "table {\n" + "  0: unionCol: required struct<1: tag_0: optional int>\n" + "}";
+        Assert.assertEquals(expectedIcebergSchema, icebergSchema.toString());
+      }
+      @Test
+      public void testOptionSchema() {
+        Schema avroSchema = SchemaBuilder.record("root")
+            .fields()
+            .name("optionCol")
+            .type()
+            .unionOf()
+            .nullType()
+            .and()
+            .intType()
+            .endUnion()
+            .nullDefault()
+            .endRecord();
+        org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema);
+        String expectedIcebergSchema = "table {\n" + "  0: optionCol: optional int\n" + "}";
+        Assert.assertEquals(expectedIcebergSchema, icebergSchema.toString());
+      }
+      @Test
+      public void testNullUnionSchema() {
+        Schema avroSchema = SchemaBuilder.record("root")
+            .fields()
+            .name("nullUnionCol")
+            .type()
+            .unionOf()
+            .nullType()
+            .endUnion()
+            .noDefault()
+            .endRecord();
+        org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema);
+        String expectedIcebergSchema = "table {\n" + "  0: nullUnionCol: optional struct<>\n" + "}";
+        Assert.assertEquals(expectedIcebergSchema, icebergSchema.toString());
+      }
+    }

-Original file line number
+Diff line change
@@ Expand Up / @@ -27,6 +27,7 @@ @@
     import org.apache.avro.Schema;
     import org.apache.avro.io.DatumReader;
     import org.apache.avro.io.Decoder;
+    import org.apache.iceberg.avro.AvroSchemaUtil;
     import org.apache.iceberg.avro.AvroSchemaWithTypeVisitor;
     import org.apache.iceberg.avro.ValueReader;
     import org.apache.iceberg.avro.ValueReaders;
@@ Expand Down Expand Up @@
         @Override
         public ValueReader<?> union(Type expected, Schema union, List<ValueReader<?>> options) {
-          return ValueReaders.union(options);
+          if (AvroSchemaUtil.isOptionSchema(union)) {
+            return ValueReaders.union(options);
+          } else {
+            return SparkValueReaders.union(options);
+          }
         }
         @Override
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up @@
         return new StructReader(readers, struct, idToConstant);
       }
+      static ValueReader<InternalRow> union(List<ValueReader<?>> readers) {
+        return new UnionReader(readers);
+      }
       private static class StringReader implements ValueReader<UTF8String> {
         private static final StringReader INSTANCE = new StringReader();
@@ Expand Down Expand Up @@
           }
         }
       }
+      static class UnionReader implements ValueReader<InternalRow> {
+        private final ValueReader[] readers;
+        private UnionReader(List<ValueReader<?>> readers) {
+          this.readers = new ValueReader[readers.size()];
+          for (int i = 0; i < this.readers.length; i += 1) {
+            this.readers[i] = readers.get(i);
+          }
+        }
+        @Override
+        public InternalRow read(Decoder decoder, Object reuse) throws IOException {
+          InternalRow struct = new GenericInternalRow(readers.length);
+          int index = decoder.readIndex();
+          Object value = this.readers[index].read(decoder, reuse);
+          for (int i = 0; i < readers.length; i += 1) {
+            struct.setNullAt(i);
+          }
+          struct.update(index, value);
+          return struct;
+        }
+      }
     }

Support reading Avro complex union types #65

Uh oh!

Support reading Avro complex union types #65

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!