TaoZex commented on code in PR #4976: URL: https://github.com/apache/seatunnel/pull/4976#discussion_r1251834151
########## seatunnel-formats/seatunnel-format-avro/src/main/java/org/apache/seatunnel/format/avro/RowToAvroConverter.java: ########## @@ -0,0 +1,234 @@ +/* + * 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.seatunnel.format.avro; + +import org.apache.seatunnel.api.table.type.ArrayType; +import org.apache.seatunnel.api.table.type.BasicType; +import org.apache.seatunnel.api.table.type.DecimalType; +import org.apache.seatunnel.api.table.type.MapType; +import org.apache.seatunnel.api.table.type.SeaTunnelDataType; +import org.apache.seatunnel.api.table.type.SeaTunnelRow; +import org.apache.seatunnel.api.table.type.SeaTunnelRowType; +import org.apache.seatunnel.format.avro.exception.AvroFormatErrorCode; +import org.apache.seatunnel.format.avro.exception.SeaTunnelAvroFormatException; + +import org.apache.avro.Conversions; +import org.apache.avro.LogicalTypes; +import org.apache.avro.Schema; +import org.apache.avro.data.TimeConversions; +import org.apache.avro.generic.GenericDatumWriter; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.generic.GenericRecordBuilder; +import org.apache.avro.io.DatumWriter; + +import java.io.Serializable; +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.List; + +public class RowToAvroConverter implements Serializable { + + private static final long serialVersionUID = -576124379280229724L; + + private final Schema schema; + private final SeaTunnelRowType rowType; + private final DatumWriter<GenericRecord> writer; + + public RowToAvroConverter(SeaTunnelRowType rowType) { + this.schema = buildAvroSchemaWithRowType(rowType); + this.rowType = rowType; + this.writer = createWriter(); + } + + private DatumWriter<GenericRecord> createWriter() { + GenericDatumWriter<GenericRecord> datumWriter = new GenericDatumWriter<>(schema); + datumWriter.getData().addLogicalTypeConversion(new Conversions.DecimalConversion()); + datumWriter.getData().addLogicalTypeConversion(new TimeConversions.DateConversion()); + datumWriter + .getData() + .addLogicalTypeConversion(new TimeConversions.LocalTimestampMillisConversion()); + return datumWriter; + } + + public Schema getSchema() { + return schema; + } + + public DatumWriter<GenericRecord> getWriter() { + return writer; + } + + public GenericRecord convertRowToGenericRecord(SeaTunnelRow element) { + GenericRecordBuilder builder = new GenericRecordBuilder(schema); + String[] fieldNames = rowType.getFieldNames(); + for (int i = 0; i < fieldNames.length; i++) { + String fieldName = rowType.getFieldName(i); + Object value = element.getField(i); + builder.set(fieldName.toLowerCase(), resolveObject(value, rowType.getFieldType(i))); + } + return builder.build(); + } + + private Schema buildAvroSchemaWithRowType(SeaTunnelRowType seaTunnelRowType) { + List<Schema.Field> fields = new ArrayList<>(); + SeaTunnelDataType<?>[] fieldTypes = seaTunnelRowType.getFieldTypes(); + String[] fieldNames = seaTunnelRowType.getFieldNames(); + for (int i = 0; i < fieldNames.length; i++) { + fields.add(generateField(fieldNames[i], fieldTypes[i])); + } + return Schema.createRecord("SeaTunnelRecord", null, null, false, fields); + } + + private Schema.Field generateField(String fieldName, SeaTunnelDataType<?> seaTunnelDataType) { + return new Schema.Field( + fieldName, + seaTunnelDataType2AvroDataType(fieldName, seaTunnelDataType), + null, + null); + } + + private Schema seaTunnelDataType2AvroDataType( + String fieldName, SeaTunnelDataType<?> seaTunnelDataType) { + + switch (seaTunnelDataType.getSqlType()) { + case STRING: + return Schema.create(Schema.Type.STRING); + case BYTES: + return Schema.create(Schema.Type.BYTES); + case TINYINT: + case SMALLINT: + case INT: + return Schema.create(Schema.Type.INT); + case BIGINT: + return Schema.create(Schema.Type.LONG); + case FLOAT: + return Schema.create(Schema.Type.FLOAT); + case DOUBLE: + return Schema.create(Schema.Type.DOUBLE); + case BOOLEAN: + return Schema.create(Schema.Type.BOOLEAN); + case MAP: + SeaTunnelDataType<?> valueType = ((MapType<?, ?>) seaTunnelDataType).getValueType(); + return Schema.createMap(seaTunnelDataType2AvroDataType(fieldName, valueType)); + case ARRAY: + BasicType<?> elementType = ((ArrayType<?, ?>) seaTunnelDataType).getElementType(); + return Schema.createArray(seaTunnelDataType2AvroDataType(fieldName, elementType)); + case ROW: + SeaTunnelDataType<?>[] fieldTypes = + ((SeaTunnelRowType) seaTunnelDataType).getFieldTypes(); + String[] fieldNames = ((SeaTunnelRowType) seaTunnelDataType).getFieldNames(); + List<Schema.Field> subField = new ArrayList<>(); + for (int i = 0; i < fieldNames.length; i++) { + subField.add(generateField(fieldNames[i], fieldTypes[i])); + } + return Schema.createRecord(fieldName, null, null, false, subField); + case DECIMAL: + int precision = ((DecimalType) seaTunnelDataType).getPrecision(); + int scale = ((DecimalType) seaTunnelDataType).getScale(); + LogicalTypes.Decimal decimal = LogicalTypes.decimal(precision, scale); + return decimal.addToSchema(Schema.create(Schema.Type.BYTES)); + case TIMESTAMP: + return LogicalTypes.timestampMillis().addToSchema(Schema.create(Schema.Type.LONG)); + case DATE: + return LogicalTypes.date().addToSchema(Schema.create(Schema.Type.INT)); + case NULL: + return Schema.create(Schema.Type.NULL); + default: + String errorMsg = + String.format( + "SeaTunnel file connector is not supported for this data type [%s]", Review Comment: ```suggestion "SeaTunnel avro format is not supported for this data type [%s]", ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
