[core] SstFileReader supports range query (#6842)

Author: steFaiz

paimon-common/src/main/java/org/apache/paimon/sst/SstFileReader.java

@@ -28,6 +28,7 @@
 import org.apache.paimon.memory.MemorySliceInput;
 import org.apache.paimon.utils.FileBasedBloomFilter;
 import org.apache.paimon.utils.MurmurHashUtils;
+import org.apache.paimon.utils.Preconditions;
 
 import javax.annotation.Nullable;
 
@@ -39,7 +40,8 @@
 import static org.apache.paimon.utils.Preconditions.checkArgument;
 
 /**
- * An SST File Reader which only serves point queries now.
+ * An SST File Reader which serves point queries and range queries. Users can call {@code
+ * createIterator} to create a file iterator and then use seek and read methods to do range queries.
  *
  * <p>Note that this class is NOT thread-safe.
  */
@@ -109,6 +111,10 @@ public byte[] lookup(byte[] key) throws IOException {
         return null;
     }
 
+    public SstFileIterator createIterator() {
+        return new SstFileIterator(indexBlock.iterator());
+    }
+
     private BlockIterator getNextBlock(BlockIterator indexBlockIterator) {
         // index block handle, point to the key, value position.
         MemorySlice blockHandle = indexBlockIterator.next().getValue();
@@ -180,4 +186,55 @@ public void close() throws IOException {
         blockCache.close();
         // do not need to close input, since it will be closed by outer classes
     }
+
+    /** An Iterator for range queries. */
+    public class SstFileIterator {
+        private final BlockIterator indexIterator;
+        private @Nullable BlockIterator seekedDataBlock = null;
+
+        SstFileIterator(BlockIterator indexBlockIterator) {
+            this.indexIterator = indexBlockIterator;
+        }
+
+        /**
+         * Seek to the position of the record whose key is exactly equal to or greater than the
+         * specified key.
+         */
+        public void seekTo(byte[] key) {
+            MemorySlice keySlice = MemorySlice.wrap(key);
+
+            indexIterator.seekTo(keySlice);
+            if (indexIterator.hasNext()) {
+                seekedDataBlock = getNextBlock(indexIterator);
+                // The index block entry key is the last key of the corresponding data block.
+                // If there is some index entry key >= targetKey, the related data block must
+                // also contain some key >= target key, which means seekedDataBlock.hasNext()
+                // must be true
+                seekedDataBlock.seekTo(keySlice);
+                Preconditions.checkState(seekedDataBlock.hasNext());
+            } else {
+                seekedDataBlock = null;
+            }
+        }
+
+        /**
+         * Read a batch of records from this SST File and move current record position to the next
+         * batch.
+         *
+         * @return current batch of records, null if reaching file end.
+         */
+        public BlockIterator readBatch() throws IOException {
+            if (seekedDataBlock != null) {
+                BlockIterator result = seekedDataBlock;
+                seekedDataBlock = null;
+                return result;
+            }
+
+            if (!indexIterator.hasNext()) {
+                return null;
+            }
+
+            return getNextBlock(indexIterator);
+        }
+    }
 }

paimon-common/src/test/java/org/apache/paimon/sst/SstFileTest.java

@@ -41,6 +41,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.Comparator;
 import java.util.List;
@@ -109,12 +110,16 @@ private void innerTestLookup() throws Exception {
 
             // 1. lookup random existing keys
             for (int i = 0; i < 100; i++) {
-                int key = random.nextInt(5000);
+                int key = random.nextInt(5000 * 2 - 1);
                 keyOut.reset();
                 keyOut.writeInt(key);
                 byte[] queried = reader.lookup(keyOut.toSlice().getHeapMemory());
-                Assertions.assertNotNull(queried);
-                Assertions.assertEquals(key, MemorySlice.wrap(queried).readInt(0));
+                if (key % 2 == 0) {
+                    Assertions.assertNotNull(queried);
+                    Assertions.assertEquals(key, MemorySlice.wrap(queried).readInt(0));
+                } else {
+                    Assertions.assertNull(queried);
+                }
             }
 
             // 2. lookup boundaries
@@ -125,16 +130,16 @@ private void innerTestLookup() throws Exception {
             Assertions.assertEquals(0, MemorySlice.wrap(queried).readInt(0));
 
             keyOut.reset();
-            keyOut.writeInt(511);
+            keyOut.writeInt(511 * 2);
             byte[] queried1 = reader.lookup(keyOut.toSlice().getHeapMemory());
             Assertions.assertNotNull(queried1);
-            Assertions.assertEquals(511, MemorySlice.wrap(queried1).readInt(0));
+            Assertions.assertEquals(511 * 2, MemorySlice.wrap(queried1).readInt(0));
 
             keyOut.reset();
-            keyOut.writeInt(4999);
+            keyOut.writeInt(4999 * 2);
             byte[] queried2 = reader.lookup(keyOut.toSlice().getHeapMemory());
             Assertions.assertNotNull(queried2);
-            Assertions.assertEquals(4999, MemorySlice.wrap(queried2).readInt(0));
+            Assertions.assertEquals(4999 * 2, MemorySlice.wrap(queried2).readInt(0));
 
             // 2. lookup key smaller than first key
             for (int i = 0; i < 100; i++) {
@@ -146,12 +151,91 @@ private void innerTestLookup() throws Exception {
             // 3. lookup key greater than last key
             for (int i = 0; i < 100; i++) {
                 keyOut.reset();
-                keyOut.writeInt(10000 + i);
+                keyOut.writeInt(10001 + i);
                 Assertions.assertNull(reader.lookup(keyOut.toSlice().getHeapMemory()));
             }
         }
     }
 
+    @TestTemplate
+    public void testScan() throws Exception {
+        int recordNum = 20000;
+        writeData(recordNum, bloomFilterEnabled);
+
+        long fileSize = fileIO.getFileSize(file);
+        try (SeekableInputStream inputStream = fileIO.newInputStream(file);
+                SstFileReader reader =
+                        new SstFileReader(
+                                Comparator.comparingInt(slice -> slice.readInt(0)),
+                                fileSize,
+                                file,
+                                inputStream,
+                                CACHE_MANAGER); ) {
+            SstFileReader.SstFileIterator fileIterator = reader.createIterator();
+
+            MemorySliceOutput keyOut = new MemorySliceOutput(4);
+
+            // 1. test full scan
+            assertScan(0, recordNum - 1, fileIterator);
+
+            // 2. test random seek and scan
+            Random random = new Random();
+            for (int i = 0; i < 1000; i++) {
+                resetIterator(fileIterator, keyOut);
+                int key = random.nextInt(recordNum * 2 - 1);
+                int targetPosition = (key + 1) / 2;
+
+                keyOut.reset();
+                keyOut.writeInt(key);
+                fileIterator.seekTo(keyOut.toSlice().getHeapMemory());
+
+                // lookup should not affect reader position
+                interleaveLookup(reader, keyOut);
+
+                assertScan(targetPosition, recordNum - 1, fileIterator);
+            }
+
+            // 3. test boundaries
+            resetIterator(fileIterator, keyOut);
+            keyOut.reset();
+            keyOut.writeInt(0);
+            fileIterator.seekTo(keyOut.toSlice().getHeapMemory());
+            assertScan(0, recordNum - 1, fileIterator);
+
+            resetIterator(fileIterator, keyOut);
+            keyOut.reset();
+            keyOut.writeInt(recordNum * 2 - 2);
+            fileIterator.seekTo(keyOut.toSlice().getHeapMemory());
+            assertScan(recordNum - 1, recordNum - 1, fileIterator);
+
+            // 4. test out of boundaries
+            resetIterator(fileIterator, keyOut);
+            keyOut.reset();
+            keyOut.writeInt(-10);
+            fileIterator.seekTo(keyOut.toSlice().getHeapMemory());
+            assertScan(0, recordNum - 1, fileIterator);
+
+            resetIterator(fileIterator, keyOut);
+            keyOut.reset();
+            keyOut.writeInt(recordNum * 2 + 10);
+            fileIterator.seekTo(keyOut.toSlice().getHeapMemory());
+            Assertions.assertNull(fileIterator.readBatch());
+        }
+    }
+
+    private void resetIterator(SstFileReader.SstFileIterator iterator, MemorySliceOutput keyOut)
+            throws IOException {
+        keyOut.reset();
+        keyOut.writeInt(-1);
+        iterator.seekTo(keyOut.toSlice().getHeapMemory());
+    }
+
+    private void interleaveLookup(SstFileReader reader, MemorySliceOutput keyOut) throws Exception {
+        keyOut.reset();
+        keyOut.writeInt(0);
+        reader.lookup(keyOut.toSlice().getHeapMemory());
+    }
+
     private void writeData(int recordCount, boolean withBloomFilter) throws Exception {
         BloomFilter.Builder bloomFilter = null;
         if (withBloomFilter) {
@@ -168,11 +252,27 @@ private void writeData(int recordCount, boolean withBloomFilter) throws Exceptio
             for (int i = 0; i < recordCount; i++) {
                 keyOut.reset();
                 valueOut.reset();
-                keyOut.writeInt(i);
-                valueOut.writeInt(i);
+                keyOut.writeInt(i * 2);
+                valueOut.writeInt(i * 2);
                 writer.put(keyOut.toSlice().getHeapMemory(), valueOut.toSlice().getHeapMemory());
             }
             LOG.info("Write {} data cost {} ms", recordCount, System.currentTimeMillis() - start);
         }
     }
+
+    private static void assertScan(
+            int startPosition, int lastPosition, SstFileReader.SstFileIterator iterator)
+            throws Exception {
+        int count = startPosition;
+        BlockIterator iter;
+        while ((iter = iterator.readBatch()) != null) {
+            while (iter.hasNext()) {
+                BlockEntry entry = iter.next();
+                Assertions.assertEquals(count * 2, entry.getKey().readInt(0));
+                Assertions.assertEquals(count * 2, entry.getValue().readInt(0));
+                count++;
+            }
+        }
+        Assertions.assertEquals(lastPosition, count - 1);
+    }
 }