diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/reader/SeekableLakeSnapshotSplitScanner.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/reader/SeekableLakeSnapshotSplitScanner.java index b0df3c17403..6a09f9d4aa1 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/reader/SeekableLakeSnapshotSplitScanner.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/lake/reader/SeekableLakeSnapshotSplitScanner.java @@ -67,18 +67,20 @@ public SeekableLakeSnapshotSplitScanner( @Nullable @Override public CloseableIterator pollBatch(Duration timeout) throws IOException { - if (currentLakeRecordIterator == null) { - updateCurrentIterator(); - } + while (true) { + if (currentLakeRecordIterator == null) { + updateCurrentIterator(); + } - // has no next record in currentIterator, update currentIterator - if (currentLakeRecordIterator != null && !currentLakeRecordIterator.hasNext()) { - updateCurrentIterator(); - } + if (currentLakeRecordIterator == null || currentLakeRecordIterator.hasNext()) { + return currentLakeRecordIterator; + } - return currentLakeRecordIterator != null && currentLakeRecordIterator.hasNext() - ? currentLakeRecordIterator - : null; + // An inner lake split may become empty after deletes or filtering. Keep looking for + // data instead of reporting the end of the whole bounded split. + currentLakeRecordIterator.close(); + currentLakeRecordIterator = null; + } } private void updateCurrentIterator() throws IOException { diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/BoundedSplitReader.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/BoundedSplitReader.java index 2bb425954bb..038b187483a 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/BoundedSplitReader.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/reader/BoundedSplitReader.java @@ -52,6 +52,7 @@ public class BoundedSplitReader implements AutoCloseable { private final BatchScanner splitScanner; private long currentReadRecordsCount; + private int currentSplitIndex; private long toSkip; private final BlockingQueue recordAndPosBatchPool; @@ -60,6 +61,7 @@ public BoundedSplitReader(BatchScanner splitScanner, final long toSkip) { this.splitScanner = splitScanner; this.toSkip = toSkip; this.currentReadRecordsCount = 0; + this.currentSplitIndex = RecordAndPos.DEFAULT_SPLIT_INDEX; this.recordAndPosBatchPool = new ArrayBlockingQueue<>(1); this.recordAndPosBatchPool.add(new RecordAndPosBatch()); } @@ -102,6 +104,7 @@ private RecordAndPosBatch pollRecordAndPosBatch() throws IOException { private CloseableIterator poll() throws IOException { CloseableIterator nextBatch = null; + long skippedRecordsCount = 0; // may skip records while (toSkip > 0) { // pool a batch of records @@ -111,13 +114,18 @@ private CloseableIterator poll() throws IOException { throw new RuntimeException( String.format( "Skip more than the number of total records, has skipped %d record(s), but remain %s record(s) to skip.", - currentReadRecordsCount, toSkip)); + skippedRecordsCount, toSkip)); } // skip while (toSkip > 0 && nextBatch.hasNext()) { nextBatch.next(); toSkip--; currentReadRecordsCount++; + skippedRecordsCount++; + } + if (!nextBatch.hasNext()) { + nextBatch.close(); + nextBatch = null; } } // if any batch remains while skipping, return the batch @@ -133,7 +141,19 @@ private CloseableIterator poll() throws IOException { private CloseableIterator pollBatch() throws IOException { CloseableIterator records = splitScanner.pollBatch(POLL_TIMEOUT); - return records == null ? null : new ScanRecordBatch(records); + if (records == null) { + return null; + } + + ScanRecordBatch batch = new ScanRecordBatch(records); + int nextSplitIndex = batch.getCurrentSplitIndex(); + if (currentSplitIndex != nextSplitIndex) { + currentSplitIndex = nextSplitIndex; + currentReadRecordsCount = 0; + } + // Keep toSkip unchanged across indexed batches. Existing checkpoints start at index 0 and + // store a global skip count, while newly emitted positions use a split-local count. + return batch; } @Override @@ -143,13 +163,15 @@ public void close() throws Exception { private static class ScanRecordBatch implements CloseableIterator { private final CloseableIterator rowIterator; - private int currentSplitIndex; + private final int currentSplitIndex; public ScanRecordBatch(CloseableIterator rowIterator) { this.rowIterator = rowIterator; if (rowIterator instanceof IndexedLakeSplitRecordIterator) { currentSplitIndex = ((IndexedLakeSplitRecordIterator) rowIterator).getCurrentLakeSplitIndex(); + } else { + currentSplitIndex = RecordAndPos.DEFAULT_SPLIT_INDEX; } } @@ -184,15 +206,17 @@ public int getCurrentSplitIndex() { private class RecordAndPosBatch implements CloseableIterator { private CloseableIterator records; + private int currentSplitIndex; private final MutableRecordAndPos recordAndPosition = new MutableRecordAndPos(); RecordAndPosBatch replace(CloseableIterator records) { this.records = records; if (records instanceof ScanRecordBatch) { - int currentSplitIndex = ((ScanRecordBatch) records).getCurrentSplitIndex(); + currentSplitIndex = ((ScanRecordBatch) records).getCurrentSplitIndex(); recordAndPosition.setRecord(null, NO_READ_RECORDS_COUNT, currentSplitIndex); } else { + currentSplitIndex = RecordAndPos.DEFAULT_SPLIT_INDEX; recordAndPosition.setRecord(null, NO_READ_RECORDS_COUNT); } return this; @@ -205,7 +229,8 @@ public boolean hasNext() { @Override public RecordAndPos next() { - recordAndPosition.setRecord(records.next(), ++currentReadRecordsCount); + recordAndPosition.setRecord( + records.next(), ++currentReadRecordsCount, currentSplitIndex); return recordAndPosition; } diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/lake/reader/SeekableLakeSnapshotSplitScannerTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/lake/reader/SeekableLakeSnapshotSplitScannerTest.java new file mode 100644 index 00000000000..73634e135d6 --- /dev/null +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/lake/reader/SeekableLakeSnapshotSplitScannerTest.java @@ -0,0 +1,96 @@ +/* + * 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.fluss.flink.lake.reader; + +import org.apache.fluss.lake.source.LakeSource; +import org.apache.fluss.lake.source.LakeSplit; +import org.apache.fluss.lake.source.RecordReader; +import org.apache.fluss.record.ChangeType; +import org.apache.fluss.record.GenericRecord; +import org.apache.fluss.record.LogRecord; +import org.apache.fluss.row.InternalRow; +import org.apache.fluss.utils.CloseableIterator; + +import org.junit.jupiter.api.Test; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.apache.fluss.testutils.DataTestUtils.row; +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** Tests for {@link SeekableLakeSnapshotSplitScanner}. */ +class SeekableLakeSnapshotSplitScannerTest { + + @Test + void testSkipEmptyLakeSplits() throws Exception { + LakeSplit firstSplit = mock(LakeSplit.class); + LakeSplit emptySplit = mock(LakeSplit.class); + LakeSplit lastSplit = mock(LakeSplit.class); + InternalRow firstRow = row(1, "first"); + InternalRow lastRow = row(2, "last"); + + Map> rowsBySplit = new HashMap<>(); + rowsBySplit.put(firstSplit, Arrays.asList(firstRow)); + rowsBySplit.put(emptySplit, new ArrayList<>()); + rowsBySplit.put(lastSplit, Arrays.asList(lastRow)); + + LakeSource lakeSource = createLakeSource(rowsBySplit); + SeekableLakeSnapshotSplitScanner scanner = + new SeekableLakeSnapshotSplitScanner( + lakeSource, Arrays.asList(firstSplit, emptySplit, lastSplit), 0); + + List actualRows = new ArrayList<>(); + CloseableIterator batch; + while ((batch = scanner.pollBatch(Duration.ZERO)) != null) { + while (batch.hasNext()) { + actualRows.add(batch.next()); + } + batch.close(); + } + scanner.close(); + + assertThat(actualRows).containsExactly(firstRow, lastRow); + } + + @SuppressWarnings("unchecked") + private LakeSource createLakeSource(Map> rowsBySplit) + throws Exception { + LakeSource lakeSource = mock(LakeSource.class); + when(lakeSource.createRecordReader(any())) + .thenAnswer( + invocation -> { + LakeSource.ReaderContext context = invocation.getArgument(0); + List records = new ArrayList<>(); + for (InternalRow row : rowsBySplit.get(context.lakeSplit())) { + records.add(new GenericRecord(0, 0, ChangeType.INSERT, row)); + } + RecordReader reader = () -> CloseableIterator.wrap(records.iterator()); + return reader; + }); + return lakeSource; + } +} diff --git a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/reader/BoundedSplitReaderTest.java b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/reader/BoundedSplitReaderTest.java index 2ecb48ea193..bc2724b0f2f 100644 --- a/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/reader/BoundedSplitReaderTest.java +++ b/fluss-flink/fluss-flink-common/src/test/java/org/apache/fluss/flink/source/reader/BoundedSplitReaderTest.java @@ -18,6 +18,10 @@ package org.apache.fluss.flink.source.reader; import org.apache.fluss.client.table.scanner.batch.BatchScanner; +import org.apache.fluss.flink.lake.reader.IndexedLakeSplitRecordIterator; +import org.apache.fluss.record.ChangeType; +import org.apache.fluss.record.GenericRecord; +import org.apache.fluss.record.LogRecord; import org.apache.fluss.row.InternalRow; import org.apache.fluss.row.ProjectedRow; import org.apache.fluss.row.indexed.IndexedRow; @@ -80,6 +84,54 @@ void testReadWithSkipOverTotalRecordsNum() { "Skip more than the number of total records, has skipped 10 record(s), but remain 1 record(s) to skip."); } + @Test + void testReadPositionAcrossIndexedBatches() throws IOException { + List> rowsBySplit = new ArrayList<>(); + rowsBySplit.add(mockRows(2)); + rowsBySplit.add(mockRows(2)); + BoundedSplitReader reader = + new BoundedSplitReader(new TestingIndexedBatchScanner(rowsBySplit, 0), 0); + + List records = collectRecords(reader); + + assertThat(records) + .extracting(RecordAndPos::getCurrentSplitIndex) + .containsExactly(0, 0, 1, 1); + assertThat(records) + .extracting(RecordAndPos::readRecordsCount) + .containsExactly(1L, 2L, 1L, 2L); + } + + @Test + void testRestoreLegacyPositionAcrossIndexedBatches() throws IOException { + List> rowsBySplit = new ArrayList<>(); + rowsBySplit.add(mockRows(2)); + rowsBySplit.add(mockRows(2)); + + // Existing checkpoints start at split 0 and store a global records-to-skip count. + BoundedSplitReader reader = + new BoundedSplitReader(new TestingIndexedBatchScanner(rowsBySplit, 0), 3); + List records = collectRecords(reader); + + assertThat(records).hasSize(1); + assertThat(records.get(0).getCurrentSplitIndex()).isEqualTo(1); + assertThat(records.get(0).readRecordsCount()).isEqualTo(2); + } + + @Test + void testRestoreSplitLocalPosition() throws IOException { + List> rowsBySplit = new ArrayList<>(); + rowsBySplit.add(mockRows(2)); + + BoundedSplitReader reader = + new BoundedSplitReader(new TestingIndexedBatchScanner(rowsBySplit, 1), 1); + List records = collectRecords(reader); + + assertThat(records).hasSize(1); + assertThat(records.get(0).getCurrentSplitIndex()).isEqualTo(1); + assertThat(records.get(0).readRecordsCount()).isEqualTo(2); + } + @Test void testSizeInBytesWithIndexedRow() throws IOException { // Use IndexedRow which implements MemoryAwareGetters @@ -164,6 +216,40 @@ public void close() throws IOException { } } + /** A testing scanner that returns one indexed batch for each inner lake split. */ + private static class TestingIndexedBatchScanner implements BatchScanner { + + private final List> rowsBySplit; + private final int firstSplitIndex; + private int nextBatchIndex; + + private TestingIndexedBatchScanner( + List> rowsBySplit, int firstSplitIndex) { + this.rowsBySplit = rowsBySplit; + this.firstSplitIndex = firstSplitIndex; + } + + @Override + @Nullable + public CloseableIterator pollBatch(Duration timeout) { + if (nextBatchIndex >= rowsBySplit.size()) { + return null; + } + + List records = new ArrayList<>(); + for (InternalRow row : rowsBySplit.get(nextBatchIndex)) { + records.add(new GenericRecord(0, 0, ChangeType.INSERT, row)); + } + return new IndexedLakeSplitRecordIterator( + CloseableIterator.wrap(records.iterator()), firstSplitIndex + nextBatchIndex++); + } + + @Override + public void close() throws IOException { + // do nothing + } + } + private List mockRows(int numRows) { List rows = new ArrayList<>(numRows); for (int i = 0; i < numRows; i++) { @@ -206,7 +292,10 @@ private List collectRecords(BoundedSplitReader reader) throws IOEx while (recordIter.hasNext()) { RecordAndPos recordAndPos = recordIter.next(); records.add( - new RecordAndPos(recordAndPos.scanRecord, recordAndPos.readRecordsCount)); + new RecordAndPos( + recordAndPos.scanRecord, + recordAndPos.readRecordsCount, + recordAndPos.getCurrentSplitIndex())); } recordIter.close(); }