diff --git a/java/core/src/java/org/apache/orc/impl/InStream.java b/java/core/src/java/org/apache/orc/impl/InStream.java index bf23a95087..252b806652 100644 --- a/java/core/src/java/org/apache/orc/impl/InStream.java +++ b/java/core/src/java/org/apache/orc/impl/InStream.java @@ -223,6 +223,15 @@ public void seek(long desired) throws IOException { if (desired == 0 && bytes == null) { return; } + // The read planner assigns zero bytes to a stream when the selected + // row groups need no data from it (e.g. predicate push-down selects + // only an all-null trailing row group, whose row group index points + // at the end of the stream). Seeking to the end of such a stream is + // a legitimate position: park at the end and read() will return -1. + if (bytes == null && desired == length) { + position = desired; + return; + } // compute the position of the desired point in file long positionFile = desired + offset; // If we are seeking inside of the current range, just reposition. @@ -250,7 +259,7 @@ public void seek(long desired) throws IOException { public String toString() { return "uncompressed stream " + name + " position: " + position + " length: " + length + " range: " + getRangeNumber(bytes, currentRange) + - " offset: " + currentRange.getOffset() + + " offset: " + (currentRange == null ? -1 : currentRange.getOffset()) + " position: " + (decrypted == null ? 0 : decrypted.position()) + " limit: " + (decrypted == null ? 0 : decrypted.limit()); } @@ -671,6 +680,17 @@ boolean seek(long desired) throws IOException { // Header already at the required position return false; } + // Same as UncompressedStream above: no bytes were read for this + // stream, so seeking to its end is a legitimate position: park at the + // end of the stream and read() will return -1. The buffered chunk is + // dropped, so the cached header position is no longer valid. + if (bytes == null && desired == length) { + position = desired; + currentCompressedStart = -1; + compressed = null; + uncompressed = null; + return true; + } long posn = desired + offset; for (DiskRangeList range = bytes; range != null; range = range.next) { if (range.getOffset() <= posn && diff --git a/java/core/src/test/org/apache/orc/impl/TestInStream.java b/java/core/src/test/org/apache/orc/impl/TestInStream.java index 596afcf304..ca58b98cb8 100644 --- a/java/core/src/test/org/apache/orc/impl/TestInStream.java +++ b/java/core/src/test/org/apache/orc/impl/TestInStream.java @@ -40,6 +40,7 @@ import static org.junit.jupiter.api.Assertions.assertNotEquals; 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.fail; public class TestInStream { @@ -1000,4 +1001,65 @@ public void testStreamResetWithoutIncreasedLength() throws IOException { byte[] inBuffer = new byte[5]; assertEquals(5, inStream.read(inBuffer)); } + + @Test + public void testUncompressedSeekToEndOfStream() throws Exception { + // ORC-2226: the row group index of an all-null trailing row group points + // at the end of the stream, and the read planner may not have read any + // bytes for the stream, so seeking to the end must be accepted. + InStream in = InStream.create("test", null, 0, 1024); + InStream.UncompressedStream stream = (InStream.UncompressedStream) in; + stream.seek(0); + stream.seek(1024); + assertEquals(-1, in.read()); + assertEquals(0, in.available()); + // seeking past the end of the stream still fails + assertThrows(IllegalArgumentException.class, () -> stream.seek(1025)); + } + + @Test + public void testCompressedSeekToEndOfStream() throws Exception { + // ORC-2226: same as above, for compressed streams. + InStream in = InStream.create("test", null, 0, 1024, + InStream.options().withCodec(new ZlibCodec()).withBufferSize(300)); + InStream.CompressedStream stream = (InStream.CompressedStream) in; + stream.seek(0); + stream.seek(1024); + assertEquals(-1, in.read()); + assertEquals(0, in.available()); + // seeking past the end of the stream still fails + assertThrows(IOException.class, () -> stream.seek(1025)); + } + + @Test + public void testCompressedSeekBackFromEndOfStream() throws Exception { + // ORC-2226: seeking to the end of the stream must not corrupt the + // cached chunk state; seeking back must return the correct data. + OutputCollector collect = new OutputCollector(); + CompressionCodec codec = new ZlibCodec(); + try (OutStream out = new OutStream("test", + new StreamOptions(300).withCodec(codec, codec.getDefaultOptions()), + collect)) { + for (int i = 0; i < 1024; ++i) { + out.write(i); + } + } + byte[] bytes = collect.buffer.get(); + InStream in = InStream.create("test", + new BufferChunk(ByteBuffer.wrap(bytes), 0), 0, bytes.length, + InStream.options().withCodec(new ZlibCodec()).withBufferSize(300)); + assertEquals(0, in.read()); + // seek to the end of the stream + PositionCollector toEnd = new PositionCollector(); + toEnd.addPosition(bytes.length); + toEnd.addPosition(0); + in.seek(toEnd); + assertEquals(-1, in.read()); + // seek back to the beginning + PositionCollector toStart = new PositionCollector(); + toStart.addPosition(0); + toStart.addPosition(0); + in.seek(toStart); + assertEquals(0, in.read()); + } } diff --git a/java/core/src/test/org/apache/orc/impl/TestRecordReaderImpl.java b/java/core/src/test/org/apache/orc/impl/TestRecordReaderImpl.java index 3c70b7284a..cd5a3be4f3 100644 --- a/java/core/src/test/org/apache/orc/impl/TestRecordReaderImpl.java +++ b/java/core/src/test/org/apache/orc/impl/TestRecordReaderImpl.java @@ -83,6 +83,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.List; +import java.util.Locale; import java.util.TimeZone; import static org.apache.orc.impl.RecordReaderUtils.MAX_BYTE_WIDTH; @@ -2817,4 +2818,121 @@ public void testDecimalIsRepeatingFlag() throws IOException { } } } + + @Test + public void testSargOnAllNullTrailingRowGroup() throws Exception { + // ORC-2226: with predicate pushdown, selecting only an all-null trailing + // row group of an uncompressed stripe makes the read planner assign zero + // bytes to the column's DATA stream; seeking it to the recorded position + // (the end of the stream) must not fail and must read all-null values. + Path path = new Path(workDir, "allNullTrailingRowGroup.orc"); + FileSystem.get(conf).delete(path, true); + TypeDescription schema = TypeDescription.fromString("struct"); + Writer writer = OrcFile.createWriter(path, OrcFile.writerOptions(conf) + .setSchema(schema) + .compress(CompressionKind.NONE) + .rowIndexStride(1024)); + VectorizedRowBatch batch = schema.createRowBatch(); + BytesColumnVector k = (BytesColumnVector) batch.cols[0]; + BytesColumnVector v = (BytesColumnVector) batch.cols[1]; + batch.size = 1024; + for (int i = 0; i < batch.size; i++) { + k.setVal(i, "a".getBytes(StandardCharsets.UTF_8)); + // distinct values so that the writer uses DIRECT encoding + v.setVal(i, String.format(Locale.ROOT, "%07dx", i).getBytes(StandardCharsets.UTF_8)); + } + writer.addRowBatch(batch); + batch.reset(); + batch.size = 512; + v.noNulls = false; + v.isRepeating = true; + v.isNull[0] = true; + for (int i = 0; i < batch.size; i++) { + k.setVal(i, "c".getBytes(StandardCharsets.UTF_8)); + } + writer.addRowBatch(batch); + writer.close(); + + Reader reader = OrcFile.createReader(path, OrcFile.readerOptions(conf)); + SearchArgument sarg = SearchArgumentFactory.newBuilder() + .equals("k", PredicateLeaf.Type.STRING, "c") + .build(); + VectorizedRowBatch readBatch = reader.getSchema().createRowBatch(); + long rowCount = 0; + try (RecordReader rows = reader.rows(reader.options() + .searchArgument(sarg, new String[] {"k"}))) { + while (rows.nextBatch(readBatch)) { + BytesColumnVector readK = (BytesColumnVector) readBatch.cols[0]; + BytesColumnVector readV = (BytesColumnVector) readBatch.cols[1]; + for (int i = 0; i < readBatch.size; i++) { + assertEquals("c", readK.toString(i)); + assertTrue(readV.isRepeating ? readV.isNull[0] : readV.isNull[i]); + } + rowCount += readBatch.size; + } + } + assertEquals(512, rowCount); + } + + @Test + public void testSargOnAllNullMiddleRowGroup() throws Exception { + // An all-null row group in the middle of a stripe is covered by the + // following row group's estimated end offset; make sure that such a + // selection keeps working. + Path path = new Path(workDir, "allNullMiddleRowGroup.orc"); + FileSystem.get(conf).delete(path, true); + TypeDescription schema = TypeDescription.fromString("struct"); + Writer writer = OrcFile.createWriter(path, OrcFile.writerOptions(conf) + .setSchema(schema) + .compress(CompressionKind.NONE) + .rowIndexStride(1024)); + VectorizedRowBatch batch = schema.createRowBatch(); + BytesColumnVector k = (BytesColumnVector) batch.cols[0]; + BytesColumnVector v = (BytesColumnVector) batch.cols[1]; + batch.size = 1024; + for (int i = 0; i < batch.size; i++) { + k.setVal(i, "a".getBytes(StandardCharsets.UTF_8)); + v.setVal(i, String.format(Locale.ROOT, "%07dx", i).getBytes(StandardCharsets.UTF_8)); + } + writer.addRowBatch(batch); + batch.reset(); + batch.size = 1024; + v.noNulls = false; + v.isRepeating = true; + v.isNull[0] = true; + for (int i = 0; i < batch.size; i++) { + k.setVal(i, "b".getBytes(StandardCharsets.UTF_8)); + } + writer.addRowBatch(batch); + batch.reset(); + batch.size = 1024; + v.noNulls = true; + v.isRepeating = false; + for (int i = 0; i < batch.size; i++) { + k.setVal(i, "d".getBytes(StandardCharsets.UTF_8)); + v.setVal(i, String.format(Locale.ROOT, "%07dy", i).getBytes(StandardCharsets.UTF_8)); + } + writer.addRowBatch(batch); + writer.close(); + + Reader reader = OrcFile.createReader(path, OrcFile.readerOptions(conf)); + SearchArgument sarg = SearchArgumentFactory.newBuilder() + .equals("k", PredicateLeaf.Type.STRING, "b") + .build(); + VectorizedRowBatch readBatch = reader.getSchema().createRowBatch(); + long rowCount = 0; + try (RecordReader rows = reader.rows(reader.options() + .searchArgument(sarg, new String[] {"k"}))) { + while (rows.nextBatch(readBatch)) { + BytesColumnVector readK = (BytesColumnVector) readBatch.cols[0]; + BytesColumnVector readV = (BytesColumnVector) readBatch.cols[1]; + for (int i = 0; i < readBatch.size; i++) { + assertEquals("b", readK.toString(i)); + assertTrue(readV.isRepeating ? readV.isNull[0] : readV.isNull[i]); + } + rowCount += readBatch.size; + } + } + assertEquals(1024, rowCount); + } }