From 34653761eea0c5a8f1a210217bd3129eee4b86ce Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Mon, 13 Jul 2026 16:38:09 +0800 Subject: [PATCH] [core] Maintain sorted primary-key index payloads --- .../BucketedSortedIndexMaintainer.java | 539 ++++++++++++++++++ .../pksorted/PkSortedBucketIndexState.java | 110 ++++ .../BucketedSortedIndexMaintainerTest.java | 539 ++++++++++++++++++ .../PkSortedBucketIndexStateTest.java | 206 +++++++ 4 files changed, 1394 insertions(+) create mode 100644 paimon-core/src/main/java/org/apache/paimon/index/pksorted/BucketedSortedIndexMaintainer.java create mode 100644 paimon-core/src/main/java/org/apache/paimon/index/pksorted/PkSortedBucketIndexState.java create mode 100644 paimon-core/src/test/java/org/apache/paimon/index/pksorted/BucketedSortedIndexMaintainerTest.java create mode 100644 paimon-core/src/test/java/org/apache/paimon/index/pksorted/PkSortedBucketIndexStateTest.java diff --git a/paimon-core/src/main/java/org/apache/paimon/index/pksorted/BucketedSortedIndexMaintainer.java b/paimon-core/src/main/java/org/apache/paimon/index/pksorted/BucketedSortedIndexMaintainer.java new file mode 100644 index 000000000000..6a887f470d41 --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/index/pksorted/BucketedSortedIndexMaintainer.java @@ -0,0 +1,539 @@ +/* + * 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.paimon.index.pksorted; + +import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceFile; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourcePolicy; +import org.apache.paimon.io.CompactIncrement; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.io.DataIncrement; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.HashSet; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.RejectedExecutionException; + +import static org.apache.paimon.utils.Preconditions.checkArgument; + +/** Maintains one bucket-local source-backed BTree or Bitmap definition. */ +public class BucketedSortedIndexMaintainer { + + private static final Logger LOG = LoggerFactory.getLogger(BucketedSortedIndexMaintainer.class); + private static final int MAX_BUILD_ATTEMPTS = 3; + private static final long INITIAL_RETRY_BACKOFF_MILLIS = 10L; + + private final int fieldId; + private final String indexType; + private final PkSortedIndexFile indexFile; + private final BuildFunction buildFunction; + private final Map activeSourceFiles = new LinkedHashMap<>(); + private final List groups = new ArrayList<>(); + private final List pendingRestoredDeletions = new ArrayList<>(); + private ExecutorService executor; + @Nullable private PendingBuild pendingBuild; + + public BucketedSortedIndexMaintainer( + int fieldId, + String indexType, + PkSortedIndexFile indexFile, + BuildFunction buildFunction, + List restoredDataFiles, + List restoredPayloads, + ExecutorService executor) { + this.fieldId = fieldId; + this.indexType = indexType; + this.indexFile = indexFile; + this.buildFunction = buildFunction; + this.executor = executor; + for (DataFileMeta dataFile : restoredDataFiles) { + if (PrimaryKeyIndexSourcePolicy.shouldRead(dataFile)) { + activeSourceFiles.put(dataFile.fileName(), dataFile); + } + } + + List definitionPayloads = new ArrayList<>(); + for (IndexFileMeta payload : restoredPayloads) { + if (indexType.equals(payload.indexType()) + && payload.globalIndexMeta() != null + && payload.globalIndexMeta().indexFieldId() == fieldId) { + definitionPayloads.add(payload); + } + } + PkSortedBucketIndexState restoredState = + PkSortedBucketIndexState.fromActivePayloads( + fieldId, indexType, sourceFiles(), definitionPayloads); + groups.addAll(restoredState.groups()); + pendingRestoredDeletions.addAll(restoredState.rejectedPayloads()); + } + + public synchronized SortedIndexCommit prepareCommit( + DataIncrement appendIncrement, + CompactIncrement compactIncrement, + boolean waitCompaction) + throws Exception { + return prepareCommit(appendIncrement, compactIncrement, waitCompaction, true); + } + + public synchronized SortedIndexCommit prepareCommit( + DataIncrement appendIncrement, + CompactIncrement compactIncrement, + boolean waitCompaction, + boolean allowBuildStart) + throws Exception { + checkArgument( + eligibleFiles(appendIncrement.newFiles()).isEmpty(), + "Append files must not be primary-key sorted index sources."); + + Map originalSourceFiles = new LinkedHashMap<>(activeSourceFiles); + List originalGroups = new ArrayList<>(groups); + List originalRestoredDeletions = new ArrayList<>(pendingRestoredDeletions); + List created = new ArrayList<>(); + try { + boolean hasCompactDataTransition = + !compactIncrement.compactBefore().isEmpty() + || !compactIncrement.compactAfter().isEmpty(); + applySourceTransition(compactIncrement); + + List removed = new ArrayList<>(pendingRestoredDeletions); + pendingRestoredDeletions.clear(); + removed.addAll(removeInactiveGroups()); + Set failedSources = new HashSet<>(); + while (true) { + Optional completed = + finishPendingBuild(waitCompaction, failedSources); + if (completed.isPresent()) { + acceptOrDelete(completed.get(), created, failedSources); + } + + if (pendingBuild == null && allowBuildStart) { + DataFileMeta uncovered = firstUncoveredSource(failedSources); + if (uncovered != null) { + PendingBuild next = new PendingBuild(uncovered); + try { + next.start(); + pendingBuild = next; + } catch (RejectedExecutionException e) { + failedSources.add(sourceIdentity(uncovered)); + LOG.warn( + "Primary-key {} index build for source file {} was rejected.", + indexType, + uncovered.fileName(), + e); + } + } + } + if (!waitCompaction || pendingBuild == null) { + break; + } + } + + boolean changed = !created.isEmpty() || !removed.isEmpty(); + Optional appendChange = + changed && !hasCompactDataTransition + ? Optional.of(new SortedIndexIncrement(created, removed)) + : Optional.empty(); + Optional compactChange = + changed && hasCompactDataTransition + ? Optional.of(new SortedIndexIncrement(created, removed)) + : Optional.empty(); + return new SortedIndexCommit(appendChange, compactChange); + } catch (Throwable failure) { + rollbackPrepareCommit( + originalSourceFiles, + originalGroups, + originalRestoredDeletions, + created, + failure); + if (failure instanceof Exception) { + throw (Exception) failure; + } + if (failure instanceof Error) { + throw (Error) failure; + } + throw new RuntimeException(failure); + } + } + + private void rollbackPrepareCommit( + Map originalSourceFiles, + List originalGroups, + List originalRestoredDeletions, + List created, + Throwable failure) { + activeSourceFiles.clear(); + activeSourceFiles.putAll(originalSourceFiles); + groups.clear(); + groups.addAll(originalGroups); + pendingRestoredDeletions.clear(); + pendingRestoredDeletions.addAll(originalRestoredDeletions); + + PendingBuild build = pendingBuild; + pendingBuild = null; + if (build != null) { + try { + build.cancel(); + } catch (Throwable cleanupFailure) { + failure.addSuppressed(cleanupFailure); + } + } + for (IndexFileMeta payload : created) { + try { + indexFile.delete(payload); + } catch (Throwable cleanupFailure) { + failure.addSuppressed(cleanupFailure); + } + } + } + + private void applySourceTransition(CompactIncrement compactIncrement) { + for (DataFileMeta before : compactIncrement.compactBefore()) { + if (!containsFile(compactIncrement.compactAfter(), before.fileName())) { + activeSourceFiles.remove(before.fileName()); + } + } + for (DataFileMeta after : compactIncrement.compactAfter()) { + if (PrimaryKeyIndexSourcePolicy.shouldRead(after)) { + activeSourceFiles.put(after.fileName(), after); + } + } + } + + private List removeInactiveGroups() { + List removed = new ArrayList<>(); + Iterator iterator = groups.iterator(); + while (iterator.hasNext()) { + PkSortedIndexGroup group = iterator.next(); + DataFileMeta active = activeSourceFiles.get(group.sourceFile().fileName()); + if (active == null || active.rowCount() != group.sourceFile().rowCount()) { + iterator.remove(); + removed.addAll(group.payloads()); + } + } + return removed; + } + + @Nullable + private DataFileMeta firstUncoveredSource(Set failedSources) { + List candidates = new ArrayList<>(activeSourceFiles.values()); + candidates.sort(Comparator.comparing(DataFileMeta::fileName)); + for (DataFileMeta candidate : candidates) { + if (!isCovered(candidate) && !failedSources.contains(sourceIdentity(candidate))) { + return candidate; + } + } + return null; + } + + private boolean isCovered(DataFileMeta candidate) { + for (PkSortedIndexGroup group : groups) { + if (group.sourceFile().fileName().equals(candidate.fileName()) + && group.sourceFile().rowCount() == candidate.rowCount()) { + return true; + } + } + return false; + } + + private Optional finishPendingBuild(boolean blocking, Set failedSources) + throws InterruptedException { + if (pendingBuild == null || (!blocking && !pendingBuild.isDone())) { + return Optional.empty(); + } + PendingBuild completed = pendingBuild; + try { + List payloads = completed.get(); + pendingBuild = null; + return Optional.of(new CompletedBuild(completed.sourceFile, payloads)); + } catch (CancellationException e) { + pendingBuild = null; + failedSources.add(sourceIdentity(completed.sourceFile)); + return Optional.empty(); + } catch (ExecutionException e) { + pendingBuild = null; + failedSources.add(sourceIdentity(completed.sourceFile)); + LOG.warn( + "Primary-key {} index build for source file {} failed after {} attempts; " + + "the source remains uncovered.", + indexType, + completed.sourceFile.fileName(), + MAX_BUILD_ATTEMPTS, + e.getCause()); + return Optional.empty(); + } + } + + private void acceptOrDelete( + CompletedBuild completed, List created, Set failedSources) { + DataFileMeta active = activeSourceFiles.get(completed.sourceFile.fileName()); + PrimaryKeyIndexSourceFile source = + new PrimaryKeyIndexSourceFile( + completed.sourceFile.fileName(), completed.sourceFile.rowCount()); + Optional group; + try { + group = PkSortedIndexGroup.create(fieldId, indexType, source, completed.payloads); + } catch (RuntimeException e) { + failedSources.add(sourceIdentity(completed.sourceFile)); + deleteGenerated(completed.payloads); + LOG.warn( + "Primary-key {} index build for source file {} produced invalid metadata.", + indexType, + completed.sourceFile.fileName(), + e); + return; + } + if (active == null + || active.rowCount() != completed.sourceFile.rowCount() + || isCovered(active) + || !group.isPresent()) { + deleteGenerated(completed.payloads); + failedSources.add(sourceIdentity(completed.sourceFile)); + return; + } + groups.add(group.get()); + created.addAll(completed.payloads); + } + + private void deleteGenerated(List payloads) { + for (IndexFileMeta payload : payloads) { + try { + indexFile.delete(payload); + } catch (RuntimeException e) { + LOG.warn("Failed to delete unpublished primary-key sorted index payload.", e); + } + } + } + + public synchronized boolean buildNotCompleted() { + return pendingBuild != null; + } + + public int fieldId() { + return fieldId; + } + + public synchronized void withExecutor(ExecutorService executor) { + checkArgument(pendingBuild == null, "Cannot replace executor during a sorted index build."); + this.executor = executor; + } + + public synchronized void close() { + PendingBuild build = pendingBuild; + pendingBuild = null; + if (build != null) { + build.cancel(); + } + } + + public synchronized PkSortedBucketIndexState state() { + return PkSortedBucketIndexState.fromActivePayloads( + fieldId, indexType, sourceFiles(), activePayloads()); + } + + private List sourceFiles() { + List sources = new ArrayList<>(); + for (DataFileMeta dataFile : activeSourceFiles.values()) { + sources.add(new PrimaryKeyIndexSourceFile(dataFile.fileName(), dataFile.rowCount())); + } + return sources; + } + + private List activePayloads() { + List payloads = new ArrayList<>(); + for (PkSortedIndexGroup group : groups) { + payloads.addAll(group.payloads()); + } + return payloads; + } + + private final class PendingBuild { + + private final DataFileMeta sourceFile; + @Nullable private List result; + @Nullable private Future> future; + private boolean cancelled; + + private PendingBuild(DataFileMeta sourceFile) { + this.sourceFile = sourceFile; + } + + private void start() { + future = + executor.submit( + () -> { + List payloads = buildWithRetries(); + synchronized (PendingBuild.this) { + if (!cancelled) { + result = payloads; + return payloads; + } + } + deleteGenerated(payloads); + throw new CancellationException(); + }); + } + + private List buildWithRetries() throws Exception { + for (int attempt = 1; ; attempt++) { + try { + return buildFunction.build(sourceFile); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new CancellationException(); + } catch (Exception e) { + if (attempt >= MAX_BUILD_ATTEMPTS) { + throw e; + } + try { + Thread.sleep(INITIAL_RETRY_BACKOFF_MILLIS << (attempt - 1)); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new CancellationException(); + } + } + } + } + + private boolean isDone() { + return future.isDone(); + } + + private List get() throws InterruptedException, ExecutionException { + return future.get(); + } + + private void cancel() { + Future> buildFuture; + List payloads; + synchronized (this) { + cancelled = true; + buildFuture = future; + payloads = result; + result = null; + } + if (buildFuture != null) { + buildFuture.cancel(true); + } + if (payloads != null) { + deleteGenerated(payloads); + } + } + } + + private static final class CompletedBuild { + + private final DataFileMeta sourceFile; + private final List payloads; + + private CompletedBuild(DataFileMeta sourceFile, List payloads) { + this.sourceFile = sourceFile; + this.payloads = payloads; + } + } + + private static List eligibleFiles(List files) { + List eligible = new ArrayList<>(); + for (DataFileMeta file : files) { + if (PrimaryKeyIndexSourcePolicy.shouldRead(file)) { + eligible.add(file); + } + } + return eligible; + } + + private static boolean containsFile(List files, String fileName) { + for (DataFileMeta file : files) { + if (file.fileName().equals(fileName)) { + return true; + } + } + return false; + } + + private static String sourceIdentity(DataFileMeta file) { + return file.fileName() + '\0' + file.rowCount(); + } + + /** Builds all rotated payloads for one physical source file. */ + @FunctionalInterface + public interface BuildFunction { + + List build(DataFileMeta sourceFile) throws Exception; + } + + /** Sorted-index changes for append and compact snapshot routing. */ + public static final class SortedIndexCommit { + + private final Optional appendIncrement; + private final Optional compactIncrement; + + private SortedIndexCommit( + Optional appendIncrement, + Optional compactIncrement) { + this.appendIncrement = appendIncrement; + this.compactIncrement = compactIncrement; + } + + public Optional appendIncrement() { + return appendIncrement; + } + + public Optional compactIncrement() { + return compactIncrement; + } + } + + /** Index-file additions and deletions emitted by one sorted definition. */ + public static final class SortedIndexIncrement { + + private final List newIndexFiles; + private final List deletedIndexFiles; + + private SortedIndexIncrement( + List newIndexFiles, List deletedIndexFiles) { + this.newIndexFiles = Collections.unmodifiableList(new ArrayList<>(newIndexFiles)); + this.deletedIndexFiles = + Collections.unmodifiableList(new ArrayList<>(deletedIndexFiles)); + } + + public List newIndexFiles() { + return newIndexFiles; + } + + public List deletedIndexFiles() { + return deletedIndexFiles; + } + } +} 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 new file mode 100644 index 000000000000..8fd6729522e7 --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/index/pksorted/PkSortedBucketIndexState.java @@ -0,0 +1,110 @@ +/* + * 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.paimon.index.pksorted; + +import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceFile; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceMeta; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +/** Immutable sorted-index state for one field and bucket. */ +public final class PkSortedBucketIndexState { + + private final List groups; + private final List coveredSourceFiles; + private final List uncoveredSourceFiles; + private final List rejectedPayloads; + + private PkSortedBucketIndexState( + List groups, + List coveredSourceFiles, + List uncoveredSourceFiles, + List rejectedPayloads) { + this.groups = Collections.unmodifiableList(groups); + this.coveredSourceFiles = Collections.unmodifiableList(coveredSourceFiles); + this.uncoveredSourceFiles = Collections.unmodifiableList(uncoveredSourceFiles); + this.rejectedPayloads = Collections.unmodifiableList(rejectedPayloads); + } + + public static PkSortedBucketIndexState fromActivePayloads( + int fieldId, + String indexType, + List activeSourceFiles, + List activePayloads) { + Map> payloadsBySource = new LinkedHashMap<>(); + List rejected = new ArrayList<>(); + for (IndexFileMeta payload : activePayloads) { + try { + PrimaryKeyIndexSourceFile sourceFile = + PrimaryKeyIndexSourceMeta.fromIndexFile(payload).sourceFile(); + payloadsBySource + .computeIfAbsent(sourceFile.fileName(), key -> new ArrayList<>()) + .add(payload); + } catch (RuntimeException ignored) { + rejected.add(payload); + } + } + + List groups = new ArrayList<>(); + List covered = new ArrayList<>(); + List uncovered = new ArrayList<>(); + for (PrimaryKeyIndexSourceFile sourceFile : activeSourceFiles) { + List payloads = payloadsBySource.remove(sourceFile.fileName()); + if (payloads == null || payloads.isEmpty()) { + uncovered.add(sourceFile); + } else { + Optional group = + PkSortedIndexGroup.create(fieldId, indexType, sourceFile, payloads); + if (group.isPresent()) { + groups.add(group.get()); + covered.add(sourceFile); + } else { + uncovered.add(sourceFile); + rejected.addAll(payloads); + } + } + } + for (List inactivePayloads : payloadsBySource.values()) { + rejected.addAll(inactivePayloads); + } + return new PkSortedBucketIndexState(groups, covered, uncovered, rejected); + } + + public List groups() { + return groups; + } + + public List coveredSourceFiles() { + return coveredSourceFiles; + } + + public List uncoveredSourceFiles() { + return uncoveredSourceFiles; + } + + public List rejectedPayloads() { + return rejectedPayloads; + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/index/pksorted/BucketedSortedIndexMaintainerTest.java b/paimon-core/src/test/java/org/apache/paimon/index/pksorted/BucketedSortedIndexMaintainerTest.java new file mode 100644 index 000000000000..167faddce86b --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/index/pksorted/BucketedSortedIndexMaintainerTest.java @@ -0,0 +1,539 @@ +/* + * 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.paimon.index.pksorted; + +import org.apache.paimon.fs.FileIO; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.index.GlobalIndexMeta; +import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.index.IndexPathFactory; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceFile; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceMeta; +import org.apache.paimon.io.CompactIncrement; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.io.DataIncrement; +import org.apache.paimon.manifest.FileSource; +import org.apache.paimon.stats.SimpleStats; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.UUID; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests bucket-local BTree/Bitmap source maintenance. */ +class BucketedSortedIndexMaintainerTest { + + @TempDir java.nio.file.Path tempPath; + private final ExecutorService executor = Executors.newSingleThreadExecutor(); + + @AfterEach + void shutdownExecutor() { + executor.shutdownNow(); + } + + @Test + void testRestoreBuildAndSourceRemoval() throws Exception { + DataFileMeta oldSource = dataFile("data-1", 3); + DataFileMeta newSource = dataFile("data-2", 3); + List oldPayloads = + Arrays.asList(payload("old-1", oldSource, 2), payload("old-2", oldSource, 1)); + List newPayloads = + Arrays.asList(payload("new-1", newSource, 2), payload("new-2", newSource, 1)); + PkSortedIndexFile indexFile = new PkSortedIndexFile(LocalFileIO.create(), pathFactory()); + BucketedSortedIndexMaintainer maintainer = + new BucketedSortedIndexMaintainer( + 7, + "btree", + indexFile, + dataFile -> { + assertThat(dataFile).isEqualTo(newSource); + return newPayloads; + }, + Collections.singletonList(oldSource), + oldPayloads, + executor); + + BucketedSortedIndexMaintainer.SortedIndexCommit commit = + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), + new CompactIncrement( + Collections.singletonList(oldSource), + Collections.singletonList(newSource), + Collections.emptyList()), + true); + + assertThat(commit.compactIncrement()).isPresent(); + assertThat(commit.compactIncrement().get().newIndexFiles()) + .containsExactlyElementsOf(newPayloads); + assertThat(commit.compactIncrement().get().deletedIndexFiles()) + .containsExactlyElementsOf(oldPayloads); + assertThat(maintainer.state().coveredSourceFiles()) + .containsExactly(new PrimaryKeyIndexSourceFile("data-2", 3)); + assertThat(maintainer.state().uncoveredSourceFiles()).isEmpty(); + } + + @Test + void testRestoreDeletesInvalidPayloadsBeforePublishingReplacement() throws Exception { + DataFileMeta source = dataFile("data-1", 3); + List invalidPayloads = + Collections.singletonList(payload("invalid", source, 2)); + List replacementPayloads = + Arrays.asList(payload("new-1", source, 2), payload("new-2", source, 1)); + PkSortedIndexFile indexFile = new PkSortedIndexFile(LocalFileIO.create(), pathFactory()); + BucketedSortedIndexMaintainer maintainer = + new BucketedSortedIndexMaintainer( + 7, + "btree", + indexFile, + dataFile -> replacementPayloads, + Collections.singletonList(source), + invalidPayloads, + executor); + + BucketedSortedIndexMaintainer.SortedIndexCommit repair = + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), CompactIncrement.emptyIncrement(), true); + + assertThat(repair.appendIncrement()).isPresent(); + assertThat(repair.appendIncrement().get().deletedIndexFiles()) + .containsExactlyElementsOf(invalidPayloads); + assertThat(repair.appendIncrement().get().newIndexFiles()) + .containsExactlyElementsOf(replacementPayloads); + + BucketedSortedIndexMaintainer restored = + new BucketedSortedIndexMaintainer( + 7, + "btree", + indexFile, + dataFile -> { + throw new AssertionError("Covered source must not be rebuilt."); + }, + Collections.singletonList(source), + replacementPayloads, + executor); + BucketedSortedIndexMaintainer.SortedIndexCommit stable = + restored.prepareCommit( + DataIncrement.emptyIncrement(), + CompactIncrement.emptyIncrement(), + false, + false); + + assertThat(stable.appendIncrement()).isEmpty(); + assertThat(stable.compactIncrement()).isEmpty(); + assertThat(restored.state().coveredSourceFiles()) + .containsExactly(new PrimaryKeyIndexSourceFile("data-1", 3)); + } + + @Test + void testBuildFailureDoesNotBlockSourceRemoval() throws Exception { + DataFileMeta oldSource = dataFile("data-1", 3); + DataFileMeta newSource = dataFile("data-2", 3); + List oldPayloads = + Arrays.asList(payload("old-1", oldSource, 2), payload("old-2", oldSource, 1)); + AtomicInteger attempts = new AtomicInteger(); + BucketedSortedIndexMaintainer maintainer = + new BucketedSortedIndexMaintainer( + 7, + "btree", + new PkSortedIndexFile(LocalFileIO.create(), pathFactory()), + dataFile -> { + attempts.incrementAndGet(); + throw new IllegalStateException("expected build failure"); + }, + Collections.singletonList(oldSource), + oldPayloads, + executor); + + BucketedSortedIndexMaintainer.SortedIndexCommit commit = + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), + new CompactIncrement( + Collections.singletonList(oldSource), + Collections.singletonList(newSource), + Collections.emptyList()), + true); + + assertThat(attempts).hasValue(3); + assertThat(commit.compactIncrement()).isPresent(); + assertThat(commit.compactIncrement().get().newIndexFiles()).isEmpty(); + assertThat(commit.compactIncrement().get().deletedIndexFiles()) + .containsExactlyElementsOf(oldPayloads); + assertThat(maintainer.state().coveredSourceFiles()).isEmpty(); + assertThat(maintainer.state().uncoveredSourceFiles()) + .containsExactly(new PrimaryKeyIndexSourceFile("data-2", 3)); + } + + @Test + void testTransientFailureRetriesAndPublishesWholeGroup() throws Exception { + DataFileMeta source = dataFile("data-1", 3); + List payloads = + Arrays.asList(payload("new-1", source, 2), payload("new-2", source, 1)); + AtomicInteger attempts = new AtomicInteger(); + BucketedSortedIndexMaintainer maintainer = + new BucketedSortedIndexMaintainer( + 7, + "btree", + new PkSortedIndexFile(LocalFileIO.create(), pathFactory()), + dataFile -> { + if (attempts.incrementAndGet() < 3) { + throw new IllegalStateException("expected transient failure"); + } + return payloads; + }, + Collections.emptyList(), + Collections.emptyList(), + executor); + + BucketedSortedIndexMaintainer.SortedIndexCommit commit = + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), compactAfter(source), true); + + assertThat(attempts).hasValue(3); + assertThat(commit.compactIncrement()).isPresent(); + assertThat(commit.compactIncrement().get().newIndexFiles()) + .containsExactlyElementsOf(payloads); + assertThat(maintainer.state().coveredSourceFiles()) + .containsExactly(new PrimaryKeyIndexSourceFile("data-1", 3)); + } + + @Test + void testNonBlockingBuildPublishesOnLaterCommit() throws Exception { + DataFileMeta source = dataFile("data-1", 3); + List payloads = + Arrays.asList(payload("new-1", source, 2), payload("new-2", source, 1)); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CountDownLatch completed = new CountDownLatch(1); + BucketedSortedIndexMaintainer maintainer = + new BucketedSortedIndexMaintainer( + 7, + "btree", + new PkSortedIndexFile(LocalFileIO.create(), pathFactory()), + dataFile -> { + started.countDown(); + release.await(); + completed.countDown(); + return payloads; + }, + Collections.emptyList(), + Collections.emptyList(), + executor); + + BucketedSortedIndexMaintainer.SortedIndexCommit first = + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), compactAfter(source), false); + assertThat(started.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(first.compactIncrement()).isEmpty(); + assertThat(maintainer.buildNotCompleted()).isTrue(); + + release.countDown(); + assertThat(completed.await(5, TimeUnit.SECONDS)).isTrue(); + BucketedSortedIndexMaintainer.SortedIndexCommit second = + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), CompactIncrement.emptyIncrement(), false); + + assertThat(second.appendIncrement()).isPresent(); + assertThat(second.appendIncrement().get().newIndexFiles()) + .containsExactlyElementsOf(payloads); + assertThat(maintainer.buildNotCompleted()).isFalse(); + } + + @Test + void testStaleCompletionIsDeletedBeforeReplacementPublishes() throws Exception { + DataFileMeta staleSource = dataFile("data-1", 3); + DataFileMeta activeSource = dataFile("data-2", 3); + List stalePayloads = + Arrays.asList( + payload("stale-1", staleSource, 2), payload("stale-2", staleSource, 1)); + List activePayloads = + Arrays.asList( + payload("active-1", activeSource, 2), payload("active-2", activeSource, 1)); + CountDownLatch staleStarted = new CountDownLatch(1); + CountDownLatch releaseStale = new CountDownLatch(1); + CountDownLatch staleCompleted = new CountDownLatch(1); + TrackingPkSortedIndexFile indexFile = + new TrackingPkSortedIndexFile(LocalFileIO.create(), pathFactory()); + BucketedSortedIndexMaintainer maintainer = + new BucketedSortedIndexMaintainer( + 7, + "btree", + indexFile, + dataFile -> { + if (dataFile.fileName().equals(staleSource.fileName())) { + staleStarted.countDown(); + releaseStale.await(); + staleCompleted.countDown(); + return stalePayloads; + } + assertThat(dataFile).isEqualTo(activeSource); + return activePayloads; + }, + Collections.emptyList(), + Collections.emptyList(), + executor); + + maintainer.prepareCommit(DataIncrement.emptyIncrement(), compactAfter(staleSource), false); + assertThat(staleStarted.await(5, TimeUnit.SECONDS)).isTrue(); + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), + new CompactIncrement( + Collections.singletonList(staleSource), + Collections.singletonList(activeSource), + Collections.emptyList()), + false); + releaseStale.countDown(); + assertThat(staleCompleted.await(5, TimeUnit.SECONDS)).isTrue(); + + BucketedSortedIndexMaintainer.SortedIndexCommit commit = + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), CompactIncrement.emptyIncrement(), true); + + assertThat(indexFile.deleted()).containsExactlyElementsOf(stalePayloads); + assertThat(commit.appendIncrement()).isPresent(); + assertThat(commit.appendIncrement().get().newIndexFiles()) + .containsExactlyElementsOf(activePayloads); + assertThat(maintainer.state().coveredSourceFiles()) + .containsExactly(new PrimaryKeyIndexSourceFile("data-2", 3)); + } + + @Test + void testRejectedSubmissionLeavesSourceUncovered() throws Exception { + DataFileMeta source = dataFile("data-1", 3); + ExecutorService rejectedExecutor = Executors.newSingleThreadExecutor(); + rejectedExecutor.shutdownNow(); + BucketedSortedIndexMaintainer maintainer = + new BucketedSortedIndexMaintainer( + 7, + "btree", + new PkSortedIndexFile(LocalFileIO.create(), pathFactory()), + dataFile -> Collections.singletonList(payload("new", source, 3)), + Collections.emptyList(), + Collections.emptyList(), + rejectedExecutor); + + BucketedSortedIndexMaintainer.SortedIndexCommit commit = + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), compactAfter(source), false); + + assertThat(commit.appendIncrement()).isEmpty(); + assertThat(commit.compactIncrement()).isEmpty(); + assertThat(maintainer.buildNotCompleted()).isFalse(); + assertThat(maintainer.state().uncoveredSourceFiles()) + .containsExactly(new PrimaryKeyIndexSourceFile("data-1", 3)); + } + + @Test + void testCloseDeletesResultCompletedAfterCancellation() throws Exception { + DataFileMeta source = dataFile("data-1", 3); + IndexFileMeta generated = payload("cancelled", source, 3); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + TrackingPkSortedIndexFile indexFile = + new TrackingPkSortedIndexFile(LocalFileIO.create(), pathFactory()); + BucketedSortedIndexMaintainer maintainer = + new BucketedSortedIndexMaintainer( + 7, + "btree", + indexFile, + dataFile -> { + started.countDown(); + boolean released = false; + while (!released) { + try { + release.await(); + released = true; + } catch (InterruptedException ignored) { + // Simulate a builder completing despite cancellation. + } + } + return Collections.singletonList(generated); + }, + Collections.emptyList(), + Collections.emptyList(), + executor); + + maintainer.prepareCommit(DataIncrement.emptyIncrement(), compactAfter(source), false); + assertThat(started.await(5, TimeUnit.SECONDS)).isTrue(); + maintainer.close(); + release.countDown(); + + assertThat(indexFile.awaitDeletion()).isTrue(); + assertThat(indexFile.deleted()).containsExactly(generated); + assertThat(maintainer.buildNotCompleted()).isFalse(); + } + + @Test + void testInterruptedCommitRollsBackStateAndCancelsBuild() throws Exception { + DataFileMeta oldSource = dataFile("data-1", 3); + DataFileMeta newSource = dataFile("data-2", 3); + List oldPayloads = + Arrays.asList(payload("old-1", oldSource, 2), payload("old-2", oldSource, 1)); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch cancelled = new CountDownLatch(1); + BucketedSortedIndexMaintainer maintainer = + new BucketedSortedIndexMaintainer( + 7, + "btree", + new PkSortedIndexFile(LocalFileIO.create(), pathFactory()), + dataFile -> { + started.countDown(); + try { + new CountDownLatch(1).await(); + throw new AssertionError("Build must be cancelled."); + } catch (InterruptedException e) { + cancelled.countDown(); + throw e; + } + }, + Collections.singletonList(oldSource), + oldPayloads, + executor); + CompactIncrement transition = + new CompactIncrement( + Collections.singletonList(oldSource), + Collections.singletonList(newSource), + Collections.emptyList()); + AtomicReference failure = new AtomicReference<>(); + Thread commitThread = + new Thread( + () -> { + try { + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), transition, true); + } catch (Throwable t) { + failure.set(t); + } + }); + + commitThread.start(); + assertThat(started.await(5, TimeUnit.SECONDS)).isTrue(); + commitThread.interrupt(); + commitThread.join(TimeUnit.SECONDS.toMillis(5)); + + assertThat(commitThread.isAlive()).isFalse(); + assertThat(failure.get()).isInstanceOf(InterruptedException.class); + assertThat(cancelled.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(maintainer.buildNotCompleted()).isFalse(); + assertThat(maintainer.state().coveredSourceFiles()) + .containsExactly(new PrimaryKeyIndexSourceFile("data-1", 3)); + assertThat(maintainer.state().uncoveredSourceFiles()).isEmpty(); + } + + private static CompactIncrement compactAfter(DataFileMeta source) { + return new CompactIncrement( + Collections.emptyList(), + Collections.singletonList(source), + Collections.emptyList()); + } + + private static IndexFileMeta payload( + String fileName, DataFileMeta sourceFile, long payloadRowCount) { + PrimaryKeyIndexSourceFile source = + new PrimaryKeyIndexSourceFile(sourceFile.fileName(), sourceFile.rowCount()); + return new IndexFileMeta( + "btree", + fileName, + 1, + payloadRowCount, + new GlobalIndexMeta( + 0, + source.rowCount() - 1, + 7, + null, + new byte[] {1}, + new PrimaryKeyIndexSourceMeta(source).serialize()), + null); + } + + private static DataFileMeta dataFile(String fileName, long rowCount) { + return DataFileMeta.forAppend( + fileName, + 100, + rowCount, + SimpleStats.EMPTY_STATS, + 0, + 1, + 1, + Collections.emptyList(), + null, + FileSource.COMPACT, + null, + null, + null, + null) + .upgrade(1); + } + + private IndexPathFactory pathFactory() { + Path directory = new Path(tempPath.toUri()); + return new IndexPathFactory() { + @Override + public Path toPath(String fileName) { + return new Path(directory, fileName); + } + + @Override + public Path newPath() { + return new Path(directory, UUID.randomUUID().toString()); + } + + @Override + public boolean isExternalPath() { + return false; + } + }; + } + + private static class TrackingPkSortedIndexFile extends PkSortedIndexFile { + + private final List deleted = new java.util.ArrayList<>(); + private final CountDownLatch deletion = new CountDownLatch(1); + + private TrackingPkSortedIndexFile(FileIO fileIO, IndexPathFactory pathFactory) { + super(fileIO, pathFactory); + } + + @Override + public synchronized void delete(IndexFileMeta file) { + deleted.add(file); + deletion.countDown(); + } + + private synchronized List deleted() { + return new java.util.ArrayList<>(deleted); + } + + private boolean awaitDeletion() throws InterruptedException { + return deletion.await(1, TimeUnit.SECONDS); + } + } +} 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 new file mode 100644 index 000000000000..4235bf745606 --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/index/pksorted/PkSortedBucketIndexStateTest.java @@ -0,0 +1,206 @@ +/* + * 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.paimon.index.pksorted; + +import org.apache.paimon.index.GlobalIndexMeta; +import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceFile; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceMeta; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link PkSortedBucketIndexState}. */ +class PkSortedBucketIndexStateTest { + + @Test + void testRotatedPayloadsFormOneCoveredGroup() { + PrimaryKeyIndexSourceFile source = new PrimaryKeyIndexSourceFile("data-1", 10); + PkSortedBucketIndexState state = + PkSortedBucketIndexState.fromActivePayloads( + 7, + "btree", + Collections.singletonList(source), + Arrays.asList( + payload("index-1", source, "btree", 7, 0, 9, 4), + payload("index-2", source, "btree", 7, 0, 9, 6))); + + assertThat(state.groups()).hasSize(1); + assertThat(state.groups().get(0).payloads()) + .extracting(IndexFileMeta::fileName) + .containsExactly("index-1", "index-2"); + assertThat(state.coveredSourceFiles()).containsExactly(source); + assertThat(state.uncoveredSourceFiles()).isEmpty(); + assertThat(state.rejectedPayloads()).isEmpty(); + } + + @Test + void testIncompletePayloadRowCountLeavesSourceUncovered() { + PrimaryKeyIndexSourceFile source = new PrimaryKeyIndexSourceFile("data-1", 10); + PkSortedBucketIndexState state = + PkSortedBucketIndexState.fromActivePayloads( + 7, + "btree", + Collections.singletonList(source), + Arrays.asList( + payload("index-1", source, "btree", 7, 0, 9, 4), + payload("index-2", source, "btree", 7, 0, 9, 5))); + + assertThat(state.groups()).isEmpty(); + assertThat(state.coveredSourceFiles()).isEmpty(); + assertThat(state.uncoveredSourceFiles()).containsExactly(source); + assertThat(state.rejectedPayloads()) + .extracting(IndexFileMeta::fileName) + .containsExactly("index-1", "index-2"); + } + + @Test + void testWrongPayloadRangeLeavesSourceUncovered() { + PrimaryKeyIndexSourceFile source = new PrimaryKeyIndexSourceFile("data-1", 10); + PkSortedBucketIndexState state = + PkSortedBucketIndexState.fromActivePayloads( + 7, + "bitmap", + Collections.singletonList(source), + Arrays.asList( + payload("index-1", source, "bitmap", 7, 0, 9, 4), + payload("index-2", source, "bitmap", 7, 0, 8, 6))); + + assertThat(state.groups()).isEmpty(); + assertThat(state.coveredSourceFiles()).isEmpty(); + assertThat(state.uncoveredSourceFiles()).containsExactly(source); + } + + @Test + void testMixedIndexTypeLeavesSourceUncovered() { + PrimaryKeyIndexSourceFile source = new PrimaryKeyIndexSourceFile("data-1", 10); + PkSortedBucketIndexState state = + PkSortedBucketIndexState.fromActivePayloads( + 7, + "btree", + Collections.singletonList(source), + Arrays.asList( + payload("index-1", source, "btree", 7, 0, 9, 4), + payload("index-2", source, "bitmap", 7, 0, 9, 6))); + + assertThat(state.groups()).isEmpty(); + assertThat(state.coveredSourceFiles()).isEmpty(); + assertThat(state.uncoveredSourceFiles()).containsExactly(source); + } + + @Test + void testMixedFieldLeavesSourceUncovered() { + PrimaryKeyIndexSourceFile source = new PrimaryKeyIndexSourceFile("data-1", 10); + PkSortedBucketIndexState state = + PkSortedBucketIndexState.fromActivePayloads( + 7, + "btree", + Collections.singletonList(source), + Arrays.asList( + payload("index-1", source, "btree", 7, 0, 9, 4), + payload("index-2", source, "btree", 8, 0, 9, 6))); + + assertThat(state.groups()).isEmpty(); + assertThat(state.coveredSourceFiles()).isEmpty(); + assertThat(state.uncoveredSourceFiles()).containsExactly(source); + } + + @Test + void testMismatchedSourceMetadataLeavesSourceUncovered() { + PrimaryKeyIndexSourceFile source = new PrimaryKeyIndexSourceFile("data-1", 10); + PrimaryKeyIndexSourceFile mismatchedSource = new PrimaryKeyIndexSourceFile("data-1", 11); + PkSortedBucketIndexState state = + PkSortedBucketIndexState.fromActivePayloads( + 7, + "btree", + Collections.singletonList(source), + Arrays.asList( + payload("index-1", source, "btree", 7, 0, 9, 4), + payload("index-2", mismatchedSource, "btree", 7, 0, 9, 6))); + + assertThat(state.groups()).isEmpty(); + assertThat(state.coveredSourceFiles()).isEmpty(); + assertThat(state.uncoveredSourceFiles()).containsExactly(source); + } + + @Test + void testDuplicatePayloadNameLeavesSourceUncovered() { + PrimaryKeyIndexSourceFile source = new PrimaryKeyIndexSourceFile("data-1", 10); + PkSortedBucketIndexState state = + PkSortedBucketIndexState.fromActivePayloads( + 7, + "btree", + Collections.singletonList(source), + Arrays.asList( + payload("index-1", source, "btree", 7, 0, 9, 5), + payload("index-1", source, "btree", 7, 0, 9, 5))); + + assertThat(state.groups()).isEmpty(); + assertThat(state.coveredSourceFiles()).isEmpty(); + assertThat(state.uncoveredSourceFiles()).containsExactly(source); + } + + @Test + void testMalformedSourceMetadataLeavesSourceUncovered() { + PrimaryKeyIndexSourceFile source = new PrimaryKeyIndexSourceFile("data-1", 10); + IndexFileMeta malformed = + new IndexFileMeta( + "btree", + "index-1", + 100, + 10, + new GlobalIndexMeta(0, 9, 7, null, new byte[] {1}, new byte[] {1}), + null); + + PkSortedBucketIndexState state = + PkSortedBucketIndexState.fromActivePayloads( + 7, + "btree", + Collections.singletonList(source), + Collections.singletonList(malformed)); + + assertThat(state.groups()).isEmpty(); + assertThat(state.coveredSourceFiles()).isEmpty(); + assertThat(state.uncoveredSourceFiles()).containsExactly(source); + assertThat(state.rejectedPayloads()).containsExactly(malformed); + } + + private static IndexFileMeta payload( + String fileName, + PrimaryKeyIndexSourceFile source, + String indexType, + int fieldId, + long rangeStart, + long rangeEnd, + long rowCount) { + byte[] sourceMeta = new PrimaryKeyIndexSourceMeta(source).serialize(); + return new IndexFileMeta( + indexType, + fileName, + 100, + rowCount, + new GlobalIndexMeta( + rangeStart, rangeEnd, fieldId, null, new byte[] {1}, sourceMeta), + null); + } +}