diff --git a/parquet-common/src/main/java/org/apache/parquet/bytes/ByteBufferInputStream.java b/parquet-common/src/main/java/org/apache/parquet/bytes/ByteBufferInputStream.java index 786d5059ff..8868433ad1 100644 --- a/parquet-common/src/main/java/org/apache/parquet/bytes/ByteBufferInputStream.java +++ b/parquet-common/src/main/java/org/apache/parquet/bytes/ByteBufferInputStream.java @@ -100,6 +100,14 @@ public void skipFully(long n) throws IOException { } } + /** + * Copies bytes from this stream into {@code out}, reading at most {@code out.remaining()} bytes. The position of + * {@code out} is advanced by the number of bytes copied and its limit is unchanged. + * + * @param out the buffer to copy bytes into + * @return the number of bytes copied, {@code 0} if {@code out} has no remaining space, or {@code -1} if this stream + * is exhausted + */ public int read(ByteBuffer out) { return delegate.read(out); } diff --git a/parquet-common/src/main/java/org/apache/parquet/bytes/SingleBufferInputStream.java b/parquet-common/src/main/java/org/apache/parquet/bytes/SingleBufferInputStream.java index 82b804bec4..ce66be2a67 100644 --- a/parquet-common/src/main/java/org/apache/parquet/bytes/SingleBufferInputStream.java +++ b/parquet-common/src/main/java/org/apache/parquet/bytes/SingleBufferInputStream.java @@ -90,6 +90,14 @@ public long skip(long n) { @Override public int read(ByteBuffer out) { + if (!out.hasRemaining()) { + return 0; + } + + if (buffer.remaining() <= 0) { + return -1; + } + int bytesToCopy; ByteBuffer copyBuffer; if (buffer.remaining() <= out.remaining()) { @@ -105,7 +113,6 @@ public int read(ByteBuffer out) { } out.put(copyBuffer); - out.flip(); return bytesToCopy; } diff --git a/parquet-common/src/test/java/org/apache/parquet/bytes/TestByteBufferInputStreamRead.java b/parquet-common/src/test/java/org/apache/parquet/bytes/TestByteBufferInputStreamRead.java new file mode 100644 index 0000000000..7960ccbc99 --- /dev/null +++ b/parquet-common/src/test/java/org/apache/parquet/bytes/TestByteBufferInputStreamRead.java @@ -0,0 +1,207 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.parquet.bytes; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.List; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +/** + * Tests {@link ByteBufferInputStream#read(ByteBuffer)} appending into the caller's destination buffer for both single + * buffer and multi buffer streams created through the public wrap entry points. + */ +public class TestByteBufferInputStreamRead { + + private static final byte PREFIX_BYTE = (byte) 0x11; + private static final byte SUFFIX_BYTE = (byte) 0x77; + private static final int PREFIX_LENGTH = 3; + private static final int DEST_CAPACITY = 16; + private static final int DEST_LIMIT = 8; + private static final int SLICE_MARGIN = 4; + + private enum OutputKind { + HEAP, + DIRECT, + SLICED + } + + private enum SourceKind { + SINGLE, + MULTI + } + + static Stream kinds() { + return Stream.of(OutputKind.values()) + .flatMap(output -> Stream.of(SourceKind.values()).map(source -> Arguments.of(output, source))); + } + + @ParameterizedTest(name = "{0} {1}") + @MethodSource("kinds") + public void partialCopyAdvancesDestinationPosition(OutputKind output, SourceKind source) { + byte[] data = data(10, 1); + int copied = DEST_LIMIT - PREFIX_LENGTH; + ByteBuffer dest = destination(output, DEST_LIMIT, PREFIX_LENGTH); + + assertRead(output + " " + source, stream(source, data, false), dest, copied); + + assertContent(dest, output, source, Arrays.copyOf(data, copied)); + } + + @ParameterizedTest(name = "{0} {1}") + @MethodSource("kinds") + public void fullSourceCopyKeepsLimitAndRoomForAppendedRead(OutputKind output, SourceKind source) { + byte[] first = data(4, 1); + byte[] second = data(3, 11); + ByteBuffer dest = destination(output, DEST_CAPACITY, PREFIX_LENGTH); + + assertRead(output + " " + source + " first read", stream(source, first, true), dest, first.length); + assertRead(output + " " + source + " second read", stream(source, second, true), dest, second.length); + + assertContent(dest, output, source, first, second); + } + + @ParameterizedTest(name = "{0} {1}") + @MethodSource("kinds") + public void readIntoExhaustedDestinationReturnsZeroWithoutMutatingIt(OutputKind output, SourceKind source) { + byte[] data = data(10, 1); + ByteBuffer dest = destination(output, DEST_LIMIT, DEST_LIMIT); + ByteBufferInputStream in = stream(source, data, false); + + assertThat(in.read(dest)).as("%s %s: bytes copied", output, source).isZero(); + assertThat(dest.position()) + .as("%s %s: destination position", output, source) + .isEqualTo(DEST_LIMIT); + assertThat(dest.limit()).as("%s %s: destination limit", output, source).isEqualTo(DEST_LIMIT); + assertThat(in.position()).as("%s %s: source position", output, source).isZero(); + assertThat(in.available()).as("%s %s: source remaining", output, source).isEqualTo(data.length); + } + + @ParameterizedTest(name = "{0} {1}") + @MethodSource("kinds") + public void readAtEndOfStreamReturnsMinusOneWithoutMutatingDestination(OutputKind output, SourceKind source) + throws Exception { + byte[] data = data(10, 1); + ByteBufferInputStream in = stream(source, data, true); + ByteBuffer dest = destination(output, DEST_LIMIT, PREFIX_LENGTH); + + assertThat(in.read(new byte[data.length])) + .as("%s %s: drained bytes", output, source) + .isEqualTo(data.length); + assertThat(in.read(dest)).as("%s %s: bytes copied", output, source).isEqualTo(-1); + assertThat(dest.position()) + .as("%s %s: destination position", output, source) + .isEqualTo(PREFIX_LENGTH); + assertThat(dest.limit()).as("%s %s: destination limit", output, source).isEqualTo(DEST_LIMIT); + assertThat(in.position()).as("%s %s: source position", output, source).isEqualTo(data.length); + assertThat(in.available()).as("%s %s: source remaining", output, source).isZero(); + } + + private static void assertRead(String description, ByteBufferInputStream in, ByteBuffer dest, int expectedCopied) { + int destStart = dest.position(); + int destLimit = dest.limit(); + long sourceStart = in.position(); + int sourceRemaining = in.available(); + + assertThat(in.read(dest)).as("%s: bytes copied", description).isEqualTo(expectedCopied); + assertThat(dest.position()).as("%s: destination position", description).isEqualTo(destStart + expectedCopied); + assertThat(dest.limit()).as("%s: destination limit", description).isEqualTo(destLimit); + assertThat(in.position()).as("%s: source position", description).isEqualTo(sourceStart + expectedCopied); + assertThat(in.available()).as("%s: source remaining", description).isEqualTo(sourceRemaining - expectedCopied); + } + + private static void assertContent(ByteBuffer dest, OutputKind output, SourceKind source, byte[]... payloads) { + byte[] expected = new byte + [PREFIX_LENGTH + + Arrays.stream(payloads) + .mapToInt(payload -> payload.length) + .sum()]; + Arrays.fill(expected, 0, PREFIX_LENGTH, PREFIX_BYTE); + int offset = PREFIX_LENGTH; + for (byte[] payload : payloads) { + System.arraycopy(payload, 0, expected, offset, payload.length); + offset += payload.length; + } + + dest.flip(); + byte[] actual = new byte[dest.remaining()]; + dest.get(actual); + assertThat(actual) + .as("%s %s: prefilled prefix followed by the copied source bytes", output, source) + .containsExactly(expected); + + ByteBuffer fullDestination = dest.duplicate(); + fullDestination.clear(); + for (int i = dest.limit(); i < dest.capacity(); i++) { + assertThat(fullDestination.get(i)) + .as("%s %s: untouched destination byte %d", output, source, i) + .isEqualTo(SUFFIX_BYTE); + } + } + + private static ByteBuffer destination(OutputKind output, int limit, int prefixLength) { + ByteBuffer dest = ByteBuffer.allocate(DEST_CAPACITY); + if (output == OutputKind.DIRECT) { + dest = ByteBuffer.allocateDirect(DEST_CAPACITY); + } else if (output == OutputKind.SLICED) { + ByteBuffer parent = ByteBuffer.allocate(DEST_CAPACITY + 2 * SLICE_MARGIN); + parent.position(SLICE_MARGIN); + parent.limit(SLICE_MARGIN + DEST_CAPACITY); + dest = parent.slice(); + } + + for (int i = 0; i < dest.capacity(); i++) { + dest.put(i, SUFFIX_BYTE); + } + for (int i = 0; i < prefixLength; i++) { + dest.put(i, PREFIX_BYTE); + } + dest.limit(limit); + dest.position(prefixLength); + + return dest; + } + + private static ByteBufferInputStream stream(SourceKind source, byte[] data, boolean viaList) { + if (source == SourceKind.SINGLE) { + ByteBuffer whole = ByteBuffer.wrap(data); + return viaList ? ByteBufferInputStream.wrap(List.of(whole)) : ByteBufferInputStream.wrap(whole); + } + + // the first buffer holds more than half of the data so that a small destination forces a partial copy + int split = data.length / 2 + 1; + ByteBuffer first = ByteBuffer.wrap(data, 0, split); + ByteBuffer second = ByteBuffer.wrap(data, split, data.length - split); + return viaList ? ByteBufferInputStream.wrap(List.of(first, second)) : ByteBufferInputStream.wrap(first, second); + } + + private static byte[] data(int length, int firstValue) { + byte[] data = new byte[length]; + for (int i = 0; i < length; i++) { + data[i] = (byte) (firstValue + i); + } + return data; + } +}