From 024ee78b8573a298cdad531e768483d6061d0c39 Mon Sep 17 00:00:00 2001 From: QuakeWang Date: Wed, 12 Aug 2026 15:07:04 +0800 Subject: [PATCH] [core] Avoid repeated sorted-index source lookups Primary-key sorted-index scans linearly searched each source group to locate every data file, making source localization quadratic in the number of files. Build a file-name lookup alongside source offsets and preserve row-count validation when creating file-local readers. Signed-off-by: QuakeWang --- .../source/PrimaryKeySortedIndexScan.java | 55 ++++++++++++------- .../source/PrimaryKeySortedIndexScanTest.java | 27 +++++---- 2 files changed, 53 insertions(+), 29 deletions(-) diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeySortedIndexScan.java b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeySortedIndexScan.java index 7bc04f5d6fdf..cb03736b7be8 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeySortedIndexScan.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeySortedIndexScan.java @@ -54,6 +54,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; import java.util.HashSet; import java.util.IdentityHashMap; import java.util.LinkedHashMap; @@ -274,8 +275,7 @@ static EvaluatedPlan evaluate( totalRowCount)); sharedReaders.put(indexGroup, reader); } - return Collections.singletonList( - fileLocalReader(file, group.get(), reader)); + return Collections.singletonList(fileLocalReader(file, reader)); }); Optional result; try { @@ -300,23 +300,20 @@ static EvaluatedPlan evaluate( } private static GlobalIndexReader fileLocalReader( - FilePlan file, PkSortedIndexGroup group, SharedGlobalIndexReader reader) { - List sourceFiles = group.sourceFiles(); - PrimaryKeyIndexSourceFile target = - new PrimaryKeyIndexSourceFile( - file.dataFile().fileName(), file.dataFile().rowCount()); - int sourceIndex = -1; - for (int i = 0; i < sourceFiles.size(); i++) { - if (sourceFiles.get(i).equals(target)) { - sourceIndex = i; - break; - } - } + FilePlan file, SharedGlobalIndexReader reader) { + DataFileMeta dataFile = file.dataFile(); + SourceLocation sourceLocation = reader.sourceLocations.get(dataFile.fileName()); checkArgument( - sourceIndex >= 0, + sourceLocation != null, "Data file %s is not covered by its sorted-index source group.", - file.dataFile().fileName()); - return new FileLocalGlobalIndexReader(reader, sourceIndex); + dataFile.fileName()); + checkArgument( + dataFile.rowCount() == sourceLocation.rowCount, + "Data file %s row count %s does not match sorted-index source row count %s.", + dataFile.fileName(), + dataFile.rowCount(), + sourceLocation.rowCount); + return new FileLocalGlobalIndexReader(reader, sourceLocation.sourceIndex); } private static long totalRowCount(List sourceFiles) { @@ -333,6 +330,17 @@ private static void rethrowIfInterrupted(RuntimeException exception) { } } + private static final class SourceLocation { + + private final long rowCount; + private final int sourceIndex; + + private SourceLocation(long rowCount, int sourceIndex) { + this.rowCount = rowCount; + this.sourceIndex = sourceIndex; + } + } + /** Shares one source-group reader and its group-global query results across source files. */ private static final class SharedGlobalIndexReader implements GlobalIndexReader { @@ -342,6 +350,7 @@ private static final class SharedGlobalIndexReader implements GlobalIndexReader CompletableFuture>, CompletableFuture>>> localizedResults; + private final Map sourceLocations; private final long[] sourceOffsets; private GlobalIndexReader reader; @@ -353,10 +362,18 @@ private SharedGlobalIndexReader( this.readerFactory = readerFactory; this.results = new ConcurrentHashMap<>(); this.localizedResults = new ConcurrentHashMap<>(); + this.sourceLocations = new HashMap<>(); this.sourceOffsets = new long[sourceFiles.size() + 1]; for (int i = 0; i < sourceFiles.size(); i++) { - sourceOffsets[i + 1] = - Math.addExact(sourceOffsets[i], sourceFiles.get(i).rowCount()); + PrimaryKeyIndexSourceFile sourceFile = sourceFiles.get(i); + checkArgument( + sourceLocations.put( + sourceFile.fileName(), + new SourceLocation(sourceFile.rowCount(), i)) + == null, + "Duplicate sorted-index source file %s.", + sourceFile.fileName()); + sourceOffsets[i + 1] = Math.addExact(sourceOffsets[i], sourceFile.rowCount()); } } diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexScanTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexScanTest.java index c704589af74a..b800fb1c61db 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexScanTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexScanTest.java @@ -205,7 +205,8 @@ void testDuplicateLevelPayloadsFallBackWithoutCreatingReader() { void testReadMergedSourceGroupInFileLocalPositions() throws IOException { DataFileMeta first = dataFile("data-1", 2); DataFileMeta second = dataFile("data-2", 3); - DataSplit split = dataSplit(11, 0, true, second, first); + DataFileMeta third = dataFile("data-3", 4); + DataSplit split = dataSplit(11, 0, true, third, first, second); PrimaryKeyIndexDefinition definition = definition( 7, @@ -216,10 +217,11 @@ void testReadMergedSourceGroupInFileLocalPositions() throws IOException { "btree-merged", Arrays.asList( new PrimaryKeyIndexSourceFile("data-1", 2), - new PrimaryKeyIndexSourceFile("data-2", 3)), + new PrimaryKeyIndexSourceFile("data-2", 3), + new PrimaryKeyIndexSourceFile("data-3", 4)), "btree", 7, - 5); + 9); PrimaryKeySortedIndexScan.Plan plan = PrimaryKeySortedIndexScan.plan( 11, @@ -232,8 +234,10 @@ void testReadMergedSourceGroupInFileLocalPositions() throws IOException { AtomicInteger queries = new AtomicInteger(); CountingRoaringNavigableMap64 groupPositions = new CountingRoaringNavigableMap64(); groupPositions.add(1); - groupPositions.add(3); + groupPositions.add(2); groupPositions.add(4); + groupPositions.add(6); + groupPositions.add(7); GlobalIndexReader reader = mock(GlobalIndexReader.class); when(reader.visitEqual(any(), eq(42))) .thenAnswer( @@ -251,23 +255,26 @@ void testReadMergedSourceGroupInFileLocalPositions() throws IOException { (ignoredFile, ignoredDefinition, payloads, totalRowCount) -> { readersCreated.incrementAndGet(); assertThat(payloads).containsExactly(mergedPayload); - assertThat(totalRowCount).isEqualTo(5); + assertThat(totalRowCount).isEqualTo(9); return reader; }); PrimaryKeySortedIndexResult result = new PrimaryKeySortedIndexResult(evaluated); assertThat(readersCreated).hasValue(1); assertThat(queries).hasValue(1); - assertThat(groupPositions.iteratedPositions()).isEqualTo(3); + assertThat(groupPositions.iteratedPositions()).isEqualTo(5); verify(reader, times(1)).close(); - assertThat(result.splits()).hasSize(2); + assertThat(result.splits()).hasSize(3); assertThat(result.splits()).allMatch(IndexedSplit.class::isInstance); - IndexedSplit secondSplit = (IndexedSplit) result.splits().get(0); - assertThat(secondSplit.dataSplit().dataFiles()).containsExactly(second); - assertThat(secondSplit.rowRanges()).containsExactly(new Range(1, 2)); + IndexedSplit thirdSplit = (IndexedSplit) result.splits().get(0); + assertThat(thirdSplit.dataSplit().dataFiles()).containsExactly(third); + assertThat(thirdSplit.rowRanges()).containsExactly(new Range(1, 2)); IndexedSplit firstSplit = (IndexedSplit) result.splits().get(1); assertThat(firstSplit.dataSplit().dataFiles()).containsExactly(first); assertThat(firstSplit.rowRanges()).containsExactly(new Range(1, 1)); + IndexedSplit secondSplit = (IndexedSplit) result.splits().get(2); + assertThat(secondSplit.dataSplit().dataFiles()).containsExactly(second); + assertThat(secondSplit.rowRanges()).containsExactly(new Range(0, 0), new Range(2, 2)); } @Test