This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new df3196146e [core][python] Reuse PK sorted indexes with retired source
files (#9460)
df3196146e is described below
commit df3196146e709c4899fade2492eb92e5f3edf8f5
Author: wangyong9999 <[email protected]>
AuthorDate: Sun Aug 30 10:11:06 2026 +0800
[core][python] Reuse PK sorted indexes with retired source files (#9460)
---
.../index/pksorted/PkSortedBucketIndexState.java | 111 ++++++++++++++++-----
.../paimon/index/pksorted/PkSortedIndexGroup.java | 5 +-
.../table/source/PrimaryKeySortedIndexScan.java | 17 +++-
.../pksorted/PkSortedBucketIndexStateTest.java | 95 +++++++++++++++++-
.../source/PrimaryKeySortedIndexScanTest.java | 99 +++++++++++++++++-
.../index/pk/primary_key_index_source_policy.py | 21 ++++
.../index/pksorted/pk_sorted_bucket_index_state.py | 63 +++++++++---
.../table/source/primary_key_full_text_scan.py | 7 +-
.../table/source/primary_key_sorted_index_scan.py | 14 ++-
.../table/source/primary_key_vector_scan.py | 7 +-
.../tests/primary_key_sorted_index_scan_test.py | 105 +++++++++++++++++++
11 files changed, 481 insertions(+), 63 deletions(-)
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 fa6a5dd10b..1933a28778 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 @@ import java.util.TreeMap;
/** 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<PrimaryKeyIndexSourceFile> activeSources;
+
+ private PayloadCandidate(
+ IndexFileMeta payload,
+ PkSortedIndexGroup group,
+ List<PrimaryKeyIndexSourceFile> activeSources) {
+ this.payload = payload;
+ this.group = group;
+ this.activeSources = activeSources;
+ }
+ }
+
private final List<PkSortedIndexGroup> groups;
private final List<PrimaryKeyIndexSourceFile> coveredSourceFiles;
private final List<PrimaryKeyIndexSourceFile> uncoveredSourceFiles;
@@ -72,55 +88,102 @@ public final class PkSortedBucketIndexState {
sources.sort(Comparator.comparing(PrimaryKeyIndexSourceFile::fileName));
}
- Map<Integer, List<IndexFileMeta>> payloadsByLevel = new TreeMap<>();
+ Map<Integer, List<PayloadCandidate>> candidatesByLevel = new
TreeMap<>();
List<IndexFileMeta> rejected = new ArrayList<>();
for (IndexFileMeta payload : activePayloads) {
try {
PrimaryKeyIndexSourceMeta sourceMeta =
PrimaryKeyIndexSourceMeta.fromIndexFile(payload);
- List<PrimaryKeyIndexSourceFile> desired =
+ List<PrimaryKeyIndexSourceFile> activeLevelSources =
sourcesByLevel.get(sourceMeta.dataLevel());
- if (desired == null ||
!desired.equals(sourceMeta.sourceFiles())) {
+ if (activeLevelSources == null) {
+ rejected.add(payload);
+ continue;
+ }
+
+ List<PrimaryKeyIndexSourceFile> payloadSources =
sourceMeta.sourceFiles();
+ List<PrimaryKeyIndexSourceFile> activeIntersection =
+ activeIntersection(activeLevelSources, payloadSources);
+ if (activeIntersection == null ||
activeIntersection.isEmpty()) {
+ rejected.add(payload);
+ continue;
+ }
+
+ Optional<PkSortedIndexGroup> 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<PkSortedIndexGroup> groups = new ArrayList<>();
- Set<Integer> coveredLevels = new HashSet<>();
- for (Map.Entry<Integer, List<IndexFileMeta>> entry :
payloadsByLevel.entrySet()) {
- List<IndexFileMeta> levelPayloads = entry.getValue();
- Optional<PkSortedIndexGroup> 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<Integer, Set<PrimaryKeyIndexSourceFile>> coveredSourcesByLevel =
new TreeMap<>();
+ for (Map.Entry<Integer, List<PayloadCandidate>> entry :
candidatesByLevel.entrySet()) {
+ List<PayloadCandidate> 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<PrimaryKeyIndexSourceFile> covered = new ArrayList<>();
List<PrimaryKeyIndexSourceFile> uncovered = new ArrayList<>();
for (Map.Entry<Integer, List<PrimaryKeyIndexSourceFile>> entry :
sourcesByLevel.entrySet()) {
- (coveredLevels.contains(entry.getKey()) ? covered :
uncovered).addAll(entry.getValue());
+ Set<PrimaryKeyIndexSourceFile> 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<PrimaryKeyIndexSourceFile> activeIntersection(
+ List<PrimaryKeyIndexSourceFile> activeSources,
+ List<PrimaryKeyIndexSourceFile> payloadSources) {
+ List<PrimaryKeyIndexSourceFile> 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<PkSortedIndexGroup> 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 54999c2162..f30047a681 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.List;
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 2bda216ce0..5003c9dbfd 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.IndexFileMeta;
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 @@ public final class PrimaryKeySortedIndexScan {
Pair<BinaryRow, Integer> bucket = bucketEntry.getKey();
List<IndexFileMeta> bucketPayloads =
payloadsByBucket.getOrDefault(bucket,
Collections.emptyList());
- Set<PrimaryKeyIndexSourceFile> activeSourceFiles = new HashSet<>();
+ Map<Integer, Set<PrimaryKeyIndexSourceFile>>
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<String, Map<Integer, PkSortedIndexGroup>> groupsBySource = new
LinkedHashMap<>();
for (PrimaryKeyIndexDefinition definition : scalarDefinitions) {
@@ -193,8 +199,11 @@ public final class PrimaryKeySortedIndexScan {
bucketEntry.getValue(),
definitionPayloads);
for (PkSortedIndexGroup group : state.groups()) {
+ Set<PrimaryKeyIndexSourceFile> 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 f7a9a60d1d..e666ee5ea0 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 @@ class PkSortedBucketIndexStateTest {
}
@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 @@ class PkSortedBucketIndexStateTest {
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 @@ class PkSortedBucketIndexStateTest {
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<PrimaryKeyIndexSourceFile> sources) {
long rowCount = 0;
for (PrimaryKeyIndexSourceFile source : sources) {
rowCount += source.rowCount();
@@ -177,4 +262,8 @@ class PkSortedBucketIndexStateTest {
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 c9c09ea7d4..25edb47eca 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 @@ class PrimaryKeySortedIndexScanTest {
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 @@ class PrimaryKeySortedIndexScanTest {
}
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 @@ class PrimaryKeySortedIndexScanTest {
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 0000000000..66ffa1e4ab
--- /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 490957149e..c541c9d798 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 dataclasses import dataclass
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 afa3a02dfc..dab5f0e66e 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.globalindex.indexed_split import IndexedSplit
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 5a2c91ad6a..865a52eebf 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.global_index_result import
GlobalIndexResult
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 ccd2535783..99156528f0 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.options import Options
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 791e319120..1eab053efd 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 @@ from pypaimon.index.pk.primary_key_index_definition import (
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 @@ class PrimaryKeySortedIndexScanTest(unittest.TestCase):
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):
