You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by lz...@apache.org on 2021/12/14 10:22:03 UTC
[flink] branch master updated: [FLINK-24565][avro] Port avro file format factory to BulkReaderFormatFactory
This is an automated email from the ASF dual-hosted git repository.
lzljs3620320 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new 388dce6 [FLINK-24565][avro] Port avro file format factory to BulkReaderFormatFactory
388dce6 is described below
commit 388dce6af221f5cbe7807d1439305555c91c7449
Author: tsreaper <ts...@gmail.com>
AuthorDate: Tue Dec 14 18:21:33 2021 +0800
[FLINK-24565][avro] Port avro file format factory to BulkReaderFormatFactory
This closes #17520
---
flink-formats/flink-avro/pom.xml | 8 +
.../flink/formats/avro/AbstractAvroBulkFormat.java | 207 ++++++++++++++++
.../flink/formats/avro/AvroFileFormatFactory.java | 68 ++++-
.../formats/avro/AvroFileSystemFormatFactory.java | 193 ---------------
.../org.apache.flink.table.factories.Factory | 1 -
.../flink/formats/avro/AvroBulkFormatITCase.java | 231 +++++++++++++++++
.../flink/formats/avro/AvroBulkFormatTest.java | 275 +++++++++++++++++++++
.../formats/avro/AvroBulkFormatTestUtils.java | 76 ++++++
.../flink/formats/avro/AvroFilesystemITCase.java | 7 +-
9 files changed, 865 insertions(+), 201 deletions(-)
diff --git a/flink-formats/flink-avro/pom.xml b/flink-formats/flink-avro/pom.xml
index d2e05ed..048f0f1 100644
--- a/flink-formats/flink-avro/pom.xml
+++ b/flink-formats/flink-avro/pom.xml
@@ -47,6 +47,7 @@ under the License.
<!-- Table ecosystem and filesystem connector -->
+ <!-- Data stream projects depending on this project might not depend on flink-table-*. -->
<dependency>
<groupId>org.apache.flink</groupId>
<artifactId>flink-table-common</artifactId>
@@ -86,6 +87,13 @@ under the License.
<dependency>
<groupId>org.apache.flink</groupId>
+ <artifactId>flink-connector-testing</artifactId>
+ <version>${project.version}</version>
+ <scope>test</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.flink</groupId>
<artifactId>flink-streaming-scala_${scala.binary.version}</artifactId>
<version>${project.version}</version>
<scope>test</scope>
diff --git a/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AbstractAvroBulkFormat.java b/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AbstractAvroBulkFormat.java
new file mode 100644
index 0000000..1221a04
--- /dev/null
+++ b/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AbstractAvroBulkFormat.java
@@ -0,0 +1,207 @@
+/*
+ * 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.flink.formats.avro;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.connector.file.src.FileSourceSplit;
+import org.apache.flink.connector.file.src.reader.BulkFormat;
+import org.apache.flink.connector.file.src.util.CheckpointedPosition;
+import org.apache.flink.connector.file.src.util.IteratorResultIterator;
+import org.apache.flink.connector.file.src.util.Pool;
+import org.apache.flink.core.fs.FileSystem;
+import org.apache.flink.core.fs.Path;
+import org.apache.flink.formats.avro.utils.FSDataInputStreamWrapper;
+
+import org.apache.avro.Schema;
+import org.apache.avro.file.DataFileReader;
+import org.apache.avro.file.SeekableInput;
+import org.apache.avro.generic.GenericDatumReader;
+import org.apache.avro.io.DatumReader;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.util.Iterator;
+
+/** Provides a {@link BulkFormat} for Avro records. */
+@Internal
+public abstract class AbstractAvroBulkFormat<A, T, SplitT extends FileSourceSplit>
+ implements BulkFormat<T, SplitT> {
+
+ private static final long serialVersionUID = 1L;
+
+ protected final Schema readerSchema;
+
+ protected AbstractAvroBulkFormat(Schema readerSchema) {
+ this.readerSchema = readerSchema;
+ }
+
+ @Override
+ public AvroReader createReader(Configuration config, SplitT split) throws IOException {
+ open(split);
+ return createReader(split);
+ }
+
+ @Override
+ public AvroReader restoreReader(Configuration config, SplitT split) throws IOException {
+ open(split);
+ return createReader(split);
+ }
+
+ @Override
+ public boolean isSplittable() {
+ return true;
+ }
+
+ private AvroReader createReader(SplitT split) throws IOException {
+ long end = split.offset() + split.length();
+ if (split.getReaderPosition().isPresent()) {
+ CheckpointedPosition position = split.getReaderPosition().get();
+ return new AvroReader(
+ split.path(),
+ split.offset(),
+ end,
+ position.getOffset(),
+ position.getRecordsAfterOffset());
+ } else {
+ return new AvroReader(split.path(), split.offset(), end, -1, 0);
+ }
+ }
+
+ protected void open(SplitT split) {}
+
+ protected abstract T convert(A record);
+
+ protected abstract A createReusedAvroRecord();
+
+ private class AvroReader implements BulkFormat.Reader<T> {
+
+ private final DataFileReader<A> reader;
+
+ private final long end;
+ private final Pool<A> pool;
+
+ private long currentBlockStart;
+ private long currentRecordsToSkip;
+
+ private AvroReader(Path path, long offset, long end, long blockStart, long recordsToSkip)
+ throws IOException {
+ A reuse = createReusedAvroRecord();
+
+ this.reader = createReaderFromPath(path);
+ if (blockStart >= 0) {
+ reader.seek(blockStart);
+ } else {
+ reader.sync(offset);
+ }
+ for (int i = 0; i < recordsToSkip; i++) {
+ reader.next(reuse);
+ }
+
+ this.end = end;
+ this.pool = new Pool<>(1);
+ this.pool.add(reuse);
+
+ this.currentBlockStart = reader.previousSync();
+ this.currentRecordsToSkip = recordsToSkip;
+ }
+
+ private DataFileReader<A> createReaderFromPath(Path path) throws IOException {
+ FileSystem fileSystem = path.getFileSystem();
+ DatumReader<A> datumReader = new GenericDatumReader<>(null, readerSchema);
+ SeekableInput in =
+ new FSDataInputStreamWrapper(
+ fileSystem.open(path), fileSystem.getFileStatus(path).getLen());
+ return (DataFileReader<A>) DataFileReader.openReader(in, datumReader);
+ }
+
+ @Nullable
+ @Override
+ public RecordIterator<T> readBatch() throws IOException {
+ A reuse;
+ try {
+ reuse = pool.pollEntry();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(
+ "Interrupted while waiting for the previous batch to be consumed", e);
+ }
+
+ if (!readNextBlock()) {
+ pool.recycler().recycle(reuse);
+ return null;
+ }
+
+ currentBlockStart = reader.previousSync();
+ Iterator<T> iterator =
+ new AvroBlockIterator(
+ reader.getBlockCount() - currentRecordsToSkip, reader, reuse);
+ long recordsToSkip = currentRecordsToSkip;
+ currentRecordsToSkip = 0;
+ return new IteratorResultIterator<>(
+ iterator,
+ currentBlockStart,
+ recordsToSkip,
+ () -> pool.recycler().recycle(reuse));
+ }
+
+ private boolean readNextBlock() throws IOException {
+ // read the next block with reader,
+ // returns true if a block is read and false if we reach the end of this split
+ return reader.hasNext() && !reader.pastSync(end);
+ }
+
+ @Override
+ public void close() throws IOException {
+ reader.close();
+ }
+ }
+
+ private class AvroBlockIterator implements Iterator<T> {
+
+ private long numRecordsRemaining;
+ private final DataFileReader<A> reader;
+ private final A reuse;
+
+ private AvroBlockIterator(long numRecordsRemaining, DataFileReader<A> reader, A reuse) {
+ this.numRecordsRemaining = numRecordsRemaining;
+ this.reader = reader;
+ this.reuse = reuse;
+ }
+
+ @Override
+ public boolean hasNext() {
+ return numRecordsRemaining > 0;
+ }
+
+ @Override
+ public T next() {
+ try {
+ numRecordsRemaining--;
+ // reader.next merely deserialize bytes in memory to java objects
+ // and will not read from file
+ return convert(reader.next(reuse));
+ } catch (IOException e) {
+ throw new RuntimeException(
+ "Encountered exception when reading from avro format file", e);
+ }
+ }
+ }
+}
diff --git a/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroFileFormatFactory.java b/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroFileFormatFactory.java
index 7b9d8c0..630d990 100644
--- a/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroFileFormatFactory.java
+++ b/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroFileFormatFactory.java
@@ -20,14 +20,21 @@ package org.apache.flink.formats.avro;
import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.serialization.BulkWriter;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.configuration.ConfigOption;
import org.apache.flink.configuration.ReadableConfig;
+import org.apache.flink.connector.file.src.FileSourceSplit;
+import org.apache.flink.connector.file.src.reader.BulkFormat;
+import org.apache.flink.connector.file.table.factories.BulkReaderFormatFactory;
import org.apache.flink.connector.file.table.factories.BulkWriterFormatFactory;
+import org.apache.flink.connector.file.table.format.BulkDecodingFormat;
import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.formats.avro.typeutils.AvroSchemaConverter;
import org.apache.flink.table.connector.ChangelogMode;
import org.apache.flink.table.connector.format.EncodingFormat;
import org.apache.flink.table.connector.sink.DynamicTableSink;
+import org.apache.flink.table.connector.source.DynamicTableSource;
+import org.apache.flink.table.data.GenericRowData;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.factories.DynamicTableFactory;
import org.apache.flink.table.types.DataType;
@@ -36,6 +43,7 @@ import org.apache.flink.table.types.logical.RowType;
import org.apache.avro.Schema;
import org.apache.avro.file.CodecFactory;
import org.apache.avro.file.DataFileWriter;
+import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericDatumWriter;
import org.apache.avro.generic.GenericRecord;
import org.apache.avro.io.DatumWriter;
@@ -49,11 +57,29 @@ import static org.apache.flink.formats.avro.AvroFormatOptions.AVRO_OUTPUT_CODEC;
/** Avro format factory for file system. */
@Internal
-public class AvroFileFormatFactory implements BulkWriterFormatFactory {
+public class AvroFileFormatFactory implements BulkReaderFormatFactory, BulkWriterFormatFactory {
public static final String IDENTIFIER = "avro";
@Override
+ public BulkDecodingFormat<RowData> createDecodingFormat(
+ DynamicTableFactory.Context context, ReadableConfig formatOptions) {
+ return new BulkDecodingFormat<RowData>() {
+ @Override
+ public BulkFormat<RowData, FileSourceSplit> createRuntimeDecoder(
+ DynamicTableSource.Context sourceContext, DataType producedDataType) {
+ return new AvroGenericRecordBulkFormat(
+ sourceContext, (RowType) producedDataType.getLogicalType().copy(false));
+ }
+
+ @Override
+ public ChangelogMode getChangelogMode() {
+ return ChangelogMode.insertOnly();
+ }
+ };
+ }
+
+ @Override
public EncodingFormat<BulkWriter.Factory<RowData>> createEncodingFormat(
DynamicTableFactory.Context context, ReadableConfig formatOptions) {
return new EncodingFormat<BulkWriter.Factory<RowData>>() {
@@ -90,6 +116,46 @@ public class AvroFileFormatFactory implements BulkWriterFormatFactory {
return options;
}
+ private static class AvroGenericRecordBulkFormat
+ extends AbstractAvroBulkFormat<GenericRecord, RowData, FileSourceSplit> {
+
+ private static final long serialVersionUID = 1L;
+
+ private final RowType producedRowType;
+ private final TypeInformation<RowData> producedTypeInfo;
+
+ private transient AvroToRowDataConverters.AvroToRowDataConverter converter;
+ private transient GenericRecord reusedAvroRecord;
+
+ public AvroGenericRecordBulkFormat(
+ DynamicTableSource.Context context, RowType producedRowType) {
+ super(AvroSchemaConverter.convertToSchema(producedRowType));
+ this.producedRowType = producedRowType;
+ this.producedTypeInfo = context.createTypeInformation(producedRowType);
+ }
+
+ @Override
+ protected void open(FileSourceSplit split) {
+ converter = AvroToRowDataConverters.createRowConverter(producedRowType);
+ reusedAvroRecord = new GenericData.Record(readerSchema);
+ }
+
+ @Override
+ protected RowData convert(GenericRecord record) {
+ return record == null ? null : (GenericRowData) converter.convert(record);
+ }
+
+ @Override
+ protected GenericRecord createReusedAvroRecord() {
+ return reusedAvroRecord;
+ }
+
+ @Override
+ public TypeInformation<RowData> getProducedType() {
+ return producedTypeInfo;
+ }
+ }
+
/**
* A {@link BulkWriter.Factory} to convert {@link RowData} to {@link GenericRecord} and wrap
* {@link AvroWriterFactory}.
diff --git a/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroFileSystemFormatFactory.java b/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroFileSystemFormatFactory.java
deleted file mode 100644
index 205ae2f..0000000
--- a/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroFileSystemFormatFactory.java
+++ /dev/null
@@ -1,193 +0,0 @@
-/*
- * 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.flink.formats.avro;
-
-import org.apache.flink.annotation.Internal;
-import org.apache.flink.api.common.io.InputFormat;
-import org.apache.flink.configuration.ConfigOption;
-import org.apache.flink.connector.file.table.factories.FileSystemFormatFactory;
-import org.apache.flink.core.fs.FileInputSplit;
-import org.apache.flink.core.fs.Path;
-import org.apache.flink.formats.avro.typeutils.AvroSchemaConverter;
-import org.apache.flink.table.data.GenericRowData;
-import org.apache.flink.table.data.RowData;
-import org.apache.flink.table.types.DataType;
-import org.apache.flink.table.types.logical.RowType;
-import org.apache.flink.table.utils.PartitionPathUtils;
-
-import org.apache.avro.Schema;
-import org.apache.avro.generic.GenericData;
-import org.apache.avro.generic.GenericRecord;
-import org.apache.avro.generic.IndexedRecord;
-
-import java.io.IOException;
-import java.util.Arrays;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Set;
-import java.util.stream.Collectors;
-
-import static org.apache.flink.formats.avro.AvroFormatOptions.AVRO_OUTPUT_CODEC;
-
-/** Avro format factory for file system. */
-@Internal
-public class AvroFileSystemFormatFactory implements FileSystemFormatFactory {
-
- public static final String IDENTIFIER = "avro";
-
- @Override
- public String factoryIdentifier() {
- return IDENTIFIER;
- }
-
- @Override
- public Set<ConfigOption<?>> requiredOptions() {
- return new HashSet<>();
- }
-
- @Override
- public Set<ConfigOption<?>> optionalOptions() {
- Set<ConfigOption<?>> options = new HashSet<>();
- options.add(AVRO_OUTPUT_CODEC);
- return options;
- }
-
- @Override
- public InputFormat<RowData, ?> createReader(ReaderContext context) {
- String[] fieldNames = context.getSchema().getFieldNames();
- List<String> projectFields =
- Arrays.stream(context.getProjectFields())
- .mapToObj(idx -> fieldNames[idx])
- .collect(Collectors.toList());
- List<String> csvFields =
- Arrays.stream(fieldNames)
- .filter(field -> !context.getPartitionKeys().contains(field))
- .collect(Collectors.toList());
-
- int[] selectFieldToProjectField =
- context.getFormatProjectFields().stream()
- .mapToInt(projectFields::indexOf)
- .toArray();
- int[] selectFieldToFormatField =
- context.getFormatProjectFields().stream().mapToInt(csvFields::indexOf).toArray();
-
- //noinspection unchecked
- return new RowDataAvroInputFormat(
- context.getPaths(),
- context.getFormatRowType(),
- context.getSchema().getFieldDataTypes(),
- context.getSchema().getFieldNames(),
- context.getProjectFields(),
- context.getPartitionKeys(),
- context.getDefaultPartName(),
- context.getPushedDownLimit(),
- selectFieldToProjectField,
- selectFieldToFormatField);
- }
-
- /**
- * InputFormat that reads avro record into {@link RowData}.
- *
- * <p>This extends {@link AvroInputFormat}, but {@link RowData} is not a avro record, so now we
- * remove generic information. TODO {@link AvroInputFormat} support type conversion.
- */
- private static class RowDataAvroInputFormat extends AvroInputFormat {
-
- private static final long serialVersionUID = 1L;
-
- private final DataType[] fieldTypes;
- private final String[] fieldNames;
- private final int[] selectFields;
- private final List<String> partitionKeys;
- private final String defaultPartValue;
- private final long limit;
- private final int[] selectFieldToProjectField;
- private final int[] selectFieldToFormatField;
- private final RowType formatRowType;
-
- private transient long emitted;
- // reuse object for per record
- private transient GenericRowData rowData;
- private transient IndexedRecord record;
- private transient AvroToRowDataConverters.AvroToRowDataConverter converter;
-
- public RowDataAvroInputFormat(
- Path[] filePaths,
- RowType formatRowType,
- DataType[] fieldTypes,
- String[] fieldNames,
- int[] selectFields,
- List<String> partitionKeys,
- String defaultPartValue,
- long limit,
- int[] selectFieldToProjectField,
- int[] selectFieldToFormatField) {
- super(filePaths[0], GenericRecord.class);
- super.setFilePaths(filePaths);
- this.formatRowType = formatRowType;
- this.fieldTypes = fieldTypes;
- this.fieldNames = fieldNames;
- this.partitionKeys = partitionKeys;
- this.defaultPartValue = defaultPartValue;
- this.selectFields = selectFields;
- this.limit = limit;
- this.emitted = 0;
- this.selectFieldToProjectField = selectFieldToProjectField;
- this.selectFieldToFormatField = selectFieldToFormatField;
- }
-
- @Override
- public void open(FileInputSplit split) throws IOException {
- super.open(split);
- Schema schema = AvroSchemaConverter.convertToSchema(formatRowType);
- record = new GenericData.Record(schema);
- rowData =
- PartitionPathUtils.fillPartitionValueForRecord(
- fieldNames,
- fieldTypes,
- selectFields,
- partitionKeys,
- currentSplit.getPath(),
- defaultPartValue);
- this.converter = AvroToRowDataConverters.createRowConverter(formatRowType);
- }
-
- @Override
- public boolean reachedEnd() throws IOException {
- return emitted >= limit || super.reachedEnd();
- }
-
- @Override
- public Object nextRecord(Object reuse) throws IOException {
- @SuppressWarnings("unchecked")
- IndexedRecord r = (IndexedRecord) super.nextRecord(record);
- if (r == null) {
- return null;
- }
- GenericRowData row = (GenericRowData) converter.convert(r);
-
- for (int i = 0; i < selectFieldToFormatField.length; i++) {
- rowData.setField(
- selectFieldToProjectField[i], row.getField(selectFieldToFormatField[i]));
- }
- emitted++;
- return rowData;
- }
- }
-}
diff --git a/flink-formats/flink-avro/src/main/resources/META-INF/services/org.apache.flink.table.factories.Factory b/flink-formats/flink-avro/src/main/resources/META-INF/services/org.apache.flink.table.factories.Factory
index 252b311..40e19c5 100644
--- a/flink-formats/flink-avro/src/main/resources/META-INF/services/org.apache.flink.table.factories.Factory
+++ b/flink-formats/flink-avro/src/main/resources/META-INF/services/org.apache.flink.table.factories.Factory
@@ -15,4 +15,3 @@
org.apache.flink.formats.avro.AvroFormatFactory
org.apache.flink.formats.avro.AvroFileFormatFactory
-org.apache.flink.formats.avro.AvroFileSystemFormatFactory
diff --git a/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroBulkFormatITCase.java b/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroBulkFormatITCase.java
new file mode 100644
index 0000000..7d13d7f
--- /dev/null
+++ b/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroBulkFormatITCase.java
@@ -0,0 +1,231 @@
+/*
+ * 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.flink.formats.avro;
+
+import org.apache.flink.api.connector.source.Boundedness;
+import org.apache.flink.api.connector.source.Source;
+import org.apache.flink.connector.file.src.FileSource;
+import org.apache.flink.connectors.test.common.TestResource;
+import org.apache.flink.connectors.test.common.environment.ClusterControllable;
+import org.apache.flink.connectors.test.common.environment.MiniClusterTestEnvironment;
+import org.apache.flink.connectors.test.common.environment.TestEnvironment;
+import org.apache.flink.connectors.test.common.external.ExternalContext;
+import org.apache.flink.connectors.test.common.external.SourceSplitDataWriter;
+import org.apache.flink.connectors.test.common.junit.annotations.ExternalContextFactory;
+import org.apache.flink.connectors.test.common.junit.annotations.ExternalSystem;
+import org.apache.flink.connectors.test.common.junit.annotations.TestEnv;
+import org.apache.flink.connectors.test.common.testsuites.SourceTestSuiteBase;
+import org.apache.flink.formats.avro.typeutils.AvroSchemaConverter;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.StringData;
+import org.apache.flink.table.data.binary.BinaryRowData;
+import org.apache.flink.table.runtime.typeutils.RowDataSerializer;
+import org.apache.flink.util.FileUtils;
+
+import org.apache.avro.Schema;
+import org.apache.avro.file.DataFileWriter;
+import org.apache.avro.generic.GenericDatumWriter;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.avro.io.DatumWriter;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Random;
+
+import static org.apache.flink.formats.avro.AvroBulkFormatTestUtils.ROW_TYPE;
+
+/** IT cases for {@link AbstractAvroBulkFormat}. */
+public class AvroBulkFormatITCase extends SourceTestSuiteBase<RowData> {
+
+ private static final RowDataSerializer SERIALIZER = new RowDataSerializer(ROW_TYPE);
+
+ @SuppressWarnings("unused")
+ @TestEnv
+ MiniClusterTestEnvironment flink = new MiniClusterTestEnvironment();
+
+ @SuppressWarnings("unused")
+ @ExternalSystem
+ EmptyExternalSystem externalSystem = new EmptyExternalSystem();
+
+ @SuppressWarnings("unused")
+ @ExternalContextFactory
+ AvroBulkFormatExternalContext.Factory oneBlockPerFile =
+ new AvroBulkFormatExternalContext.Factory(1);
+
+ @SuppressWarnings("unused")
+ @ExternalContextFactory
+ AvroBulkFormatExternalContext.Factory twoBlocksPerFile =
+ new AvroBulkFormatExternalContext.Factory(2);
+
+ @Override
+ public void testTaskManagerFailure(
+ TestEnvironment testEnv,
+ ExternalContext<RowData> externalContext,
+ ClusterControllable controller)
+ throws Exception {
+ // this test needs an unbounded source so we currently ignore this,
+ // add this test back once an unbounded file source is added
+ }
+
+ /** An empty {@link TestResource} which does nothing. */
+ public static class EmptyExternalSystem implements TestResource {
+
+ @Override
+ public void startUp() throws Exception {}
+
+ @Override
+ public void tearDown() throws Exception {}
+ }
+
+ /**
+ * {@link ExternalContext} for the IT case. It prepares avro test files and creates avro
+ * sources.
+ */
+ public static class AvroBulkFormatExternalContext implements ExternalContext<RowData> {
+
+ private final Path tmpDir;
+ private final int blocksPerFile;
+ private int index;
+
+ private AvroBulkFormatExternalContext(int blocksPerFile) {
+ try {
+ this.tmpDir = Files.createTempDirectory("avro-bulk-format-it-case");
+ } catch (IOException e) {
+ throw new RuntimeException(
+ "Encountered exception when creating temp directory for tests", e);
+ }
+ this.blocksPerFile = blocksPerFile;
+ this.index = 0;
+ }
+
+ @Override
+ public Source<RowData, ?, ?> createSource(Boundedness boundedness) {
+ AvroBulkFormatTestUtils.TestingAvroBulkFormat format =
+ new AvroBulkFormatTestUtils.TestingAvroBulkFormat();
+ return FileSource.forBulkFileFormat(
+ format, org.apache.flink.core.fs.Path.fromLocalFile(tmpDir.toFile()))
+ .build();
+ }
+
+ @Override
+ public SourceSplitDataWriter<RowData> createSourceSplitDataWriter() {
+ File file = Paths.get(tmpDir.toString(), String.valueOf(index)).toFile();
+ AvroBulkFormatSourceSplitDataWriter writer;
+ try {
+ file.createNewFile();
+ writer = new AvroBulkFormatSourceSplitDataWriter(new FileOutputStream(file));
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ index++;
+ return writer;
+ }
+
+ @Override
+ public List<RowData> generateTestData(int splitIndex, long seed) {
+ Random random = new Random(seed);
+ List<RowData> data = new ArrayList<>();
+ for (int i = 0; i < blocksPerFile; i++) {
+ data.add(getBinaryRow(randomString(4, random), randomString(8, random)));
+ data.add(getBinaryRow(randomString(16, random), null));
+ data.add(getBinaryRow(randomString(32, random), randomString(1024, random)));
+ }
+ return data;
+ }
+
+ @Override
+ public void close() throws Exception {
+ FileUtils.deleteDirectory(tmpDir.toFile());
+ }
+
+ private StringData randomString(int len, Random random) {
+ StringBuilder builder = new StringBuilder();
+ for (int i = 0; i < len; i++) {
+ builder.append((char) (random.nextInt(26) + 'a'));
+ }
+ return StringData.fromString(builder.toString());
+ }
+
+ private RowData getBinaryRow(Object... values) {
+ GenericRowData genericRowData = new GenericRowData(values.length);
+ for (int i = 0; i < values.length; i++) {
+ genericRowData.setField(i, values[i]);
+ }
+ BinaryRowData binaryRowData = SERIALIZER.toBinaryRow(genericRowData);
+ return SERIALIZER.copy(binaryRowData);
+ }
+
+ /** Factory to create {@link AvroBulkFormatExternalContext}. */
+ public static class Factory implements ExternalContext.Factory<RowData> {
+
+ private final int blocksPerFile;
+
+ public Factory(int blocksPerFile) {
+ this.blocksPerFile = blocksPerFile;
+ }
+
+ @Override
+ public ExternalContext<RowData> createExternalContext() {
+ return new AvroBulkFormatExternalContext(blocksPerFile);
+ }
+ }
+ }
+
+ private static class AvroBulkFormatSourceSplitDataWriter
+ implements SourceSplitDataWriter<RowData> {
+
+ private final Schema schema;
+ private final RowDataToAvroConverters.RowDataToAvroConverter converter;
+ private final DataFileWriter<GenericRecord> dataFileWriter;
+
+ private AvroBulkFormatSourceSplitDataWriter(FileOutputStream out) throws IOException {
+ schema = AvroSchemaConverter.convertToSchema(ROW_TYPE);
+ converter = RowDataToAvroConverters.createConverter(ROW_TYPE);
+
+ DatumWriter<GenericRecord> datumWriter = new GenericDatumWriter<>(schema);
+ dataFileWriter = new DataFileWriter<>(datumWriter);
+ dataFileWriter.create(schema, out);
+ dataFileWriter.setSyncInterval(128);
+ }
+
+ @Override
+ public void writeRecords(Collection<RowData> records) {
+ for (RowData rowData : records) {
+ try {
+ dataFileWriter.append((GenericRecord) converter.convert(schema, rowData));
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ }
+
+ @Override
+ public void close() throws Exception {
+ dataFileWriter.close();
+ }
+ }
+}
diff --git a/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroBulkFormatTest.java b/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroBulkFormatTest.java
new file mode 100644
index 0000000..9080efa
--- /dev/null
+++ b/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroBulkFormatTest.java
@@ -0,0 +1,275 @@
+/*
+ * 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.flink.formats.avro;
+
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.connector.file.src.FileSourceSplit;
+import org.apache.flink.connector.file.src.reader.BulkFormat;
+import org.apache.flink.connector.file.src.util.CheckpointedPosition;
+import org.apache.flink.connector.file.src.util.RecordAndPosition;
+import org.apache.flink.core.fs.Path;
+import org.apache.flink.formats.avro.typeutils.AvroSchemaConverter;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.StringData;
+import org.apache.flink.util.FileUtils;
+import org.apache.flink.util.StringUtils;
+
+import org.apache.avro.Schema;
+import org.apache.avro.file.DataFileConstants;
+import org.apache.avro.file.DataFileWriter;
+import org.apache.avro.generic.GenericDatumWriter;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.avro.io.DatumWriter;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.UUID;
+
+import static org.apache.flink.formats.avro.AvroBulkFormatTestUtils.ROW_TYPE;
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link AbstractAvroBulkFormat}. */
+public class AvroBulkFormatTest {
+
+ private static final List<RowData> TEST_DATA =
+ Arrays.asList(
+ // -------- batch 0, block start 186 --------
+ GenericRowData.of(
+ StringData.fromString("AvroBulk"), StringData.fromString("FormatTest")),
+ GenericRowData.of(
+ StringData.fromString("Apache"), StringData.fromString("Flink")),
+ GenericRowData.of(
+ StringData.fromString(
+ "永和九年，岁在癸丑，暮春之初，会于会稽山阴之兰亭，修禊事也。群贤毕至，少"
+ + "长咸集。此地有崇山峻岭，茂林修竹，又有清流激湍，映带左右。引"
+ + "以为流觞曲水，列坐其次。虽无丝竹管弦之盛，一觞一咏，亦足以畅"
+ + "叙幽情。"),
+ StringData.fromString("")),
+ // -------- batch 1, block start 547 --------
+ GenericRowData.of(
+ StringData.fromString("File"), StringData.fromString("Format")),
+ GenericRowData.of(
+ null,
+ StringData.fromString(
+ "This is a string with English, 中文 and even 🍎🍌🍑🥝🍍🥭🍐")),
+ // -------- batch 2, block start 659 --------
+ GenericRowData.of(
+ StringData.fromString("block with"),
+ StringData.fromString("only one record"))
+ // -------- file length 706 --------
+ );
+ private static final List<Integer> BLOCK_STARTS = Arrays.asList(186, 547, 659);
+
+ private File tmpFile;
+
+ @BeforeEach
+ public void before() throws IOException {
+ tmpFile = Files.createTempFile("avro-bulk-format-test", ".avro").toFile();
+ tmpFile.createNewFile();
+ FileOutputStream out = new FileOutputStream(tmpFile);
+
+ Schema schema = AvroSchemaConverter.convertToSchema(ROW_TYPE);
+ RowDataToAvroConverters.RowDataToAvroConverter converter =
+ RowDataToAvroConverters.createConverter(ROW_TYPE);
+
+ DatumWriter<GenericRecord> datumWriter = new GenericDatumWriter<>(schema);
+ DataFileWriter<GenericRecord> dataFileWriter = new DataFileWriter<>(datumWriter);
+ dataFileWriter.create(schema, out);
+ dataFileWriter.setSyncInterval(64);
+
+ for (RowData rowData : TEST_DATA) {
+ dataFileWriter.append((GenericRecord) converter.convert(schema, rowData));
+ }
+
+ dataFileWriter.close();
+ }
+
+ @AfterEach
+ public void after() throws IOException {
+ FileUtils.deleteFileOrDirectory(tmpFile);
+ }
+
+ @Test
+ public void testReadWholeFileWithOneSplit() throws IOException {
+ AvroBulkFormatTestUtils.TestingAvroBulkFormat bulkFormat =
+ new AvroBulkFormatTestUtils.TestingAvroBulkFormat();
+ assertSplit(
+ bulkFormat,
+ Collections.singletonList(
+ new SplitInfo(
+ 0,
+ tmpFile.length(),
+ Arrays.asList(
+ new BatchInfo(0, 3),
+ new BatchInfo(3, 5),
+ new BatchInfo(5, 6)))));
+ }
+
+ @Test
+ public void testReadWholeFileWithMultipleSplits() throws IOException {
+ AvroBulkFormatTestUtils.TestingAvroBulkFormat bulkFormat =
+ new AvroBulkFormatTestUtils.TestingAvroBulkFormat();
+ long splitLength = tmpFile.length() / 3;
+ assertSplit(
+ bulkFormat,
+ Arrays.asList(
+ new SplitInfo(
+ 0, splitLength, Collections.singletonList(new BatchInfo(0, 3))),
+ new SplitInfo(splitLength, splitLength * 2, Collections.emptyList()),
+ new SplitInfo(
+ splitLength * 2,
+ tmpFile.length(),
+ Arrays.asList(new BatchInfo(3, 5), new BatchInfo(5, 6)))));
+ }
+
+ @Test
+ public void testSplitsAtCriticalLocations() throws IOException {
+ AvroBulkFormatTestUtils.TestingAvroBulkFormat bulkFormat =
+ new AvroBulkFormatTestUtils.TestingAvroBulkFormat();
+ assertSplit(
+ bulkFormat,
+ Arrays.asList(
+ // ends just before the new block
+ new SplitInfo(
+ BLOCK_STARTS.get(0) - DataFileConstants.SYNC_SIZE,
+ BLOCK_STARTS.get(1) - DataFileConstants.SYNC_SIZE,
+ Collections.singletonList(new BatchInfo(0, 3))),
+ // ends just at the beginning of new block
+ new SplitInfo(
+ BLOCK_STARTS.get(1) - DataFileConstants.SYNC_SIZE,
+ BLOCK_STARTS.get(2) - DataFileConstants.SYNC_SIZE + 1,
+ Arrays.asList(new BatchInfo(3, 5), new BatchInfo(5, 6)))));
+ }
+
+ @Test
+ public void testRestoreReader() throws IOException {
+ AvroBulkFormatTestUtils.TestingAvroBulkFormat bulkFormat =
+ new AvroBulkFormatTestUtils.TestingAvroBulkFormat();
+ long splitLength = tmpFile.length() / 3;
+ String splitId = UUID.randomUUID().toString();
+
+ FileSourceSplit split =
+ new FileSourceSplit(
+ splitId, new Path(tmpFile.toString()), splitLength * 2, tmpFile.length());
+ BulkFormat.Reader<RowData> reader = bulkFormat.createReader(new Configuration(), split);
+ long offset1 = assertBatch(reader, new BatchInfo(3, 5));
+ assertBatch(reader, new BatchInfo(5, 6));
+ assertThat(reader.readBatch()).isNull();
+ reader.close();
+
+ split =
+ new FileSourceSplit(
+ splitId,
+ new Path(tmpFile.toString()),
+ splitLength * 2,
+ tmpFile.length(),
+ StringUtils.EMPTY_STRING_ARRAY,
+ new CheckpointedPosition(offset1, 1));
+ reader = bulkFormat.restoreReader(new Configuration(), split);
+ long offset2 = assertBatch(reader, new BatchInfo(3, 5), 1);
+ assertBatch(reader, new BatchInfo(5, 6));
+ assertThat(reader.readBatch()).isNull();
+ reader.close();
+
+ assertThat(offset2).isEqualTo(offset1);
+ }
+
+ private void assertSplit(
+ AvroBulkFormatTestUtils.TestingAvroBulkFormat bulkFormat, List<SplitInfo> splitInfos)
+ throws IOException {
+ for (SplitInfo splitInfo : splitInfos) {
+ FileSourceSplit split =
+ new FileSourceSplit(
+ UUID.randomUUID().toString(),
+ new Path(tmpFile.toString()),
+ splitInfo.start,
+ splitInfo.end - splitInfo.start);
+ BulkFormat.Reader<RowData> reader = bulkFormat.createReader(new Configuration(), split);
+ List<Long> offsets = new ArrayList<>();
+ for (BatchInfo batch : splitInfo.batches) {
+ offsets.add(assertBatch(reader, batch));
+ }
+ assertThat(reader.readBatch()).isNull();
+ for (int j = 1; j < offsets.size(); j++) {
+ assertThat(offsets.get(j - 1) < offsets.get(j)).isTrue();
+ }
+ reader.close();
+ }
+ }
+
+ private long assertBatch(BulkFormat.Reader<RowData> reader, BatchInfo batchInfo)
+ throws IOException {
+ return assertBatch(reader, batchInfo, 0);
+ }
+
+ private long assertBatch(
+ BulkFormat.Reader<RowData> reader, BatchInfo batchInfo, int initialSkipCount)
+ throws IOException {
+ long ret = -1;
+ int skipCount = initialSkipCount;
+ BulkFormat.RecordIterator<RowData> iterator = reader.readBatch();
+ for (RecordAndPosition<RowData> recordAndPos = iterator.next();
+ recordAndPos != null;
+ recordAndPos = iterator.next()) {
+ if (ret == -1) {
+ ret = recordAndPos.getOffset();
+ }
+ assertThat(recordAndPos.getRecord())
+ .isEqualTo(TEST_DATA.get(batchInfo.start + skipCount));
+ assertThat(recordAndPos.getOffset()).isEqualTo(ret);
+ skipCount++;
+ assertThat(recordAndPos.getRecordSkipCount()).isEqualTo(skipCount);
+ }
+ assertThat(skipCount).isEqualTo(batchInfo.end - batchInfo.start);
+ iterator.releaseBatch();
+ return ret;
+ }
+
+ private static class SplitInfo {
+ private final long start;
+ private final long end;
+ private final List<BatchInfo> batches;
+
+ private SplitInfo(long start, long end, List<BatchInfo> batches) {
+ this.start = start;
+ this.end = end;
+ this.batches = batches;
+ }
+ }
+
+ private static class BatchInfo {
+ private final int start;
+ private final int end;
+
+ private BatchInfo(int start, int end) {
+ this.start = start;
+ this.end = end;
+ }
+ }
+}
diff --git a/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroBulkFormatTestUtils.java b/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroBulkFormatTestUtils.java
new file mode 100644
index 0000000..a688ff7
--- /dev/null
+++ b/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroBulkFormatTestUtils.java
@@ -0,0 +1,76 @@
+/*
+ * 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.flink.formats.avro;
+
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.connector.file.src.FileSourceSplit;
+import org.apache.flink.formats.avro.typeutils.AvroSchemaConverter;
+import org.apache.flink.table.api.DataTypes;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.runtime.typeutils.InternalTypeInfo;
+import org.apache.flink.table.types.logical.LogicalType;
+import org.apache.flink.table.types.logical.RowType;
+
+import org.apache.avro.generic.GenericData;
+import org.apache.avro.generic.GenericRecord;
+
+/** Testing utils for tests related to {@link AbstractAvroBulkFormat}. */
+public class AvroBulkFormatTestUtils {
+
+ public static final RowType ROW_TYPE =
+ RowType.of(
+ false,
+ new LogicalType[] {
+ DataTypes.STRING().getLogicalType(), DataTypes.STRING().getLogicalType()
+ },
+ new String[] {"a", "b"});
+
+ /** {@link AbstractAvroBulkFormat} for tests. */
+ public static class TestingAvroBulkFormat
+ extends AbstractAvroBulkFormat<GenericRecord, RowData, FileSourceSplit> {
+
+ private transient GenericRecord reusedAvroRecord;
+ private transient AvroToRowDataConverters.AvroToRowDataConverter converter;
+
+ protected TestingAvroBulkFormat() {
+ super(AvroSchemaConverter.convertToSchema(ROW_TYPE));
+ }
+
+ @Override
+ protected void open(FileSourceSplit split) {
+ reusedAvroRecord = new GenericData.Record(readerSchema);
+ converter = AvroToRowDataConverters.createRowConverter(ROW_TYPE);
+ }
+
+ @Override
+ protected RowData convert(GenericRecord record) {
+ return record == null ? null : (RowData) converter.convert(record);
+ }
+
+ @Override
+ protected GenericRecord createReusedAvroRecord() {
+ return reusedAvroRecord;
+ }
+
+ @Override
+ public TypeInformation<RowData> getProducedType() {
+ return InternalTypeInfo.of(ROW_TYPE);
+ }
+ }
+}
diff --git a/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroFilesystemITCase.java b/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroFilesystemITCase.java
index b89c2be..f240024 100644
--- a/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroFilesystemITCase.java
+++ b/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroFilesystemITCase.java
@@ -30,7 +30,7 @@ import java.util.List;
import static org.apache.avro.file.DataFileConstants.NULL_CODEC;
-/** ITCase to test avro format for {@link AvroFileSystemFormatFactory} in batch mode. */
+/** ITCase to test avro format for {@link AvroFileFormatFactory} in batch mode. */
@RunWith(Parameterized.class)
public class AvroFilesystemITCase extends BatchFileSystemITCaseBase {
@@ -46,11 +46,6 @@ public class AvroFilesystemITCase extends BatchFileSystemITCaseBase {
}
@Override
- public boolean supportsReadingMetadata() {
- return false;
- }
-
- @Override
public String[] formatProperties() {
List<String> ret = new ArrayList<>();
ret.add("'format'='avro'");