Flink: Add Avro value reader, writer implementations (apache#1153)

core/src/main/java/org/apache/iceberg/data/avro/DataReader.java

@@ -59,7 +59,7 @@ public static <D> DataReader<D> create(org.apache.iceberg.Schema expectedSchema,
   private Schema fileSchema = null;
 
   @SuppressWarnings("unchecked")
-  private DataReader(org.apache.iceberg.Schema expectedSchema, Schema readSchema, Map<Integer, ?> idToConstant) {
+  protected DataReader(org.apache.iceberg.Schema expectedSchema, Schema readSchema, Map<Integer, ?> idToConstant) {
     this.readSchema = readSchema;
     this.reader = (ValueReader<T>) AvroSchemaWithTypeVisitor
         .visit(expectedSchema, readSchema, new ReadBuilder(idToConstant));
@@ -102,7 +102,12 @@ private ResolvingDecoder newResolver() {
     }
   }
 
-  private static class ReadBuilder extends AvroSchemaWithTypeVisitor<ValueReader<?>> {
+  protected ValueReader<?> createStructReader(Types.StructType struct,
+                                              List<ValueReader<?>> fields, Map<Integer, ?> idToConstant) {
+    return GenericReaders.struct(struct, fields, idToConstant);
+  }
+
+  private class ReadBuilder extends AvroSchemaWithTypeVisitor<ValueReader<?>> {
     private final Map<Integer, ?> idToConstant;
 
     private ReadBuilder(Map<Integer, ?> idToConstant) {
@@ -112,7 +117,7 @@ private ReadBuilder(Map<Integer, ?> idToConstant) {
     @Override
     public ValueReader<?> record(Types.StructType struct, Schema record,
                                  List<String> names, List<ValueReader<?>> fields) {
-      return GenericReaders.struct(struct, fields, idToConstant);
+      return createStructReader(struct, fields, idToConstant);
     }
 
     @Override

core/src/main/java/org/apache/iceberg/data/avro/DataWriter.java

@@ -40,7 +40,7 @@ public static <D> DataWriter<D> create(Schema schema) {
     return new DataWriter<>(schema);
   }
 
-  private DataWriter(Schema schema) {
+  protected DataWriter(Schema schema) {
     setSchema(schema);
   }
 
@@ -55,13 +55,14 @@ public void write(T datum, Encoder out) throws IOException {
     writer.write(datum, out);
   }
 
-  private static class WriteBuilder extends AvroSchemaVisitor<ValueWriter<?>> {
-    private WriteBuilder() {
-    }
+  protected ValueWriter<?> createStructWriter(List<ValueWriter<?>> fields) {
+    return GenericWriters.struct(fields);
+  }
 
+  private class WriteBuilder extends AvroSchemaVisitor<ValueWriter<?>> {
     @Override
     public ValueWriter<?> record(Schema record, List<String> names, List<ValueWriter<?>> fields) {
-      return GenericWriters.struct(fields);
+      return createStructWriter(fields);
     }
 
     @Override

flink/src/main/java/org/apache/iceberg/flink/data/FlinkAvroReader.java

@@ -0,0 +1,72 @@
+/*
+ * 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.flink.data;
+
+import java.util.List;
+import java.util.Map;
+import org.apache.avro.Schema;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.avro.ValueReader;
+import org.apache.iceberg.avro.ValueReaders;
+import org.apache.iceberg.data.avro.DataReader;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.types.Types;
+
+public class FlinkAvroReader extends DataReader<Row> {
+
+  public FlinkAvroReader(org.apache.iceberg.Schema expectedSchema, Schema readSchema) {
+    super(expectedSchema, readSchema, ImmutableMap.of());
+  }
+
+  @Override
+  protected ValueReader<?> createStructReader(Types.StructType struct,
+                                              List<ValueReader<?>> fields,
+                                              Map<Integer, ?> idToConstant) {
+    return new RowReader(fields, struct, idToConstant);
+  }
+
+  private static class RowReader extends ValueReaders.StructReader<Row> {
+    private final Types.StructType structType;
+
+    private RowReader(List<ValueReader<?>> readers, Types.StructType struct, Map<Integer, ?> idToConstant) {
+      super(readers, struct, idToConstant);
+      this.structType = struct;
+    }
+
+    @Override
+    protected Row reuseOrCreate(Object reuse) {
+      if (reuse instanceof Row) {
+        return (Row) reuse;
+      } else {
+        return new Row(structType.fields().size());
+      }
+    }
+
+    @Override
+    protected Object get(Row row, int pos) {
+      return row.getField(pos);
+    }
+
+    @Override
+    protected void set(Row row, int pos, Object value) {
+      row.setField(pos, value);
+    }
+  }
+}

flink/src/main/java/org/apache/iceberg/flink/data/FlinkAvroWriter.java

@@ -0,0 +1,49 @@
+/*
+ * 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.flink.data;
+
+import java.util.List;
+import org.apache.avro.Schema;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.avro.ValueWriter;
+import org.apache.iceberg.avro.ValueWriters;
+import org.apache.iceberg.data.avro.DataWriter;
+
+public class FlinkAvroWriter extends DataWriter<Row> {
+  public FlinkAvroWriter(Schema schema) {
+    super(schema);
+  }
+
+  @Override
+  protected ValueWriter<?> createStructWriter(List<ValueWriter<?>> fields) {
+    return new RowWriter(fields);
+  }
+
+  private static class RowWriter extends ValueWriters.StructWriter<Row> {
+    private RowWriter(List<ValueWriter<?>> writers) {
+      super(writers);
+    }
+
+    @Override
+    protected Object get(Row struct, int pos) {
+      return struct.getField(pos);
+    }
+  }
+}

flink/src/test/java/org/apache/iceberg/flink/data/RandomData.java

@@ -33,10 +33,40 @@
 import org.apache.iceberg.types.Types;
 import org.apache.iceberg.util.RandomUtil;
 
+import static org.apache.iceberg.types.Types.NestedField.optional;
+import static org.apache.iceberg.types.Types.NestedField.required;
+
 class RandomData {
   private RandomData() {
   }
 
+  static final Schema COMPLEX_SCHEMA = new Schema(
+      required(1, "roots", Types.LongType.get()),
+      optional(3, "lime", Types.ListType.ofRequired(4, Types.DoubleType.get())),
+      required(5, "strict", Types.StructType.of(
+          required(9, "tangerine", Types.StringType.get()),
+          optional(6, "hopeful", Types.StructType.of(
+              required(7, "steel", Types.FloatType.get()),
+              required(8, "lantern", Types.DateType.get())
+          )),
+          optional(10, "vehement", Types.LongType.get())
+      )),
+      optional(11, "metamorphosis", Types.MapType.ofRequired(12, 13,
+          Types.StringType.get(), Types.TimestampType.withZone())),
+      required(14, "winter", Types.ListType.ofOptional(15, Types.StructType.of(
+          optional(16, "beet", Types.DoubleType.get()),
+          required(17, "stamp", Types.FloatType.get()),
+          optional(18, "wheeze", Types.StringType.get())
+      ))),
+      optional(19, "renovate", Types.MapType.ofRequired(20, 21,
+          Types.StringType.get(), Types.StructType.of(
+              optional(22, "jumpy", Types.DoubleType.get()),
+              required(23, "koala", Types.IntegerType.get()),
+              required(24, "couch rope", Types.IntegerType.get())
+          ))),
+      optional(2, "slide", Types.StringType.get())
+  );
+
   private static Iterable<Row> generateData(Schema schema, int numRecords, Supplier<RandomRowGenerator> supplier) {
     return () -> new Iterator<Row>() {
       private final RandomRowGenerator generator = supplier.get();

flink/src/test/java/org/apache/iceberg/flink/data/TestFlinkAvroReaderWriter.java

@@ -0,0 +1,73 @@
+/*
+ * 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.flink.data;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.Iterator;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.Files;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.avro.Avro;
+import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.io.FileAppender;
+import org.junit.Assert;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import static org.apache.iceberg.flink.data.RandomData.COMPLEX_SCHEMA;
+
+public class TestFlinkAvroReaderWriter {
+  private static final int NUM_RECORDS = 20_000;
+
+  @Rule
+  public TemporaryFolder temp = new TemporaryFolder();
+
+  private void testCorrectness(Schema schema, int numRecords, Iterable<Row> iterable) throws IOException {
+    File testFile = temp.newFile();
+    Assert.assertTrue("Delete should succeed", testFile.delete());
+
+    try (FileAppender<Row> writer = Avro.write(Files.localOutput(testFile))
+        .schema(schema)
+        .createWriterFunc(FlinkAvroWriter::new)
+        .build()) {
+      writer.addAll(iterable);
+    }
+
+    try (CloseableIterable<Row> reader = Avro.read(Files.localInput(testFile))
+        .project(schema)
+        .createReaderFunc(FlinkAvroReader::new)
+        .build()) {
+      Iterator<Row> expected = iterable.iterator();
+      Iterator<Row> rows = reader.iterator();
+      for (int i = 0; i < numRecords; i += 1) {
+        Assert.assertTrue("Should have expected number of rows", rows.hasNext());
+        Assert.assertEquals(expected.next(), rows.next());
+      }
+      Assert.assertFalse("Should not have extra rows", rows.hasNext());
+    }
+  }
+
+  @Test
+  public void testNormalData() throws IOException {
+    testCorrectness(COMPLEX_SCHEMA, NUM_RECORDS, RandomData.generate(COMPLEX_SCHEMA, NUM_RECORDS, 19982));
+  }
+}

flink/src/test/java/org/apache/iceberg/flink/data/TestFlinkParquetReaderWriter.java

@@ -28,48 +28,19 @@
 import org.apache.iceberg.io.CloseableIterable;
 import org.apache.iceberg.io.FileAppender;
 import org.apache.iceberg.parquet.Parquet;
-import org.apache.iceberg.types.Types;
 import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
 
-import static org.apache.iceberg.types.Types.NestedField.optional;
-import static org.apache.iceberg.types.Types.NestedField.required;
+import static org.apache.iceberg.flink.data.RandomData.COMPLEX_SCHEMA;
 
 public class TestFlinkParquetReaderWriter {
   private static final int NUM_RECORDS = 20_000;
 
   @Rule
   public TemporaryFolder temp = new TemporaryFolder();
 
-  private static final Schema COMPLEX_SCHEMA = new Schema(
-      required(1, "roots", Types.LongType.get()),
-      optional(3, "lime", Types.ListType.ofRequired(4, Types.DoubleType.get())),
-      required(5, "strict", Types.StructType.of(
-          required(9, "tangerine", Types.StringType.get()),
-          optional(6, "hopeful", Types.StructType.of(
-              required(7, "steel", Types.FloatType.get()),
-              required(8, "lantern", Types.DateType.get())
-          )),
-          optional(10, "vehement", Types.LongType.get())
-      )),
-      optional(11, "metamorphosis", Types.MapType.ofRequired(12, 13,
-          Types.StringType.get(), Types.TimestampType.withZone())),
-      required(14, "winter", Types.ListType.ofOptional(15, Types.StructType.of(
-          optional(16, "beet", Types.DoubleType.get()),
-          required(17, "stamp", Types.FloatType.get()),
-          optional(18, "wheeze", Types.StringType.get())
-      ))),
-      optional(19, "renovate", Types.MapType.ofRequired(20, 21,
-          Types.StringType.get(), Types.StructType.of(
-              optional(22, "jumpy", Types.DoubleType.get()),
-              required(23, "koala", Types.IntegerType.get()),
-              required(24, "couch rope", Types.IntegerType.get())
-          ))),
-      optional(2, "slide", Types.StringType.get())
-  );
-
   private void testCorrectness(Schema schema, int numRecords, Iterable<Row> iterable) throws IOException {
     File testFile = temp.newFile();
     Assert.assertTrue("Delete should succeed", testFile.delete());