ORC-2226: Allow seeking to the end of stream in InStream.seek

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 &&

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());
+  }
 }

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<k:string,v:string>");
+    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<k:string,v:string>");
+    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);
+  }
 }