From 7d6c28112004c27a3628f73471f06689af59ce75 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Thu, 13 Aug 2026 19:12:15 -0700 Subject: [PATCH] GH-3719: Fix vectored read allocation limits and fallback safety --- .../parquet/hadoop/ParquetFileReader.java | 192 ++-- .../parquet/hadoop/TestDataPageChecksums.java | 89 ++ .../TestParquetFileReaderVectoredIO.java | 881 ++++++++++++++++++ 3 files changed, 1094 insertions(+), 68 deletions(-) create mode 100644 parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileReaderVectoredIO.java diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileReader.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileReader.java index 9af4b4ac60..52e639bf81 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileReader.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileReader.java @@ -36,6 +36,7 @@ import java.io.Closeable; import java.io.IOException; import java.io.InputStream; +import java.io.InterruptedIOException; import java.io.SequenceInputStream; import java.nio.ByteBuffer; import java.util.ArrayList; @@ -1293,14 +1294,14 @@ public ColumnChunkPageReadStore readFilteredRowGroup( private void readAllPartsVectoredOrNormal(List allParts, ChunkListBuilder builder) throws IOException { - if (shouldUseVectoredIo(allParts)) { + if (shouldUseVectoredIo()) { try { readVectored(allParts, builder); return; } catch (IllegalArgumentException | UnsupportedOperationException e) { - // Either the arguments are wrong or somehow this is being invoked against - // a hadoop release which doesn't have the API and yet somehow it got here. - LOG.warn("readVectored() failed; falling back to normal IO against {}", f, e); + // At this point only range preparation can have failed; exceptions from the + // vectored call itself are wrapped below because reads may already be active. + LOG.warn("Preparing vectored reads failed; falling back to normal IO against {}", f, e); } } for (ConsecutivePartList consecutiveChunks : allParts) { @@ -1315,37 +1316,14 @@ private void readAllPartsVectoredOrNormal(List allParts, Ch *
    *
  1. The option is enabled
    2. *
  2. The Hadoop version supports vectored IO
    3. - *
  3. The part lengths are all valid for vectored IO
    4. *
  4. The stream implementation explicitly supports the API; for other streams the classic * API is always used.
    5. *
  5. The allocator is not direct. This is to avoid HADOOP-19101 surfacing. *
- * @param allParts all parts to read. * @return true or false. */ - private boolean shouldUseVectoredIo(final List allParts) { - return options.useHadoopVectoredIo() - && f.readVectoredAvailable(options.getAllocator()) - && arePartsValidForVectoredIo(allParts); - } - - /** - * Validate the parts for vectored IO. - * Vectored IO doesn't support reading ranges of size greater than - * Integer.MAX_VALUE. - * @param allParts all parts to read. - * @return true or false. - */ - private boolean arePartsValidForVectoredIo(List allParts) { - for (ConsecutivePartList consecutivePart : allParts) { - if (consecutivePart.length >= Integer.MAX_VALUE) { - LOG.debug( - "Part length {} greater than Integer.MAX_VALUE thus disabling vectored IO", - consecutivePart.length); - return false; - } - } - return true; + private boolean shouldUseVectoredIo() { + return options.useHadoopVectoredIo() && f.readVectoredAvailable(options.getAllocator()); } /** @@ -1357,32 +1335,100 @@ private boolean arePartsValidForVectoredIo(List allParts) { * If directly implemented by a Filesystem then it is likely to be a more efficient * operation such as a scatter-gather read (native IO) or set of parallel * GET requests against an object store. + * The allocation limit applies to filesystem buffers; decoders can still require a + * contiguous buffer for an individual logical value larger than that limit. * @param allParts all parts to be read. * @param builder used to build chunk list to read the pages for the different columns. - * @throws IOException any IOE. - * @throws IllegalArgumentException arguments are invalid. - * @throws UnsupportedOperationException if the filesystem does not support vectored IO. + * @throws IOException if submitting or consuming the vectored reads fails. + * @throws IllegalArgumentException if range preparation fails before any reads are submitted. */ private void readVectored(List allParts, ChunkListBuilder builder) throws IOException { - + final int maximumAllocation = options.getMaxAllocationSize(); + Preconditions.checkArgument(maximumAllocation > 0, "Invalid maximum allocation size %s", maximumAllocation); + final long fileLength = file.getLength(); List ranges = new ArrayList<>(allParts.size()); + List partRangeCounts = new ArrayList<>(allParts.size()); long totalSize = 0; for (ConsecutivePartList consecutiveChunks : allParts) { final long len = consecutiveChunks.length; - Preconditions.checkArgument( - len < Integer.MAX_VALUE, - "Invalid length %s for vectored read operation. It must be less than max integer value.", - len); - ranges.add(new ParquetFileRange(consecutiveChunks.offset, (int) len)); + final long start = consecutiveChunks.offset; + if (start < 0 || len < 0 || start > fileLength || len > fileLength - start) { + throw new IOException(String.format( + "Invalid vectored read range (offset %d, length %d) for file length %d", + start, len, fileLength)); + } + final int firstRange = ranges.size(); + long remaining = len; + long offset = start; + do { + int rangeLength = (int) Math.min(remaining, maximumAllocation); + ranges.add(new ParquetFileRange(offset, rangeLength)); + offset += rangeLength; + remaining -= rangeLength; + } while (remaining > 0); + partRangeCounts.add(ranges.size() - firstRange); totalSize += len; } LOG.debug("Reading {} bytes of data with vectored IO in {} ranges", totalSize, ranges.size()); - // Request a vectored read; - f.readVectored(ranges, options.getAllocator()); - int k = 0; - for (ConsecutivePartList consecutivePart : allParts) { - ParquetFileRange currRange = ranges.get(k++); - consecutivePart.readFromVectoredRange(currRange, builder); + final long readStart = System.nanoTime(); + try { + // Even a synchronous failure can occur after some reads have been scheduled, + // so falling back to normal IO is unsafe once this call has been entered. + f.readVectored(ranges, options.getAllocator()); + int firstRange = 0; + for (int partIndex = 0; partIndex < allParts.size(); partIndex++) { + int endRange = firstRange + partRangeCounts.get(partIndex); + allParts.get(partIndex).readFromVectoredRanges(ranges.subList(firstRange, endRange), builder); + firstRange = endRange; + } + } catch (IllegalArgumentException | UnsupportedOperationException e) { + IOException failure = + new IOException("Vectored read failed after asynchronous reads may have been submitted", e); + awaitRemainingVectoredReads(ranges, readStart, failure); + throw failure; + } catch (IOException | RuntimeException e) { + awaitRemainingVectoredReads(ranges, readStart, e); + throw e; + } + } + + /** + * Wait for submitted reads with published futures to finish before their stream can be + * closed. Cancelling result futures does not stop all Hadoop backends from continuing IO. + */ + private void awaitRemainingVectoredReads(List ranges, long readStart, Throwable failure) { + if (Thread.currentThread().isInterrupted() + || failure instanceof InterruptedIOException && failure.getCause() instanceof InterruptedException) { + return; + } + + final long timeoutNanos = TimeUnit.SECONDS.toNanos(HADOOP_VECTORED_READ_TIMEOUT_SECONDS); + for (ParquetFileRange range : ranges) { + Future future = range.getDataReadFuture(); + if (future == null || future.isDone()) { + continue; + } + + long remainingNanos = Math.max(timeoutNanos - (System.nanoTime() - readStart), 0L); + try { + FutureIO.awaitFuture(future, remainingNanos, TimeUnit.NANOSECONDS); + } catch (InterruptedIOException e) { + if (failure != e) { + failure.addSuppressed(e); + } + if (e.getCause() instanceof InterruptedException) { + Thread.currentThread().interrupt(); + return; + } + } catch (TimeoutException e) { + failure.addSuppressed(e); + LOG.warn("Timed out waiting for vectored read {} after another read failed", range, e); + return; + } catch (IOException | RuntimeException e) { + if (failure != e) { + failure.addSuppressed(e); + } + } } } @@ -1951,10 +1997,12 @@ protected PageHeader readPageHeader(BlockCipher.Decryptor blockDecryptor, byte[] * Calculate checksum of input bytes, throw decoding exception if it does not match the provided * reference crc */ - private void verifyCrc(int referenceCrc, BytesInput bytes, String exceptionMsg) { + private void verifyCrc(int referenceCrc, String exceptionMsg, BytesInput... inputs) throws IOException { crc.reset(); - try (ByteBufferReleaser releaser = crcAllocator.getReleaser()) { - crc.update(bytes.toByteBuffer(releaser)); + for (BytesInput input : inputs) { + for (ByteBuffer buffer : input.toInputStream().remainingBuffers()) { + crc.update(buffer); + } } if (crc.getValue() != ((long) referenceCrc & 0xffffffffL)) { throw new ParquetDecodingException(exceptionMsg); @@ -2021,8 +2069,8 @@ public ColumnChunkPageReader readAllPages( if (options.usePageChecksumVerification() && pageHeader.isSetCrc()) { verifyCrc( pageHeader.getCrc(), - pageBytes, - "could not verify dictionary page integrity, CRC checksum verification failed"); + "could not verify dictionary page integrity, CRC checksum verification failed", + pageBytes); } DictionaryPageHeader dicHeader = pageHeader.getDictionary_page_header(); dictionaryPage = new DictionaryPage( @@ -2041,8 +2089,8 @@ public ColumnChunkPageReader readAllPages( if (options.usePageChecksumVerification() && pageHeader.isSetCrc()) { verifyCrc( pageHeader.getCrc(), - pageBytes, - "could not verify page integrity, CRC checksum verification failed"); + "could not verify page integrity, CRC checksum verification failed", + pageBytes); } DataPageV1 dataPageV1 = new DataPageV1( pageBytes, @@ -2072,11 +2120,12 @@ public ColumnChunkPageReader readAllPages( this.readAsBytesInput(dataHeaderV2.getDefinition_levels_byte_length()); final BytesInput values = this.readAsBytesInput(dataSize); if (options.usePageChecksumVerification() && pageHeader.isSetCrc()) { - pageBytes = BytesInput.concat(repetitionLevels, definitionLevels, values); verifyCrc( pageHeader.getCrc(), - pageBytes, - "could not verify page integrity, CRC checksum verification failed"); + "could not verify page integrity, CRC checksum verification failed", + repetitionLevels, + definitionLevels, + values); } DataPageV2 dataPageV2 = new DataPageV2( dataHeaderV2.getNum_rows(), @@ -2343,32 +2392,39 @@ private void setReadMetrics(long startNs, long len) { } /** - * Populate data in a parquet file range from a vectored range; will block for up - * to {@link #HADOOP_VECTORED_READ_TIMEOUT_SECONDS} seconds. - * @param currRange range to populated. + * Populate data in a parquet file range from one or more bounded vectored ranges; together + * they may block for up to {@link #HADOOP_VECTORED_READ_TIMEOUT_SECONDS} seconds. + * @param ranges bounded ranges containing this part. * @param builder used to build chunk list to read the pages for the different columns. * @throws IOException if there is an error while reading from the stream, including a timeout. */ - public void readFromVectoredRange(ParquetFileRange currRange, ChunkListBuilder builder) throws IOException { - ByteBuffer buffer; + public void readFromVectoredRanges(List ranges, ChunkListBuilder builder) throws IOException { + List buffers = new ArrayList<>(ranges.size()); + ParquetFileRange currentRange = null; final long timeoutSeconds = HADOOP_VECTORED_READ_TIMEOUT_SECONDS; + final long timeoutNanos = TimeUnit.SECONDS.toNanos(timeoutSeconds); long readStart = System.nanoTime(); try { - LOG.debug( - "Waiting for vectored read to finish for range {} with timeout {} seconds", - currRange, - timeoutSeconds); - buffer = FutureIO.awaitFuture(currRange.getDataReadFuture(), timeoutSeconds, TimeUnit.SECONDS); - setReadMetrics(readStart, currRange.getLength()); + for (ParquetFileRange range : ranges) { + currentRange = range; + LOG.debug( + "Waiting for vectored read to finish for range {} with timeout {} seconds", + range, + timeoutSeconds); + long remainingNanos = Math.max(timeoutNanos - (System.nanoTime() - readStart), 0L); + buffers.add(FutureIO.awaitFuture(range.getDataReadFuture(), remainingNanos, TimeUnit.NANOSECONDS)); + } + setReadMetrics(readStart, length); // report in a counter the data we just scanned - BenchmarkCounter.incrementBytesRead(currRange.getLength()); + BenchmarkCounter.incrementBytesRead(length); } catch (TimeoutException e) { String error = String.format( - "Timeout while fetching result for %s with time limit %d seconds", currRange, timeoutSeconds); + "Timeout while fetching result for %s with time limit %d seconds", + currentRange, timeoutSeconds); LOG.error(error, e); throw new IOException(error, e); } - ByteBufferInputStream stream = ByteBufferInputStream.wrap(buffer); + ByteBufferInputStream stream = ByteBufferInputStream.wrap(buffers); for (ChunkDescriptor descriptor : chunks) { builder.add(descriptor, stream.sliceBuffers(descriptor.size), f); } diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestDataPageChecksums.java b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestDataPageChecksums.java index da3f9248b2..faa2bb967b 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestDataPageChecksums.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestDataPageChecksums.java @@ -24,16 +24,19 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; import java.io.IOException; +import java.nio.ByteBuffer; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Random; +import java.util.concurrent.CompletableFuture; import java.util.zip.CRC32; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.Path; import org.apache.parquet.HadoopReadOptions; import org.apache.parquet.ParquetReadOptions; +import org.apache.parquet.bytes.ByteBufferAllocator; import org.apache.parquet.bytes.BytesInput; import org.apache.parquet.bytes.HeapByteBufferAllocator; import org.apache.parquet.bytes.TrackingByteBufferAllocator; @@ -58,9 +61,11 @@ import org.apache.parquet.hadoop.metadata.ParquetMetadata; import org.apache.parquet.hadoop.util.HadoopInputFile; import org.apache.parquet.hadoop.util.HadoopOutputFile; +import org.apache.parquet.io.DelegatingSeekableInputStream; import org.apache.parquet.io.InputFile; import org.apache.parquet.io.OutputFile; import org.apache.parquet.io.ParquetDecodingException; +import org.apache.parquet.io.ParquetFileRange; import org.apache.parquet.io.PositionOutputStream; import org.apache.parquet.io.SeekableInputStream; import org.apache.parquet.schema.MessageType; @@ -447,6 +452,90 @@ public void testWriteOnVerifyOnV2() throws IOException { testWriteOnVerifyOn(ParquetProperties.WriterVersion.PARQUET_2_0); } + private void testVectoredChecksumsRespectAllocationLimit(ParquetProperties.WriterVersion version) + throws IOException { + Configuration conf = new Configuration(); + conf.setBoolean(ParquetOutputFormat.PAGE_WRITE_CHECKSUM_ENABLED, true); + Path path = writeSimpleParquetFile(conf, CompressionCodecName.UNCOMPRESSED, version); + InputFile inputFile = HadoopInputFile.fromPath(path, conf); + final int allocationLimit = 64 * 1024; + final int[] vectorRangeCount = {0}; + + ByteBufferAllocator boundedAllocator = new HeapByteBufferAllocator() { + @Override + public ByteBuffer allocate(int size) { + assertThat(size) + .as("Checksum verification must not allocate above the %s-byte limit", allocationLimit) + .isLessThanOrEqualTo(allocationLimit); + return super.allocate(size); + } + }; + + SeekableInputStream delegate = inputFile.newStream(); + SeekableInputStream vectoredStream = new DelegatingSeekableInputStream(delegate) { + @Override + public long getPos() throws IOException { + return delegate.getPos(); + } + + @Override + public void seek(long position) throws IOException { + delegate.seek(position); + } + + @Override + public boolean readVectoredAvailable(ByteBufferAllocator allocator) { + return true; + } + + @Override + public void readVectored(List ranges, ByteBufferAllocator allocator) throws IOException { + long originalPosition = delegate.getPos(); + try { + for (ParquetFileRange range : ranges) { + vectorRangeCount[0]++; + ByteBuffer buffer = allocator.allocate(range.getLength()); + delegate.seek(range.getOffset()); + delegate.readFully(buffer); + buffer.flip(); + range.setDataReadFuture(CompletableFuture.completedFuture(buffer)); + } + } finally { + delegate.seek(originalPosition); + } + } + }; + + ParquetReadOptions options = ParquetReadOptions.builder() + .withUseHadoopVectoredIo(true) + .withAllocator(boundedAllocator) + .withMaxAllocationInBytes(allocationLimit) + .withPageChecksumVerification(true) + .build(); + + try (ParquetFileReader reader = ParquetFileReader.open(inputFile, options, vectoredStream); + PageReadStore pages = reader.readNextRowGroup()) { + assertCorrectContent(getPageBytes(readNextPage(colADesc, pages)), colAPage1Bytes); + assertCorrectContent(getPageBytes(readNextPage(colADesc, pages)), colAPage2Bytes); + assertCorrectContent(getPageBytes(readNextPage(colBDesc, pages)), colBPage1Bytes); + assertCorrectContent(getPageBytes(readNextPage(colBDesc, pages)), colBPage2Bytes); + } + + assertThat(vectorRangeCount[0]) + .as("Expected each checksummed page to span multiple vectored ranges") + .isGreaterThan(4); + } + + @Test + public void testVectoredChecksumsRespectAllocationLimitV1() throws IOException { + testVectoredChecksumsRespectAllocationLimit(ParquetProperties.WriterVersion.PARQUET_1_0); + } + + @Test + public void testVectoredChecksumsRespectAllocationLimitV2() throws IOException { + testVectoredChecksumsRespectAllocationLimit(ParquetProperties.WriterVersion.PARQUET_2_0); + } + /** * Test whether corruption in the page content is detected by checksum verification */ diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileReaderVectoredIO.java b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileReaderVectoredIO.java new file mode 100644 index 0000000000..570495bf56 --- /dev/null +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileReaderVectoredIO.java @@ -0,0 +1,881 @@ +/* + * 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.parquet.hadoop; + +import static org.apache.parquet.filter2.predicate.FilterApi.and; +import static org.apache.parquet.filter2.predicate.FilterApi.gtEq; +import static org.apache.parquet.filter2.predicate.FilterApi.intColumn; +import static org.apache.parquet.filter2.predicate.FilterApi.ltEq; +import static org.apache.parquet.filter2.predicate.FilterApi.or; +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.io.IOException; +import java.net.SocketTimeoutException; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.PrimitiveIterator; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.Path; +import org.apache.parquet.ParquetReadOptions; +import org.apache.parquet.bytes.ByteBufferAllocator; +import org.apache.parquet.bytes.HeapByteBufferAllocator; +import org.apache.parquet.column.page.PageReadStore; +import org.apache.parquet.example.data.Group; +import org.apache.parquet.example.data.simple.SimpleGroupFactory; +import org.apache.parquet.example.data.simple.convert.GroupRecordConverter; +import org.apache.parquet.filter2.compat.FilterCompat; +import org.apache.parquet.filter2.predicate.FilterPredicate; +import org.apache.parquet.hadoop.example.ExampleParquetWriter; +import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; +import org.apache.parquet.hadoop.util.HadoopInputFile; +import org.apache.parquet.internal.column.columnindex.OffsetIndex; +import org.apache.parquet.io.ColumnIOFactory; +import org.apache.parquet.io.DelegatingSeekableInputStream; +import org.apache.parquet.io.MessageColumnIO; +import org.apache.parquet.io.ParquetFileRange; +import org.apache.parquet.io.RecordReader; +import org.apache.parquet.io.SeekableInputStream; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.MessageTypeParser; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +public class TestParquetFileReaderVectoredIO { + private static final int ROW_COUNT = 128; + private static final int OTHER_COLUMN_BASE = 10000; + private static final MessageType SCHEMA = MessageTypeParser.parseMessageType( + "message test { required int32 id; required binary padding (UTF8); required int32 other; }"); + private static final MessageType PROJECTED_SCHEMA = new MessageType( + SCHEMA.getName(), SCHEMA.getFields().get(0), SCHEMA.getFields().get(2)); + private static final MessageType ID_ONLY_SCHEMA = + new MessageType(SCHEMA.getName(), SCHEMA.getFields().get(0)); + + @TempDir + private java.nio.file.Path tempDir; + + private Path path; + private HadoopInputFile inputFile; + + @BeforeEach + public void writeTestFile() throws IOException { + path = new Path(tempDir.resolve("vectored.parquet").toUri()); + Configuration configuration = new Configuration(); + try (ParquetWriter writer = ExampleParquetWriter.builder(path) + .withConf(configuration) + .withType(SCHEMA) + .withWriteMode(ParquetFileWriter.Mode.OVERWRITE) + .withRowGroupSize(256 * 1024) + .withPageSize(128) + .withPageRowCountLimit(8) + .withDictionaryEncoding(false) + .build()) { + SimpleGroupFactory groups = new SimpleGroupFactory(SCHEMA); + for (int row = 0; row < ROW_COUNT; row++) { + writer.write(groups.newGroup() + .append("id", row) + .append("padding", "padding_" + row) + .append("other", OTHER_COLUMN_BASE + row)); + } + } + inputFile = HadoopInputFile.fromPath(path, configuration); + } + + @Test + public void testSplitsAdjacentColumnsAtMaximumAllocation() throws Exception { + List columns = ParquetFileReader.readFooter(new Configuration(), path) + .getBlocks() + .get(0) + .getColumns(); + int maximumAllocation = 0; + long totalColumnBytes = 0; + for (ColumnChunkMetaData column : columns) { + maximumAllocation = Math.max(maximumAllocation, Math.toIntExact(column.getTotalSize())); + totalColumnBytes += column.getTotalSize(); + } + assertThat(totalColumnBytes > maximumAllocation).isTrue(); + + RecordingAllocator allocator = new RecordingAllocator(); + RecordingSeekableInputStream stream = newStream(FailureMode.NONE); + try (ParquetFileReader reader = + ParquetFileReader.open(inputFile, readOptions(allocator, maximumAllocation, false), stream)) { + allocator.reset(); + try (PageReadStore pages = reader.readNextRowGroup()) { + assertRows(pages, SCHEMA, false); + } + } + + assertEquals(1, stream.vectorCalls); + assertThat(stream.rangeLengths.size() > 1).isTrue(); + for (int rangeLength : stream.rangeLengths) { + assertThat(rangeLength <= maximumAllocation).isTrue(); + } + assertThat(allocator.maximumAllocation <= maximumAllocation).isTrue(); + } + + @Test + public void testSplitsSingleOversizedColumnIntoBoundedVectoredRanges() throws Exception { + int maximumAllocation = 128; + ColumnChunkMetaData column = ParquetFileReader.readFooter(new Configuration(), path) + .getBlocks() + .get(0) + .getColumns() + .get(0); + assertThat(column.getTotalSize() > maximumAllocation).isTrue(); + + RecordingAllocator allocator = new RecordingAllocator(); + RecordingSeekableInputStream stream = newStream(FailureMode.NONE); + try (ParquetFileReader reader = + ParquetFileReader.open(inputFile, readOptions(allocator, maximumAllocation, false), stream)) { + allocator.reset(); + reader.setRequestedSchema(ID_ONLY_SCHEMA); + try (PageReadStore pages = reader.readNextRowGroup()) { + assertRows(pages, ID_ONLY_SCHEMA, false); + } + } + + assertEquals(1, stream.vectorCalls); + assertEquals((column.getTotalSize() + maximumAllocation - 1) / maximumAllocation, stream.rangeLengths.size()); + long nextOffset = column.getStartingPos(); + for (int rangeIndex = 0; rangeIndex < stream.rangeLengths.size(); rangeIndex++) { + assertEquals(nextOffset, stream.rangeOffsets.get(rangeIndex).longValue()); + int rangeLength = stream.rangeLengths.get(rangeIndex); + assertThat(rangeLength <= maximumAllocation).isTrue(); + nextOffset += rangeLength; + } + assertEquals(column.getStartingPos() + column.getTotalSize(), nextOffset); + assertThat(allocator.maximumAllocation <= maximumAllocation).isTrue(); + } + + @Test + public void testSplitsProductionSizedColumnIntoEightMegabyteVectoredRanges() throws Exception { + int maximumAllocation = 8 * 1024 * 1024; + Path largePath = new Path(tempDir.resolve("large-vectored.parquet").toUri()); + Configuration configuration = new Configuration(); + char[] paddingCharacters = new char[136000]; + Arrays.fill(paddingCharacters, 'x'); + String padding = new String(paddingCharacters); + try (ParquetWriter writer = ExampleParquetWriter.builder(largePath) + .withConf(configuration) + .withType(SCHEMA) + .withWriteMode(ParquetFileWriter.Mode.OVERWRITE) + .withRowGroupSize(32 * 1024 * 1024) + .withPageSize(256 * 1024) + .withPageRowCountLimit(2) + .withDictionaryEncoding(false) + .build()) { + SimpleGroupFactory groups = new SimpleGroupFactory(SCHEMA); + for (int row = 0; row < ROW_COUNT; row++) { + writer.write(groups.newGroup() + .append("id", row) + .append("padding", padding) + .append("other", OTHER_COLUMN_BASE + row)); + } + } + + HadoopInputFile largeInputFile = HadoopInputFile.fromPath(largePath, configuration); + ColumnChunkMetaData paddingColumn = ParquetFileReader.readFooter(configuration, largePath) + .getBlocks() + .get(0) + .getColumns() + .get(1); + assertThat(paddingColumn.getTotalSize() > 2L * maximumAllocation).isTrue(); + + RecordingAllocator allocator = new RecordingAllocator(); + RecordingSeekableInputStream stream = + new RecordingSeekableInputStream(largeInputFile.newStream(), FailureMode.NONE); + try (ParquetFileReader reader = + ParquetFileReader.open(largeInputFile, readOptions(allocator, maximumAllocation, false), stream)) { + allocator.reset(); + stream.resetOrdinaryReads(); + try (PageReadStore pages = reader.readNextRowGroup()) { + assertRows(pages, SCHEMA, false, ROW_COUNT, padding); + } + assertEquals(0, stream.normalSeekCalls); + assertEquals(0, stream.normalReadCalls); + } + + assertEquals(1, stream.vectorCalls); + assertThat(stream.rangeLengths.size() >= 3).isTrue(); + for (int rangeLength : stream.rangeLengths) { + assertThat(rangeLength <= maximumAllocation).isTrue(); + } + assertThat(allocator.maximumAllocation <= maximumAllocation).isTrue(); + } + + @Test + public void testFilteredVectoredRangesRespectMaximumAllocation() throws Exception { + int maximumAllocation = 512; + RecordingAllocator allocator = new RecordingAllocator(); + RecordingSeekableInputStream stream = newStream(FailureMode.NONE); + try (ParquetFileReader reader = + ParquetFileReader.open(inputFile, readOptions(allocator, maximumAllocation, true), stream)) { + allocator.reset(); + reader.setRequestedSchema(PROJECTED_SCHEMA); + try (PageReadStore pages = reader.readNextFilteredRowGroup()) { + assertRows(pages, PROJECTED_SCHEMA, true); + } + } + + assertEquals(1, stream.vectorCalls); + assertThat(allocator.maximumAllocation <= maximumAllocation).isTrue(); + } + + @Test + public void testSplitsOversizedFilteredPageIntoBoundedVectoredRanges() throws Exception { + int maximumAllocation = 32; + RecordingAllocator allocator = new RecordingAllocator(); + RecordingSeekableInputStream stream = newStream(FailureMode.NONE); + try (ParquetFileReader reader = + ParquetFileReader.open(inputFile, readOptions(allocator, maximumAllocation, true), stream)) { + ColumnChunkMetaData column = + reader.getFooter().getBlocks().get(0).getColumns().get(0); + OffsetIndex offsetIndex = reader.readOffsetIndex(column); + boolean hasOversizedPage = false; + for (int page = 0; page < offsetIndex.getPageCount(); page++) { + hasOversizedPage |= offsetIndex.getCompressedPageSize(page) > maximumAllocation; + } + assertThat(hasOversizedPage).isTrue(); + + allocator.reset(); + reader.setRequestedSchema(PROJECTED_SCHEMA); + try (PageReadStore pages = reader.readNextFilteredRowGroup()) { + assertRows(pages, PROJECTED_SCHEMA, true); + } + } + + assertEquals(1, stream.vectorCalls); + assertThat(stream.rangeLengths.size() > 2).isTrue(); + for (int rangeLength : stream.rangeLengths) { + assertThat(rangeLength <= maximumAllocation).isTrue(); + } + assertThat(allocator.maximumAllocation <= maximumAllocation).isTrue(); + } + + @Test + public void testSplitsAdjacentFilteredPageRangesAtMaximumAllocation() throws Exception { + int maximumAllocation = 1416; + RecordingAllocator allocator = new RecordingAllocator(); + RecordingSeekableInputStream stream = newStream(FailureMode.NONE); + FilterPredicate predicate = or(ltEq(intColumn("id"), 99), gtEq(intColumn("id"), 120)); + ParquetReadOptions options = ParquetReadOptions.builder() + .withUseHadoopVectoredIo(true) + .withAllocator(allocator) + .withMaxAllocationInBytes(maximumAllocation) + .useColumnIndexFilter(true) + .withRecordFilter(FilterCompat.get(predicate)) + .build(); + + try (ParquetFileReader reader = ParquetFileReader.open(inputFile, options, stream)) { + allocator.reset(); + try (PageReadStore pages = reader.readNextFilteredRowGroup()) { + assertRows(pages, SCHEMA, true, 112); + } + } + + assertEquals(1, stream.vectorCalls); + assertThat(stream.rangeLengths.size() > 1).isTrue(); + for (int rangeLength : stream.rangeLengths) { + assertThat(rangeLength <= maximumAllocation).isTrue(); + } + assertThat(allocator.maximumAllocation <= maximumAllocation).isTrue(); + } + + @Test + public void testLeavesNonVectoredReadsUnchanged() throws Exception { + int maximumAllocation = 128; + RecordingAllocator allocator = new RecordingAllocator(); + RecordingSeekableInputStream stream = newStream(FailureMode.NONE); + ParquetReadOptions options = ParquetReadOptions.builder() + .withUseHadoopVectoredIo(false) + .withAllocator(allocator) + .withMaxAllocationInBytes(maximumAllocation) + .build(); + + try (ParquetFileReader reader = ParquetFileReader.open(inputFile, options, stream)) { + allocator.reset(); + stream.resetOrdinaryReads(); + try (PageReadStore pages = reader.readNextRowGroup()) { + assertRows(pages, SCHEMA, false); + } + assertEquals(1, stream.normalSeekCalls); + } + + assertEquals(0, stream.vectorCalls); + assertThat(allocator.maximumAllocation <= maximumAllocation).isTrue(); + } + + @Test + public void testUnsupportedBackendPreservesContiguousOrdinaryRead() throws Exception { + int maximumAllocation = 128; + RecordingAllocator allocator = new RecordingAllocator(); + RecordingSeekableInputStream stream = newStream(FailureMode.UNSUPPORTED_BACKEND); + try (ParquetFileReader reader = + ParquetFileReader.open(inputFile, readOptions(allocator, maximumAllocation, false), stream)) { + allocator.reset(); + stream.resetOrdinaryReads(); + try (PageReadStore pages = reader.readNextRowGroup()) { + assertRows(pages, SCHEMA, false); + } + assertEquals(1, stream.normalSeekCalls); + } + + assertEquals(0, stream.vectorCalls); + assertThat(allocator.maximumAllocation <= maximumAllocation).isTrue(); + } + + @Test + public void testUnsupportedBackendPreservesContiguousFilteredOrdinaryRead() throws Exception { + int maximumAllocation = 32; + RecordingAllocator allocator = new RecordingAllocator(); + RecordingSeekableInputStream stream = newStream(FailureMode.UNSUPPORTED_BACKEND); + try (ParquetFileReader reader = + ParquetFileReader.open(inputFile, readOptions(allocator, maximumAllocation, true), stream)) { + reader.setRequestedSchema(ID_ONLY_SCHEMA); + preloadFilteredIndexes(reader, ID_ONLY_SCHEMA); + allocator.reset(); + stream.resetOrdinaryReads(); + try (PageReadStore pages = reader.readNextFilteredRowGroup()) { + assertRows(pages, ID_ONLY_SCHEMA, true); + } + // The predicate selects two disjoint page spans; each span should require exactly one seek. + assertEquals(2, stream.normalSeekCalls); + } + + assertEquals(0, stream.vectorCalls); + assertThat(allocator.maximumAllocation <= maximumAllocation).isTrue(); + } + + @Test + public void testFailsFastAfterPartiallyConsumingVectoredData() throws Exception { + assertFailsAfterVectoredSubmission(FailureMode.SECOND_ILLEGAL_ARGUMENT, IllegalArgumentException.class); + assertFailsAfterVectoredSubmission(FailureMode.SECOND_UNSUPPORTED, UnsupportedOperationException.class); + } + + @Test + public void testFailsFastWhenFirstVectoredRangeFails() throws Exception { + assertFailsAfterVectoredSubmission(FailureMode.FIRST_ILLEGAL_ARGUMENT, IllegalArgumentException.class); + assertFailsAfterVectoredSubmission(FailureMode.FIRST_UNSUPPORTED, UnsupportedOperationException.class); + } + + @Test + public void testFailsFastWhenFirstVectoredRangeFailsWhileSiblingRemainsPending() throws Exception { + assertFailsAfterVectoredSubmission( + FailureMode.FIRST_ILLEGAL_ARGUMENT_PENDING_SIBLING, IllegalArgumentException.class); + assertFailsAfterVectoredSubmission( + FailureMode.FIRST_UNSUPPORTED_PENDING_SIBLING, UnsupportedOperationException.class); + } + + @Test + public void testDrainsPendingVectoredReadsBeforeClosingBackendThatDoesNotCancelThem() throws Exception { + assertDrainsPendingVectoredReadsBeforeClosing( + FailureMode.FIRST_IO_EXCEPTION_PENDING_SIBLING_NO_CLOSE_CANCELLATION, IOException.class); + } + + @Test + public void testDrainsPendingVectoredReadsAfterSocketTimeout() throws Exception { + assertDrainsPendingVectoredReadsBeforeClosing( + FailureMode.FIRST_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION, SocketTimeoutException.class); + } + + @Test + public void testContinuesDrainingVectoredReadsAfterSiblingSocketTimeout() throws Exception { + assertDrainsPendingVectoredReadsBeforeClosing( + FailureMode.FIRST_IO_EXCEPTION_THEN_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION, + IOException.class); + } + + private void assertDrainsPendingVectoredReadsBeforeClosing( + FailureMode failureMode, Class failureClass) throws Exception { + RecordingAllocator allocator = new RecordingAllocator(); + RecordingSeekableInputStream stream = newStream(failureMode); + int maximumAllocation = + failureMode == FailureMode.FIRST_IO_EXCEPTION_THEN_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION + ? 128 + : 4096; + try (ParquetFileReader reader = + ParquetFileReader.open(inputFile, readOptions(allocator, maximumAllocation, true), stream)) { + reader.setRequestedSchema(PROJECTED_SCHEMA); + preloadFilteredIndexes(reader, PROJECTED_SCHEMA); + stream.resetOrdinaryReads(); + + CompletableFuture readAttempt = CompletableFuture.supplyAsync(() -> { + try { + reader.readNextFilteredRowGroup(); + throw new AssertionError("Expected the first vectored range to fail"); + } catch (IOException failure) { + assertThat(Thread.currentThread().isInterrupted()) + .as("A sibling socket timeout must not interrupt the scan thread") + .isFalse(); + return failure; + } + }); + + try { + CompletableFuture.anyOf(stream.pendingDrainStarted, readAttempt).get(10, TimeUnit.SECONDS); + assertThat(readAttempt.isDone()) + .as("The original failure must wait for unfinished sibling reads") + .isFalse(); + assertThat(stream.pendingFutureCount() > 0).isTrue(); + } finally { + stream.allowPendingPhysicalReads.complete(null); + } + + IOException failure = readAttempt.get(10, TimeUnit.SECONDS); + assertThat(failureClass.isInstance(failure)).isTrue(); + assertThat(failure.getMessage().contains("injected asynchronous vectored")) + .isTrue(); + if (failureMode + == FailureMode.FIRST_IO_EXCEPTION_THEN_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION) { + assertThat(stream.rangeLengths.size() >= 3).isTrue(); + assertEquals(1, failure.getSuppressed().length); + assertThat(failure.getSuppressed()[0] instanceof SocketTimeoutException) + .isTrue(); + } + assertEquals(0, stream.pendingFutureCount()); + assertEquals(0, stream.normalSeekCalls); + assertEquals(0, stream.normalReadCalls); + } + + assertEquals(1, stream.vectorCalls); + assertEquals(0, stream.postClosePhysicalReads.get()); + assertThat(stream.completedPendingPhysicalReads.get() > 0).isTrue(); + } + + @Test + public void testFailsFastWhenSplitColumnFailsBeforeBuilderIsPopulated() throws Exception { + assertFailsAfterSplitColumnSubmission(FailureMode.SECOND_ILLEGAL_ARGUMENT, IllegalArgumentException.class); + assertFailsAfterSplitColumnSubmission(FailureMode.SECOND_UNSUPPORTED, UnsupportedOperationException.class); + } + + @Test + public void testFailsFastWhenOversizedColumnRangeFails() throws Exception { + assertFailsAfterConsumingColumnBeforeSplitRange( + FailureMode.SECOND_ILLEGAL_ARGUMENT, IllegalArgumentException.class); + assertFailsAfterConsumingColumnBeforeSplitRange( + FailureMode.SECOND_UNSUPPORTED, UnsupportedOperationException.class); + } + + @Test + public void testFailsFastWhenVectoredSubmissionFails() throws Exception { + assertFailsAfterVectoredSubmission(FailureMode.SUBMISSION_ILLEGAL_ARGUMENT, IllegalArgumentException.class); + assertFailsAfterVectoredSubmission(FailureMode.SUBMISSION_UNSUPPORTED, UnsupportedOperationException.class); + } + + @Test + public void testFailsFastWhenVectoredSubmissionFailsAfterSchedulingPendingRead() throws Exception { + assertFailsAfterVectoredSubmission( + FailureMode.PARTIAL_SUBMISSION_ILLEGAL_ARGUMENT, IllegalArgumentException.class); + assertFailsAfterVectoredSubmission( + FailureMode.PARTIAL_SUBMISSION_UNSUPPORTED, UnsupportedOperationException.class); + } + + private void assertFailsAfterVectoredSubmission(FailureMode failureMode, Class causeType) throws Exception { + RecordingAllocator allocator = new RecordingAllocator(); + RecordingSeekableInputStream stream = newStream(failureMode); + if (failureMode.hasPublishedPendingRead()) { + stream.allowPendingPhysicalReads.complete(null); + } + try (ParquetFileReader reader = ParquetFileReader.open(inputFile, readOptions(allocator, 4096, true), stream)) { + reader.setRequestedSchema(PROJECTED_SCHEMA); + preloadFilteredIndexes(reader, PROJECTED_SCHEMA); + stream.resetOrdinaryReads(); + IOException failure = assertThrows(IOException.class, reader::readNextFilteredRowGroup); + assertThat(failure.getMessage().contains("asynchronous reads may have been submitted")) + .isTrue(); + assertThat(causeType.isInstance(failure.getCause())).isTrue(); + assertEquals(0, stream.normalSeekCalls); + assertEquals(0, stream.normalReadCalls); + if (failureMode.hasPendingRead()) { + if (failureMode.hasPublishedPendingRead()) { + assertEquals(0, stream.pendingFutureCount()); + } else { + assertThat(stream.pendingFutureCount() > 0).isTrue(); + } + } + } + assertEquals(1, stream.vectorCalls); + assertEquals(0, stream.pendingFutureCount()); + } + + private void assertFailsAfterSplitColumnSubmission(FailureMode failureMode, Class causeType) throws Exception { + int maximumAllocation = 128; + RecordingAllocator allocator = new RecordingAllocator(); + RecordingSeekableInputStream stream = newStream(failureMode); + try (ParquetFileReader reader = + ParquetFileReader.open(inputFile, readOptions(allocator, maximumAllocation, false), stream)) { + allocator.reset(); + reader.setRequestedSchema(ID_ONLY_SCHEMA); + stream.resetOrdinaryReads(); + IOException failure = assertThrows(IOException.class, reader::readNextRowGroup); + assertThat(failure.getMessage().contains("asynchronous reads may have been submitted")) + .isTrue(); + assertThat(causeType.isInstance(failure.getCause())).isTrue(); + assertEquals(0, stream.normalSeekCalls); + assertEquals(0, stream.normalReadCalls); + } + + assertEquals(1, stream.vectorCalls); + assertThat(stream.rangeLengths.size() > 1).isTrue(); + assertThat(allocator.maximumAllocation <= maximumAllocation).isTrue(); + } + + private void assertFailsAfterConsumingColumnBeforeSplitRange(FailureMode failureMode, Class causeType) + throws Exception { + List columns = ParquetFileReader.readFooter(new Configuration(), path) + .getBlocks() + .get(0) + .getColumns(); + int maximumAllocation = Math.toIntExact(columns.get(0).getTotalSize()); + assertThat(columns.get(1).getTotalSize() > maximumAllocation).isTrue(); + + RecordingAllocator allocator = new RecordingAllocator(); + RecordingSeekableInputStream stream = newStream(failureMode); + try (ParquetFileReader reader = + ParquetFileReader.open(inputFile, readOptions(allocator, maximumAllocation, false), stream)) { + stream.resetOrdinaryReads(); + IOException failure = assertThrows(IOException.class, reader::readNextRowGroup); + assertThat(failure.getMessage().contains("asynchronous reads may have been submitted")) + .isTrue(); + assertThat(causeType.isInstance(failure.getCause())).isTrue(); + assertEquals(0, stream.normalSeekCalls); + assertEquals(0, stream.normalReadCalls); + } + assertEquals(1, stream.vectorCalls); + } + + private static void preloadFilteredIndexes(ParquetFileReader reader, MessageType projection) { + reader.getFilteredRecordCount(); + for (ColumnChunkMetaData column : reader.getFooter().getBlocks().get(0).getColumns()) { + if (projection.containsField(column.getPath().toDotString())) { + reader.getColumnIndexStore(0).getOffsetIndex(column.getPath()); + } + } + } + + private RecordingSeekableInputStream newStream(FailureMode failureMode) throws IOException { + return new RecordingSeekableInputStream(inputFile.newStream(), failureMode); + } + + private static ParquetReadOptions readOptions( + RecordingAllocator allocator, int maximumAllocation, boolean filterPages) { + ParquetReadOptions.Builder builder = ParquetReadOptions.builder() + .withUseHadoopVectoredIo(true) + .withAllocator(allocator) + .withMaxAllocationInBytes(maximumAllocation); + if (filterPages) { + FilterPredicate predicate = + or(ltEq(intColumn("id"), 99), and(gtEq(intColumn("id"), 108), ltEq(intColumn("id"), 115))); + builder.useColumnIndexFilter(true).withRecordFilter(FilterCompat.get(predicate)); + } + return builder.build(); + } + + private static void assertRows(PageReadStore pages, MessageType projection, boolean filtered) { + assertRows(pages, projection, filtered, filtered ? 108 : ROW_COUNT); + } + + private static void assertRows(PageReadStore pages, MessageType projection, boolean filtered, long expectedRows) { + assertRows(pages, projection, filtered, expectedRows, null); + } + + private static void assertRows( + PageReadStore pages, MessageType projection, boolean filtered, long expectedRows, String expectedPadding) { + assertEquals(expectedRows, pages.getRowCount()); + + PrimitiveIterator.OfLong rowIndexes = filtered ? pages.getRowIndexes().get() : null; + MessageColumnIO columns = new ColumnIOFactory().getColumnIO(projection, SCHEMA); + RecordReader records = columns.getRecordReader(pages, new GroupRecordConverter(projection)); + for (long row = 0; row < expectedRows; row++) { + long expectedIndex = filtered ? rowIndexes.nextLong() : row; + Group record = records.read(); + assertEquals(expectedIndex, record.getInteger("id", 0)); + if (projection.containsField("padding")) { + assertEquals( + expectedPadding == null ? "padding_" + expectedIndex : expectedPadding, + record.getString("padding", 0)); + } + if (projection.containsField("other")) { + assertEquals(OTHER_COLUMN_BASE + expectedIndex, record.getInteger("other", 0)); + } + } + } + + private enum FailureMode { + NONE, + UNSUPPORTED_BACKEND, + FIRST_ILLEGAL_ARGUMENT, + FIRST_UNSUPPORTED, + FIRST_ILLEGAL_ARGUMENT_PENDING_SIBLING, + FIRST_UNSUPPORTED_PENDING_SIBLING, + FIRST_IO_EXCEPTION_PENDING_SIBLING_NO_CLOSE_CANCELLATION, + FIRST_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION, + FIRST_IO_EXCEPTION_THEN_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION, + SECOND_ILLEGAL_ARGUMENT, + SECOND_UNSUPPORTED, + SUBMISSION_ILLEGAL_ARGUMENT, + SUBMISSION_UNSUPPORTED, + PARTIAL_SUBMISSION_ILLEGAL_ARGUMENT, + PARTIAL_SUBMISSION_UNSUPPORTED; + + private boolean hasPendingRead() { + return hasPublishedPendingRead() + || this == PARTIAL_SUBMISSION_ILLEGAL_ARGUMENT + || this == PARTIAL_SUBMISSION_UNSUPPORTED; + } + + private boolean hasPublishedPendingRead() { + return this == FIRST_ILLEGAL_ARGUMENT_PENDING_SIBLING + || this == FIRST_UNSUPPORTED_PENDING_SIBLING + || this == FIRST_IO_EXCEPTION_PENDING_SIBLING_NO_CLOSE_CANCELLATION + || this == FIRST_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION + || this == FIRST_IO_EXCEPTION_THEN_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION; + } + } + + private static final class RecordingAllocator extends HeapByteBufferAllocator { + private int maximumAllocation; + + @Override + public ByteBuffer allocate(int size) { + maximumAllocation = Math.max(maximumAllocation, size); + return super.allocate(size); + } + + private void reset() { + maximumAllocation = 0; + } + } + + private static final class RecordingSeekableInputStream extends DelegatingSeekableInputStream { + private final SeekableInputStream delegate; + private final FailureMode failureMode; + private final List rangeOffsets = new ArrayList<>(); + private final List rangeLengths = new ArrayList<>(); + private final List> pendingFutures = new ArrayList<>(); + private final CompletableFuture allowPendingPhysicalReads = new CompletableFuture<>(); + private final CompletableFuture pendingDrainStarted = new CompletableFuture<>(); + private final CompletableFuture socketTimeoutDrainStarted = new CompletableFuture<>(); + private final AtomicInteger completedPendingPhysicalReads = new AtomicInteger(); + private final AtomicInteger postClosePhysicalReads = new AtomicInteger(); + private volatile boolean closed; + private int vectorCalls; + private int normalSeekCalls; + private int normalReadCalls; + + private RecordingSeekableInputStream(SeekableInputStream delegate, FailureMode failureMode) { + super(delegate); + this.delegate = delegate; + this.failureMode = failureMode; + } + + @Override + public long getPos() throws IOException { + return delegate.getPos(); + } + + @Override + public void seek(long newPos) throws IOException { + normalSeekCalls++; + delegate.seek(newPos); + } + + @Override + public void readFully(ByteBuffer buffer) throws IOException { + normalReadCalls++; + delegate.readFully(buffer); + } + + @Override + public boolean readVectoredAvailable(ByteBufferAllocator allocator) { + return failureMode != FailureMode.UNSUPPORTED_BACKEND; + } + + @Override + public void readVectored(List ranges, ByteBufferAllocator allocator) throws IOException { + vectorCalls++; + if (failureMode == FailureMode.SUBMISSION_ILLEGAL_ARGUMENT) { + throw new IllegalArgumentException("injected vectored submission failure"); + } + if (failureMode == FailureMode.SUBMISSION_UNSUPPORTED) { + throw new UnsupportedOperationException("injected vectored submission failure"); + } + if (failureMode == FailureMode.PARTIAL_SUBMISSION_ILLEGAL_ARGUMENT + || failureMode == FailureMode.PARTIAL_SUBMISSION_UNSUPPORTED) { + // Hadoop's bridge does not publish backend futures until submission returns successfully. + pendingFutures.add(new CompletableFuture<>()); + throw failure(); + } + + long originalPosition = delegate.getPos(); + try { + for (int index = 0; index < ranges.size(); index++) { + ParquetFileRange range = ranges.get(index); + rangeOffsets.add(range.getOffset()); + rangeLengths.add(range.getLength()); + boolean socketTimeoutSibling = index == 1 + && failureMode + == FailureMode + .FIRST_IO_EXCEPTION_THEN_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION; + CompletableFuture future = hasPendingSibling(index) + ? new PendingPhysicalReadFuture( + socketTimeoutSibling ? socketTimeoutDrainStarted : pendingDrainStarted, + socketTimeoutSibling + ? new SocketTimeoutException("injected sibling socket timeout") + : null) + : new CompletableFuture<>(); + if (shouldFail(index)) { + if (failureMode == FailureMode.FIRST_IO_EXCEPTION_PENDING_SIBLING_NO_CLOSE_CANCELLATION + || failureMode + == FailureMode + .FIRST_IO_EXCEPTION_THEN_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION) { + future.completeExceptionally(new IOException("injected asynchronous vectored IO failure")); + } else if (failureMode + == FailureMode.FIRST_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION) { + future.completeExceptionally( + new SocketTimeoutException("injected asynchronous vectored socket timeout")); + } else { + future.completeExceptionally(failure()); + } + } else if (hasPendingSibling(index)) { + pendingFutures.add(future); + if (!socketTimeoutSibling) { + ByteBuffer buffer = allocator.allocate(range.getLength()); + CompletableFuture.runAsync(() -> { + allowPendingPhysicalReads.join(); + if (closed) { + postClosePhysicalReads.incrementAndGet(); + } + completedPendingPhysicalReads.incrementAndGet(); + future.complete(buffer); + }); + } + } else { + ByteBuffer buffer = allocator.allocate(range.getLength()); + delegate.seek(range.getOffset()); + delegate.readFully(buffer); + buffer.flip(); + future.complete(buffer); + } + range.setDataReadFuture(future); + } + } finally { + delegate.seek(originalPosition); + } + } + + private boolean shouldFail(int index) { + return index == 0 + && (failureMode == FailureMode.FIRST_ILLEGAL_ARGUMENT + || failureMode == FailureMode.FIRST_UNSUPPORTED + || failureMode == FailureMode.FIRST_ILLEGAL_ARGUMENT_PENDING_SIBLING + || failureMode == FailureMode.FIRST_UNSUPPORTED_PENDING_SIBLING + || failureMode + == FailureMode.FIRST_IO_EXCEPTION_PENDING_SIBLING_NO_CLOSE_CANCELLATION + || failureMode + == FailureMode.FIRST_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION + || failureMode + == FailureMode + .FIRST_IO_EXCEPTION_THEN_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION) + || index == 1 + && (failureMode == FailureMode.SECOND_ILLEGAL_ARGUMENT + || failureMode == FailureMode.SECOND_UNSUPPORTED); + } + + private boolean hasPendingSibling(int index) { + return index > 0 + && (failureMode == FailureMode.FIRST_ILLEGAL_ARGUMENT_PENDING_SIBLING + || failureMode == FailureMode.FIRST_UNSUPPORTED_PENDING_SIBLING + || failureMode == FailureMode.FIRST_IO_EXCEPTION_PENDING_SIBLING_NO_CLOSE_CANCELLATION + || failureMode == FailureMode.FIRST_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION + || failureMode + == FailureMode + .FIRST_IO_EXCEPTION_THEN_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION); + } + + private RuntimeException failure() { + if (failureMode == FailureMode.FIRST_UNSUPPORTED + || failureMode == FailureMode.FIRST_UNSUPPORTED_PENDING_SIBLING + || failureMode == FailureMode.SECOND_UNSUPPORTED + || failureMode == FailureMode.PARTIAL_SUBMISSION_UNSUPPORTED) { + return new UnsupportedOperationException("injected asynchronous vectored failure"); + } + return new IllegalArgumentException("injected asynchronous vectored failure"); + } + + private void resetOrdinaryReads() { + normalSeekCalls = 0; + normalReadCalls = 0; + } + + private int pendingFutureCount() { + int count = 0; + for (CompletableFuture pendingFuture : pendingFutures) { + if (!pendingFuture.isDone()) { + count++; + } + } + return count; + } + + @Override + public void close() throws IOException { + closed = true; + if (failureMode != FailureMode.FIRST_IO_EXCEPTION_PENDING_SIBLING_NO_CLOSE_CANCELLATION + && failureMode != FailureMode.FIRST_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION + && failureMode + != FailureMode + .FIRST_IO_EXCEPTION_THEN_SOCKET_TIMEOUT_PENDING_SIBLING_NO_CLOSE_CANCELLATION) { + for (CompletableFuture pendingFuture : pendingFutures) { + pendingFuture.cancel(false); + } + } + super.close(); + } + } + + private static final class PendingPhysicalReadFuture extends CompletableFuture { + private final CompletableFuture drainStarted; + private final IOException failure; + + private PendingPhysicalReadFuture(CompletableFuture drainStarted, IOException failure) { + this.drainStarted = drainStarted; + this.failure = failure; + } + + @Override + public ByteBuffer get(long timeout, TimeUnit unit) + throws InterruptedException, ExecutionException, TimeoutException { + drainStarted.complete(null); + if (failure != null) { + completeExceptionally(failure); + } + return super.get(timeout, unit); + } + } +}