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..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 @@ -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; @@ -113,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; @@ -772,6 +774,12 @@ 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 boolean vectoredIoDisabled; + private ExecutorService vectoredReadExecutor; + private int currentBlock = 0; private ColumnChunkPageReadStore currentRowGroup = null; private DictionaryPageReader nextDictionaryReader = null; @@ -1190,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; } /** @@ -1293,14 +1305,16 @@ public ColumnChunkPageReadStore readFilteredRowGroup( private void readAllPartsVectoredOrNormal(List allParts, ChunkListBuilder builder) throws IOException { - if (shouldUseVectoredIo(allParts)) { + if (shouldUseVectoredIo()) { try { - readVectored(allParts, builder); - return; + if (prepareVectoredReadFileLength()) { + 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) { @@ -1310,42 +1324,35 @@ 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 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. + *

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(final List allParts) { - return options.useHadoopVectoredIo() - && f.readVectoredAvailable(options.getAllocator()) - && arePartsValidForVectoredIo(allParts); + private boolean shouldUseVectoredIo() { + return !vectoredIoDisabled && options.useHadoopVectoredIo() && f.readVectoredAvailable(options.getAllocator()); } - /** - * 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; + 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; } - return true; } /** @@ -1357,32 +1364,153 @@ 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. + * 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. * @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 = vectoredReadFileLength; 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); + 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 + // 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, 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); + if (operation.submissionSucceeded()) { + awaitRemainingVectoredReads(ranges, operation, failure); + } + abortVectoredRead(operation, failure); + throw failure; + } 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, VectoredReadOperation operation, Throwable failure) { + if (Thread.currentThread().isInterrupted() + || failure instanceof InterruptedIOException && failure.getCause() instanceof InterruptedException) { + return; + } + + for (ParquetFileRange range : ranges) { + Future future = range.getDataReadFuture(); + if (future == null || future.isDone()) { + continue; + } + + try { + FutureIO.awaitFuture(future, operation.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); + } + } } } @@ -1464,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) @@ -1860,6 +2002,9 @@ public void close() throws IOException { f.close(); } } finally { + if (vectoredReadExecutor != null) { + vectoredReadExecutor.shutdownNow(); + } AutoCloseables.uncheckedClose(nextDictionaryReader, crcAllocator); options.getCodecFactory().release(); } @@ -1951,10 +2096,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 +2168,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 +2188,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 +2219,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 +2491,41 @@ 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. + * @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 readFromVectoredRange(ParquetFileRange currRange, ChunkListBuilder builder) throws IOException { - ByteBuffer buffer; + 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; 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); + buffers.add(FutureIO.awaitFuture( + range.getDataReadFuture(), operation.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/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 7720f7fe94..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 @@ -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(); @@ -155,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. *

@@ -198,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/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..41d65591c9 --- /dev/null +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileReaderVectoredIO.java @@ -0,0 +1,1034 @@ +/* + * 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.io.InterruptedIOException; +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.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; +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 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 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) + .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 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); + } + + @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); + IOException retryFailure = assertThrows(IOException.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()); + } + } + } + stream.closeCompleted.get(10, TimeUnit.SECONDS); + 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 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; + } + + @Override + public long getLength() throws IOException { + lengthCalls++; + if (lengthCalls > successfulLengthCalls) { + throw lengthFailure; + } + 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; + + @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 CompletableFuture closeCompleted = 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); + } + } + try { + super.close(); + } finally { + closeCompleted.complete(null); + } + } + } + + 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); + } + } +} 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. */