diff --git a/paimon-core/src/main/java/org/apache/paimon/index/pksorted/PkSortedBucketIndexState.java b/paimon-core/src/main/java/org/apache/paimon/index/pksorted/PkSortedBucketIndexState.java index fa6a5dd10b06..1933a28778ef 100644 --- a/paimon-core/src/main/java/org/apache/paimon/index/pksorted/PkSortedBucketIndexState.java +++ b/paimon-core/src/main/java/org/apache/paimon/index/pksorted/PkSortedBucketIndexState.java @@ -37,6 +37,22 @@ /** Immutable sorted-index state for one field and bucket. */ public final class PkSortedBucketIndexState { + private static final class PayloadCandidate { + + private final IndexFileMeta payload; + private final PkSortedIndexGroup group; + private final List activeSources; + + private PayloadCandidate( + IndexFileMeta payload, + PkSortedIndexGroup group, + List activeSources) { + this.payload = payload; + this.group = group; + this.activeSources = activeSources; + } + } + private final List groups; private final List coveredSourceFiles; private final List uncoveredSourceFiles; @@ -72,55 +88,102 @@ public static PkSortedBucketIndexState fromActiveDataFiles( sources.sort(Comparator.comparing(PrimaryKeyIndexSourceFile::fileName)); } - Map> payloadsByLevel = new TreeMap<>(); + Map> candidatesByLevel = new TreeMap<>(); List rejected = new ArrayList<>(); for (IndexFileMeta payload : activePayloads) { try { PrimaryKeyIndexSourceMeta sourceMeta = PrimaryKeyIndexSourceMeta.fromIndexFile(payload); - List desired = + List activeLevelSources = sourcesByLevel.get(sourceMeta.dataLevel()); - if (desired == null || !desired.equals(sourceMeta.sourceFiles())) { + if (activeLevelSources == null) { + rejected.add(payload); + continue; + } + + List payloadSources = sourceMeta.sourceFiles(); + List activeIntersection = + activeIntersection(activeLevelSources, payloadSources); + if (activeIntersection == null || activeIntersection.isEmpty()) { + rejected.add(payload); + continue; + } + + Optional group = + PkSortedIndexGroup.create( + fieldId, + indexType, + payloadSources, + Collections.singletonList(payload)); + if (!group.isPresent()) { rejected.add(payload); - } else { - payloadsByLevel - .computeIfAbsent(sourceMeta.dataLevel(), ignored -> new ArrayList<>()) - .add(payload); + continue; } + candidatesByLevel + .computeIfAbsent(sourceMeta.dataLevel(), ignored -> new ArrayList<>()) + .add(new PayloadCandidate(payload, group.get(), activeIntersection)); } catch (RuntimeException ignored) { rejected.add(payload); } } List groups = new ArrayList<>(); - Set coveredLevels = new HashSet<>(); - for (Map.Entry> entry : payloadsByLevel.entrySet()) { - List levelPayloads = entry.getValue(); - Optional group = - levelPayloads.size() == 1 - ? PkSortedIndexGroup.create( - fieldId, - indexType, - sourcesByLevel.get(entry.getKey()), - levelPayloads) - : Optional.empty(); - if (group.isPresent()) { - groups.add(group.get()); - coveredLevels.add(entry.getKey()); - } else { - rejected.addAll(levelPayloads); + Map> coveredSourcesByLevel = new TreeMap<>(); + for (Map.Entry> entry : candidatesByLevel.entrySet()) { + List levelCandidates = entry.getValue(); + if (levelCandidates.size() != 1) { + for (PayloadCandidate candidate : levelCandidates) { + rejected.add(candidate.payload); + } + continue; } + PayloadCandidate candidate = levelCandidates.get(0); + groups.add(candidate.group); + coveredSourcesByLevel + .computeIfAbsent(entry.getKey(), ignored -> new HashSet<>()) + .addAll(candidate.activeSources); } List covered = new ArrayList<>(); List uncovered = new ArrayList<>(); for (Map.Entry> entry : sourcesByLevel.entrySet()) { - (coveredLevels.contains(entry.getKey()) ? covered : uncovered).addAll(entry.getValue()); + Set coveredSources = + coveredSourcesByLevel.getOrDefault(entry.getKey(), Collections.emptySet()); + for (PrimaryKeyIndexSourceFile source : entry.getValue()) { + (coveredSources.contains(source) ? covered : uncovered).add(source); + } } return new PkSortedBucketIndexState(groups, covered, uncovered, rejected); } + private static List activeIntersection( + List activeSources, + List payloadSources) { + List intersection = new ArrayList<>(); + int activeSourceIndex = 0; + for (int i = 0; i < payloadSources.size(); i++) { + PrimaryKeyIndexSourceFile source = payloadSources.get(i); + if (i > 0 && payloadSources.get(i - 1).fileName().compareTo(source.fileName()) >= 0) { + return null; + } + while (activeSourceIndex < activeSources.size() + && activeSources.get(activeSourceIndex).fileName().compareTo(source.fileName()) + < 0) { + activeSourceIndex++; + } + if (activeSourceIndex == activeSources.size() + || !activeSources.get(activeSourceIndex).fileName().equals(source.fileName())) { + continue; + } + if (activeSources.get(activeSourceIndex).rowCount() != source.rowCount()) { + return null; + } + intersection.add(source); + } + return intersection; + } + public List groups() { return groups; } diff --git a/paimon-core/src/main/java/org/apache/paimon/index/pksorted/PkSortedIndexGroup.java b/paimon-core/src/main/java/org/apache/paimon/index/pksorted/PkSortedIndexGroup.java index 54999c216247..f30047a68119 100644 --- a/paimon-core/src/main/java/org/apache/paimon/index/pksorted/PkSortedIndexGroup.java +++ b/paimon-core/src/main/java/org/apache/paimon/index/pksorted/PkSortedIndexGroup.java @@ -30,7 +30,10 @@ import java.util.Optional; import java.util.Set; -/** The single payload which indexes one complete data level. */ +/** + * One validated payload which indexes an immutable source group at one data level. + * Snapshot-specific active coverage is validated by {@link PkSortedBucketIndexState}. + */ public final class PkSortedIndexGroup { private final int dataLevel; 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 2bda216ce0f8..5003c9dbfd92 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 @@ -32,6 +32,7 @@ import org.apache.paimon.index.IndexPathFactory; import org.apache.paimon.index.pk.PrimaryKeyIndexDefinition; import org.apache.paimon.index.pk.PrimaryKeyIndexSourceFile; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourcePolicy; import org.apache.paimon.index.pksorted.PkSortedBucketIndexState; import org.apache.paimon.index.pksorted.PkSortedIndexGroup; import org.apache.paimon.io.DataFileMeta; @@ -169,10 +170,15 @@ static Plan plan( Pair bucket = bucketEntry.getKey(); List bucketPayloads = payloadsByBucket.getOrDefault(bucket, Collections.emptyList()); - Set activeSourceFiles = new HashSet<>(); + Map> activeSourceFilesByLevel = new HashMap<>(); for (DataFileMeta dataFile : bucketEntry.getValue()) { - activeSourceFiles.add( - new PrimaryKeyIndexSourceFile(dataFile.fileName(), dataFile.rowCount())); + if (PrimaryKeyIndexSourcePolicy.shouldRead(dataFile)) { + activeSourceFilesByLevel + .computeIfAbsent(dataFile.level(), ignored -> new HashSet<>()) + .add( + new PrimaryKeyIndexSourceFile( + dataFile.fileName(), dataFile.rowCount())); + } } Map> groupsBySource = new LinkedHashMap<>(); for (PrimaryKeyIndexDefinition definition : scalarDefinitions) { @@ -193,8 +199,11 @@ static Plan plan( bucketEntry.getValue(), definitionPayloads); for (PkSortedIndexGroup group : state.groups()) { + Set activeGroupSources = + activeSourceFilesByLevel.getOrDefault( + group.dataLevel(), Collections.emptySet()); for (PrimaryKeyIndexSourceFile sourceFile : group.sourceFiles()) { - if (!activeSourceFiles.contains(sourceFile)) { + if (!activeGroupSources.contains(sourceFile)) { continue; } groupsBySource diff --git a/paimon-core/src/test/java/org/apache/paimon/index/pksorted/PkSortedBucketIndexStateTest.java b/paimon-core/src/test/java/org/apache/paimon/index/pksorted/PkSortedBucketIndexStateTest.java index f7a9a60d1dc3..e666ee5ea0ad 100644 --- a/paimon-core/src/test/java/org/apache/paimon/index/pksorted/PkSortedBucketIndexStateTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/index/pksorted/PkSortedBucketIndexStateTest.java @@ -61,7 +61,7 @@ void testAcceptsOnePayloadForCompleteLevel() { } @Test - void testRejectsPartialLevelPayload() { + void testAcceptsPayloadSubsetAndLeavesNewSourceUncovered() { DataFileMeta first = dataFile("data-a", 3, 2); DataFileMeta second = dataFile("data-b", 7, 2); IndexFileMeta partial = payload("partial", 2, first); @@ -73,9 +73,89 @@ void testRejectsPartialLevelPayload() { Arrays.asList(first, second), Collections.singletonList(partial)); + assertThat(state.groups()).hasSize(1); + assertThat(state.coveredSourceFiles()).containsExactly(sourceFile(first)); + assertThat(state.uncoveredSourceFiles()).containsExactly(sourceFile(second)); + assertThat(state.rejectedPayloads()).isEmpty(); + } + + @Test + void testRetainsRetiredSourcesAndCoversOnlyActiveIntersection() { + DataFileMeta retired = dataFile("data-a", 3, 2); + DataFileMeta active = dataFile("data-b", 7, 2); + DataFileMeta newlyActive = dataFile("data-c", 5, 2); + IndexFileMeta payload = payload("index", 2, retired, active); + + PkSortedBucketIndexState state = + PkSortedBucketIndexState.fromActiveDataFiles( + 7, + "btree", + Arrays.asList(active, newlyActive), + Collections.singletonList(payload)); + + assertThat(state.groups()).hasSize(1); + assertThat(state.groups().get(0).sourceFiles()) + .containsExactly(sourceFile(retired), sourceFile(active)); + assertThat(state.coveredSourceFiles()).containsExactly(sourceFile(active)); + assertThat(state.uncoveredSourceFiles()).containsExactly(sourceFile(newlyActive)); + assertThat(state.rejectedPayloads()).isEmpty(); + } + + @Test + void testRejectsPayloadWithoutActiveSource() { + DataFileMeta retired = dataFile("data-a", 3, 2); + DataFileMeta active = dataFile("data-b", 7, 2); + IndexFileMeta payload = payload("index", 2, retired); + + PkSortedBucketIndexState state = + PkSortedBucketIndexState.fromActiveDataFiles( + 7, + "btree", + Collections.singletonList(active), + Collections.singletonList(payload)); + + assertThat(state.groups()).isEmpty(); + assertThat(state.coveredSourceFiles()).isEmpty(); + assertThat(state.uncoveredSourceFiles()).containsExactly(sourceFile(active)); + assertThat(state.rejectedPayloads()).containsExactly(payload); + } + + @Test + void testRejectsMismatchedActiveSourceRowCount() { + DataFileMeta active = dataFile("data", 3, 2); + DataFileMeta stale = dataFile("data", 4, 2); + IndexFileMeta payload = payload("index", 2, stale); + + PkSortedBucketIndexState state = + PkSortedBucketIndexState.fromActiveDataFiles( + 7, + "btree", + Collections.singletonList(active), + Collections.singletonList(payload)); + assertThat(state.groups()).isEmpty(); - assertThat(state.uncoveredSourceFiles()).hasSize(2); - assertThat(state.rejectedPayloads()).containsExactly(partial); + assertThat(state.uncoveredSourceFiles()).containsExactly(sourceFile(active)); + assertThat(state.rejectedPayloads()).containsExactly(payload); + } + + @Test + void testRejectsMisorderedPayloadSources() { + DataFileMeta first = dataFile("data-a", 3, 2); + DataFileMeta second = dataFile("data-b", 7, 2); + IndexFileMeta payload = + payload("index", 2, Arrays.asList(sourceFile(second), sourceFile(first))); + + PkSortedBucketIndexState state = + PkSortedBucketIndexState.fromActiveDataFiles( + 7, + "btree", + Arrays.asList(first, second), + Collections.singletonList(payload)); + + assertThat(state.groups()).isEmpty(); + assertThat(state.uncoveredSourceFiles()) + .containsExactly(sourceFile(first), sourceFile(second)); + assertThat(state.rejectedPayloads()).containsExactly(payload); } @Test @@ -159,6 +239,11 @@ private static IndexFileMeta payload(String name, int level, DataFileMeta... fil new PrimaryKeyIndexSourceFile( file.fileName(), file.rowCount())) .collect(java.util.stream.Collectors.toList()); + return payload(name, level, sources); + } + + private static IndexFileMeta payload( + String name, int level, List sources) { long rowCount = 0; for (PrimaryKeyIndexSourceFile source : sources) { rowCount += source.rowCount(); @@ -177,4 +262,8 @@ private static IndexFileMeta payload(String name, int level, DataFileMeta... fil new PrimaryKeyIndexSourceMeta(level, sources).serialize()), null); } + + private static PrimaryKeyIndexSourceFile sourceFile(DataFileMeta file) { + return new PrimaryKeyIndexSourceFile(file.fileName(), file.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 c9c09ea7d498..25edb47ecaf7 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 @@ -363,6 +363,99 @@ void testReadMergedSourceGroupInFileLocalPositions() throws IOException { assertThat(secondSplit.rowRanges()).containsExactly(new Range(0, 0), new Range(2, 2)); } + @Test + void testRetiredSourceOffsetsAndNewSourceFallback() throws IOException { + DataFileMeta retired = dataFile("data-1", 2); + DataFileMeta active = dataFile("data-2", 3); + DataFileMeta newlyActive = dataFile("data-3", 4); + DataSplit split = dataSplit(11, 0, active, newlyActive); + PrimaryKeyIndexDefinition definition = + definition( + 7, + BTreeGlobalIndexerFactory.IDENTIFIER, + PrimaryKeyIndexDefinition.Family.BTREE); + IndexFileMeta payload = + payload( + "btree-retired-source", + Arrays.asList( + new PrimaryKeyIndexSourceFile( + retired.fileName(), retired.rowCount()), + new PrimaryKeyIndexSourceFile( + active.fileName(), active.rowCount())), + "btree", + 7, + 5); + PrimaryKeySortedIndexScan.Plan plan = + PrimaryKeySortedIndexScan.plan( + 11, + Collections.singletonList(split), + Collections.singletonList(definition), + Collections.singletonList(payloadEntry(0, payload))); + + assertThat(plan.files()).hasSize(2); + assertThat(plan.files().get(0).group(7)).isPresent(); + assertThat(plan.files().get(0).group(7).get().sourceFiles()) + .extracting(PrimaryKeyIndexSourceFile::fileName) + .containsExactly("data-1", "data-2"); + assertThat(plan.files().get(1).group(7)).isEmpty(); + + RowType rowType = RowType.of(new DataField(7, "f7", DataTypes.INT())); + Predicate predicate = new PredicateBuilder(rowType).equal(0, 42); + GlobalIndexReader reader = readerWithPositions(2, 4); + AtomicInteger readersCreated = new AtomicInteger(); + + PrimaryKeySortedIndexScan.EvaluatedPlan evaluated = + PrimaryKeySortedIndexScan.evaluate( + plan, + rowType, + predicate, + Collections.singletonList(definition), + (ignoredFile, ignoredDefinition, payloads, totalRowCount) -> { + readersCreated.incrementAndGet(); + assertThat(payloads).containsExactly(payload); + assertThat(totalRowCount).isEqualTo(5); + return reader; + }); + + assertThat(readersCreated).hasValue(1); + assertThat(evaluated.files().get(0).result()).isPresent(); + assertThat(evaluated.files().get(0).result().get().results()).containsExactly(0L, 2L); + assertThat(evaluated.files().get(1).result()).isEmpty(); + verify(reader).close(); + } + + @Test + void testRetiredSourceAtAnotherLevelDoesNotInheritGroup() { + DataFileMeta active = dataFile("data-1", 2, 1); + DataFileMeta movedToAnotherLevel = dataFile("data-2", 3, 2); + DataSplit split = dataSplit(11, 0, active, movedToAnotherLevel); + PrimaryKeyIndexDefinition definition = + definition( + 7, + BTreeGlobalIndexerFactory.IDENTIFIER, + PrimaryKeyIndexDefinition.Family.BTREE); + IndexFileMeta payload = + payload( + "btree-level-1", + Arrays.asList( + new PrimaryKeyIndexSourceFile("data-1", 2), + new PrimaryKeyIndexSourceFile("data-2", 3)), + "btree", + 7, + 5); + + PrimaryKeySortedIndexScan.Plan plan = + PrimaryKeySortedIndexScan.plan( + 11, + Collections.singletonList(split), + Collections.singletonList(definition), + Collections.singletonList(payloadEntry(0, payload))); + + assertThat(plan.files()).hasSize(2); + assertThat(plan.files().get(0).group(7)).isPresent(); + assertThat(plan.files().get(1).group(7)).isEmpty(); + } + @Test void testArrayContainsIsCachedAndLocalized() throws IOException { DataFileMeta first = dataFile("data-1", 2); @@ -645,6 +738,10 @@ private static DataSplit dataSplit( } private static DataFileMeta dataFile(String fileName, long rowCount) { + return dataFile(fileName, rowCount, 1); + } + + private static DataFileMeta dataFile(String fileName, long rowCount, int level) { return DataFileMeta.forAppend( fileName, 100, @@ -660,7 +757,7 @@ private static DataFileMeta dataFile(String fileName, long rowCount) { null, null, null) - .upgrade(1); + .upgrade(level); } private static IndexManifestEntry payloadEntry(int bucket, IndexFileMeta payload) { diff --git a/paimon-python/pypaimon/index/pk/primary_key_index_source_policy.py b/paimon-python/pypaimon/index/pk/primary_key_index_source_policy.py new file mode 100644 index 000000000000..66ffa1e4ab07 --- /dev/null +++ b/paimon-python/pypaimon/index/pk/primary_key_index_source_policy.py @@ -0,0 +1,21 @@ +# 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. + + +def should_read(data_file): + """Match Java PrimaryKeyIndexSourcePolicy.shouldRead.""" + # FileSource.COMPACT = 1. + return data_file.file_source == 1 and data_file.level > 0 diff --git a/paimon-python/pypaimon/index/pksorted/pk_sorted_bucket_index_state.py b/paimon-python/pypaimon/index/pksorted/pk_sorted_bucket_index_state.py index 490957149efd..c541c9d798ea 100644 --- a/paimon-python/pypaimon/index/pksorted/pk_sorted_bucket_index_state.py +++ b/paimon-python/pypaimon/index/pksorted/pk_sorted_bucket_index_state.py @@ -18,6 +18,7 @@ from pypaimon.index.pk.primary_key_index_source_file import PrimaryKeyIndexSourceFile from pypaimon.index.pk.primary_key_index_source_meta import PrimaryKeyIndexSourceMeta +from pypaimon.index.pk.primary_key_index_source_policy import should_read from pypaimon.index.pksorted.pk_sorted_index_group import PkSortedIndexGroup @@ -32,40 +33,70 @@ class PkSortedBucketIndexState: def from_active_data_files(field_id, index_type, active_data_files, active_payloads): sources_by_level = {} for data_file in active_data_files: - if data_file.file_source != 1 or data_file.level <= 0: + if not should_read(data_file): continue sources_by_level.setdefault(data_file.level, []).append( PrimaryKeyIndexSourceFile(data_file.file_name, data_file.row_count)) for sources in sources_by_level.values(): sources.sort(key=lambda source: source.file_name) - payloads_by_level = {} + candidates_by_level = {} rejected = [] for payload in active_payloads: try: source_meta = PrimaryKeyIndexSourceMeta.from_index_file(payload) - desired = sources_by_level.get(source_meta.data_level) - if desired is None or tuple(desired) != tuple(source_meta.source_files): + active_level_sources = sources_by_level.get(source_meta.data_level) + if active_level_sources is None: rejected.append(payload) - else: - payloads_by_level.setdefault(source_meta.data_level, []).append(payload) + continue + active_intersection = _active_intersection( + active_level_sources, source_meta.source_files) + if not active_intersection: + rejected.append(payload) + continue + group = PkSortedIndexGroup.create( + field_id, index_type, source_meta.source_files, [payload]) + if group is None: + rejected.append(payload) + continue + candidates_by_level.setdefault(source_meta.data_level, []).append( + (payload, group, active_intersection)) except (TypeError, ValueError): rejected.append(payload) groups = [] - covered_levels = set() - for level, payloads in sorted(payloads_by_level.items()): - group = PkSortedIndexGroup.create( - field_id, index_type, sources_by_level[level], payloads) - if group is None: - rejected.extend(payloads) - else: - groups.append(group) - covered_levels.add(level) + covered_sources_by_level = {} + for level, candidates in sorted(candidates_by_level.items()): + if len(candidates) != 1: + rejected.extend(candidate[0] for candidate in candidates) + continue + _, group, active_intersection = candidates[0] + groups.append(group) + covered_sources_by_level[level] = set(active_intersection) covered = [] uncovered = [] for level, sources in sorted(sources_by_level.items()): - (covered if level in covered_levels else uncovered).extend(sources) + covered_sources = covered_sources_by_level.get(level, set()) + for source in sources: + (covered if source in covered_sources else uncovered).append(source) return PkSortedBucketIndexState( tuple(groups), tuple(covered), tuple(uncovered), tuple(rejected)) + + +def _active_intersection(active_sources, payload_sources): + intersection = [] + active_source_index = 0 + for index, source in enumerate(payload_sources): + if index > 0 and payload_sources[index - 1].file_name >= source.file_name: + return None + while (active_source_index < len(active_sources) + and active_sources[active_source_index].file_name < source.file_name): + active_source_index += 1 + if (active_source_index == len(active_sources) + or active_sources[active_source_index].file_name != source.file_name): + continue + if active_sources[active_source_index].row_count != source.row_count: + return None + intersection.append(source) + return intersection diff --git a/paimon-python/pypaimon/table/source/primary_key_full_text_scan.py b/paimon-python/pypaimon/table/source/primary_key_full_text_scan.py index afa3a02dfc26..dab5f0e66ea2 100644 --- a/paimon-python/pypaimon/table/source/primary_key_full_text_scan.py +++ b/paimon-python/pypaimon/table/source/primary_key_full_text_scan.py @@ -23,6 +23,8 @@ from pypaimon.index.index_file_handler import IndexFileHandler from pypaimon.index.pk.primary_key_index_source_meta import ( PrimaryKeyIndexSourceMeta) +from pypaimon.index.pk.primary_key_index_source_policy import ( + should_read as _should_read_source) from pypaimon.read.query_auth_split import QueryAuthSplit from pypaimon.read.split import DataSplit from pypaimon.snapshot.time_travel_util import TimeTravelUtil @@ -184,11 +186,6 @@ def _current_payloads(active_files, active_payloads): return current, covered -def _should_read_source(data_file): - # FileSource.COMPACT = 1. Match Java PrimaryKeyIndexSourcePolicy. - return data_file.file_source == 1 and data_file.level > 0 - - class PrimaryKeyFullTextScanPlan(FullTextScanPlan): def __init__(self, snapshot_id, splits): super().__init__(splits) diff --git a/paimon-python/pypaimon/table/source/primary_key_sorted_index_scan.py b/paimon-python/pypaimon/table/source/primary_key_sorted_index_scan.py index 5a2c91ad6ad9..865a52eebf0a 100644 --- a/paimon-python/pypaimon/table/source/primary_key_sorted_index_scan.py +++ b/paimon-python/pypaimon/table/source/primary_key_sorted_index_scan.py @@ -24,6 +24,7 @@ from pypaimon.globalindex.data_evolution_global_index_scanner import _create_inner_readers from pypaimon.common.options.core_options import CoreOptions from pypaimon.index.pk.primary_key_index_source_file import PrimaryKeyIndexSourceFile +from pypaimon.index.pk.primary_key_index_source_policy import should_read from pypaimon.index.pksorted.pk_sorted_bucket_index_state import PkSortedBucketIndexState from pypaimon.utils.roaring_bitmap import RoaringBitmap64 @@ -77,9 +78,12 @@ def plan(snapshot_id, data_splits, definitions, index_entries): groups_by_bucket = {} for bucket, data_files in data_files_by_bucket.items(): payloads = payloads_by_bucket.get(bucket, []) - active_sources = { - PrimaryKeyIndexSourceFile(f.file_name, f.row_count) for f in data_files - } + active_sources_by_level = {} + for data_file in data_files: + if should_read(data_file): + active_sources_by_level.setdefault(data_file.level, set()).add( + PrimaryKeyIndexSourceFile( + data_file.file_name, data_file.row_count)) by_source = {} for definition in definitions: definition_payloads = [ @@ -92,8 +96,10 @@ def plan(snapshot_id, data_splits, definitions, index_entries): definition.field_id, definition.index_type, data_files, definition_payloads) for group in state.groups: + active_group_sources = active_sources_by_level.get( + group.data_level, set()) for source in group.source_files: - if source in active_sources: + if source in active_group_sources: by_source.setdefault(source.file_name, {})[definition.field_id] = group except Exception as exc: LOG.warning("Failed to plan primary-key sorted index for field %s: %s", diff --git a/paimon-python/pypaimon/table/source/primary_key_vector_scan.py b/paimon-python/pypaimon/table/source/primary_key_vector_scan.py index ccd25357839e..99156528f061 100644 --- a/paimon-python/pypaimon/table/source/primary_key_vector_scan.py +++ b/paimon-python/pypaimon/table/source/primary_key_vector_scan.py @@ -21,6 +21,8 @@ from pypaimon.common.options.core_options import CoreOptions from pypaimon.index.index_file_handler import IndexFileHandler from pypaimon.index.pk.primary_key_index_source_meta import PrimaryKeyIndexSourceMeta +from pypaimon.index.pk.primary_key_index_source_policy import ( + should_read as _should_read_source) from pypaimon.read.query_auth_split import QueryAuthSplit from pypaimon.read.split import DataSplit from pypaimon.globalindex.indexed_split import IndexedSplit @@ -187,11 +189,6 @@ def _bucket_splits(source_splits, entries): return result -def _should_read_source(data_file): - # FileSource.COMPACT = 1. Match Java PrimaryKeyIndexSourcePolicy. - return data_file.file_source == 1 and data_file.level > 0 - - def _residual_row_ranges(table, predicate, split, candidate_ranges): """Evaluate the residual predicate on physical rows before ANN search.""" from pypaimon.read.push_down_utils import ( diff --git a/paimon-python/pypaimon/tests/primary_key_sorted_index_scan_test.py b/paimon-python/pypaimon/tests/primary_key_sorted_index_scan_test.py index 791e319120d6..1eab053efd26 100644 --- a/paimon-python/pypaimon/tests/primary_key_sorted_index_scan_test.py +++ b/paimon-python/pypaimon/tests/primary_key_sorted_index_scan_test.py @@ -28,6 +28,7 @@ PrimaryKeyIndexDefinition, PrimaryKeyIndexFamily) from pypaimon.index.pk.primary_key_index_source_file import PrimaryKeyIndexSourceFile from pypaimon.index.pk.primary_key_index_source_meta import PrimaryKeyIndexSourceMeta +from pypaimon.index.pksorted.pk_sorted_bucket_index_state import PkSortedBucketIndexState from pypaimon.manifest.index_manifest_entry import IndexManifestEntry from pypaimon.read.split import DataSplit from pypaimon.schema.data_types import AtomicType, DataField @@ -86,6 +87,110 @@ def reader_factory(*ignored): self.assertEqual([1], evaluated.files[0].result.results().to_list()) self.assertEqual([1], evaluated.files[1].result.results().to_list()) + def test_retired_source_offsets_and_new_source_falls_back(self): + field = DataField(3, "value", AtomicType("INT")) + definition = PrimaryKeyIndexDefinition( + "value", 3, "btree", Options.from_none(), PrimaryKeyIndexFamily.BTREE) + files = [ + SimpleNamespace(file_name="b", row_count=3, level=1, file_source=1), + SimpleNamespace(file_name="c", row_count=4, level=1, file_source=1), + ] + partition = GenericRow([], []) + split = DataSplit(files, partition, 0, raw_convertible=True) + source_meta = PrimaryKeyIndexSourceMeta( + 1, [PrimaryKeyIndexSourceFile("a", 2), + PrimaryKeyIndexSourceFile("b", 3)]).serialize() + payload = IndexFileMeta( + "btree", "index", 1, 5, + global_index_meta=GlobalIndexMeta(0, 4, 3, source_meta=source_meta)) + planned = scan.plan( + 9, [split], [definition], + [IndexManifestEntry(0, partition, 0, payload)]) + + self.assertEqual((PrimaryKeyIndexSourceFile("a", 2), + PrimaryKeyIndexSourceFile("b", 3)), + planned.files[0].groups[3].source_files) + self.assertNotIn(3, planned.files[1].groups) + + created = [] + + def reader_factory(*ignored): + created.append(_Reader([Range(2, 2), Range(4, 4)])) + return created[-1] + + evaluated = scan.evaluate( + planned, [field], PredicateBuilder([field]).equal("value", 1), + [definition], reader_factory) + + self.assertEqual(1, len(created)) + self.assertEqual([0, 2], evaluated.files[0].result.results().to_list()) + self.assertIsNone(evaluated.files[1].result) + + def test_source_at_another_level_does_not_inherit_group(self): + definition = PrimaryKeyIndexDefinition( + "value", 3, "btree", Options.from_none(), PrimaryKeyIndexFamily.BTREE) + files = [ + SimpleNamespace(file_name="a", row_count=2, level=1, file_source=1), + SimpleNamespace(file_name="b", row_count=3, level=2, file_source=1), + ] + partition = GenericRow([], []) + split = DataSplit(files, partition, 0, raw_convertible=True) + source_meta = PrimaryKeyIndexSourceMeta( + 1, [PrimaryKeyIndexSourceFile("a", 2), + PrimaryKeyIndexSourceFile("b", 3)]).serialize() + payload = IndexFileMeta( + "btree", "index", 1, 5, + global_index_meta=GlobalIndexMeta(0, 4, 3, source_meta=source_meta)) + + planned = scan.plan( + 9, [split], [definition], + [IndexManifestEntry(0, partition, 0, payload)]) + + self.assertIn(3, planned.files[0].groups) + self.assertNotIn(3, planned.files[1].groups) + + def test_bucket_state_keeps_invalid_payloads_uncovered(self): + active = SimpleNamespace( + file_name="a", row_count=2, level=1, file_source=1) + + def payload(name, level, sources, source_meta=None): + row_count = sum(source.row_count for source in sources) + serialized = source_meta + if serialized is None: + serialized = PrimaryKeyIndexSourceMeta(level, sources).serialize() + return IndexFileMeta( + "btree", name, 1, row_count, + global_index_meta=GlobalIndexMeta( + 0, row_count - 1, 3, source_meta=serialized)) + + invalid_payloads = [ + payload("wrong-level", 2, [PrimaryKeyIndexSourceFile("a", 2)]), + payload("no-active", 1, [PrimaryKeyIndexSourceFile("b", 2)]), + payload("wrong-row-count", 1, [PrimaryKeyIndexSourceFile("a", 3)]), + payload( + "misordered", 1, + [PrimaryKeyIndexSourceFile("b", 1), + PrimaryKeyIndexSourceFile("a", 2)]), + payload( + "malformed", 1, [PrimaryKeyIndexSourceFile("a", 2)], b"\x00"), + ] + for invalid in invalid_payloads: + with self.subTest(payload=invalid.file_name): + state = PkSortedBucketIndexState.from_active_data_files( + 3, "btree", [active], [invalid]) + self.assertFalse(state.groups) + self.assertEqual( + (PrimaryKeyIndexSourceFile("a", 2),), + state.uncovered_source_files) + self.assertEqual((invalid,), state.rejected_payloads) + + first = payload("first", 1, [PrimaryKeyIndexSourceFile("a", 2)]) + second = payload("second", 1, [PrimaryKeyIndexSourceFile("a", 2)]) + duplicate_state = PkSortedBucketIndexState.from_active_data_files( + 3, "btree", [active], [first, second]) + self.assertFalse(duplicate_state.groups) + self.assertEqual((first, second), duplicate_state.rejected_payloads) + def test_shared_result_is_partitioned_only_once(self): class CountingResult(GlobalIndexResult): def __init__(self):