From 449688f5c222e0ea4b5fe0995fdffe015f31794a Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 23 Jul 2020 12:11:39 +0800 Subject: [PATCH 1/2] avro: Abstract AvroWithPartnerSchemaVisitor --- .../avro/AvroSchemaWithTypeVisitor.java | 118 +++-------- .../avro/AvroWithPartnerSchemaVisitor.java | 188 ++++++++++++++++++ .../apache/iceberg/data/avro/DataReader.java | 12 +- .../data/AvroWithSparkSchemaVisitor.java | 140 ++++--------- .../iceberg/spark/data/SparkAvroReader.java | 13 +- .../iceberg/spark/data/SparkAvroWriter.java | 23 +-- 6 files changed, 274 insertions(+), 220 deletions(-) create mode 100644 core/src/main/java/org/apache/iceberg/avro/AvroWithPartnerSchemaVisitor.java diff --git a/core/src/main/java/org/apache/iceberg/avro/AvroSchemaWithTypeVisitor.java b/core/src/main/java/org/apache/iceberg/avro/AvroSchemaWithTypeVisitor.java index e6f1c6eb5097..4a46dc3bc140 100644 --- a/core/src/main/java/org/apache/iceberg/avro/AvroSchemaWithTypeVisitor.java +++ b/core/src/main/java/org/apache/iceberg/avro/AvroSchemaWithTypeVisitor.java @@ -19,117 +19,57 @@ package org.apache.iceberg.avro; -import java.util.Deque; -import java.util.List; import org.apache.avro.Schema; -import org.apache.iceberg.relocated.com.google.common.base.Preconditions; -import org.apache.iceberg.relocated.com.google.common.collect.Lists; import org.apache.iceberg.types.Type; import org.apache.iceberg.types.Types; -public abstract class AvroSchemaWithTypeVisitor { +/** + * Avro {@link Schema} with expected Iceberg schema visitor. See {@link #schemaEvolution}, this class is for schema + * evolution reading. The avro schema can evolve into the expected Iceberg type. + */ +public abstract class AvroSchemaWithTypeVisitor extends AvroWithPartnerSchemaVisitor { public static T visit(org.apache.iceberg.Schema iSchema, Schema schema, AvroSchemaWithTypeVisitor visitor) { return visit(iSchema.asStruct(), schema, visitor); } - public static T visit(Type iType, Schema schema, AvroSchemaWithTypeVisitor visitor) { - switch (schema.getType()) { - case RECORD: - return visitRecord(iType != null ? iType.asStructType() : null, schema, visitor); - - case UNION: - return visitUnion(iType, schema, visitor); - - case ARRAY: - return visitArray(iType, schema, visitor); - - case MAP: - Types.MapType map = iType != null ? iType.asMapType() : null; - return visitor.map(map, schema, - visit(map != null ? map.valueType() : null, schema.getValueType(), visitor)); - - default: - return visitor.primitive(iType != null ? iType.asPrimitiveType() : null, schema); - } - } - - private static T visitRecord(Types.StructType struct, Schema record, AvroSchemaWithTypeVisitor visitor) { - // check to make sure this hasn't been visited before - String name = record.getFullName(); - Preconditions.checkState(!visitor.recordLevels.contains(name), - "Cannot process recursive Avro record %s", name); - - visitor.recordLevels.push(name); - - List fields = record.getFields(); - List names = Lists.newArrayListWithExpectedSize(fields.size()); - List results = Lists.newArrayListWithExpectedSize(fields.size()); - for (Schema.Field field : fields) { - int fieldId = AvroSchemaUtil.getFieldId(field); - Types.NestedField iField = struct != null ? struct.field(fieldId) : null; - names.add(field.name()); - results.add(visit(iField != null ? iField.type() : null, field.schema(), visitor)); - } - - visitor.recordLevels.pop(); - - return visitor.record(struct, record, names, results); - } - - private static T visitUnion(Type type, Schema union, AvroSchemaWithTypeVisitor visitor) { - List types = union.getTypes(); - List options = Lists.newArrayListWithExpectedSize(types.size()); - for (Schema branch : types) { - if (branch.getType() == Schema.Type.NULL) { - options.add(visit((Type) null, branch, visitor)); - } else { - options.add(visit(type, branch, visitor)); - } - } - return visitor.union(type, union, options); + @Override + public boolean schemaEvolution() { + return true; } - private static T visitArray(Type type, Schema array, AvroSchemaWithTypeVisitor visitor) { - if (array.getLogicalType() instanceof LogicalMap || (type != null && type.isMapType())) { - Preconditions.checkState( - AvroSchemaUtil.isKeyValueSchema(array.getElementType()), - "Cannot visit invalid logical map type: %s", array); - Types.MapType map = type != null ? type.asMapType() : null; - List keyValueFields = array.getElementType().getFields(); - return visitor.map(map, array, - visit(map != null ? map.keyType() : null, keyValueFields.get(0).schema(), visitor), - visit(map != null ? map.valueType() : null, keyValueFields.get(1).schema(), visitor)); - - } else { - Types.ListType list = type != null ? type.asListType() : null; - return visitor.array(list, array, - visit(list != null ? list.elementType() : null, array.getElementType(), visitor)); - } + @Override + public boolean isMapType(Type type) { + return type != null && type.isMapType(); } - private Deque recordLevels = Lists.newLinkedList(); - - public T record(Types.StructType iStruct, Schema record, List names, List fields) { - return null; + @Override + public boolean isValidMapKey(Type type) { + return type == null || type instanceof Types.StringType; } - public T union(Type iType, Schema union, List options) { - return null; + @Override + public Type arrayElementType(Type arrayType) { + return arrayType == null ? null : arrayType.asListType().elementType(); } - public T array(Types.ListType iList, Schema array, T element) { - return null; + @Override + public Type mapKeyType(Type mapType) { + return mapType == null ? null : mapType.asMapType().keyType(); } - public T map(Types.MapType iMap, Schema map, T key, T value) { - return null; + @Override + public Type mapValueType(Type mapType) { + return mapType == null ? null : mapType.asMapType().valueType(); } - public T map(Types.MapType iMap, Schema map, T value) { - return null; + @Override + public Type structFieldTypeById(Type structType, int id) { + Types.NestedField field = structType == null ? null : structType.asStructType().field(id); + return field == null ? null : field.type(); } - public T primitive(Type.PrimitiveType iPrimitive, Schema primitive) { + @Override + public Type nullType() { return null; } } diff --git a/core/src/main/java/org/apache/iceberg/avro/AvroWithPartnerSchemaVisitor.java b/core/src/main/java/org/apache/iceberg/avro/AvroWithPartnerSchemaVisitor.java new file mode 100644 index 000000000000..ecbf5b575442 --- /dev/null +++ b/core/src/main/java/org/apache/iceberg/avro/AvroWithPartnerSchemaVisitor.java @@ -0,0 +1,188 @@ +/* + * 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 java.util.Deque; +import java.util.List; +import org.apache.avro.Schema; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; + +/** + * A abstract avro schema visitor with partner type. This class is for both reading and writing: + * - For reading, the avro schema could evolve into the partner type. (schema evolution) + * - For writing, the avro schema should be consistent with partner type. + * + * @param

Partner type. + * @param Return T. + */ +public abstract class AvroWithPartnerSchemaVisitor { + + public static T visit(P partner, Schema schema, AvroWithPartnerSchemaVisitor visitor) { + switch (schema.getType()) { + case RECORD: + return visitRecord(partner, schema, visitor); + + case UNION: + return visitUnion(partner, schema, visitor); + + case ARRAY: + return visitArray(partner, schema, visitor); + + case MAP: + P keyType = visitor.mapKeyType(partner); + Preconditions.checkArgument( + visitor.isValidMapKey(keyType), + "Invalid map: %s is not a string", keyType); + return visitor.map(partner, schema, visit(visitor.mapValueType(partner), schema.getValueType(), visitor)); + + default: + return visitor.primitive(partner, schema); + } + } + + // ---------------------------------- Static helpers --------------------------------------------- + + private static T visitRecord(P struct, Schema record, AvroWithPartnerSchemaVisitor visitor) { + // check to make sure this hasn't been visited before + String name = record.getFullName(); + Preconditions.checkState(!visitor.recordLevels.contains(name), + "Cannot process recursive Avro record %s", name); + List fields = record.getFields(); + visitor.recordLevels.push(name); + + List names = Lists.newArrayListWithExpectedSize(fields.size()); + List results = Lists.newArrayListWithExpectedSize(fields.size()); + + if (visitor.schemaEvolution()) { + for (Schema.Field field : fields) { + int fieldId = AvroSchemaUtil.getFieldId(field); + names.add(field.name()); + results.add(visit(visitor.structFieldTypeById(struct, fieldId), field.schema(), visitor)); + } + } else { + String[] fieldNames = visitor.structFieldNames(struct); + P[] fieldTypes = visitor.structFieldTypes(struct); + Preconditions.checkArgument(fieldTypes.length == fields.size(), + "Structs do not match: %s != %s", struct, record); + for (int i = 0; i < fieldTypes.length; i += 1) { + String fieldName = fieldNames[i]; + Schema.Field field = fields.get(i); + Preconditions.checkArgument(AvroSchemaUtil.makeCompatibleName(fieldName).equals(field.name()), + "Structs do not match: field %s != %s", fieldName, field.name()); + results.add(visit(fieldTypes[i], field.schema(), visitor)); + } + } + + visitor.recordLevels.pop(); + + return visitor.record(struct, record, names, results); + } + + private static T visitUnion(P type, Schema union, AvroWithPartnerSchemaVisitor visitor) { + List types = union.getTypes(); + Preconditions.checkArgument(AvroSchemaUtil.isOptionSchema(union), + "Cannot visit non-option union: %s", union); + List options = Lists.newArrayListWithExpectedSize(types.size()); + for (Schema branch : types) { + if (branch.getType() == Schema.Type.NULL) { + options.add(visit(visitor.nullType(), branch, visitor)); + } else { + options.add(visit(type, branch, visitor)); + } + } + return visitor.union(type, union, options); + } + + private static T visitArray(P type, Schema array, AvroWithPartnerSchemaVisitor visitor) { + if (array.getLogicalType() instanceof LogicalMap || visitor.isMapType(type)) { + Preconditions.checkState( + AvroSchemaUtil.isKeyValueSchema(array.getElementType()), + "Cannot visit invalid logical map type: %s", array); + List keyValueFields = array.getElementType().getFields(); + return visitor.map(type, array, + visit(visitor.mapKeyType(type), keyValueFields.get(0).schema(), visitor), + visit(visitor.mapValueType(type), keyValueFields.get(1).schema(), visitor)); + + } else { + return visitor.array(type, array, visit(visitor.arrayElementType(type), array.getElementType(), visitor)); + } + } + + /** + * Just for checking state. + */ + private Deque recordLevels = Lists.newLinkedList(); + + // ---------------------------------- Partner type methods --------------------------------------------- + + public boolean schemaEvolution() { + return false; + } + + public abstract boolean isMapType(P type); + + public abstract boolean isValidMapKey(P type); + + public abstract P arrayElementType(P arrayType); + + public abstract P mapKeyType(P mapType); + public abstract P mapValueType(P mapType); + + public String[] structFieldNames(P structType) { + throw new UnsupportedOperationException(); + } + + public P[] structFieldTypes(P structType) { + throw new UnsupportedOperationException(); + } + + public P structFieldTypeById(P structType, int id) { + throw new UnsupportedOperationException(); + } + + public abstract P nullType(); + + // ---------------------------------- Type visitors --------------------------------------------- + + public T record(P struct, Schema record, List names, List fields) { + return null; + } + + public T union(P type, Schema union, List options) { + return null; + } + + public T array(P sArray, Schema array, T element) { + return null; + } + + public T map(P sMap, Schema map, T key, T value) { + return null; + } + + public T map(P sMap, Schema map, T value) { + return null; + } + + public T primitive(P type, Schema primitive) { + return null; + } +} diff --git a/core/src/main/java/org/apache/iceberg/data/avro/DataReader.java b/core/src/main/java/org/apache/iceberg/data/avro/DataReader.java index 7c8cabc9c1fb..7f9d486b0919 100644 --- a/core/src/main/java/org/apache/iceberg/data/avro/DataReader.java +++ b/core/src/main/java/org/apache/iceberg/data/avro/DataReader.java @@ -89,9 +89,9 @@ private ReadBuilder(Map idToConstant) { } @Override - public ValueReader record(Types.StructType struct, Schema record, + public ValueReader record(Type struct, Schema record, List names, List> fields) { - return createStructReader(struct, fields, idToConstant); + return createStructReader(struct.asStructType(), fields, idToConstant); } @Override @@ -100,22 +100,22 @@ public ValueReader union(Type ignored, Schema union, List> opt } @Override - public ValueReader array(Types.ListType ignored, Schema array, ValueReader elementReader) { + public ValueReader array(Type ignored, Schema array, ValueReader elementReader) { return ValueReaders.array(elementReader); } @Override - public ValueReader map(Types.MapType iMap, Schema map, ValueReader keyReader, ValueReader valueReader) { + public ValueReader map(Type iMap, Schema map, ValueReader keyReader, ValueReader valueReader) { return ValueReaders.arrayMap(keyReader, valueReader); } @Override - public ValueReader map(Types.MapType ignored, Schema map, ValueReader valueReader) { + public ValueReader map(Type ignored, Schema map, ValueReader valueReader) { return ValueReaders.map(ValueReaders.strings(), valueReader); } @Override - public ValueReader primitive(Type.PrimitiveType ignored, Schema primitive) { + public ValueReader primitive(Type ignored, Schema primitive) { LogicalType logicalType = primitive.getLogicalType(); if (logicalType != null) { switch (logicalType.getName()) { diff --git a/spark/src/main/java/org/apache/iceberg/spark/data/AvroWithSparkSchemaVisitor.java b/spark/src/main/java/org/apache/iceberg/spark/data/AvroWithSparkSchemaVisitor.java index 9e3cb805a65e..b2af9fbe3e84 100644 --- a/spark/src/main/java/org/apache/iceberg/spark/data/AvroWithSparkSchemaVisitor.java +++ b/spark/src/main/java/org/apache/iceberg/spark/data/AvroWithSparkSchemaVisitor.java @@ -19,13 +19,10 @@ package org.apache.iceberg.spark.data; -import java.util.Deque; -import java.util.List; +import java.util.Arrays; import org.apache.avro.Schema; -import org.apache.iceberg.avro.AvroSchemaUtil; -import org.apache.iceberg.avro.LogicalMap; +import org.apache.iceberg.avro.AvroWithPartnerSchemaVisitor; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; -import org.apache.iceberg.relocated.com.google.common.collect.Lists; import org.apache.spark.sql.types.ArrayType; import org.apache.spark.sql.types.DataType; import org.apache.spark.sql.types.DataTypes; @@ -34,119 +31,54 @@ import org.apache.spark.sql.types.StructField; import org.apache.spark.sql.types.StructType; -public abstract class AvroWithSparkSchemaVisitor { - public static T visit(StructType struct, Schema schema, AvroWithSparkSchemaVisitor visitor) { - return visitRecord(struct, schema, visitor); - } - - public static T visit(DataType type, Schema schema, AvroWithSparkSchemaVisitor visitor) { - switch (schema.getType()) { - case RECORD: - Preconditions.checkArgument(type instanceof StructType, "Invalid struct: %s is not a struct", type); - return visitRecord((StructType) type, schema, visitor); - - case UNION: - return visitUnion(type, schema, visitor); - - case ARRAY: - return visitArray(type, schema, visitor); - - case MAP: - Preconditions.checkArgument(type instanceof MapType, "Invalid map: %s is not a map", type); - MapType map = (MapType) type; - Preconditions.checkArgument(map.keyType() instanceof StringType, - "Invalid map: %s is not a string", map.keyType()); - return visitor.map(map, schema, visit(map.valueType(), schema.getValueType(), visitor)); - - default: - return visitor.primitive(type, schema); - } - } - - private static T visitRecord(StructType struct, Schema record, AvroWithSparkSchemaVisitor visitor) { - // check to make sure this hasn't been visited before - String name = record.getFullName(); - Preconditions.checkState(!visitor.recordLevels.contains(name), - "Cannot process recursive Avro record %s", name); - StructField[] sFields = struct.fields(); - List fields = record.getFields(); - Preconditions.checkArgument(sFields.length == fields.size(), - "Structs do not match: %s != %s", struct, record); - - visitor.recordLevels.push(name); - - List names = Lists.newArrayListWithExpectedSize(fields.size()); - List results = Lists.newArrayListWithExpectedSize(fields.size()); - for (int i = 0; i < sFields.length; i += 1) { - StructField sField = sFields[i]; - Schema.Field field = fields.get(i); - Preconditions.checkArgument(AvroSchemaUtil.makeCompatibleName(sField.name()).equals(field.name()), - "Structs do not match: field %s != %s", sField.name(), field.name()); - results.add(visit(sField.dataType(), field.schema(), visitor)); - } - - visitor.recordLevels.pop(); - - return visitor.record(struct, record, names, results); - } +/** + * Avro {@link Schema} with Spark {@link DataType} visitor. This class is for writing. The avro schema should be + * consistent with the spark data type. + */ +public abstract class AvroWithSparkSchemaVisitor extends AvroWithPartnerSchemaVisitor { - private static T visitUnion(DataType type, Schema union, AvroWithSparkSchemaVisitor visitor) { - List types = union.getTypes(); - Preconditions.checkArgument(AvroSchemaUtil.isOptionSchema(union), - "Cannot visit non-option union: %s", union); - List options = Lists.newArrayListWithExpectedSize(types.size()); - for (Schema branch : types) { - if (branch.getType() == Schema.Type.NULL) { - options.add(visit(DataTypes.NullType, branch, visitor)); - } else { - options.add(visit(type, branch, visitor)); - } - } - return visitor.union(type, union, options); + @Override + public boolean isValidMapKey(DataType dataType) { + return dataType instanceof StringType; } - private static T visitArray(DataType type, Schema array, AvroWithSparkSchemaVisitor visitor) { - if (array.getLogicalType() instanceof LogicalMap || type instanceof MapType) { - Preconditions.checkState( - AvroSchemaUtil.isKeyValueSchema(array.getElementType()), - "Cannot visit invalid logical map type: %s", array); - Preconditions.checkArgument(type instanceof MapType, "Invalid map: %s is not a map", type); - MapType map = (MapType) type; - List keyValueFields = array.getElementType().getFields(); - return visitor.map(map, array, - visit(map.keyType(), keyValueFields.get(0).schema(), visitor), - visit(map.valueType(), keyValueFields.get(1).schema(), visitor)); - - } else { - Preconditions.checkArgument(type instanceof ArrayType, "Invalid array: %s is not an array", type); - ArrayType list = (ArrayType) type; - return visitor.array(list, array, visit(list.elementType(), array.getElementType(), visitor)); - } + @Override + public boolean isMapType(DataType dataType) { + return dataType instanceof MapType; } - private Deque recordLevels = Lists.newLinkedList(); - - public T record(StructType struct, Schema record, List names, List fields) { - return null; + @Override + public DataType arrayElementType(DataType arrayType) { + Preconditions.checkArgument(arrayType instanceof ArrayType, "Invalid array: %s is not an array", arrayType); + return ((ArrayType) arrayType).elementType(); } - public T union(DataType type, Schema union, List options) { - return null; + @Override + public DataType mapKeyType(DataType mapType) { + Preconditions.checkArgument(mapType instanceof MapType, "Invalid map: %s is not a map", mapType); + return ((MapType) mapType).keyType(); } - public T array(ArrayType sArray, Schema array, T element) { - return null; + @Override + public DataType mapValueType(DataType mapType) { + Preconditions.checkArgument(mapType instanceof MapType, "Invalid map: %s is not a map", mapType); + return ((MapType) mapType).valueType(); } - public T map(MapType sMap, Schema map, T key, T value) { - return null; + @Override + public String[] structFieldNames(DataType structType) { + Preconditions.checkArgument(structType instanceof StructType, "Invalid struct: %s is not a struct", structType); + return ((StructType) structType).fieldNames(); } - public T map(MapType sMap, Schema map, T value) { - return null; + @Override + public DataType[] structFieldTypes(DataType structType) { + Preconditions.checkArgument(structType instanceof StructType, "Invalid struct: %s is not a struct", structType); + return Arrays.stream(((StructType) structType).fields()).map(StructField::dataType).toArray(DataType[]::new); } - public T primitive(DataType type, Schema primitive) { - return null; + @Override + public DataType nullType() { + return DataTypes.NullType; } } diff --git a/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroReader.java b/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroReader.java index 46c594e56a6c..928ff3f72e19 100644 --- a/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroReader.java +++ b/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroReader.java @@ -33,7 +33,6 @@ import org.apache.iceberg.data.avro.DecoderResolver; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.apache.iceberg.types.Type; -import org.apache.iceberg.types.Types; import org.apache.spark.sql.catalyst.InternalRow; @@ -72,9 +71,9 @@ private ReadBuilder(Map idToConstant) { } @Override - public ValueReader record(Types.StructType expected, Schema record, List names, + public ValueReader record(Type expected, Schema record, List names, List> fields) { - return SparkValueReaders.struct(fields, expected, idToConstant); + return SparkValueReaders.struct(fields, expected.asStructType(), idToConstant); } @Override @@ -83,23 +82,23 @@ public ValueReader union(Type expected, Schema union, List> op } @Override - public ValueReader array(Types.ListType expected, Schema array, ValueReader elementReader) { + public ValueReader array(Type expected, Schema array, ValueReader elementReader) { return SparkValueReaders.array(elementReader); } @Override - public ValueReader map(Types.MapType expected, Schema map, + public ValueReader map(Type expected, Schema map, ValueReader keyReader, ValueReader valueReader) { return SparkValueReaders.arrayMap(keyReader, valueReader); } @Override - public ValueReader map(Types.MapType expected, Schema map, ValueReader valueReader) { + public ValueReader map(Type expected, Schema map, ValueReader valueReader) { return SparkValueReaders.map(SparkValueReaders.strings(), valueReader); } @Override - public ValueReader primitive(Type.PrimitiveType expected, Schema primitive) { + public ValueReader primitive(Type expected, Schema primitive) { LogicalType logicalType = primitive.getLogicalType(); if (logicalType != null) { switch (logicalType.getName()) { diff --git a/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroWriter.java b/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroWriter.java index 2f82cc12f469..d892bd34a501 100644 --- a/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroWriter.java +++ b/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroWriter.java @@ -20,9 +20,8 @@ package org.apache.iceberg.spark.data; import java.io.IOException; +import java.util.Arrays; import java.util.List; -import java.util.stream.Collectors; -import java.util.stream.Stream; import org.apache.avro.LogicalType; import org.apache.avro.LogicalTypes; import org.apache.avro.Schema; @@ -32,12 +31,9 @@ import org.apache.iceberg.avro.ValueWriters; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; import org.apache.spark.sql.catalyst.InternalRow; -import org.apache.spark.sql.types.ArrayType; import org.apache.spark.sql.types.ByteType; import org.apache.spark.sql.types.DataType; -import org.apache.spark.sql.types.MapType; import org.apache.spark.sql.types.ShortType; -import org.apache.spark.sql.types.StructField; import org.apache.spark.sql.types.StructType; public class SparkAvroWriter implements DatumWriter { @@ -62,9 +58,8 @@ public void write(InternalRow datum, Encoder out) throws IOException { private static class WriteBuilder extends AvroWithSparkSchemaVisitor> { @Override - public ValueWriter record(StructType struct, Schema record, List names, List> fields) { - List types = Stream.of(struct.fields()).map(StructField::dataType).collect(Collectors.toList()); - return SparkValueWriters.struct(fields, types); + public ValueWriter record(DataType struct, Schema record, List names, List> fields) { + return SparkValueWriters.struct(fields, Arrays.asList(structFieldTypes(struct))); } @Override @@ -81,18 +76,18 @@ public ValueWriter union(DataType type, Schema union, List> op } @Override - public ValueWriter array(ArrayType sArray, Schema array, ValueWriter elementWriter) { - return SparkValueWriters.array(elementWriter, sArray.elementType()); + public ValueWriter array(DataType sArray, Schema array, ValueWriter elementWriter) { + return SparkValueWriters.array(elementWriter, arrayElementType(sArray)); } @Override - public ValueWriter map(MapType sMap, Schema map, ValueWriter valueReader) { - return SparkValueWriters.map(SparkValueWriters.strings(), sMap.keyType(), valueReader, sMap.valueType()); + public ValueWriter map(DataType sMap, Schema map, ValueWriter valueReader) { + return SparkValueWriters.map(SparkValueWriters.strings(), mapKeyType(sMap), valueReader, mapValueType(sMap)); } @Override - public ValueWriter map(MapType sMap, Schema map, ValueWriter keyWriter, ValueWriter valueWriter) { - return SparkValueWriters.arrayMap(keyWriter, sMap.keyType(), valueWriter, sMap.valueType()); + public ValueWriter map(DataType sMap, Schema map, ValueWriter keyWriter, ValueWriter valueWriter) { + return SparkValueWriters.arrayMap(keyWriter, mapKeyType(sMap), valueWriter, mapValueType(sMap)); } @Override From 24cf7cd1a630821506706e0d06de6acc48866db5 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Tue, 28 Jul 2020 15:09:37 +0800 Subject: [PATCH 2/2] Address Ryan's comments --- .../avro/AvroSchemaWithTypeVisitor.java | 118 +++++++++++++----- ...=> AvroWithPartnerByStructureVisitor.java} | 74 ++++------- .../apache/iceberg/data/avro/DataReader.java | 12 +- .../data/AvroWithSparkSchemaVisitor.java | 38 +++--- .../iceberg/spark/data/SparkAvroReader.java | 13 +- .../iceberg/spark/data/SparkAvroWriter.java | 6 +- 6 files changed, 145 insertions(+), 116 deletions(-) rename core/src/main/java/org/apache/iceberg/avro/{AvroWithPartnerSchemaVisitor.java => AvroWithPartnerByStructureVisitor.java} (68%) diff --git a/core/src/main/java/org/apache/iceberg/avro/AvroSchemaWithTypeVisitor.java b/core/src/main/java/org/apache/iceberg/avro/AvroSchemaWithTypeVisitor.java index 4a46dc3bc140..e6f1c6eb5097 100644 --- a/core/src/main/java/org/apache/iceberg/avro/AvroSchemaWithTypeVisitor.java +++ b/core/src/main/java/org/apache/iceberg/avro/AvroSchemaWithTypeVisitor.java @@ -19,57 +19,117 @@ package org.apache.iceberg.avro; +import java.util.Deque; +import java.util.List; import org.apache.avro.Schema; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; import org.apache.iceberg.types.Type; import org.apache.iceberg.types.Types; -/** - * Avro {@link Schema} with expected Iceberg schema visitor. See {@link #schemaEvolution}, this class is for schema - * evolution reading. The avro schema can evolve into the expected Iceberg type. - */ -public abstract class AvroSchemaWithTypeVisitor extends AvroWithPartnerSchemaVisitor { +public abstract class AvroSchemaWithTypeVisitor { public static T visit(org.apache.iceberg.Schema iSchema, Schema schema, AvroSchemaWithTypeVisitor visitor) { return visit(iSchema.asStruct(), schema, visitor); } - @Override - public boolean schemaEvolution() { - return true; + public static T visit(Type iType, Schema schema, AvroSchemaWithTypeVisitor visitor) { + switch (schema.getType()) { + case RECORD: + return visitRecord(iType != null ? iType.asStructType() : null, schema, visitor); + + case UNION: + return visitUnion(iType, schema, visitor); + + case ARRAY: + return visitArray(iType, schema, visitor); + + case MAP: + Types.MapType map = iType != null ? iType.asMapType() : null; + return visitor.map(map, schema, + visit(map != null ? map.valueType() : null, schema.getValueType(), visitor)); + + default: + return visitor.primitive(iType != null ? iType.asPrimitiveType() : null, schema); + } + } + + private static T visitRecord(Types.StructType struct, Schema record, AvroSchemaWithTypeVisitor visitor) { + // check to make sure this hasn't been visited before + String name = record.getFullName(); + Preconditions.checkState(!visitor.recordLevels.contains(name), + "Cannot process recursive Avro record %s", name); + + visitor.recordLevels.push(name); + + List fields = record.getFields(); + List names = Lists.newArrayListWithExpectedSize(fields.size()); + List results = Lists.newArrayListWithExpectedSize(fields.size()); + for (Schema.Field field : fields) { + int fieldId = AvroSchemaUtil.getFieldId(field); + Types.NestedField iField = struct != null ? struct.field(fieldId) : null; + names.add(field.name()); + results.add(visit(iField != null ? iField.type() : null, field.schema(), visitor)); + } + + visitor.recordLevels.pop(); + + return visitor.record(struct, record, names, results); + } + + private static T visitUnion(Type type, Schema union, AvroSchemaWithTypeVisitor visitor) { + List types = union.getTypes(); + List options = Lists.newArrayListWithExpectedSize(types.size()); + for (Schema branch : types) { + if (branch.getType() == Schema.Type.NULL) { + options.add(visit((Type) null, branch, visitor)); + } else { + options.add(visit(type, branch, visitor)); + } + } + return visitor.union(type, union, options); } - @Override - public boolean isMapType(Type type) { - return type != null && type.isMapType(); + private static T visitArray(Type type, Schema array, AvroSchemaWithTypeVisitor visitor) { + if (array.getLogicalType() instanceof LogicalMap || (type != null && type.isMapType())) { + Preconditions.checkState( + AvroSchemaUtil.isKeyValueSchema(array.getElementType()), + "Cannot visit invalid logical map type: %s", array); + Types.MapType map = type != null ? type.asMapType() : null; + List keyValueFields = array.getElementType().getFields(); + return visitor.map(map, array, + visit(map != null ? map.keyType() : null, keyValueFields.get(0).schema(), visitor), + visit(map != null ? map.valueType() : null, keyValueFields.get(1).schema(), visitor)); + + } else { + Types.ListType list = type != null ? type.asListType() : null; + return visitor.array(list, array, + visit(list != null ? list.elementType() : null, array.getElementType(), visitor)); + } } - @Override - public boolean isValidMapKey(Type type) { - return type == null || type instanceof Types.StringType; + private Deque recordLevels = Lists.newLinkedList(); + + public T record(Types.StructType iStruct, Schema record, List names, List fields) { + return null; } - @Override - public Type arrayElementType(Type arrayType) { - return arrayType == null ? null : arrayType.asListType().elementType(); + public T union(Type iType, Schema union, List options) { + return null; } - @Override - public Type mapKeyType(Type mapType) { - return mapType == null ? null : mapType.asMapType().keyType(); + public T array(Types.ListType iList, Schema array, T element) { + return null; } - @Override - public Type mapValueType(Type mapType) { - return mapType == null ? null : mapType.asMapType().valueType(); + public T map(Types.MapType iMap, Schema map, T key, T value) { + return null; } - @Override - public Type structFieldTypeById(Type structType, int id) { - Types.NestedField field = structType == null ? null : structType.asStructType().field(id); - return field == null ? null : field.type(); + public T map(Types.MapType iMap, Schema map, T value) { + return null; } - @Override - public Type nullType() { + public T primitive(Type.PrimitiveType iPrimitive, Schema primitive) { return null; } } diff --git a/core/src/main/java/org/apache/iceberg/avro/AvroWithPartnerSchemaVisitor.java b/core/src/main/java/org/apache/iceberg/avro/AvroWithPartnerByStructureVisitor.java similarity index 68% rename from core/src/main/java/org/apache/iceberg/avro/AvroWithPartnerSchemaVisitor.java rename to core/src/main/java/org/apache/iceberg/avro/AvroWithPartnerByStructureVisitor.java index ecbf5b575442..2b28f54aaaaa 100644 --- a/core/src/main/java/org/apache/iceberg/avro/AvroWithPartnerSchemaVisitor.java +++ b/core/src/main/java/org/apache/iceberg/avro/AvroWithPartnerByStructureVisitor.java @@ -24,18 +24,18 @@ import org.apache.avro.Schema; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; import org.apache.iceberg.relocated.com.google.common.collect.Lists; +import org.apache.iceberg.util.Pair; /** - * A abstract avro schema visitor with partner type. This class is for both reading and writing: - * - For reading, the avro schema could evolve into the partner type. (schema evolution) - * - For writing, the avro schema should be consistent with partner type. + * A abstract avro schema visitor with partner type. The visitor rely on the structure matching exactly and are + * guaranteed that because both schemas are derived from the same Iceberg schema. * * @param

Partner type. * @param Return T. */ -public abstract class AvroWithPartnerSchemaVisitor { +public abstract class AvroWithPartnerByStructureVisitor { - public static T visit(P partner, Schema schema, AvroWithPartnerSchemaVisitor visitor) { + public static T visit(P partner, Schema schema, AvroWithPartnerByStructureVisitor visitor) { switch (schema.getType()) { case RECORD: return visitRecord(partner, schema, visitor); @@ -49,7 +49,7 @@ public static T visit(P partner, Schema schema, AvroWithPartnerSchemaVisi case MAP: P keyType = visitor.mapKeyType(partner); Preconditions.checkArgument( - visitor.isValidMapKey(keyType), + visitor.isStringType(keyType), "Invalid map: %s is not a string", keyType); return visitor.map(partner, schema, visit(visitor.mapValueType(partner), schema.getValueType(), visitor)); @@ -60,35 +60,25 @@ public static T visit(P partner, Schema schema, AvroWithPartnerSchemaVisi // ---------------------------------- Static helpers --------------------------------------------- - private static T visitRecord(P struct, Schema record, AvroWithPartnerSchemaVisitor visitor) { + private static T visitRecord(P struct, Schema record, AvroWithPartnerByStructureVisitor visitor) { // check to make sure this hasn't been visited before String name = record.getFullName(); Preconditions.checkState(!visitor.recordLevels.contains(name), "Cannot process recursive Avro record %s", name); List fields = record.getFields(); + visitor.recordLevels.push(name); List names = Lists.newArrayListWithExpectedSize(fields.size()); List results = Lists.newArrayListWithExpectedSize(fields.size()); - - if (visitor.schemaEvolution()) { - for (Schema.Field field : fields) { - int fieldId = AvroSchemaUtil.getFieldId(field); - names.add(field.name()); - results.add(visit(visitor.structFieldTypeById(struct, fieldId), field.schema(), visitor)); - } - } else { - String[] fieldNames = visitor.structFieldNames(struct); - P[] fieldTypes = visitor.structFieldTypes(struct); - Preconditions.checkArgument(fieldTypes.length == fields.size(), - "Structs do not match: %s != %s", struct, record); - for (int i = 0; i < fieldTypes.length; i += 1) { - String fieldName = fieldNames[i]; - Schema.Field field = fields.get(i); - Preconditions.checkArgument(AvroSchemaUtil.makeCompatibleName(fieldName).equals(field.name()), - "Structs do not match: field %s != %s", fieldName, field.name()); - results.add(visit(fieldTypes[i], field.schema(), visitor)); - } + for (int i = 0; i < fields.size(); i += 1) { + Pair nameAndType = visitor.fieldNameAndType(struct, i); + String fieldName = nameAndType.first(); + Schema.Field field = fields.get(i); + Preconditions.checkArgument(AvroSchemaUtil.makeCompatibleName(fieldName).equals(field.name()), + "Structs do not match: field %s != %s", fieldName, field.name()); + results.add(visit(nameAndType.second(), field.schema(), visitor)); + names.add(fieldName); } visitor.recordLevels.pop(); @@ -96,7 +86,7 @@ private static T visitRecord(P struct, Schema record, AvroWithPartnerSche return visitor.record(struct, record, names, results); } - private static T visitUnion(P type, Schema union, AvroWithPartnerSchemaVisitor visitor) { + private static T visitUnion(P type, Schema union, AvroWithPartnerByStructureVisitor visitor) { List types = union.getTypes(); Preconditions.checkArgument(AvroSchemaUtil.isOptionSchema(union), "Cannot visit non-option union: %s", union); @@ -111,7 +101,7 @@ private static T visitUnion(P type, Schema union, AvroWithPartnerSchemaVi return visitor.union(type, union, options); } - private static T visitArray(P type, Schema array, AvroWithPartnerSchemaVisitor visitor) { + private static T visitArray(P type, Schema array, AvroWithPartnerByStructureVisitor visitor) { if (array.getLogicalType() instanceof LogicalMap || visitor.isMapType(type)) { Preconditions.checkState( AvroSchemaUtil.isKeyValueSchema(array.getElementType()), @@ -133,32 +123,18 @@ private static T visitArray(P type, Schema array, AvroWithPartnerSchemaVi // ---------------------------------- Partner type methods --------------------------------------------- - public boolean schemaEvolution() { - return false; - } - - public abstract boolean isMapType(P type); - - public abstract boolean isValidMapKey(P type); - - public abstract P arrayElementType(P arrayType); + protected abstract boolean isMapType(P type); - public abstract P mapKeyType(P mapType); - public abstract P mapValueType(P mapType); + protected abstract boolean isStringType(P type); - public String[] structFieldNames(P structType) { - throw new UnsupportedOperationException(); - } + protected abstract P arrayElementType(P arrayType); - public P[] structFieldTypes(P structType) { - throw new UnsupportedOperationException(); - } + protected abstract P mapKeyType(P mapType); + protected abstract P mapValueType(P mapType); - public P structFieldTypeById(P structType, int id) { - throw new UnsupportedOperationException(); - } + protected abstract Pair fieldNameAndType(P structType, int pos); - public abstract P nullType(); + protected abstract P nullType(); // ---------------------------------- Type visitors --------------------------------------------- diff --git a/core/src/main/java/org/apache/iceberg/data/avro/DataReader.java b/core/src/main/java/org/apache/iceberg/data/avro/DataReader.java index 7f9d486b0919..7c8cabc9c1fb 100644 --- a/core/src/main/java/org/apache/iceberg/data/avro/DataReader.java +++ b/core/src/main/java/org/apache/iceberg/data/avro/DataReader.java @@ -89,9 +89,9 @@ private ReadBuilder(Map idToConstant) { } @Override - public ValueReader record(Type struct, Schema record, + public ValueReader record(Types.StructType struct, Schema record, List names, List> fields) { - return createStructReader(struct.asStructType(), fields, idToConstant); + return createStructReader(struct, fields, idToConstant); } @Override @@ -100,22 +100,22 @@ public ValueReader union(Type ignored, Schema union, List> opt } @Override - public ValueReader array(Type ignored, Schema array, ValueReader elementReader) { + public ValueReader array(Types.ListType ignored, Schema array, ValueReader elementReader) { return ValueReaders.array(elementReader); } @Override - public ValueReader map(Type iMap, Schema map, ValueReader keyReader, ValueReader valueReader) { + public ValueReader map(Types.MapType iMap, Schema map, ValueReader keyReader, ValueReader valueReader) { return ValueReaders.arrayMap(keyReader, valueReader); } @Override - public ValueReader map(Type ignored, Schema map, ValueReader valueReader) { + public ValueReader map(Types.MapType ignored, Schema map, ValueReader valueReader) { return ValueReaders.map(ValueReaders.strings(), valueReader); } @Override - public ValueReader primitive(Type ignored, Schema primitive) { + public ValueReader primitive(Type.PrimitiveType ignored, Schema primitive) { LogicalType logicalType = primitive.getLogicalType(); if (logicalType != null) { switch (logicalType.getName()) { diff --git a/spark/src/main/java/org/apache/iceberg/spark/data/AvroWithSparkSchemaVisitor.java b/spark/src/main/java/org/apache/iceberg/spark/data/AvroWithSparkSchemaVisitor.java index b2af9fbe3e84..40ed05b4ce65 100644 --- a/spark/src/main/java/org/apache/iceberg/spark/data/AvroWithSparkSchemaVisitor.java +++ b/spark/src/main/java/org/apache/iceberg/spark/data/AvroWithSparkSchemaVisitor.java @@ -19,10 +19,9 @@ package org.apache.iceberg.spark.data; -import java.util.Arrays; -import org.apache.avro.Schema; -import org.apache.iceberg.avro.AvroWithPartnerSchemaVisitor; +import org.apache.iceberg.avro.AvroWithPartnerByStructureVisitor; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.util.Pair; import org.apache.spark.sql.types.ArrayType; import org.apache.spark.sql.types.DataType; import org.apache.spark.sql.types.DataTypes; @@ -31,54 +30,45 @@ import org.apache.spark.sql.types.StructField; import org.apache.spark.sql.types.StructType; -/** - * Avro {@link Schema} with Spark {@link DataType} visitor. This class is for writing. The avro schema should be - * consistent with the spark data type. - */ -public abstract class AvroWithSparkSchemaVisitor extends AvroWithPartnerSchemaVisitor { +public abstract class AvroWithSparkSchemaVisitor extends AvroWithPartnerByStructureVisitor { @Override - public boolean isValidMapKey(DataType dataType) { + protected boolean isStringType(DataType dataType) { return dataType instanceof StringType; } @Override - public boolean isMapType(DataType dataType) { + protected boolean isMapType(DataType dataType) { return dataType instanceof MapType; } @Override - public DataType arrayElementType(DataType arrayType) { + protected DataType arrayElementType(DataType arrayType) { Preconditions.checkArgument(arrayType instanceof ArrayType, "Invalid array: %s is not an array", arrayType); return ((ArrayType) arrayType).elementType(); } @Override - public DataType mapKeyType(DataType mapType) { - Preconditions.checkArgument(mapType instanceof MapType, "Invalid map: %s is not a map", mapType); + protected DataType mapKeyType(DataType mapType) { + Preconditions.checkArgument(isMapType(mapType), "Invalid map: %s is not a map", mapType); return ((MapType) mapType).keyType(); } @Override - public DataType mapValueType(DataType mapType) { - Preconditions.checkArgument(mapType instanceof MapType, "Invalid map: %s is not a map", mapType); + protected DataType mapValueType(DataType mapType) { + Preconditions.checkArgument(isMapType(mapType), "Invalid map: %s is not a map", mapType); return ((MapType) mapType).valueType(); } @Override - public String[] structFieldNames(DataType structType) { - Preconditions.checkArgument(structType instanceof StructType, "Invalid struct: %s is not a struct", structType); - return ((StructType) structType).fieldNames(); - } - - @Override - public DataType[] structFieldTypes(DataType structType) { + protected Pair fieldNameAndType(DataType structType, int pos) { Preconditions.checkArgument(structType instanceof StructType, "Invalid struct: %s is not a struct", structType); - return Arrays.stream(((StructType) structType).fields()).map(StructField::dataType).toArray(DataType[]::new); + StructField field = ((StructType) structType).apply(pos); + return Pair.of(field.name(), field.dataType()); } @Override - public DataType nullType() { + protected DataType nullType() { return DataTypes.NullType; } } diff --git a/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroReader.java b/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroReader.java index 928ff3f72e19..46c594e56a6c 100644 --- a/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroReader.java +++ b/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroReader.java @@ -33,6 +33,7 @@ import org.apache.iceberg.data.avro.DecoderResolver; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.Types; import org.apache.spark.sql.catalyst.InternalRow; @@ -71,9 +72,9 @@ private ReadBuilder(Map idToConstant) { } @Override - public ValueReader record(Type expected, Schema record, List names, + public ValueReader record(Types.StructType expected, Schema record, List names, List> fields) { - return SparkValueReaders.struct(fields, expected.asStructType(), idToConstant); + return SparkValueReaders.struct(fields, expected, idToConstant); } @Override @@ -82,23 +83,23 @@ public ValueReader union(Type expected, Schema union, List> op } @Override - public ValueReader array(Type expected, Schema array, ValueReader elementReader) { + public ValueReader array(Types.ListType expected, Schema array, ValueReader elementReader) { return SparkValueReaders.array(elementReader); } @Override - public ValueReader map(Type expected, Schema map, + public ValueReader map(Types.MapType expected, Schema map, ValueReader keyReader, ValueReader valueReader) { return SparkValueReaders.arrayMap(keyReader, valueReader); } @Override - public ValueReader map(Type expected, Schema map, ValueReader valueReader) { + public ValueReader map(Types.MapType expected, Schema map, ValueReader valueReader) { return SparkValueReaders.map(SparkValueReaders.strings(), valueReader); } @Override - public ValueReader primitive(Type expected, Schema primitive) { + public ValueReader primitive(Type.PrimitiveType expected, Schema primitive) { LogicalType logicalType = primitive.getLogicalType(); if (logicalType != null) { switch (logicalType.getName()) { diff --git a/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroWriter.java b/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroWriter.java index d892bd34a501..b1625d7df9b1 100644 --- a/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroWriter.java +++ b/spark/src/main/java/org/apache/iceberg/spark/data/SparkAvroWriter.java @@ -20,8 +20,9 @@ package org.apache.iceberg.spark.data; import java.io.IOException; -import java.util.Arrays; import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.IntStream; import org.apache.avro.LogicalType; import org.apache.avro.LogicalTypes; import org.apache.avro.Schema; @@ -59,7 +60,8 @@ public void write(InternalRow datum, Encoder out) throws IOException { private static class WriteBuilder extends AvroWithSparkSchemaVisitor> { @Override public ValueWriter record(DataType struct, Schema record, List names, List> fields) { - return SparkValueWriters.struct(fields, Arrays.asList(structFieldTypes(struct))); + return SparkValueWriters.struct(fields, IntStream.range(0, names.size()) + .mapToObj(i -> fieldNameAndType(struct, i).second()).collect(Collectors.toList())); } @Override