HIVE/AVRO: Handle all union options coersion to single type

trinodb · Jan 24, 2024 · 43c20d6 · 43c20d6
1 parent d92ad42
commit 43c20d6
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 3 deletions.
diff --git a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/avro/AvroPageDataReader.java b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/avro/AvroPageDataReader.java
@@ -193,11 +193,12 @@ private static BlockBuildingDecoder createBlockBuildingDecoderForAction(Resolver
             case RECORD -> new RowBlockBuildingDecoder(action, typeManager);
             case ENUM -> new EnumBlockBuildingDecoder((Resolver.EnumAdjust) action);
             case WRITER_UNION -> {
-                if (isSimpleNullableUnion(action.reader)) {
-                    yield new WriterUnionBlockBuildingDecoder((Resolver.WriterUnion) action, typeManager);
+                if (action.reader.getType() == Schema.Type.UNION && !isSimpleNullableUnion(action.reader)) {
+                    yield new WriterUnionCoercedIntoRowBlockBuildingDecoder((Resolver.WriterUnion) action, typeManager);
                 }
                 else {
-                    yield new WriterUnionCoercedIntoRowBlockBuildingDecoder((Resolver.WriterUnion) action, typeManager);
+                    // reading a union with non-union or nullable union, optimistically try to create the reader, will fail at read time with any underlying issues
+                    yield new WriterUnionBlockBuildingDecoder((Resolver.WriterUnion) action, typeManager);
                 }
             }
             case READER_UNION -> {

diff --git a/...ts/src/test/java/io/trino/hive/formats/avro/TestAvroPageDataReaderWithoutTypeManager.java b/...ts/src/test/java/io/trino/hive/formats/avro/TestAvroPageDataReaderWithoutTypeManager.java
@@ -370,4 +370,37 @@ public void testRead3UnionWith2UnionDataWith2Union()
             }
         }
     }
+
+    @Test
+    public void testReadUnionWithNonUnionAllCoercions()
+            throws IOException, AvroTypeException
+    {
+        Schema nonUnion = Schema.create(Schema.Type.STRING);
+        Schema union = Schema.createUnion(Schema.create(Schema.Type.STRING), Schema.create(Schema.Type.BYTES));
+
+        Schema nonUnionRecord = SchemaBuilder.builder()
+                .record("aRecord")
+                .fields()
+                .name("aField")
+                .type(nonUnion)
+                .noDefault()
+                .endRecord();
+
+        Schema unionRecord = SchemaBuilder.builder()
+                .record("aRecord")
+                .fields()
+                .name("aField")
+                .type(union)
+                .noDefault()
+                .endRecord();
+
+        TrinoInputFile inputFile = createWrittenFileWithSchema(1000, unionRecord);
+
+        //read the file with the non-union schema and ensure that no error thrown
+        try (AvroFileReader avroFileReader = new AvroFileReader(inputFile, nonUnionRecord, NoOpAvroTypeManager.INSTANCE)) {
+            while (avroFileReader.hasNext()) {
+                assertThat(avroFileReader.next()).isNotNull();
+            }
+        }
+    }
 }