[avro] Unwrap AvroRuntimeException and throw the real IOException (#4179)

Author: tsreaper

paimon-common/src/main/java/org/apache/paimon/utils/IteratorResultIterator.java

@@ -25,18 +25,18 @@
 
 import javax.annotation.Nullable;
 
-import java.util.Iterator;
+import java.io.IOException;
 
 /** A simple {@link RecordReader.RecordIterator} that returns the elements of an iterator. */
 public final class IteratorResultIterator extends RecyclableIterator<InternalRow>
         implements FileRecordIterator<InternalRow> {
 
-    private final Iterator<InternalRow> records;
+    private final IteratorWithException<InternalRow, IOException> records;
     private final Path filePath;
     private long nextFilePos;
 
     public IteratorResultIterator(
-            final Iterator<InternalRow> records,
+            final IteratorWithException<InternalRow, IOException> records,
             final @Nullable Runnable recycler,
             final Path filePath,
             long pos) {
@@ -48,7 +48,7 @@ public IteratorResultIterator(
 
     @Nullable
     @Override
-    public InternalRow next() {
+    public InternalRow next() throws IOException {
         if (records.hasNext()) {
             nextFilePos++;
             return records.next();

paimon-common/src/main/java/org/apache/paimon/utils/IteratorWithException.java

@@ -0,0 +1,27 @@
+/*
+ * 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.paimon.utils;
+
+/** An iterator which might throws exception. */
+public interface IteratorWithException<V, E extends Exception> {
+
+    boolean hasNext() throws E;
+
+    V next() throws E;
+}

paimon-format/src/main/java/org/apache/paimon/format/avro/AvroBulkFormat.java

@@ -26,16 +26,18 @@
 import org.apache.paimon.types.RowType;
 import org.apache.paimon.utils.IOUtils;
 import org.apache.paimon.utils.IteratorResultIterator;
+import org.apache.paimon.utils.IteratorWithException;
 import org.apache.paimon.utils.Pool;
 
+import org.apache.avro.AvroRuntimeException;
 import org.apache.avro.file.DataFileReader;
 import org.apache.avro.file.SeekableInput;
 import org.apache.avro.io.DatumReader;
 
 import javax.annotation.Nullable;
 
 import java.io.IOException;
-import java.util.Iterator;
+import java.util.function.Supplier;
 
 /** Provides a {@link FormatReaderFactory} for Avro records. */
 public class AvroBulkFormat implements FormatReaderFactory {
@@ -105,15 +107,16 @@ public RecordIterator<InternalRow> readBatch() throws IOException {
 
             long rowPosition = currentRowPosition;
             currentRowPosition += reader.getBlockCount();
-            Iterator<InternalRow> iterator = new AvroBlockIterator(reader.getBlockCount(), reader);
+            IteratorWithException<InternalRow, IOException> iterator =
+                    new AvroBlockIterator(reader.getBlockCount(), reader);
             return new IteratorResultIterator(
                     iterator, () -> pool.recycler().recycle(ticket), filePath, rowPosition);
         }
 
         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);
+            return replaceAvroRuntimeException(reader::hasNext) && !reader.pastSync(end);
         }
 
         @Override
@@ -122,7 +125,8 @@ public void close() throws IOException {
         }
     }
 
-    private static class AvroBlockIterator implements Iterator<InternalRow> {
+    private static class AvroBlockIterator
+            implements IteratorWithException<InternalRow, IOException> {
 
         private long numRecordsRemaining;
         private final DataFileReader<InternalRow> reader;
@@ -138,17 +142,23 @@ public boolean hasNext() {
         }
 
         @Override
-        public InternalRow next() {
-            try {
-                numRecordsRemaining--;
-                // reader.next merely deserialize bytes in memory to java objects
-                // and will not read from file
-                // Do not reuse object, manifest file assumes no object reuse
-                return reader.next(null);
-            } catch (IOException e) {
-                throw new RuntimeException(
-                        "Encountered exception when reading from avro format file", e);
+        public InternalRow next() throws IOException {
+            numRecordsRemaining--;
+            // reader.next merely deserialize bytes in memory to java objects
+            // and will not read from file
+            // Do not reuse object, manifest file assumes no object reuse
+            return replaceAvroRuntimeException(reader::next);
+        }
+    }
+
+    private static <T> T replaceAvroRuntimeException(Supplier<T> supplier) throws IOException {
+        try {
+            return supplier.get();
+        } catch (AvroRuntimeException e) {
+            if (e.getCause() != null && e.getCause() instanceof IOException) {
+                throw (IOException) e.getCause();
             }
+            throw e;
         }
     }
 }

paimon-format/src/test/java/org/apache/paimon/format/avro/AvroFileFormatTest.java

@@ -24,8 +24,10 @@
 import org.apache.paimon.format.FileFormatFactory.FormatContext;
 import org.apache.paimon.format.FormatReaderContext;
 import org.apache.paimon.format.FormatWriter;
+import org.apache.paimon.fs.FileIO;
 import org.apache.paimon.fs.Path;
 import org.apache.paimon.fs.PositionOutputStream;
+import org.apache.paimon.fs.SeekableInputStream;
 import org.apache.paimon.fs.local.LocalFileIO;
 import org.apache.paimon.options.Options;
 import org.apache.paimon.reader.RecordReader;
@@ -37,11 +39,15 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
 
+import java.io.File;
+import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.UUID;
+import java.util.concurrent.ThreadLocalRandom;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /** Test for avro file format. */
 public class AvroFileFormatTest {
@@ -122,10 +128,74 @@ void testReadRowPosition() throws IOException {
         try (RecordReader<InternalRow> reader =
                 format.createReaderFactory(rowType)
                         .createReader(
-                                new FormatReaderContext(
-                                        fileIO, file, fileIO.getFileSize(file))); ) {
+                                new FormatReaderContext(fileIO, file, fileIO.getFileSize(file)))) {
             reader.forEachRemainingWithPosition(
                     (rowPosition, row) -> assertThat(row.getInt(0) == rowPosition).isTrue());
         }
     }
+
+    @Test
+    void testGetRealIOException() throws IOException {
+        RowType rowType = DataTypes.ROW(DataTypes.INT().notNull());
+        FileFormat format = new AvroFileFormat(new FormatContext(new Options(), 16, 16));
+
+        LocalFileIO localFileIO = LocalFileIO.create();
+        Path file = new Path(new Path(tempPath.toUri()), UUID.randomUUID().toString());
+        try (PositionOutputStream out = localFileIO.newOutputStream(file, false)) {
+            FormatWriter writer = format.createWriterFactory(rowType).create(out, "zstd");
+            ThreadLocalRandom random = ThreadLocalRandom.current();
+            // magic number tested by hand
+            for (int i = 0; i < 100000; i++) {
+                writer.addElement(GenericRow.of(random.nextInt()));
+            }
+            writer.close();
+        }
+
+        FileIO failingFileIO =
+                new LocalFileIO() {
+
+                    @Override
+                    public SeekableInputStream newInputStream(Path path) throws IOException {
+                        return new FailingInputStream(toFile(path));
+                    }
+
+                    class FailingInputStream extends LocalFileIO.LocalSeekableInputStream {
+
+                        private int cnt;
+
+                        public FailingInputStream(File file) throws FileNotFoundException {
+                            super(file);
+                            cnt = 0;
+                        }
+
+                        @Override
+                        public int read() throws IOException {
+                            checkException();
+                            return super.read();
+                        }
+
+                        @Override
+                        public int read(byte[] b, int off, int len) throws IOException {
+                            checkException();
+                            return super.read(b, off, len);
+                        }
+
+                        private void checkException() throws IOException {
+                            cnt++;
+                            // magic number tested by hand
+                            if (cnt == 200) {
+                                throw new IOException("Artificial exception");
+                            }
+                        }
+                    }
+                };
+        RecordReader<InternalRow> reader =
+                format.createReaderFactory(rowType)
+                        .createReader(
+                                new FormatReaderContext(
+                                        failingFileIO, file, failingFileIO.getFileSize(file)));
+        assertThatThrownBy(() -> reader.forEachRemaining(row -> {}))
+                .isInstanceOf(IOException.class)
+                .hasMessageContaining("Artificial exception");
+    }
 }