From 7d6c28112004c27a3628f73471f06689af59ce75 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Thu, 13 Aug 2026 19:12:15 -0700 Subject: [PATCH 1/3] 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); + } + } +} From ad623d6f4a44934ee2ccb98ab28bfee7242b0885 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Mon, 17 Aug 2026 19:43:59 -0700 Subject: [PATCH 2/3] Address vectored-read review feedback --- .../parquet/hadoop/ParquetFileReader.java | 28 +++--- .../util/wrapped/io/VectorIoBridge.java | 3 +- .../TestParquetFileReaderVectoredIO.java | 88 ++++++++++++++++++- 3 files changed, 105 insertions(+), 14 deletions(-) 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 52e639bf81..931d211105 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 @@ -773,6 +773,10 @@ public static ParquetFileReader open( // not final. in some cases, this may be lazily loaded for backward-compat. private ParquetMetadata footer; + // Some InputFile implementations fetch remote metadata for getLength(). Cache the + // vectored range-validation length lazily so ordinary reads need no extra lookup. + private long vectoredReadFileLength = -1; + private int currentBlock = 0; private ColumnChunkPageReadStore currentRowGroup = null; private DictionaryPageReader nextDictionaryReader = null; @@ -1311,15 +1315,10 @@ private void readAllPartsVectoredOrNormal(List allParts, Ch /** * Should the read use vectored IO? - *

- * This returns true if all necessary conditions are met: - *

    - *
  1. The option is enabled
    2. - *
  2. The Hadoop version supports vectored IO
    3. - *
  3. The stream implementation explicitly supports the API; for other streams the classic - * API is always used.
    4. - *
  4. The allocator is not direct. This is to avoid HADOOP-19101 surfacing. - *
+ *

The option must be enabled and the stream's availability probe must accept the + * allocator. For Hadoop streams, that probe checks runtime API availability and excludes + * direct allocators to avoid HADOOP-19101. It does not guarantee that a particular + * vectored-read request will be accepted. * @return true or false. */ private boolean shouldUseVectoredIo() { @@ -1345,7 +1344,10 @@ private boolean shouldUseVectoredIo() { 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(); + if (vectoredReadFileLength < 0) { + vectoredReadFileLength = file.getLength(); + } + final long fileLength = vectoredReadFileLength; List ranges = new ArrayList<>(allParts.size()); List partRangeCounts = new ArrayList<>(allParts.size()); long totalSize = 0; @@ -1372,8 +1374,9 @@ private void readVectored(List allParts, ChunkListBuilder b LOG.debug("Reading {} bytes of data with vectored IO in {} ranges", totalSize, ranges.size()); 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. + // Even a synchronous rejection can follow partial submission. The Hadoop bridge + // publishes futures only after submission returns, so missing futures do not prove + // that no reads started. Once this call is entered, normal-read fallback is unsafe. f.readVectored(ranges, options.getAllocator()); int firstRange = 0; for (int partIndex = 0; partIndex < allParts.size(); partIndex++) { @@ -1382,6 +1385,7 @@ private void readVectored(List allParts, ChunkListBuilder b firstRange = endRange; } } catch (IllegalArgumentException | UnsupportedOperationException e) { + // Consumption may also have populated the builder. Do not replay those chunks. IOException failure = new IOException("Vectored read failed after asynchronous reads may have been submitted", e); awaitRemainingVectoredReads(ranges, readStart, failure); diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/util/wrapped/io/VectorIoBridge.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/util/wrapped/io/VectorIoBridge.java index 7720f7fe94..be5cc80251 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/util/wrapped/io/VectorIoBridge.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/util/wrapped/io/VectorIoBridge.java @@ -124,7 +124,8 @@ private VectorIoBridge() { * @param stream input stream to query. * @param allocator allocator to be used. * - * @return true if the stream declares the capability is available. + * @return true if the runtime API and allocator are supported; an individual request + * may still be rejected. */ public boolean readVectoredAvailable(final FSDataInputStream stream, final ByteBufferAllocator allocator) { return available() && !allocator.isDirect(); 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 index 570495bf56..4261250aab 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileReaderVectoredIO.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileReaderVectoredIO.java @@ -52,10 +52,12 @@ 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.metadata.ParquetMetadata; 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.InputFile; import org.apache.parquet.io.MessageColumnIO; import org.apache.parquet.io.ParquetFileRange; import org.apache.parquet.io.RecordReader; @@ -106,6 +108,64 @@ public void writeTestFile() throws IOException { inputFile = HadoopInputFile.fromPath(path, configuration); } + @Test + public void testCachesFileLengthAcrossVectoredReads() throws Exception { + assertCachesFileLengthAcrossVectoredReads(false); + } + + @Test + public void testCachesFileLengthWithSuppliedFooter() throws Exception { + assertCachesFileLengthAcrossVectoredReads(true); + } + + private void assertCachesFileLengthAcrossVectoredReads(boolean supplyFooter) throws Exception { + ParquetMetadata footer = ParquetFileReader.readFooter(new Configuration(), path); + CountingInputFile countingFile = new CountingInputFile(inputFile); + RecordingSeekableInputStream stream = newStream(FailureMode.NONE); + ParquetReadOptions options = readOptions(new RecordingAllocator(), 128, true); + + try (ParquetFileReader reader = supplyFooter + ? ParquetFileReader.open(countingFile, footer, options, stream) + : ParquetFileReader.open(countingFile, options, stream)) { + // Reading the footer needs the length, but a supplied footer does not. + int initialLengthCalls = supplyFooter ? 0 : 1; + assertEquals(initialLengthCalls, countingFile.lengthCalls); + reader.setRequestedSchema(PROJECTED_SCHEMA); + + try (PageReadStore pages = reader.readRowGroup(0)) { + assertRows(pages, PROJECTED_SCHEMA, false); + } + assertEquals(initialLengthCalls + 1, countingFile.lengthCalls); + + try (PageReadStore pages = reader.readFilteredRowGroup(0)) { + assertRows(pages, PROJECTED_SCHEMA, true); + } + try (PageReadStore pages = reader.readRowGroup(0)) { + assertRows(pages, PROJECTED_SCHEMA, false); + } + assertEquals(initialLengthCalls + 1, countingFile.lengthCalls); + } + + assertEquals(3, stream.vectorCalls); + } + + @Test + public void testSuppliedFooterAvoidsLengthLookupForOrdinaryReads() throws Exception { + ParquetMetadata footer = ParquetFileReader.readFooter(new Configuration(), path); + CountingInputFile countingFile = new CountingInputFile(inputFile); + RecordingSeekableInputStream stream = newStream(FailureMode.NONE); + ParquetReadOptions options = + ParquetReadOptions.builder().withUseHadoopVectoredIo(false).build(); + + try (ParquetFileReader reader = ParquetFileReader.open(countingFile, footer, options, stream)) { + try (PageReadStore pages = reader.readRowGroup(0)) { + assertRows(pages, SCHEMA, false); + } + assertEquals(0, countingFile.lengthCalls); + } + assertEquals(0, stream.vectorCalls); + } + @Test public void testSplitsAdjacentColumnsAtMaximumAllocation() throws Exception { List columns = ParquetFileReader.readFooter(new Configuration(), path) @@ -482,7 +542,8 @@ public void testFailsFastWhenOversizedColumnRangeFails() throws Exception { } @Test - public void testFailsFastWhenVectoredSubmissionFails() throws Exception { + public void testFailsFastWhenVectoredSubmissionRejectsBeforeStartingReads() throws Exception { + // A synchronous rejection cannot be distinguished from the partial-submission case below. assertFailsAfterVectoredSubmission(FailureMode.SUBMISSION_ILLEGAL_ARGUMENT, IllegalArgumentException.class); assertFailsAfterVectoredSubmission(FailureMode.SUBMISSION_UNSUPPORTED, UnsupportedOperationException.class); } @@ -658,6 +719,31 @@ private boolean hasPublishedPendingRead() { } } + private static final class CountingInputFile implements InputFile { + private final InputFile delegate; + private int lengthCalls; + + private CountingInputFile(InputFile delegate) { + this.delegate = delegate; + } + + @Override + public long getLength() throws IOException { + lengthCalls++; + return delegate.getLength(); + } + + @Override + public SeekableInputStream newStream() throws IOException { + return delegate.newStream(); + } + + @Override + public String toString() { + return delegate.toString(); + } + } + private static final class RecordingAllocator extends HeapByteBufferAllocator { private int maximumAllocation; From f82eaa428528849cf019c99410fe879c4d641875 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sat, 29 Aug 2026 17:17:46 +0000 Subject: [PATCH 3/3] GH-3719: Fix vectored read ownership and submission deadlines Track original filesystem allocations, including checksum buffers and buffers backing returned slices, and release them with the row group. Clean up the row group when reading or parsing fails. Include submission and every requested range in a single read deadline. Invalidate failed readers and defer stream cleanup until submission exits, without delaying interruption or recycling buffers still used by IO. Publish available futures after partial submission failures and fall back to ordinary reads when a pre-submission file-length lookup fails. Add regression coverage for buffer ownership, submission failures, timeouts, interruption, and metadata fallback. --- .../parquet/hadoop/ParquetFileReader.java | 159 ++++-- .../hadoop/VectoredReadBufferAllocator.java | 110 ++++ .../parquet/hadoop/VectoredReadOperation.java | 179 +++++++ .../util/wrapped/io/VectorIoBridge.java | 53 +- .../TestParquetFileReaderVectoredIO.java | 73 ++- ...estParquetFileReaderVectoredOwnership.java | 314 ++++++++++++ .../TestVectoredReadBufferAllocator.java | 256 ++++++++++ .../hadoop/TestVectoredReadOperation.java | 472 ++++++++++++++++++ .../util/wrapped/io/TestVectorIoBridge.java | 82 +++ 9 files changed, 1649 insertions(+), 49 deletions(-) create mode 100644 parquet-hadoop/src/main/java/org/apache/parquet/hadoop/VectoredReadBufferAllocator.java create mode 100644 parquet-hadoop/src/main/java/org/apache/parquet/hadoop/VectoredReadOperation.java create mode 100644 parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileReaderVectoredOwnership.java create mode 100644 parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestVectoredReadBufferAllocator.java create mode 100644 parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestVectoredReadOperation.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 931d211105..4c4e0d8995 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 @@ -114,6 +114,7 @@ import org.apache.parquet.internal.filter2.columnindex.ColumnIndexFilter; import org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore; import org.apache.parquet.internal.hadoop.metadata.IndexReference; +import org.apache.parquet.io.DelegatingSeekableInputStream; import org.apache.parquet.io.InputFile; import org.apache.parquet.io.ParquetDecodingException; import org.apache.parquet.io.ParquetFileRange; @@ -776,6 +777,8 @@ public static ParquetFileReader open( // Some InputFile implementations fetch remote metadata for getLength(). Cache the // vectored range-validation length lazily so ordinary reads need no extra lookup. private long vectoredReadFileLength = -1; + private boolean vectoredIoDisabled; + private ExecutorService vectoredReadExecutor; private int currentBlock = 0; private ColumnChunkPageReadStore currentRowGroup = null; @@ -1195,13 +1198,17 @@ private ColumnChunkPageReadStore internalReadRowGroup(int blockIndex) throws IOE } // actually read all the chunks ChunkListBuilder builder = new ChunkListBuilder(block.getRowCount()); - readAllPartsVectoredOrNormal(allParts, builder); rowGroup.setReleaser(builder.releaser); - for (Chunk chunk : builder.build()) { - readChunkPages(chunk, block, rowGroup); + try { + readAllPartsVectoredOrNormal(allParts, builder); + for (Chunk chunk : builder.build()) { + readChunkPages(chunk, block, rowGroup); + } + return rowGroup; + } catch (IOException | RuntimeException | Error failure) { + closeRowGroupAfterFailure(rowGroup, failure); + throw failure; } - - return rowGroup; } /** @@ -1300,8 +1307,10 @@ private void readAllPartsVectoredOrNormal(List allParts, Ch if (shouldUseVectoredIo()) { try { - readVectored(allParts, builder); - return; + if (prepareVectoredReadFileLength()) { + readVectored(allParts, builder); + return; + } } catch (IllegalArgumentException | UnsupportedOperationException 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. @@ -1322,7 +1331,28 @@ private void readAllPartsVectoredOrNormal(List allParts, Ch * @return true or false. */ private boolean shouldUseVectoredIo() { - return options.useHadoopVectoredIo() && f.readVectoredAvailable(options.getAllocator()); + return !vectoredIoDisabled && options.useHadoopVectoredIo() && f.readVectoredAvailable(options.getAllocator()); + } + + private boolean prepareVectoredReadFileLength() throws IOException { + if (vectoredReadFileLength >= 0) { + return true; + } + try { + vectoredReadFileLength = file.getLength(); + return true; + } catch (IOException failure) { + if (Thread.currentThread().isInterrupted() + || failure instanceof InterruptedIOException + && failure.getCause() instanceof InterruptedException) { + throw failure; + } + // No requests have been submitted and the builder is still empty. An already + // open stream can remain readable when a remote metadata lookup fails. + vectoredIoDisabled = true; + LOG.warn("Cannot determine file length for vectored IO; using normal IO against {}", f, failure); + return false; + } } /** @@ -1334,6 +1364,9 @@ private boolean shouldUseVectoredIo() { * 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. + * Submission and all requested ranges share the vectored-read timeout. Failed + * operations invalidate the stream; cleanup waits for submission to exit before + * closing it, even if the backend does not respond promptly to interruption. * 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. @@ -1344,9 +1377,6 @@ private boolean shouldUseVectoredIo() { private void readVectored(List allParts, ChunkListBuilder builder) throws IOException { final int maximumAllocation = options.getMaxAllocationSize(); Preconditions.checkArgument(maximumAllocation > 0, "Invalid maximum allocation size %s", maximumAllocation); - if (vectoredReadFileLength < 0) { - vectoredReadFileLength = file.getLength(); - } final long fileLength = vectoredReadFileLength; List ranges = new ArrayList<>(allParts.size()); List partRangeCounts = new ArrayList<>(allParts.size()); @@ -1372,50 +1402,98 @@ private void readVectored(List allParts, ChunkListBuilder b totalSize += len; } LOG.debug("Reading {} bytes of data with vectored IO in {} ranges", totalSize, ranges.size()); - final long readStart = System.nanoTime(); + if (vectoredReadExecutor == null) { + vectoredReadExecutor = Executors.newSingleThreadExecutor(task -> { + Thread thread = new Thread(task, "parquet-vectored-read"); + thread.setDaemon(true); + return thread; + }); + } + VectoredReadOperation operation = new VectoredReadOperation( + f, + ranges, + options.getAllocator(), + vectoredReadExecutor, + HADOOP_VECTORED_READ_TIMEOUT_SECONDS, + TimeUnit.SECONDS); try { // Even a synchronous rejection can follow partial submission. The Hadoop bridge - // publishes futures only after submission returns, so missing futures do not prove - // that no reads started. Once this call is entered, normal-read fallback is unsafe. - f.readVectored(ranges, options.getAllocator()); + // may expose futures for reads which were never scheduled. Once submission is + // attempted, neither missing futures nor a synchronous error permit replay. + operation.awaitSubmission(); 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); + allParts.get(partIndex) + .readFromVectoredRanges(ranges.subList(firstRange, endRange), builder, operation); firstRange = endRange; } + operation.transferTo(builder.releaser); + } catch (TimeoutException e) { + IOException failure = new IOException("Timed out submitting vectored reads", e); + abortVectoredRead(operation, failure); + throw failure; } catch (IllegalArgumentException | UnsupportedOperationException e) { // Consumption may also have populated the builder. Do not replay those chunks. IOException failure = new IOException("Vectored read failed after asynchronous reads may have been submitted", e); - awaitRemainingVectoredReads(ranges, readStart, failure); + if (operation.submissionSucceeded()) { + awaitRemainingVectoredReads(ranges, operation, failure); + } + abortVectoredRead(operation, failure); throw failure; - } catch (IOException | RuntimeException e) { - awaitRemainingVectoredReads(ranges, readStart, e); + } catch (IOException | RuntimeException | Error e) { + if (operation.submissionSucceeded()) { + awaitRemainingVectoredReads(ranges, operation, e); + } + abortVectoredRead(operation, e); throw e; } } + private void abortVectoredRead(VectoredReadOperation operation, Throwable failure) { + // Deferred cleanup now owns the original stream and executor. Prevent a later + // call (including an attempted ordinary read) from racing that cleanup. + f = + new DelegatingSeekableInputStream(new InputStream() { + @Override + public int read() throws IOException { + throw new IOException("Cannot reuse a reader after a vectored read failure", failure); + } + }) { + @Override + public long getPos() throws IOException { + throw new IOException("Cannot reuse a reader after a vectored read failure", failure); + } + + @Override + public void seek(long position) throws IOException { + throw new IOException("Cannot reuse a reader after a vectored read failure", failure); + } + }; + vectoredReadExecutor = null; + operation.abort(failure); + } + /** * 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) { + private void awaitRemainingVectoredReads( + List ranges, VectoredReadOperation operation, 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); + FutureIO.awaitFuture(future, operation.remainingNanos(), TimeUnit.NANOSECONDS); } catch (InterruptedIOException e) { if (failure != e) { failure.addSuppressed(e); @@ -1514,13 +1592,27 @@ private ColumnChunkPageReadStore internalReadFilteredRowGroup( } } } - readAllPartsVectoredOrNormal(allParts, builder); rowGroup.setReleaser(builder.releaser); - for (Chunk chunk : builder.build()) { - readChunkPages(chunk, block, rowGroup); + try { + readAllPartsVectoredOrNormal(allParts, builder); + for (Chunk chunk : builder.build()) { + readChunkPages(chunk, block, rowGroup); + } + return rowGroup; + } catch (IOException | RuntimeException | Error failure) { + closeRowGroupAfterFailure(rowGroup, failure); + throw failure; } + } - return rowGroup; + private static void closeRowGroupAfterFailure(ColumnChunkPageReadStore rowGroup, Throwable failure) { + try { + rowGroup.close(); + } catch (RuntimeException closeFailure) { + if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } + } } private void readChunkPages(Chunk chunk, BlockMetaData block, ColumnChunkPageReadStore rowGroup) @@ -1910,6 +2002,9 @@ public void close() throws IOException { f.close(); } } finally { + if (vectoredReadExecutor != null) { + vectoredReadExecutor.shutdownNow(); + } AutoCloseables.uncheckedClose(nextDictionaryReader, crcAllocator); options.getCodecFactory().release(); } @@ -2400,13 +2495,15 @@ private void setReadMetrics(long startNs, long len) { * 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. + * @param operation owns the allocations and the deadline shared with submission and other parts. * @throws IOException if there is an error while reading from the stream, including a timeout. */ - public void readFromVectoredRanges(List ranges, ChunkListBuilder builder) throws IOException { + public void readFromVectoredRanges( + List ranges, ChunkListBuilder builder, VectoredReadOperation operation) + 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 { for (ParquetFileRange range : ranges) { @@ -2415,8 +2512,8 @@ public void readFromVectoredRanges(List ranges, ChunkListBuild "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)); + buffers.add(FutureIO.awaitFuture( + range.getDataReadFuture(), operation.remainingNanos(), TimeUnit.NANOSECONDS)); } setReadMetrics(readStart, length); // report in a counter the data we just scanned diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/VectoredReadBufferAllocator.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/VectoredReadBufferAllocator.java new file mode 100644 index 0000000000..af93d0506d --- /dev/null +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/VectoredReadBufferAllocator.java @@ -0,0 +1,110 @@ +/* + * 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 java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.IdentityHashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import org.apache.parquet.bytes.ByteBufferAllocator; +import org.apache.parquet.bytes.ByteBufferReleaser; +import org.apache.parquet.util.AutoCloseables; + +/** + * Owns the original allocations for one vectored read. Filesystems may return slices of these buffers and allocate + * additional buffers for checksums, so the range results do not identify all the allocations that need releasing. + * + *

The caller must wait for submission and all reads to finish before transferring or closing this owner. Stopping + * allocations, cancelling a result future, or closing the stream does not establish that the filesystem has stopped + * using a previously allocated buffer. + */ +final class VectoredReadBufferAllocator implements ByteBufferAllocator, AutoCloseable { + private final ByteBufferAllocator allocator; + private final Map buffers = new IdentityHashMap<>(); + private volatile boolean acceptingAllocations = true; + + VectoredReadBufferAllocator(ByteBufferAllocator allocator) { + this.allocator = Objects.requireNonNull(allocator, "allocator"); + } + + @Override + public synchronized ByteBuffer allocate(int size) { + if (!acceptingAllocations) { + throw new IllegalStateException("Vectored read is no longer accepting buffer allocations"); + } + ByteBuffer buffer = Objects.requireNonNull(allocator.allocate(size), "allocated buffer"); + // stopAllocating() may run while the delegate is allocating. The already accepted allocation still belongs + // to this owner and must remain available to the filesystem until the caller establishes completion. + buffers.put(buffer, Boolean.TRUE); + return buffer; + } + + @Override + public synchronized void release(ByteBuffer buffer) { + if (buffers.remove(buffer) == null) { + throw new IllegalArgumentException("Buffer is not owned by this vectored read"); + } + allocator.release(buffer); + } + + @Override + public boolean isDirect() { + return allocator.isDirect(); + } + + /** Rejects new allocations without waiting for a potentially blocked delegate allocation or releasing buffers. */ + void stopAllocating() { + acceptingAllocations = false; + } + + /** + * Transfers original buffers to a releaser associated with the delegate allocator. This is only valid after + * successful submission and completion of every range. The owner must not have been aborted or transferred. + */ + synchronized void transferTo(ByteBufferReleaser releaser) { + Objects.requireNonNull(releaser, "releaser"); + if (!acceptingAllocations) { + throw new IllegalStateException("Vectored read buffer ownership has already been stopped or transferred"); + } + acceptingAllocations = false; + Iterator iterator = buffers.keySet().iterator(); + while (iterator.hasNext()) { + releaser.releaseLater(iterator.next()); + iterator.remove(); + } + } + + /** + * Releases the still-owned originals exactly once. The caller must establish that no backend work can use them; + * this method itself does not wait for asynchronous IO. Buffers already transferred to a row group are unaffected. + */ + @Override + public synchronized void close() { + acceptingAllocations = false; + List releases = new ArrayList<>(buffers.size()); + for (ByteBuffer buffer : buffers.keySet()) { + releases.add(() -> allocator.release(buffer)); + } + buffers.clear(); + AutoCloseables.uncheckedClose(releases); + } +} diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/VectoredReadOperation.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/VectoredReadOperation.java new file mode 100644 index 0000000000..42618a2a65 --- /dev/null +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/VectoredReadOperation.java @@ -0,0 +1,179 @@ +/* + * 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 java.io.IOException; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import org.apache.parquet.bytes.ByteBufferAllocator; +import org.apache.parquet.bytes.ByteBufferReleaser; +import org.apache.parquet.hadoop.util.wrapped.io.FutureIO; +import org.apache.parquet.io.ParquetFileRange; +import org.apache.parquet.io.SeekableInputStream; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Owns one vectored submission and its allocations until the reader can take ownership. + * The executor must have a single worker: queued failure cleanup must not overtake a + * submission which is still running after its Future has been cancelled. + */ +final class VectoredReadOperation { + private static final Logger LOG = LoggerFactory.getLogger(VectoredReadOperation.class); + + private final SeekableInputStream stream; + private final List ranges; + private final VectoredReadBufferAllocator allocator; + private final ExecutorService executor; + private final long timeoutNanos; + private final long readStart = System.nanoTime(); + private Future submission; + private volatile boolean submissionSucceeded; + private boolean aborted; + private boolean releaseRegistered; + + VectoredReadOperation( + SeekableInputStream stream, + List ranges, + ByteBufferAllocator allocator, + ExecutorService executor, + long timeout, + TimeUnit unit) { + this.stream = stream; + this.ranges = ranges; + this.allocator = new VectoredReadBufferAllocator(allocator); + this.executor = executor; + this.timeoutNanos = unit.toNanos(timeout); + } + + void awaitSubmission() throws IOException, TimeoutException { + submission = executor.submit(() -> { + stream.readVectored(ranges, allocator); + submissionSucceeded = true; + return null; + }); + FutureIO.awaitFuture(submission, remainingNanos(), TimeUnit.NANOSECONDS); + } + + boolean submissionSucceeded() { + return submissionSucceeded; + } + + long remainingNanos() { + return Math.max(timeoutNanos - (System.nanoTime() - readStart), 0L); + } + + void transferTo(ByteBufferReleaser releaser) { + if (!submissionSucceeded || aborted) { + throw new IllegalStateException("Cannot transfer buffers from an unsuccessful vectored read"); + } + for (ParquetFileRange range : ranges) { + CompletableFuture future = range.getDataReadFuture(); + if (future == null || !future.isDone() || future.isCompletedExceptionally()) { + throw new IllegalStateException("Cannot transfer buffers before all vectored reads succeed"); + } + } + allocator.transferTo(releaser); + } + + /** + * Stop the caller's wait without treating interruption as proof that backend IO stopped. + * Once aborted, the reader must not reuse the stream or this executor. + */ + void abort(Throwable failure) { + if (aborted) { + return; + } + aborted = true; + allocator.stopAllocating(); + if (submission != null) { + submission.cancel(true); + } + if (submissionSucceeded) { + // All futures are published. If the failed read and its siblings already + // finished, reclaim their buffers before the caller closes its allocator. + releaseWhenReadsFinish(failure); + } + try { + // Future.cancel(true) may return while the callable is still running. A task on + // the same single worker cannot close the stream until that callable has exited. + executor.execute(() -> { + releaseWhenReadsFinish(failure); + try { + stream.close(); + } catch (IOException | RuntimeException closeFailure) { + if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } + LOG.warn("Failed to close a stream after a vectored read failure", closeFailure); + } + }); + } catch (RejectedExecutionException cleanupFailure) { + // The reader owns this executor and must not shut it down before queuing cleanup. + // Do not mask the read error or recycle buffers whose IO lifetime is now unknown. + failure.addSuppressed(cleanupFailure); + LOG.warn("Could not schedule cleanup after a vectored read failure", cleanupFailure); + } finally { + // shutdownNow would discard the queued cleanup or interrupt its close operation. + executor.shutdown(); + } + } + + private void releaseWhenReadsFinish(Throwable failure) { + if (releaseRegistered) { + return; + } + releaseRegistered = true; + CompletableFuture[] futures = new CompletableFuture[ranges.size()]; + for (int i = 0; i < ranges.size(); i++) { + CompletableFuture future = ranges.get(i).getDataReadFuture(); + if (future == null) { + // A backend may have started IO without publishing its result. Neither closing + // its stream nor cancelling a future establishes that pooled memory is reusable. + LOG.debug("Retaining allocations after a vectored submission with unpublished reads"); + return; + } + futures[i] = future; + } + // A failed submission may also publish a future for work never submitted. Such a + // future need not complete: retain its ownership without blocking a cleanup worker. + CompletableFuture.allOf(futures).whenComplete((ignored, readFailure) -> { + for (CompletableFuture future : futures) { + if (future.isCancelled()) { + LOG.debug("Retaining allocations after a cancelled vectored result with unknown IO lifetime"); + return; + } + } + try { + allocator.close(); + } catch (RuntimeException releaseFailure) { + if (failure != releaseFailure) { + failure.addSuppressed(releaseFailure); + } + LOG.warn("Failed to release buffers after a vectored read failure", releaseFailure); + } + }); + } +} diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/util/wrapped/io/VectorIoBridge.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/util/wrapped/io/VectorIoBridge.java index be5cc80251..ce26477e78 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/util/wrapped/io/VectorIoBridge.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/util/wrapped/io/VectorIoBridge.java @@ -156,13 +156,15 @@ private void checkAvailable() { * The default iterates through the ranges to read each synchronously, but * the intent is that FSDataInputStream subclasses can make more efficient * readers. - * The {@link ParquetFileRange} parameters all have their - * data read futures set to the range reads of the associated - * operations; callers must await these to complete. + * After a successful return, the {@link ParquetFileRange} parameters all have + * their data read futures set to the associated operations; callers can await + * these futures for the read results. *

- * As a result of the call, each range will have FileRange.setData(CompletableFuture) - * called with a future that when complete will have a ByteBuffer with the - * data from the file's range. + * If submission throws, futures already assigned by Hadoop are still published. + * A null future does not prove that no work started. A non-null future may have + * been created for a read that was never submitted and may never complete after + * a submission failure. Callers must handle that failure without waiting for + * every such future. *

* The position returned by getPos() after readVectored() is undefined. *

@@ -199,15 +201,36 @@ public void readVectoredRanges( // Setting the parquet range as a reference. List fileRanges = sorted.stream().map(rangeBridge::toFileRange).collect(Collectors.toList()); - readWrappedRanges(stream, fileRanges, allocator::allocate); - - // copy back the completable futures from the scheduled - // vector reads to the ParquetFileRange entries passed in. - fileRanges.forEach(fileRange -> { - // toFileRange() sets up this back reference - ParquetFileRange parquetFileRange = (ParquetFileRange) fileRange.getReference(); - parquetFileRange.setDataReadFuture(fileRange.getData()); - }); + Throwable submissionFailure = null; + try { + readWrappedRanges(stream, fileRanges, allocator::allocate); + } catch (IOException | RuntimeException | Error failure) { + submissionFailure = failure; + throw failure; + } finally { + publishReadFutures(fileRanges, submissionFailure); + } + } + + /** + * Publish available futures even after partial submission, without replacing + * the submission failure if a future cannot be retrieved. + */ + static void publishReadFutures(List fileRanges, Throwable submissionFailure) { + for (FileRangeBridge.WrappedFileRange fileRange : fileRanges) { + try { + // toFileRange() sets up this back reference. + ParquetFileRange parquetFileRange = (ParquetFileRange) fileRange.getReference(); + parquetFileRange.setDataReadFuture(fileRange.getData()); + } catch (RuntimeException | Error publicationFailure) { + if (submissionFailure == null) { + throw publicationFailure; + } + if (submissionFailure != publicationFailure) { + submissionFailure.addSuppressed(publicationFailure); + } + } + } } /** 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 index 4261250aab..41d65591c9 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileReaderVectoredIO.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileReaderVectoredIO.java @@ -28,6 +28,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import java.io.IOException; +import java.io.InterruptedIOException; import java.net.SocketTimeoutException; import java.nio.ByteBuffer; import java.util.ArrayList; @@ -166,6 +167,58 @@ public void testSuppliedFooterAvoidsLengthLookupForOrdinaryReads() throws Except assertEquals(0, stream.vectorCalls); } + @Test + public void testLengthLookupFailureFallsBackWithSuppliedFooter() throws Exception { + assertLengthLookupFailureFallsBack(true); + } + + @Test + public void testLengthLookupFailureFallsBackWithoutSuppliedFooter() throws Exception { + assertLengthLookupFailureFallsBack(false); + } + + private void assertLengthLookupFailureFallsBack(boolean supplyFooter) throws Exception { + ParquetMetadata footer = ParquetFileReader.readFooter(new Configuration(), path); + CountingInputFile countingFile = new CountingInputFile(inputFile); + countingFile.successfulLengthCalls = supplyFooter ? 0 : 1; + RecordingSeekableInputStream stream = newStream(FailureMode.NONE); + ParquetReadOptions options = readOptions(new RecordingAllocator(), 128, true); + try (ParquetFileReader reader = supplyFooter + ? ParquetFileReader.open(countingFile, footer, options, stream) + : ParquetFileReader.open(countingFile, options, stream)) { + reader.setRequestedSchema(PROJECTED_SCHEMA); + try (PageReadStore pages = reader.readRowGroup(0)) { + assertRows(pages, PROJECTED_SCHEMA, false); + } + try (PageReadStore pages = reader.readFilteredRowGroup(0)) { + assertRows(pages, PROJECTED_SCHEMA, true); + } + // Once this reader has selected ordinary IO, it must not repeat a failed + // metadata lookup on each subsequent row-group read. + assertEquals(countingFile.successfulLengthCalls + 1, countingFile.lengthCalls); + assertEquals(0, stream.vectorCalls); + } + } + + @Test + public void testInterruptedLengthLookupDoesNotFallBack() throws Exception { + ParquetMetadata footer = ParquetFileReader.readFooter(new Configuration(), path); + CountingInputFile countingFile = new CountingInputFile(inputFile); + countingFile.successfulLengthCalls = 0; + countingFile.lengthFailure = new InterruptedIOException("interrupted metadata lookup"); + countingFile.lengthFailure.initCause(new InterruptedException("cancelled read")); + RecordingSeekableInputStream stream = newStream(FailureMode.NONE); + try (ParquetFileReader reader = ParquetFileReader.open( + countingFile, footer, readOptions(new RecordingAllocator(), 128, false), stream)) { + stream.resetOrdinaryReads(); + IOException failure = assertThrows(InterruptedIOException.class, () -> reader.readRowGroup(0)); + assertThat(failure).isSameAs(countingFile.lengthFailure); + assertEquals(0, stream.vectorCalls); + assertEquals(0, stream.normalSeekCalls); + assertEquals(0, stream.normalReadCalls); + } + } + @Test public void testSplitsAdjacentColumnsAtMaximumAllocation() throws Exception { List columns = ParquetFileReader.readFooter(new Configuration(), path) @@ -572,14 +625,18 @@ private void assertFailsAfterVectoredSubmission(FailureMode failureMode, Class reader.readRowGroup(0)); + assertThat(retryFailure).hasMessageContaining("Cannot reuse a reader after a vectored read failure"); + assertEquals(1, stream.vectorCalls); + assertEquals(0, stream.normalSeekCalls); + assertEquals(0, stream.normalReadCalls); if (failureMode.hasPendingRead()) { if (failureMode.hasPublishedPendingRead()) { assertEquals(0, stream.pendingFutureCount()); - } else { - assertThat(stream.pendingFutureCount() > 0).isTrue(); } } } + stream.closeCompleted.get(10, TimeUnit.SECONDS); assertEquals(1, stream.vectorCalls); assertEquals(0, stream.pendingFutureCount()); } @@ -722,6 +779,8 @@ private boolean hasPublishedPendingRead() { private static final class CountingInputFile implements InputFile { private final InputFile delegate; private int lengthCalls; + private int successfulLengthCalls = Integer.MAX_VALUE; + private IOException lengthFailure = new IOException("injected metadata lookup failure"); private CountingInputFile(InputFile delegate) { this.delegate = delegate; @@ -730,6 +789,9 @@ private CountingInputFile(InputFile delegate) { @Override public long getLength() throws IOException { lengthCalls++; + if (lengthCalls > successfulLengthCalls) { + throw lengthFailure; + } return delegate.getLength(); } @@ -767,6 +829,7 @@ private static final class RecordingSeekableInputStream extends DelegatingSeekab private final CompletableFuture allowPendingPhysicalReads = new CompletableFuture<>(); private final CompletableFuture pendingDrainStarted = new CompletableFuture<>(); private final CompletableFuture socketTimeoutDrainStarted = new CompletableFuture<>(); + private final CompletableFuture closeCompleted = new CompletableFuture<>(); private final AtomicInteger completedPendingPhysicalReads = new AtomicInteger(); private final AtomicInteger postClosePhysicalReads = new AtomicInteger(); private volatile boolean closed; @@ -941,7 +1004,11 @@ public void close() throws IOException { pendingFuture.cancel(false); } } - super.close(); + try { + super.close(); + } finally { + closeCompleted.complete(null); + } } } diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileReaderVectoredOwnership.java b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileReaderVectoredOwnership.java new file mode 100644 index 0000000000..389244308e --- /dev/null +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileReaderVectoredOwnership.java @@ -0,0 +1,314 @@ +/* + * 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.intColumn; +import static org.apache.parquet.filter2.predicate.FilterApi.ltEq; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.PrimitiveIterator; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +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.bytes.TrackingByteBufferAllocator; +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.hadoop.example.ExampleParquetWriter; +import org.apache.parquet.hadoop.metadata.ParquetMetadata; +import org.apache.parquet.hadoop.util.HadoopInputFile; +import org.apache.parquet.io.ColumnIOFactory; +import org.apache.parquet.io.DelegatingSeekableInputStream; +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; + +class TestParquetFileReaderVectoredOwnership { + private static final int ROWS = 128; + private static final MessageType SCHEMA = + MessageTypeParser.parseMessageType("message test { required int32 id; required binary padding (UTF8); }"); + + @TempDir + private java.nio.file.Path tempDir; + + private HadoopInputFile inputFile; + private ParquetMetadata footer; + + @BeforeEach + void writeFile() throws IOException { + Configuration conf = new Configuration(); + Path path = new Path(tempDir.resolve("ownership.parquet").toUri()); + try (ParquetWriter writer = ExampleParquetWriter.builder(path) + .withConf(conf) + .withType(SCHEMA) + .withWriteMode(ParquetFileWriter.Mode.OVERWRITE) + .withRowGroupSize(256 * 1024) + .withPageSize(128) + .withPageRowCountLimit(8) + .withDictionaryEncoding(false) + .build()) { + SimpleGroupFactory factory = new SimpleGroupFactory(SCHEMA); + for (int row = 0; row < ROWS; row++) { + writer.write(factory.newGroup().append("id", row).append("padding", "padding_" + row)); + } + } + inputFile = HadoopInputFile.fromPath(path, conf); + footer = ParquetFileReader.readFooter(conf, path); + } + + @Test + void testReleasesOriginalVectoredBuffersWhenRowGroupCloses() throws Exception { + assertSuccessfulReadReleasesBuffers(ReadMode.ORIGINALS); + } + + @Test + void testReleasesMergedOriginalAndChecksumBufferWhenResultsAreSlices() throws Exception { + assertSuccessfulReadReleasesBuffers(ReadMode.SLICES); + } + + @Test + void testReleasesSlicedBuffersForFilteredRowGroup() throws Exception { + CountingAllocator delegate = new CountingAllocator(); + try (TrackingByteBufferAllocator tracking = TrackingByteBufferAllocator.wrap(delegate)) { + ParquetReadOptions readOptions = ParquetReadOptions.builder() + .withAllocator(tracking) + .withUseHadoopVectoredIo(true) + .withMaxAllocationInBytes(128) + .useColumnIndexFilter(true) + .withRecordFilter(FilterCompat.get(ltEq(intColumn("id"), 63))) + .build(); + try (ParquetFileReader reader = ParquetFileReader.open( + inputFile, footer, readOptions, new OwnedStream(inputFile.newStream(), ReadMode.SLICES))) { + try (PageReadStore pages = reader.readFilteredRowGroup(0)) { + assertTrue(pages.getRowCount() > 0 && pages.getRowCount() < ROWS); + assertTrue(pages.getRowIndexes().isPresent()); + PrimitiveIterator.OfLong indexes = pages.getRowIndexes().get(); + RecordReader records = new ColumnIOFactory() + .getColumnIO(SCHEMA) + .getRecordReader(pages, new GroupRecordConverter(SCHEMA)); + for (long row = 0; row < pages.getRowCount(); row++) { + long expectedIndex = indexes.nextLong(); + Group record = records.read(); + assertEquals(expectedIndex, record.getInteger("id", 0)); + assertEquals("padding_" + expectedIndex, record.getString("padding", 0)); + } + } + assertEquals(delegate.allocations.get(), delegate.releases.get()); + } + } + } + + private void assertSuccessfulReadReleasesBuffers(ReadMode mode) throws Exception { + CountingAllocator delegate = new CountingAllocator(); + try (TrackingByteBufferAllocator tracking = TrackingByteBufferAllocator.wrap(delegate)) { + OwnedStream stream = new OwnedStream(inputFile.newStream(), mode); + try (ParquetFileReader reader = ParquetFileReader.open(inputFile, footer, options(tracking), stream)) { + try (PageReadStore pages = reader.readRowGroup(0)) { + assertEquals(ROWS, pages.getRowCount()); + RecordReader records = new ColumnIOFactory() + .getColumnIO(SCHEMA) + .getRecordReader(pages, new GroupRecordConverter(SCHEMA)); + for (int row = 0; row < ROWS; row++) { + Group record = records.read(); + assertEquals(row, record.getInteger("id", 0)); + assertEquals("padding_" + row, record.getString("padding", 0)); + } + assertEquals(0, delegate.releases.get()); + } + assertTrue(stream.rangeCount > 1); + if (mode == ReadMode.SLICES) { + assertEquals(2, delegate.allocations.get()); + } else { + assertEquals(stream.rangeCount, delegate.allocations.get()); + } + assertEquals(delegate.allocations.get(), delegate.releases.get()); + } + } + } + + @Test + void testReleasesTransferredBuffersWhenPageHeaderParsingFails() throws Exception { + CountingAllocator delegate = new CountingAllocator(); + try (TrackingByteBufferAllocator tracking = TrackingByteBufferAllocator.wrap(delegate)) { + try (ParquetFileReader reader = ParquetFileReader.open( + inputFile, footer, options(tracking), new OwnedStream(inputFile.newStream(), ReadMode.CORRUPT))) { + assertThrows(IOException.class, () -> reader.readRowGroup(0)); + assertEquals(delegate.allocations.get(), delegate.releases.get()); + } + assertTrue(delegate.allocations.get() > 0); + assertEquals(delegate.allocations.get(), delegate.releases.get()); + } + } + + @Test + void testReleasesFailedAndSuccessfulAllocationsAfterAllRangesFinish() throws Exception { + CountingAllocator delegate = new CountingAllocator(); + try (TrackingByteBufferAllocator tracking = TrackingByteBufferAllocator.wrap(delegate)) { + OwnedStream stream = new OwnedStream(inputFile.newStream(), ReadMode.FAILED_RANGE); + try (ParquetFileReader reader = ParquetFileReader.open(inputFile, footer, options(tracking), stream)) { + try { + assertThrows(IOException.class, () -> reader.readRowGroup(0)); + // All range futures already finished. Reclamation must not wait for deferred stream closure, + // otherwise an outer allocator.close() can race the delayed release. + assertEquals(delegate.allocations.get(), delegate.releases.get()); + } finally { + stream.allowClose.countDown(); + } + } + assertTrue(delegate.allocations.get() > 1); + assertEquals(delegate.allocations.get(), delegate.releases.get()); + } + } + + private static ParquetReadOptions options(ByteBufferAllocator allocator) { + return ParquetReadOptions.builder() + .withAllocator(allocator) + .withUseHadoopVectoredIo(true) + .withMaxAllocationInBytes(128) + .build(); + } + + private enum ReadMode { + ORIGINALS, + SLICES, + CORRUPT, + FAILED_RANGE + } + + private static final class CountingAllocator extends HeapByteBufferAllocator { + private final AtomicInteger allocations = new AtomicInteger(); + private final AtomicInteger releases = new AtomicInteger(); + + @Override + public ByteBuffer allocate(int size) { + allocations.incrementAndGet(); + return super.allocate(size); + } + + @Override + public void release(ByteBuffer buffer) { + releases.incrementAndGet(); + } + } + + /** Supports vectored IO independently of the Hadoop version used to run the test. */ + private static final class OwnedStream extends DelegatingSeekableInputStream { + private final SeekableInputStream delegate; + private final ReadMode mode; + private final CountDownLatch allowClose = new CountDownLatch(1); + private int rangeCount; + + OwnedStream(SeekableInputStream delegate, ReadMode mode) { + super(delegate); + this.delegate = delegate; + this.mode = mode; + } + + @Override + public long getPos() throws IOException { + return delegate.getPos(); + } + + @Override + public void seek(long newPos) throws IOException { + delegate.seek(newPos); + } + + @Override + public boolean readVectoredAvailable(ByteBufferAllocator allocator) { + return true; + } + + @Override + public void close() throws IOException { + try { + if (mode == ReadMode.FAILED_RANGE) { + try { + if (!allowClose.await(10, TimeUnit.SECONDS)) { + throw new IOException("Test did not unblock stream closure"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted while waiting for test stream closure", e); + } + } + } finally { + super.close(); + } + } + + @Override + public void readVectored(List ranges, ByteBufferAllocator allocator) throws IOException { + rangeCount = ranges.size(); + ByteBuffer merged = null; + if (mode == ReadMode.SLICES) { + // Model filesystem coalescing and checksum verification: several result slices share one original, + // while an additional checksum allocation does not appear in any returned range. + int totalLength = + ranges.stream().mapToInt(ParquetFileRange::getLength).sum(); + merged = allocator.allocate(totalLength); + allocator.allocate(8); + } + for (int index = 0; index < ranges.size(); index++) { + ParquetFileRange range = ranges.get(index); + ByteBuffer buffer; + if (merged == null) { + buffer = allocator.allocate(range.getLength()); + } else { + buffer = merged.slice(); + buffer.limit(range.getLength()); + buffer = buffer.slice(); + merged.position(merged.position() + range.getLength()); + } + if (mode == ReadMode.CORRUPT) { + buffer.position(range.getLength()); // Zero-filled bytes are not a valid page header. + } else { + delegate.seek(range.getOffset()); + delegate.readFully(buffer); + } + buffer.flip(); + CompletableFuture future = new CompletableFuture<>(); + if (mode == ReadMode.FAILED_RANGE && index == 0) { + future.completeExceptionally(new IOException("injected range failure after allocation")); + } else { + future.complete(buffer); + } + range.setDataReadFuture(future); + } + } + } +} diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestVectoredReadBufferAllocator.java b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestVectoredReadBufferAllocator.java new file mode 100644 index 0000000000..0f3dfcb815 --- /dev/null +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestVectoredReadBufferAllocator.java @@ -0,0 +1,256 @@ +/* + * 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.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.parquet.bytes.ByteBufferAllocator; +import org.apache.parquet.bytes.ByteBufferReleaser; +import org.apache.parquet.bytes.HeapByteBufferAllocator; +import org.apache.parquet.bytes.TrackingByteBufferAllocator; +import org.apache.parquet.util.AutoCloseables; +import org.junit.jupiter.api.Test; + +class TestVectoredReadBufferAllocator { + @Test + void testTransfersOriginalBuffersAndAdditionalChecksumAllocations() { + CountingAllocator delegate = new CountingAllocator(); + try (TrackingByteBufferAllocator tracking = TrackingByteBufferAllocator.wrap(delegate); + VectoredReadBufferAllocator owner = new VectoredReadBufferAllocator(tracking); + ByteBufferReleaser rowGroup = new ByteBufferReleaser(tracking)) { + ByteBuffer first = owner.allocate(32); + ByteBuffer second = owner.allocate(32); + owner.allocate(8); // A filesystem checksum allocation does not appear in the range results. + assertEquals(first, second); // Content equality must not collapse distinct original allocations. + ByteBuffer firstResult = first.slice(); + ByteBuffer secondResult = second.slice(); + assertNotSame(first, firstResult); + assertNotSame(second, secondResult); + + owner.transferTo(rowGroup); + owner.close(); + assertEquals(0, delegate.releases.get()); + firstResult.put(0, (byte) 37); + assertEquals(37, first.get(0)); + assertThrows(IllegalStateException.class, () -> owner.allocate(1)); + assertThrows(IllegalStateException.class, () -> owner.transferTo(rowGroup)); + + rowGroup.close(); + assertEquals(3, delegate.releases.get()); + } + assertEquals(3, delegate.releases.get()); + } + + @Test + void testStopAllocatingDoesNotReleaseBuffersStillUsedByIo() { + CountingAllocator delegate = new CountingAllocator(); + try (TrackingByteBufferAllocator tracking = TrackingByteBufferAllocator.wrap(delegate); + VectoredReadBufferAllocator owner = new VectoredReadBufferAllocator(tracking); + ByteBufferReleaser rowGroup = new ByteBufferReleaser(tracking)) { + ByteBuffer buffer = owner.allocate(32); + owner.stopAllocating(); + owner.stopAllocating(); + assertThrows(IllegalStateException.class, () -> owner.allocate(16)); + assertThrows(IllegalStateException.class, () -> owner.transferTo(rowGroup)); + assertEquals(1, delegate.allocations.get()); + assertEquals(0, delegate.releases.get()); + buffer.putInt(0, 1234); // An already accepted read may still finish after abort. + assertEquals(1234, buffer.getInt(0)); + + owner.close(); // The caller has now established that the read finished. + owner.close(); + assertEquals(1, delegate.releases.get()); + } + } + + @Test + void testBackendReleasedOriginalIsNotReleasedAgain() { + CountingAllocator delegate = new CountingAllocator(); + try (TrackingByteBufferAllocator tracking = TrackingByteBufferAllocator.wrap(delegate); + VectoredReadBufferAllocator owner = new VectoredReadBufferAllocator(tracking); + ByteBufferReleaser rowGroup = new ByteBufferReleaser(tracking)) { + ByteBuffer original = owner.allocate(32); + owner.allocate(32); + assertThrows(IllegalArgumentException.class, () -> owner.release(original.slice())); + assertEquals(0, delegate.releases.get()); + owner.release(original); + assertEquals(1, delegate.releases.get()); + assertThrows(IllegalArgumentException.class, () -> owner.release(original)); + + owner.transferTo(rowGroup); + owner.close(); + rowGroup.close(); + assertEquals(2, delegate.releases.get()); + } + } + + @Test + void testAbortDoesNotWaitForBlockedDelegateAllocation() throws Exception { + CountDownLatch allocationStarted = new CountDownLatch(1); + CountDownLatch allowAllocation = new CountDownLatch(1); + AtomicInteger releases = new AtomicInteger(); + ByteBufferAllocator delegate = new HeapByteBufferAllocator() { + @Override + public ByteBuffer allocate(int size) { + allocationStarted.countDown(); + try { + if (!allowAllocation.await(10, TimeUnit.SECONDS)) { + throw new IllegalStateException("Test did not unblock allocation"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(e); + } + return super.allocate(size); + } + + @Override + public void release(ByteBuffer buffer) { + releases.incrementAndGet(); + } + }; + try (TrackingByteBufferAllocator tracking = TrackingByteBufferAllocator.wrap(delegate); + VectoredReadBufferAllocator owner = new VectoredReadBufferAllocator(tracking)) { + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future allocation = executor.submit(() -> owner.allocate(32)); + assertTrue(allocationStarted.await(10, TimeUnit.SECONDS)); + executor.submit(owner::stopAllocating).get(5, TimeUnit.SECONDS); + assertFalse(allocation.isDone()); + assertEquals(0, releases.get()); + + allowAllocation.countDown(); + ByteBuffer buffer = allocation.get(10, TimeUnit.SECONDS); + buffer.putInt(0, 1234); + assertEquals(1234, buffer.getInt(0)); + assertEquals(0, releases.get()); + assertThrows(IllegalStateException.class, () -> owner.allocate(32)); + owner.close(); + assertEquals(1, releases.get()); + } finally { + allowAllocation.countDown(); + owner.stopAllocating(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + } + + @Test + void testConcurrentAllocationsAndAbortRetainEveryAcceptedOriginal() throws Exception { + CountingAllocator delegate = new CountingAllocator(); + try (TrackingByteBufferAllocator tracking = TrackingByteBufferAllocator.wrap(delegate); + VectoredReadBufferAllocator owner = new VectoredReadBufferAllocator(tracking)) { + CountDownLatch start = new CountDownLatch(1); + CountDownLatch firstAllocation = new CountDownLatch(1); + ExecutorService executor = Executors.newFixedThreadPool(5); + try { + List> workers = new ArrayList<>(); + for (int worker = 0; worker < 4; worker++) { + workers.add(executor.submit(() -> { + assertTrue(start.await(10, TimeUnit.SECONDS)); + for (int allocation = 0; allocation < 100; allocation++) { + try { + owner.allocate(32); + firstAllocation.countDown(); + } catch (IllegalStateException stopped) { + break; + } + } + return null; + })); + } + Future abort = executor.submit(() -> { + assertTrue(firstAllocation.await(10, TimeUnit.SECONDS)); + owner.stopAllocating(); + return null; + }); + start.countDown(); + for (Future worker : workers) { + worker.get(10, TimeUnit.SECONDS); + } + abort.get(10, TimeUnit.SECONDS); + assertTrue(delegate.allocations.get() > 0); + assertEquals(0, delegate.releases.get()); + assertThrows(IllegalStateException.class, () -> owner.allocate(32)); + owner.close(); + assertEquals(delegate.allocations.get(), delegate.releases.get()); + } finally { + start.countDown(); + owner.stopAllocating(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + } + + @Test + void testCloseAttemptsEveryReleaseEvenWhenOneFails() { + AtomicInteger releases = new AtomicInteger(); + RuntimeException failure = new IllegalStateException("injected release failure"); + ByteBufferAllocator delegate = new HeapByteBufferAllocator() { + @Override + public void release(ByteBuffer buffer) { + if (releases.incrementAndGet() == 1) { + throw failure; + } + } + }; + VectoredReadBufferAllocator owner = new VectoredReadBufferAllocator(delegate); + owner.allocate(8); + owner.allocate(8); + owner.allocate(8); + AutoCloseables.ParquetCloseResourceException error = + assertThrows(AutoCloseables.ParquetCloseResourceException.class, owner::close); + assertSame(failure, error.getCause()); + assertEquals(3, releases.get()); + owner.close(); + assertEquals(3, releases.get()); + } + + private static final class CountingAllocator extends HeapByteBufferAllocator { + private final AtomicInteger allocations = new AtomicInteger(); + private final AtomicInteger releases = new AtomicInteger(); + + @Override + public ByteBuffer allocate(int size) { + allocations.incrementAndGet(); + return super.allocate(size); + } + + @Override + public void release(ByteBuffer buffer) { + releases.incrementAndGet(); + } + } +} diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestVectoredReadOperation.java b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestVectoredReadOperation.java new file mode 100644 index 0000000000..7c8cc8cffa --- /dev/null +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestVectoredReadOperation.java @@ -0,0 +1,472 @@ +/* + * 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.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InterruptedIOException; +import java.nio.ByteBuffer; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.parquet.bytes.ByteBufferAllocator; +import org.apache.parquet.bytes.ByteBufferReleaser; +import org.apache.parquet.io.DelegatingSeekableInputStream; +import org.apache.parquet.io.ParquetFileRange; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +@Timeout(10) +public class TestVectoredReadOperation { + private final List executors = new ArrayList<>(); + + @AfterEach + public void stopExecutors() throws InterruptedException { + for (ExecutorService executor : executors) { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS), "Submission worker did not stop"); + } + } + + @Test + public void testTimeoutIncludesBlockedSubmission() throws Exception { + RecordingAllocator allocator = new RecordingAllocator(); + CountDownLatch readsPublished = new CountDownLatch(1); + CountDownLatch finishSubmission = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + TestStream stream = new TestStream((ranges, buffers) -> { + ByteBuffer buffer = buffers.allocate(8); + CompletableFuture read = new CompletableFuture<>(); + ranges.get(0).setDataReadFuture(read); + readsPublished.countDown(); + try { + if (!finishSubmission.await(5, TimeUnit.SECONDS)) { + throw new IOException("Test submission was not released"); + } + read.complete(buffer); + } catch (InterruptedException e) { + interrupted.countDown(); + InterruptedIOException failure = new InterruptedIOException("Submission interrupted"); + failure.initCause(e); + read.completeExceptionally(failure); + throw failure; + } + }); + ExecutorService executor = newExecutor(); + VectoredReadOperation operation = + new VectoredReadOperation(stream, ranges(1), allocator, executor, 50, TimeUnit.MILLISECONDS); + try { + TimeoutException failure = awaitSubmissionTimeout(operation); + await(readsPublished); + assertFalse(operation.submissionSucceeded()); + assertEquals(0L, operation.remainingNanos()); + operation.abort(failure); + + await(interrupted); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + assertEquals(1, stream.closes.get()); + assertFalse(stream.closedDuringSubmission); + assertEquals(1, allocator.released.size()); + } finally { + finishSubmission.countDown(); + } + } + + @Test + public void testCancelledSubmissionDoesNotMeanBackendHasStopped() throws Exception { + RecordingAllocator allocator = new RecordingAllocator(); + CountDownLatch readsPublished = new CountDownLatch(1); + CountDownLatch finishSubmission = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + TestStream stream = new TestStream((ranges, buffers) -> { + ByteBuffer buffer = buffers.allocate(8); + CompletableFuture read = new CompletableFuture<>(); + ranges.get(0).setDataReadFuture(read); + readsPublished.countDown(); + awaitIgnoringInterrupts(finishSubmission, interrupted); + buffer.put(0, (byte) 37); + read.complete(buffer); + }); + ExecutorService executor = newExecutor(); + VectoredReadOperation operation = + new VectoredReadOperation(stream, ranges(1), allocator, executor, 50, TimeUnit.MILLISECONDS); + try { + TimeoutException failure = awaitSubmissionTimeout(operation); + await(readsPublished); + assertTimeoutPreemptively(Duration.ofSeconds(2), () -> operation.abort(failure)); + await(interrupted); + + assertTrue(stream.submitting); + assertEquals(0, stream.closes.get()); + assertTrue(allocator.released.isEmpty()); + + finishSubmission.countDown(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + assertFalse(stream.closedDuringSubmission); + assertEquals(1, stream.closes.get()); + assertEquals(1, allocator.released.size()); + assertEquals((byte) 37, allocator.released.get(0).get(0)); + } finally { + finishSubmission.countDown(); + } + } + + @Test + public void testAbortRejectsAllocationsRequestedByLateSubmission() throws Exception { + RecordingAllocator allocator = new RecordingAllocator(); + CountDownLatch readsPublished = new CountDownLatch(1); + CountDownLatch finishSubmission = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + AtomicReference allocationFailure = new AtomicReference<>(); + TestStream stream = new TestStream((ranges, buffers) -> { + CompletableFuture read = new CompletableFuture<>(); + ranges.get(0).setDataReadFuture(read); + readsPublished.countDown(); + awaitIgnoringInterrupts(finishSubmission, interrupted); + try { + read.complete(buffers.allocate(8)); + } catch (RuntimeException e) { + allocationFailure.set(e); + read.completeExceptionally(e); + } + }); + ExecutorService executor = newExecutor(); + VectoredReadOperation operation = + new VectoredReadOperation(stream, ranges(1), allocator, executor, 50, TimeUnit.MILLISECONDS); + try { + TimeoutException failure = awaitSubmissionTimeout(operation); + await(readsPublished); + operation.abort(failure); + await(interrupted); + assertEquals(0, stream.closes.get()); + + finishSubmission.countDown(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + assertNotNull(allocationFailure.get()); + assertTrue(allocator.allocated.isEmpty()); + assertTrue(allocator.released.isEmpty()); + assertEquals(1, stream.closes.get()); + assertFalse(stream.closedDuringSubmission); + } finally { + finishSubmission.countDown(); + } + } + + @Test + public void testAbortDoesNotBlockBehindAnAllocationInProgress() throws Exception { + CountDownLatch allocationStarted = new CountDownLatch(1); + CountDownLatch finishAllocation = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + RecordingAllocator allocator = new RecordingAllocator() { + @Override + public ByteBuffer allocate(int size) { + allocationStarted.countDown(); + try { + awaitIgnoringInterrupts(finishAllocation, interrupted); + } catch (IOException e) { + throw new IllegalStateException(e); + } + return super.allocate(size); + } + }; + TestStream stream = new TestStream((ranges, buffers) -> { + CompletableFuture read = new CompletableFuture<>(); + ranges.get(0).setDataReadFuture(read); + try { + read.complete(buffers.allocate(8)); + } catch (RuntimeException e) { + read.completeExceptionally(e); + } + }); + ExecutorService executor = newExecutor(); + VectoredReadOperation operation = + new VectoredReadOperation(stream, ranges(1), allocator, executor, 50, TimeUnit.MILLISECONDS); + try { + TimeoutException failure = awaitSubmissionTimeout(operation); + await(allocationStarted); + assertTimeoutPreemptively(Duration.ofSeconds(2), () -> operation.abort(failure)); + await(interrupted); + assertTrue(stream.submitting); + assertEquals(0, stream.closes.get()); + + finishAllocation.countDown(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + assertEquals(1, allocator.allocated.size()); + assertEquals(1, allocator.released.size()); + assertSame(allocator.allocated.get(0), allocator.released.get(0)); + assertEquals(1, stream.closes.get()); + assertFalse(stream.closedDuringSubmission); + } finally { + finishAllocation.countDown(); + } + } + + @Test + public void testSubmissionUsesTheConsumptionDeadline() throws Exception { + RecordingAllocator allocator = new RecordingAllocator(); + CountDownLatch submissionStarted = new CountDownLatch(1); + CountDownLatch finishSubmission = new CountDownLatch(1); + TestStream stream = new TestStream((ranges, buffers) -> { + submissionStarted.countDown(); + awaitIgnoringInterrupts(finishSubmission, new CountDownLatch(1)); + ranges.get(0).setDataReadFuture(CompletableFuture.completedFuture(buffers.allocate(8))); + }); + ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(); + executors.add(scheduler); + VectoredReadOperation operation = + new VectoredReadOperation(stream, ranges(1), allocator, newExecutor(), 5, TimeUnit.SECONDS); + try (ByteBufferReleaser releaser = new ByteBufferReleaser(allocator)) { + long remainingBeforeSubmission = operation.remainingNanos(); + scheduler.schedule(finishSubmission::countDown, 150, TimeUnit.MILLISECONDS); + operation.awaitSubmission(); + await(submissionStarted); + + assertTrue(operation.submissionSucceeded()); + assertTrue(operation.remainingNanos() <= remainingBeforeSubmission - TimeUnit.MILLISECONDS.toNanos(100)); + operation.transferTo(releaser); + } finally { + finishSubmission.countDown(); + } + assertEquals(1, allocator.released.size()); + assertEquals(0, stream.closes.get()); + } + + @Test + public void testSuccessTransfersOriginalBuffersInsteadOfFutureViews() throws Exception { + RecordingAllocator allocator = new RecordingAllocator(); + AtomicReference futureView = new AtomicReference<>(); + TestStream stream = new TestStream((ranges, buffers) -> { + ByteBuffer original = buffers.allocate(8); + ByteBuffer view = original.slice(); + futureView.set(view); + ranges.get(0).setDataReadFuture(CompletableFuture.completedFuture(view)); + }); + VectoredReadOperation operation = + new VectoredReadOperation(stream, ranges(1), allocator, newExecutor(), 5, TimeUnit.SECONDS); + try (ByteBufferReleaser releaser = new ByteBufferReleaser(allocator)) { + operation.awaitSubmission(); + operation.transferTo(releaser); + assertTrue(allocator.released.isEmpty()); + assertFalse(futureView.get() == allocator.allocated.get(0)); + } + assertEquals(1, allocator.released.size()); + assertSame(allocator.allocated.get(0), allocator.released.get(0)); + assertEquals(0, stream.closes.get()); + } + + @Test + public void testCancelledReadFutureDoesNotPermitBufferRelease() throws Exception { + RecordingAllocator allocator = new RecordingAllocator(); + TestStream stream = new TestStream((ranges, buffers) -> { + buffers.allocate(8); + CompletableFuture read = new CompletableFuture<>(); + ranges.get(0).setDataReadFuture(read); + read.cancel(false); + }); + ExecutorService executor = newExecutor(); + VectoredReadOperation operation = + new VectoredReadOperation(stream, ranges(1), allocator, executor, 5, TimeUnit.SECONDS); + operation.awaitSubmission(); + operation.abort(new IOException("Read was cancelled without stopping backend IO")); + + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + assertEquals(1, stream.closes.get()); + assertEquals(1, allocator.allocated.size()); + assertTrue(allocator.released.isEmpty()); + } + + @Test + public void testMissingReadFutureDoesNotPermitBufferRelease() throws Exception { + RecordingAllocator allocator = new RecordingAllocator(); + TestStream stream = new TestStream((ranges, buffers) -> { + ByteBuffer original = buffers.allocate(8); + ranges.get(0).setDataReadFuture(CompletableFuture.completedFuture(original)); + throw new IOException("Rejected before publishing the second range future"); + }); + ExecutorService executor = newExecutor(); + VectoredReadOperation operation = + new VectoredReadOperation(stream, ranges(2), allocator, executor, 5, TimeUnit.SECONDS); + IOException failure = assertThrows(IOException.class, operation::awaitSubmission); + assertFalse(operation.submissionSucceeded()); + operation.abort(failure); + + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + assertEquals(1, stream.closes.get()); + assertEquals(1, allocator.allocated.size()); + assertTrue(allocator.released.isEmpty()); + } + + @Test + public void testPendingReadReleasesBuffersWhenItActuallyFinishes() throws Exception { + RecordingAllocator allocator = new RecordingAllocator(); + CompletableFuture read = new CompletableFuture<>(); + TestStream stream = new TestStream((ranges, buffers) -> { + buffers.allocate(8); + ranges.get(0).setDataReadFuture(read); + }); + ExecutorService executor = newExecutor(); + VectoredReadOperation operation = + new VectoredReadOperation(stream, ranges(1), allocator, executor, 5, TimeUnit.SECONDS); + operation.awaitSubmission(); + operation.abort(new IOException("Another read failed")); + + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + assertEquals(1, stream.closes.get()); + assertTrue(allocator.released.isEmpty()); + assertFalse(read.isCancelled()); + + ByteBuffer original = allocator.allocated.get(0); + original.put(0, (byte) 41); + read.complete(original); + await(allocator.bufferReleased); + assertEquals(1, allocator.released.size()); + assertSame(original, allocator.released.get(0)); + } + + private ExecutorService newExecutor() { + ExecutorService executor = Executors.newSingleThreadExecutor(runnable -> { + Thread thread = new Thread(runnable, "vectored-read-test"); + thread.setDaemon(true); + return thread; + }); + executors.add(executor); + return executor; + } + + private static List ranges(int count) { + List ranges = new ArrayList<>(); + for (int i = 0; i < count; i++) { + ranges.add(new ParquetFileRange(i * 8L, 8)); + } + return ranges; + } + + private static TimeoutException awaitSubmissionTimeout(VectoredReadOperation operation) { + return assertTimeoutPreemptively( + Duration.ofSeconds(2), () -> assertThrows(TimeoutException.class, operation::awaitSubmission)); + } + + private static void await(CountDownLatch latch) throws InterruptedException { + assertTrue(latch.await(5, TimeUnit.SECONDS), "Backend did not reach the expected state"); + } + + private static void awaitIgnoringInterrupts(CountDownLatch latch, CountDownLatch interrupted) throws IOException { + while (true) { + try { + if (!latch.await(5, TimeUnit.SECONDS)) { + throw new IOException("Test backend was not released"); + } + return; + } catch (InterruptedException e) { + interrupted.countDown(); + } + } + } + + @FunctionalInterface + private interface Submission { + void submit(List ranges, ByteBufferAllocator allocator) throws IOException; + } + + private static class TestStream extends DelegatingSeekableInputStream { + private final Submission submission; + private final AtomicInteger closes = new AtomicInteger(); + private volatile boolean submitting; + private volatile boolean closedDuringSubmission; + + private TestStream(Submission submission) { + super(new ByteArrayInputStream(new byte[0])); + this.submission = submission; + } + + @Override + public void readVectored(List ranges, ByteBufferAllocator allocator) throws IOException { + submitting = true; + try { + submission.submit(ranges, allocator); + } finally { + submitting = false; + } + } + + @Override + public void close() throws IOException { + closedDuringSubmission |= submitting; + closes.incrementAndGet(); + super.close(); + } + + @Override + public long getPos() { + return 0; + } + + @Override + public void seek(long position) {} + } + + private static class RecordingAllocator implements ByteBufferAllocator { + private final List allocated = new CopyOnWriteArrayList<>(); + private final List released = new CopyOnWriteArrayList<>(); + private final CountDownLatch bufferReleased = new CountDownLatch(1); + + @Override + public ByteBuffer allocate(int size) { + ByteBuffer buffer = ByteBuffer.allocate(size); + allocated.add(buffer); + return buffer; + } + + @Override + public synchronized void release(ByteBuffer buffer) { + if (allocated.stream().noneMatch(original -> original == buffer)) { + throw new IllegalArgumentException("Released a view instead of the original buffer"); + } + if (released.stream().anyMatch(original -> original == buffer)) { + throw new IllegalStateException("Released the same buffer twice"); + } + released.add(buffer); + bufferReleased.countDown(); + } + + @Override + public boolean isDirect() { + return false; + } + } +} diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/util/wrapped/io/TestVectorIoBridge.java b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/util/wrapped/io/TestVectorIoBridge.java index 8a710d2ee2..4840db73d2 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/util/wrapped/io/TestVectorIoBridge.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/util/wrapped/io/TestVectorIoBridge.java @@ -29,10 +29,12 @@ import java.util.concurrent.CompletableFuture; 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.FSDataInputStream; import org.apache.hadoop.fs.FSDataOutputStream; import org.apache.hadoop.fs.FileSystem; +import org.apache.hadoop.fs.LocalFileSystem; import org.apache.hadoop.fs.Path; import org.apache.hadoop.io.ByteBufferPool; import org.apache.hadoop.io.ElasticByteBufferPool; @@ -205,6 +207,86 @@ public void testVectoredReadMultipleRanges() throws Exception { } } + @Test + public void testPublishesFuturesAfterPartialSubmissionFailure() throws Exception { + List fileRanges = getSampleNonOverlappingRanges(); + AtomicInteger allocations = new AtomicInteger(); + IllegalArgumentException failure = new IllegalArgumentException("second allocation failed"); + ByteBufferAllocator rejectingAllocator = new ByteBufferAllocator() { + @Override + public ByteBuffer allocate(int size) { + if (allocations.incrementAndGet() == 2) { + throw failure; + } + return allocate.allocate(size); + } + + @Override + public void release(ByteBuffer buffer) { + allocate.release(buffer); + } + + @Override + public boolean isDirect() { + return false; + } + }; + + // RawLocalFileSystem queues the first native read before allocating the second + // buffer. Its futures are created before any reads are submitted. + try (FSDataInputStream in = + ((LocalFileSystem) getFileSystem()).getRawFileSystem().open(testFilePath)) { + assertThatThrownBy(() -> vectorIOBridge.readVectoredRanges(in, fileRanges, rejectingAllocator)) + .isSameAs(failure); + assertThat(allocations.get()).isEqualTo(2); + assertThat(fileRanges.get(0).getDataReadFuture()).isNotNull(); + assertThat(fileRanges.get(1).getDataReadFuture()).isNotNull(); + // This future belongs to a read that was never submitted. Waiting for every + // published future after a submission failure would not finish. + assertThat(fileRanges.get(1).getDataReadFuture().isDone()).isFalse(); + + ByteBuffer buffer = FutureIO.awaitFuture(fileRanges.get(0).getDataReadFuture(), 5, TimeUnit.SECONDS); + try { + assertDatasetEquals(0, "partially submitted vectored read", buffer, 100, DATASET); + } finally { + allocate.release(buffer); + } + } + } + + @Test + public void testFuturePublicationPreservesSubmissionFailure() { + FileRangeBridge rangeBridge = FileRangeBridge.instance(); + ParquetFileRange failedRange = range(0, 100); + RuntimeException publicationFailure = new IllegalStateException("future lookup failed"); + FileRangeBridge.WrappedFileRange failed = rangeBridge.new WrappedFileRange(new Object()) { + @Override + public Object getReference() { + return failedRange; + } + + @Override + public CompletableFuture getData() { + throw publicationFailure; + } + }; + ParquetFileRange pendingRange = range(110, 50); + FileRangeBridge.WrappedFileRange pending = rangeBridge.toFileRange(pendingRange); + CompletableFuture pendingRead = new CompletableFuture<>(); + pending.setData(pendingRead); + ParquetFileRange unassignedRange = range(200, 50); + FileRangeBridge.WrappedFileRange unassigned = rangeBridge.toFileRange(unassignedRange); + IOException submissionFailure = new IOException("submission failed"); + + VectorIoBridge.publishReadFutures(List.of(failed, pending, unassigned), submissionFailure); + + assertThat(submissionFailure.getSuppressed()).containsExactly(publicationFailure); + assertThat(pendingRange.getDataReadFuture()).isSameAs(pendingRead); + assertThat(unassignedRange.getDataReadFuture()).isNull(); + assertThatThrownBy(() -> VectorIoBridge.publishReadFutures(List.of(failed), null)) + .isSameAs(publicationFailure); + } + /** * VectorIO and readFully() can coexist. */