diff --git a/paimon-core/src/main/java/org/apache/paimon/KeyValueFileStore.java b/paimon-core/src/main/java/org/apache/paimon/KeyValueFileStore.java index 8e0d4da7e275..7531f91a77ff 100644 --- a/paimon-core/src/main/java/org/apache/paimon/KeyValueFileStore.java +++ b/paimon-core/src/main/java/org/apache/paimon/KeyValueFileStore.java @@ -25,7 +25,6 @@ import org.apache.paimon.fs.FileIO; import org.apache.paimon.index.DynamicBucketIndexMaintainer; import org.apache.paimon.index.pk.BucketedPrimaryKeyIndexMaintainer; -import org.apache.paimon.index.pkvector.BucketedVectorIndexMaintainer; import org.apache.paimon.io.KeyValueFileReaderFactory; import org.apache.paimon.mergetree.compact.MergeFunctionFactory; import org.apache.paimon.operation.AbstractFileStoreWrite; @@ -40,7 +39,6 @@ import org.apache.paimon.schema.TableSchema; import org.apache.paimon.table.BucketMode; import org.apache.paimon.table.CatalogEnvironment; -import org.apache.paimon.types.DataField; import org.apache.paimon.types.RowType; import org.apache.paimon.utils.KeyComparatorSupplier; import org.apache.paimon.utils.UserDefinedSeqComparator; @@ -173,19 +171,12 @@ public AbstractFileStoreWrite newWrite(String commitUser, @Nullable In dvMaintainerFactory = BucketedDvMaintainer.factory(newIndexFileHandler()); } BucketedPrimaryKeyIndexMaintainer.Factory primaryKeyIndexMaintainerFactory = null; - if (options.primaryKeyVectorIndexEnabled()) { - String vectorColumn = options.primaryKeyVectorIndexColumn(); - DataField vectorField = schema.nameToFieldMap().get(vectorColumn); + if (options.primaryKeyVectorIndexEnabled() + || !options.primaryKeyBTreeIndexColumns().isEmpty() + || !options.primaryKeyBitmapIndexColumns().isEmpty()) { primaryKeyIndexMaintainerFactory = - BucketedPrimaryKeyIndexMaintainer.Factory.ofVector( - new BucketedVectorIndexMaintainer.Factory( - newIndexFileHandler(), - newReaderFactoryBuilder(), - vectorField, - options.primaryKeyVectorDistanceMetric(vectorColumn), - options.primaryKeyVectorIndexType(vectorColumn)) - .withIndexOptions( - options.primaryKeyVectorIndexOptions(vectorColumn))); + BucketedPrimaryKeyIndexMaintainer.Factory.create( + newIndexFileHandler(), newReaderFactoryBuilder(), schema); } return new KeyValueFileStoreWrite( fileIO, diff --git a/paimon-core/src/main/java/org/apache/paimon/globalindex/GlobalIndexCoverage.java b/paimon-core/src/main/java/org/apache/paimon/globalindex/GlobalIndexCoverage.java index 5191d0f942a3..6353b6616e2f 100644 --- a/paimon-core/src/main/java/org/apache/paimon/globalindex/GlobalIndexCoverage.java +++ b/paimon-core/src/main/java/org/apache/paimon/globalindex/GlobalIndexCoverage.java @@ -64,6 +64,9 @@ public GlobalIndexCoverage( this.coverageByField = new HashMap<>(); for (IndexFileMeta indexFile : indexFiles) { GlobalIndexMeta meta = checkNotNull(indexFile.globalIndexMeta()); + if (meta.sourceMeta() != null) { + continue; + } Range range = new Range(meta.rowRangeStart(), meta.rowRangeEnd()); addCoverage(meta.indexFieldId(), range); if (meta.extraFieldIds() != null) { diff --git a/paimon-core/src/main/java/org/apache/paimon/globalindex/GlobalIndexScanner.java b/paimon-core/src/main/java/org/apache/paimon/globalindex/GlobalIndexScanner.java index 7edb68365b6b..1385970b2461 100644 --- a/paimon-core/src/main/java/org/apache/paimon/globalindex/GlobalIndexScanner.java +++ b/paimon-core/src/main/java/org/apache/paimon/globalindex/GlobalIndexScanner.java @@ -187,7 +187,8 @@ public static Optional create( FileStoreTable table, @Nullable PartitionPredicate partitionFilter, Collection indexFiles) { - if (indexFiles.isEmpty()) { + List ordinaryIndexFiles = ordinaryGlobalIndexFiles(indexFiles); + if (ordinaryIndexFiles.isEmpty()) { return Optional.empty(); } return Optional.of( @@ -199,7 +200,7 @@ public static Optional create( table.rowType(), table.fileIO(), table.store().pathFactory().globalIndexFileFactory(), - indexFiles)); + ordinaryIndexFiles)); } public static Optional create( @@ -241,7 +242,7 @@ private static Filter indexFileFilter( return false; } GlobalIndexMeta globalIndex = entry.indexFile().globalIndexMeta(); - if (globalIndex == null) { + if (globalIndex == null || globalIndex.sourceMeta() != null) { return false; } // Collect indexes whose primary column is filtered, and also multi-column @@ -261,6 +262,17 @@ private static Filter indexFileFilter( return indexFileFilter; } + private static List ordinaryGlobalIndexFiles( + Collection indexFiles) { + return indexFiles.stream() + .filter( + indexFile -> { + GlobalIndexMeta meta = indexFile.globalIndexMeta(); + return meta != null && meta.sourceMeta() == null; + }) + .collect(Collectors.toList()); + } + public Optional scan(Predicate predicate) { return globalIndexEvaluator.evaluate(predicate); } diff --git a/paimon-core/src/main/java/org/apache/paimon/index/IndexFileHandler.java b/paimon-core/src/main/java/org/apache/paimon/index/IndexFileHandler.java index 799aa3ef66de..9fafedb7aa4e 100644 --- a/paimon-core/src/main/java/org/apache/paimon/index/IndexFileHandler.java +++ b/paimon-core/src/main/java/org/apache/paimon/index/IndexFileHandler.java @@ -25,6 +25,7 @@ import org.apache.paimon.deletionvectors.DeletionVectorsIndexFile; import org.apache.paimon.fs.FileIO; import org.apache.paimon.fs.Path; +import org.apache.paimon.index.pksorted.PkSortedIndexFile; import org.apache.paimon.index.pkvector.PkVectorAnnSegmentFile; import org.apache.paimon.manifest.IndexManifestEntry; import org.apache.paimon.manifest.IndexManifestEntrySerializer; @@ -89,6 +90,10 @@ public PkVectorAnnSegmentFile pkVectorAnnSegment(BinaryRow partition, int bucket return new PkVectorAnnSegmentFile(fileIO, pathFactories.get(partition, bucket)); } + public PkSortedIndexFile pkSortedIndex(BinaryRow partition, int bucket) { + return new PkSortedIndexFile(fileIO, pathFactories.get(partition, bucket)); + } + public Optional scanHashIndex( Snapshot snapshot, BinaryRow partition, int bucket) { List result = scan(snapshot, HASH_INDEX, partition, bucket); diff --git a/paimon-core/src/main/java/org/apache/paimon/index/pk/BucketedPrimaryKeyIndexMaintainer.java b/paimon-core/src/main/java/org/apache/paimon/index/pk/BucketedPrimaryKeyIndexMaintainer.java index 702a1dc9cffb..5a8403c4106b 100644 --- a/paimon-core/src/main/java/org/apache/paimon/index/pk/BucketedPrimaryKeyIndexMaintainer.java +++ b/paimon-core/src/main/java/org/apache/paimon/index/pk/BucketedPrimaryKeyIndexMaintainer.java @@ -18,44 +18,169 @@ package org.apache.paimon.index.pk; +import org.apache.paimon.CoreOptions; import org.apache.paimon.data.BinaryRow; import org.apache.paimon.disk.IOManager; import org.apache.paimon.index.IndexFileHandler; import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.index.pksorted.BucketedSortedIndexMaintainer; +import org.apache.paimon.index.pksorted.PkSortedDataFileReader; +import org.apache.paimon.index.pksorted.PkSortedIndexBuilder; +import org.apache.paimon.index.pksorted.PkSortedIndexFile; import org.apache.paimon.index.pkvector.BucketedVectorIndexMaintainer; import org.apache.paimon.io.CompactIncrement; import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.io.DataIncrement; +import org.apache.paimon.io.KeyValueFileReaderFactory; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.types.DataField; import javax.annotation.Nullable; +import java.util.ArrayList; import java.util.Collections; +import java.util.Comparator; import java.util.List; +import java.util.Map; import java.util.concurrent.ExecutorService; -/** Coordinates source-backed primary-key indexes for one bucket. */ +import static org.apache.paimon.utils.Preconditions.checkArgument; + +/** Coordinates all source-backed primary-key indexes for one bucket. */ public final class BucketedPrimaryKeyIndexMaintainer { - private final BucketedVectorIndexMaintainer vectorMaintainer; + @Nullable private final BucketedVectorIndexMaintainer vectorMaintainer; + private final List sortedMaintainers; - private BucketedPrimaryKeyIndexMaintainer(BucketedVectorIndexMaintainer vectorMaintainer) { + private BucketedPrimaryKeyIndexMaintainer( + @Nullable BucketedVectorIndexMaintainer vectorMaintainer, + List sortedMaintainers) { this.vectorMaintainer = vectorMaintainer; + List sorted = new ArrayList<>(sortedMaintainers); + sorted.sort(Comparator.comparingInt(BucketedSortedIndexMaintainer::fieldId)); + this.sortedMaintainers = Collections.unmodifiableList(sorted); } public static BucketedPrimaryKeyIndexMaintainer ofVector( BucketedVectorIndexMaintainer vectorMaintainer) { - return new BucketedPrimaryKeyIndexMaintainer(vectorMaintainer); + return new BucketedPrimaryKeyIndexMaintainer(vectorMaintainer, Collections.emptyList()); + } + + public static BucketedPrimaryKeyIndexMaintainer of( + BucketedVectorIndexMaintainer vectorMaintainer, + List sortedMaintainers) { + return new BucketedPrimaryKeyIndexMaintainer(vectorMaintainer, sortedMaintainers); + } + + public static BucketedPrimaryKeyIndexMaintainer ofSorted( + List sortedMaintainers) { + return new BucketedPrimaryKeyIndexMaintainer(null, sortedMaintainers); } - public void prepareCommit( + public synchronized void prepareCommit( DataIncrement appendIncrement, CompactIncrement compactIncrement, boolean waitCompaction) throws Exception { - BucketedVectorIndexMaintainer.VectorIndexCommit vectorCommit = - vectorMaintainer.prepareCommit(appendIncrement, compactIncrement, waitCompaction); - vectorCommit - .appendIncrement() + BucketedVectorIndexMaintainer.VectorIndexCommit vectorCommit = null; + List sortedCommits = new ArrayList<>(); + try { + if (vectorMaintainer != null) { + vectorCommit = + vectorMaintainer.prepareCommit( + appendIncrement, compactIncrement, waitCompaction); + } + + if (waitCompaction) { + prepareSortedBlocking(appendIncrement, compactIncrement, sortedCommits); + } else { + prepareSortedNonBlocking(appendIncrement, compactIncrement, sortedCommits); + } + + if (vectorCommit != null) { + mergeVectorCommit(appendIncrement, compactIncrement, vectorCommit); + } + for (BucketedSortedIndexMaintainer.SortedIndexCommit sortedCommit : sortedCommits) { + mergeSortedCommit(appendIncrement, compactIncrement, sortedCommit); + } + } catch (Throwable failure) { + for (int i = sortedCommits.size() - 1; i >= 0; i--) { + sortedCommits.get(i).abort(failure); + } + if (vectorCommit != null) { + vectorCommit.abort(failure); + } + if (failure instanceof Exception) { + throw (Exception) failure; + } + if (failure instanceof Error) { + throw (Error) failure; + } + throw new RuntimeException(failure); + } + } + + private void prepareSortedBlocking( + DataIncrement appendIncrement, + CompactIncrement compactIncrement, + List commits) + throws Exception { + BucketedSortedIndexMaintainer active = activeSortedMaintainer(); + for (BucketedSortedIndexMaintainer maintainer : sortedMaintainers) { + commits.add( + maintainer.prepareCommit( + appendIncrement, compactIncrement, true, maintainer == active)); + } + for (BucketedSortedIndexMaintainer maintainer : sortedMaintainers) { + if (maintainer != active) { + commits.add( + maintainer.prepareCommit(appendIncrement, compactIncrement, true, true)); + } + } + } + + private void prepareSortedNonBlocking( + DataIncrement appendIncrement, + CompactIncrement compactIncrement, + List commits) + throws Exception { + BucketedSortedIndexMaintainer active = activeSortedMaintainer(); + for (BucketedSortedIndexMaintainer maintainer : sortedMaintainers) { + commits.add( + maintainer.prepareCommit( + appendIncrement, compactIncrement, false, maintainer == active)); + } + if (activeSortedMaintainer() != null) { + return; + } + + for (BucketedSortedIndexMaintainer maintainer : sortedMaintainers) { + if (maintainer == active && !maintainer.state().uncoveredSourceFiles().isEmpty()) { + continue; + } + if (!maintainer.state().uncoveredSourceFiles().isEmpty()) { + commits.add( + maintainer.prepareCommit(appendIncrement, compactIncrement, false, true)); + break; + } + } + } + + @Nullable + private BucketedSortedIndexMaintainer activeSortedMaintainer() { + for (BucketedSortedIndexMaintainer maintainer : sortedMaintainers) { + if (maintainer.buildNotCompleted()) { + return maintainer; + } + } + return null; + } + + private static void mergeSortedCommit( + DataIncrement appendIncrement, + CompactIncrement compactIncrement, + BucketedSortedIndexMaintainer.SortedIndexCommit commit) { + commit.appendIncrement() .ifPresent( increment -> applyIndexIncrement( @@ -63,8 +188,29 @@ public void prepareCommit( appendIncrement.deletedIndexFiles(), increment.newIndexFiles(), increment.deletedIndexFiles())); - vectorCommit - .compactIncrement() + commit.compactIncrement() + .ifPresent( + increment -> + applyIndexIncrement( + compactIncrement.newIndexFiles(), + compactIncrement.deletedIndexFiles(), + increment.newIndexFiles(), + increment.deletedIndexFiles())); + } + + private static void mergeVectorCommit( + DataIncrement appendIncrement, + CompactIncrement compactIncrement, + BucketedVectorIndexMaintainer.VectorIndexCommit commit) { + commit.appendIncrement() + .ifPresent( + increment -> + applyIndexIncrement( + appendIncrement.newIndexFiles(), + appendIncrement.deletedIndexFiles(), + increment.newIndexFiles(), + increment.deletedIndexFiles())); + commit.compactIncrement() .ifPresent( increment -> applyIndexIncrement( @@ -84,32 +230,103 @@ private static void applyIndexIncrement( } public boolean buildNotCompleted() { - return vectorMaintainer.buildNotCompleted(); + if (vectorMaintainer != null && vectorMaintainer.buildNotCompleted()) { + return true; + } + return activeSortedMaintainer() != null; } public void withExecutor(ExecutorService executor) { - vectorMaintainer.withExecutor(executor); + if (vectorMaintainer != null) { + vectorMaintainer.withExecutor(executor); + } + for (BucketedSortedIndexMaintainer maintainer : sortedMaintainers) { + maintainer.withExecutor(executor); + } } public void close() { - vectorMaintainer.close(); + if (vectorMaintainer != null) { + vectorMaintainer.close(); + } + for (BucketedSortedIndexMaintainer maintainer : sortedMaintainers) { + maintainer.close(); + } } - /** Factory to restore configured source-backed primary-key indexes for a bucket. */ + /** Factory to restore all configured source-backed indexes for a bucket. */ public static final class Factory { - private final BucketedVectorIndexMaintainer.Factory vectorFactory; + private final IndexFileHandler handler; + @Nullable private final BucketedVectorIndexMaintainer.Factory vectorFactory; + private final List sortedFactories; - private Factory(BucketedVectorIndexMaintainer.Factory vectorFactory) { + private Factory( + IndexFileHandler handler, + @Nullable BucketedVectorIndexMaintainer.Factory vectorFactory, + List sortedFactories) { + this.handler = handler; this.vectorFactory = vectorFactory; + this.sortedFactories = Collections.unmodifiableList(new ArrayList<>(sortedFactories)); } public static Factory ofVector(BucketedVectorIndexMaintainer.Factory vectorFactory) { - return new Factory(vectorFactory); + return new Factory( + vectorFactory.indexFileHandler(), vectorFactory, Collections.emptyList()); + } + + public static Factory create( + IndexFileHandler handler, + KeyValueFileReaderFactory.Builder readerFactoryBuilder, + TableSchema schema) { + CoreOptions coreOptions = new CoreOptions(schema.options()); + Map fields = schema.nameToFieldMap(); + BucketedVectorIndexMaintainer.Factory vectorFactory = null; + List sortedFactories = new ArrayList<>(); + for (PrimaryKeyIndexDefinition definition : + PrimaryKeyIndexDefinitions.create(schema).definitions()) { + DataField field = fields.get(definition.column()); + checkArgument( + field != null, + "Primary-key index column '%s' does not exist.", + definition.column()); + switch (definition.family()) { + case VECTOR: + checkArgument( + vectorFactory == null, + "Only one primary-key vector index is supported."); + vectorFactory = + new BucketedVectorIndexMaintainer.Factory( + handler, + readerFactoryBuilder, + field, + coreOptions.primaryKeyVectorDistanceMetric( + definition.column()), + definition.indexType()) + .withIndexOptions(definition.options()); + break; + case BTREE: + case BITMAP: + sortedFactories.add( + new SortedDefinitionFactory( + readerFactoryBuilder, + field, + definition.indexType(), + definition.options())); + break; + default: + throw new IllegalArgumentException( + "Unsupported primary-key index family " + definition.family()); + } + } + checkArgument( + vectorFactory != null || !sortedFactories.isEmpty(), + "No primary-key index definition is configured."); + return new Factory(handler, vectorFactory, sortedFactories); } public IndexFileHandler indexFileHandler() { - return vectorFactory.indexFileHandler(); + return handler; } public BucketedPrimaryKeyIndexMaintainer create( @@ -118,12 +335,7 @@ public BucketedPrimaryKeyIndexMaintainer create( @Nullable List restoredDataFiles, @Nullable List restoredPayloads, ExecutorService executor) { - List dataFiles = - restoredDataFiles == null ? Collections.emptyList() : restoredDataFiles; - List payloads = - restoredPayloads == null ? Collections.emptyList() : restoredPayloads; - return BucketedPrimaryKeyIndexMaintainer.ofVector( - vectorFactory.create(partition, bucket, dataFiles, payloads, executor)); + return create(partition, bucket, restoredDataFiles, restoredPayloads, executor, null); } public BucketedPrimaryKeyIndexMaintainer create( @@ -133,7 +345,70 @@ public BucketedPrimaryKeyIndexMaintainer create( @Nullable List restoredPayloads, ExecutorService executor, @Nullable IOManager ioManager) { - return create(partition, bucket, restoredDataFiles, restoredPayloads, executor); + List dataFiles = + restoredDataFiles == null ? Collections.emptyList() : restoredDataFiles; + List payloads = + restoredPayloads == null ? Collections.emptyList() : restoredPayloads; + BucketedVectorIndexMaintainer vector = + vectorFactory == null + ? null + : vectorFactory.create( + partition, bucket, dataFiles, payloads, executor); + List sorted = new ArrayList<>(); + for (SortedDefinitionFactory factory : sortedFactories) { + sorted.add( + factory.create( + handler, partition, bucket, dataFiles, payloads, executor, + ioManager)); + } + return new BucketedPrimaryKeyIndexMaintainer(vector, sorted); + } + + private static final class SortedDefinitionFactory { + + private final KeyValueFileReaderFactory.Builder readerFactoryBuilder; + private final DataField field; + private final String indexType; + private final org.apache.paimon.options.Options options; + + private SortedDefinitionFactory( + KeyValueFileReaderFactory.Builder readerFactoryBuilder, + DataField field, + String indexType, + org.apache.paimon.options.Options options) { + this.readerFactoryBuilder = readerFactoryBuilder; + this.field = field; + this.indexType = indexType; + this.options = options; + } + + private BucketedSortedIndexMaintainer create( + IndexFileHandler handler, + BinaryRow partition, + int bucket, + List restoredDataFiles, + List restoredPayloads, + ExecutorService executor, + @Nullable IOManager ioManager) { + PkSortedIndexFile indexFile = handler.pkSortedIndex(partition, bucket); + PkSortedIndexBuilder builder = + new PkSortedIndexBuilder( + new PkSortedDataFileReader.Factory( + readerFactoryBuilder, partition, bucket, field), + indexFile, + field, + indexType, + options, + ioManager); + return new BucketedSortedIndexMaintainer( + field.id(), + indexType, + indexFile, + builder::build, + restoredDataFiles, + restoredPayloads, + executor); + } } } } 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 index 6a887f470d41..44695b5f54bd 100644 --- 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 @@ -169,7 +169,16 @@ public synchronized SortedIndexCommit prepareCommit( changed && hasCompactDataTransition ? Optional.of(new SortedIndexIncrement(created, removed)) : Optional.empty(); - return new SortedIndexCommit(appendChange, compactChange); + return new SortedIndexCommit( + appendChange, + compactChange, + failure -> + rollbackPrepareCommit( + originalSourceFiles, + originalGroups, + originalRestoredDeletions, + created, + failure)); } catch (Throwable failure) { rollbackPrepareCommit( originalSourceFiles, @@ -187,7 +196,7 @@ public synchronized SortedIndexCommit prepareCommit( } } - private void rollbackPrepareCommit( + private synchronized void rollbackPrepareCommit( Map originalSourceFiles, List originalGroups, List originalRestoredDeletions, @@ -498,12 +507,15 @@ public static final class SortedIndexCommit { private final Optional appendIncrement; private final Optional compactIncrement; + private final AbortAction abortAction; private SortedIndexCommit( Optional appendIncrement, - Optional compactIncrement) { + Optional compactIncrement, + AbortAction abortAction) { this.appendIncrement = appendIncrement; this.compactIncrement = compactIncrement; + this.abortAction = abortAction; } public Optional appendIncrement() { @@ -513,6 +525,16 @@ public Optional appendIncrement() { public Optional compactIncrement() { return compactIncrement; } + + public void abort(Throwable failure) { + abortAction.abort(failure); + } + } + + @FunctionalInterface + private interface AbortAction { + + void abort(Throwable failure); } /** Index-file additions and deletions emitted by one sorted definition. */ diff --git a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainer.java b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainer.java index 3b106f79c8b8..1e002aa78752 100644 --- a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainer.java +++ b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainer.java @@ -93,8 +93,17 @@ public class BucketedVectorIndexMaintainer { coreOptions.primaryKeyVectorIndexCompactionLevelFanout(), coreOptions.primaryKeyVectorIndexCompactionStaleRatioThreshold()); this.executor = executor; + + List definitionPayloads = new ArrayList<>(); + for (IndexFileMeta payload : restoredPayloads) { + if (algorithm.equals(payload.indexType()) + && payload.globalIndexMeta() != null + && payload.globalIndexMeta().indexFieldId() == vectorFieldId) { + definitionPayloads.add(payload); + } + } PkVectorBucketIndexState restoredState = - new PkVectorBucketIndexState(vectorFieldId, algorithm, restoredPayloads); + new PkVectorBucketIndexState(vectorFieldId, algorithm, definitionPayloads); this.annSegments = new ArrayList<>(restoredState.annSegments()); this.activeSourceFiles = new LinkedHashMap<>(); for (DataFileMeta file : restoredDataFiles) { @@ -183,7 +192,12 @@ public synchronized VectorIndexCommit prepareCommit( !indexChanged || !hasCompactDataTransition ? Optional.empty() : Optional.of(new VectorIndexIncrement(created, removed)); - return new VectorIndexCommit(appendChange, compactChange); + return new VectorIndexCommit( + appendChange, + compactChange, + failure -> + rollbackPrepareCommit( + originalSegments, originalSourceFiles, generated, failure)); } catch (Throwable failure) { rollbackPrepareCommit(originalSegments, originalSourceFiles, generated, failure); if (failure instanceof Exception) { @@ -203,7 +217,7 @@ private void startPendingBuild( pendingBuild = build; } - private void rollbackPrepareCommit( + private synchronized void rollbackPrepareCommit( List originalSegments, Map originalSourceFiles, List generated, @@ -474,12 +488,15 @@ public static class VectorIndexCommit { private final Optional appendIncrement; private final Optional compactIncrement; + private final AbortAction abortAction; private VectorIndexCommit( Optional appendIncrement, - Optional compactIncrement) { + Optional compactIncrement, + AbortAction abortAction) { this.appendIncrement = appendIncrement; this.compactIncrement = compactIncrement; + this.abortAction = abortAction; } public Optional appendIncrement() { @@ -489,6 +506,16 @@ public Optional appendIncrement() { public Optional compactIncrement() { return compactIncrement; } + + public void abort(Throwable failure) { + abortAction.abort(failure); + } + } + + @FunctionalInterface + private interface AbortAction { + + void abort(Throwable failure); } /** Index-file additions and deletions emitted by one bucket state update. */ diff --git a/paimon-core/src/test/java/org/apache/paimon/index/pk/BucketedPrimaryKeyIndexMaintainerTest.java b/paimon-core/src/test/java/org/apache/paimon/index/pk/BucketedPrimaryKeyIndexMaintainerTest.java new file mode 100644 index 000000000000..c2acfda21bff --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/index/pk/BucketedPrimaryKeyIndexMaintainerTest.java @@ -0,0 +1,351 @@ +/* + * 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.pk; + +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.pksorted.BucketedSortedIndexMaintainer; +import org.apache.paimon.index.pksorted.PkSortedIndexFile; +import org.apache.paimon.index.pkvector.BucketedVectorIndexMaintainer; +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.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Optional; +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; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** Tests for {@link BucketedPrimaryKeyIndexMaintainer}. */ +class BucketedPrimaryKeyIndexMaintainerTest { + + @TempDir java.nio.file.Path tempPath; + private final ExecutorService buildExecutor = Executors.newFixedThreadPool(2); + + @AfterEach + void shutdownExecutor() { + buildExecutor.shutdownNow(); + } + + @Test + void testDelegatesVectorLifecycle() { + BucketedVectorIndexMaintainer vector = mock(BucketedVectorIndexMaintainer.class); + ExecutorService executor = mock(ExecutorService.class); + when(vector.buildNotCompleted()).thenReturn(true); + BucketedPrimaryKeyIndexMaintainer maintainer = + BucketedPrimaryKeyIndexMaintainer.ofVector(vector); + + assertThat(maintainer.buildNotCompleted()).isTrue(); + maintainer.withExecutor(executor); + maintainer.close(); + + verify(vector).withExecutor(executor); + verify(vector).close(); + } + + @Test + void testScalarFailureDoesNotSuppressOtherDefinitionsOrVector() throws Exception { + DataFileMeta source = dataFile("data-1", 3); + IndexFileMeta vectorPayload = + new IndexFileMeta("vector", "vector", 1, 3, (GlobalIndexMeta) null, null); + BucketedVectorIndexMaintainer vector = mock(BucketedVectorIndexMaintainer.class); + BucketedVectorIndexMaintainer.VectorIndexCommit vectorCommit = + mock(BucketedVectorIndexMaintainer.VectorIndexCommit.class); + BucketedVectorIndexMaintainer.VectorIndexIncrement vectorIncrement = + mock(BucketedVectorIndexMaintainer.VectorIndexIncrement.class); + when(vector.prepareCommit(any(), any(), eq(true))).thenReturn(vectorCommit); + when(vectorCommit.appendIncrement()).thenReturn(Optional.empty()); + when(vectorCommit.compactIncrement()).thenReturn(Optional.of(vectorIncrement)); + when(vectorIncrement.newIndexFiles()).thenReturn(Collections.singletonList(vectorPayload)); + when(vectorIncrement.deletedIndexFiles()).thenReturn(Collections.emptyList()); + + List buildOrder = Collections.synchronizedList(new ArrayList<>()); + BucketedSortedIndexMaintainer btree = + sortedMaintainer( + 7, + "btree", + source, + dataFile -> { + buildOrder.add("btree"); + throw new IllegalStateException("expected BTree failure"); + }); + IndexFileMeta bitmapPayload = payload("bitmap", source, 3, 8, "bitmap"); + BucketedSortedIndexMaintainer bitmap = + sortedMaintainer( + 8, + "bitmap", + source, + dataFile -> { + buildOrder.add("bitmap"); + return Collections.singletonList(bitmapPayload); + }); + BucketedPrimaryKeyIndexMaintainer maintainer = + BucketedPrimaryKeyIndexMaintainer.of(vector, Arrays.asList(bitmap, btree)); + CompactIncrement compactIncrement = compactAfter(source); + + maintainer.prepareCommit(DataIncrement.emptyIncrement(), compactIncrement, true); + + assertThat(buildOrder).containsExactly("btree", "btree", "btree", "bitmap"); + assertThat(compactIncrement.newIndexFiles()).containsExactly(vectorPayload, bitmapPayload); + assertThat(compactIncrement.deletedIndexFiles()).isEmpty(); + assertThat(maintainer.buildNotCompleted()).isFalse(); + } + + @Test + void testScalarDefinitionsNeverBuildConcurrently() throws Exception { + DataFileMeta source = dataFile("data-1", 3); + AtomicInteger activeBuilds = new AtomicInteger(); + AtomicInteger peakBuilds = new AtomicInteger(); + CountDownLatch firstStarted = new CountDownLatch(1); + CountDownLatch releaseFirst = new CountDownLatch(1); + CountDownLatch secondStarted = new CountDownLatch(1); + BucketedSortedIndexMaintainer first = + sortedMaintainer( + 7, + "btree", + source, + dataFile -> { + enterBuild(activeBuilds, peakBuilds); + firstStarted.countDown(); + releaseFirst.await(); + activeBuilds.decrementAndGet(); + return Collections.singletonList( + payload("btree", source, 3, 7, "btree")); + }); + BucketedSortedIndexMaintainer second = + sortedMaintainer( + 8, + "bitmap", + source, + dataFile -> { + enterBuild(activeBuilds, peakBuilds); + secondStarted.countDown(); + activeBuilds.decrementAndGet(); + return Collections.singletonList( + payload("bitmap", source, 3, 8, "bitmap")); + }); + BucketedPrimaryKeyIndexMaintainer maintainer = + BucketedPrimaryKeyIndexMaintainer.ofSorted(Arrays.asList(second, first)); + + maintainer.prepareCommit(DataIncrement.emptyIncrement(), compactAfter(source), false); + assertThat(firstStarted.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(secondStarted.getCount()).isEqualTo(1); + assertThat(maintainer.buildNotCompleted()).isTrue(); + + releaseFirst.countDown(); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (secondStarted.getCount() > 0 && System.nanoTime() < deadline) { + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), CompactIncrement.emptyIncrement(), false); + secondStarted.await(10, TimeUnit.MILLISECONDS); + } + assertThat(secondStarted.getCount()).isZero(); + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), CompactIncrement.emptyIncrement(), true); + + assertThat(peakBuilds).hasValue(1); + assertThat(maintainer.buildNotCompleted()).isFalse(); + } + + @Test + void testBlockingPrepareRollsBackAllDefinitionsAfterInterruption() throws Exception { + DataFileMeta source = dataFile("data-1", 3); + AtomicInteger firstBuilds = new AtomicInteger(); + AtomicInteger secondBuilds = new AtomicInteger(); + CountDownLatch secondStarted = new CountDownLatch(1); + CountDownLatch blockSecond = new CountDownLatch(1); + BucketedSortedIndexMaintainer first = + sortedMaintainer( + 7, + "btree", + source, + dataFile -> { + int attempt = firstBuilds.incrementAndGet(); + return Collections.singletonList( + payload("btree-" + attempt, source, 3, 7, "btree")); + }); + BucketedSortedIndexMaintainer second = + sortedMaintainer( + 8, + "bitmap", + source, + dataFile -> { + int attempt = secondBuilds.incrementAndGet(); + if (attempt == 1) { + secondStarted.countDown(); + blockSecond.await(); + } + return Collections.singletonList( + payload("bitmap-" + attempt, source, 3, 8, "bitmap")); + }); + BucketedPrimaryKeyIndexMaintainer maintainer = + BucketedPrimaryKeyIndexMaintainer.ofSorted(Arrays.asList(second, first)); + CompactIncrement failedIncrement = compactAfter(source); + AtomicReference failure = new AtomicReference<>(); + Thread commitThread = + new Thread( + () -> { + try { + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), failedIncrement, true); + } catch (Throwable t) { + failure.set(t); + } + }); + + commitThread.start(); + assertThat(secondStarted.await(5, TimeUnit.SECONDS)).isTrue(); + commitThread.interrupt(); + commitThread.join(TimeUnit.SECONDS.toMillis(5)); + + assertThat(commitThread.isAlive()).isFalse(); + assertThat(failure.get()).isInstanceOf(InterruptedException.class); + assertThat(failedIncrement.newIndexFiles()).isEmpty(); + assertThat(failedIncrement.deletedIndexFiles()).isEmpty(); + assertThat(first.state().groups()).isEmpty(); + assertThat(second.state().groups()).isEmpty(); + + CompactIncrement retryIncrement = compactAfter(source); + maintainer.prepareCommit(DataIncrement.emptyIncrement(), retryIncrement, true); + + assertThat(firstBuilds).hasValue(2); + assertThat(secondBuilds).hasValue(2); + assertThat(retryIncrement.newIndexFiles()) + .extracting(IndexFileMeta::fileName) + .containsExactly("btree-2", "bitmap-2"); + } + + private BucketedSortedIndexMaintainer sortedMaintainer( + int fieldId, + String indexType, + DataFileMeta source, + BucketedSortedIndexMaintainer.BuildFunction buildFunction) { + return new BucketedSortedIndexMaintainer( + fieldId, + indexType, + new PkSortedIndexFile(LocalFileIO.create(), pathFactory()), + buildFunction, + Collections.emptyList(), + Collections.emptyList(), + buildExecutor); + } + + private static void enterBuild(AtomicInteger activeBuilds, AtomicInteger peakBuilds) { + int active = activeBuilds.incrementAndGet(); + while (true) { + int peak = peakBuilds.get(); + if (active <= peak || peakBuilds.compareAndSet(peak, active)) { + return; + } + } + } + + 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, + int fieldId, + String indexType) { + PrimaryKeyIndexSourceFile source = + new PrimaryKeyIndexSourceFile(sourceFile.fileName(), sourceFile.rowCount()); + return new IndexFileMeta( + indexType, + fileName, + 1, + payloadRowCount, + new GlobalIndexMeta( + 0, + source.rowCount() - 1, + fieldId, + 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; + } + }; + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/index/pksorted/PrimaryKeySortedIndexMaintenanceTest.java b/paimon-core/src/test/java/org/apache/paimon/index/pksorted/PrimaryKeySortedIndexMaintenanceTest.java new file mode 100644 index 000000000000..8ed3b4234f78 --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/index/pksorted/PrimaryKeySortedIndexMaintenanceTest.java @@ -0,0 +1,396 @@ +/* + * 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.CoreOptions; +import org.apache.paimon.KeyValue; +import org.apache.paimon.Snapshot; +import org.apache.paimon.TestFileStore; +import org.apache.paimon.TestKeyValueGenerator; +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.data.BinaryVector; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.disk.IOManager; +import org.apache.paimon.disk.IOManagerImpl; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceFile; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceMeta; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourcePolicy; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.manifest.ManifestCommittable; +import org.apache.paimon.memory.HeapMemorySegmentPool; +import org.apache.paimon.memory.MemoryOwner; +import org.apache.paimon.mergetree.compact.DeduplicateMergeFunction; +import org.apache.paimon.mergetree.compact.LookupMergeFunction; +import org.apache.paimon.operation.AbstractFileStoreWrite; +import org.apache.paimon.operation.FileStoreCommit; +import org.apache.paimon.schema.Schema; +import org.apache.paimon.schema.SchemaManager; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.sink.CommitMessage; +import org.apache.paimon.table.sink.CommitMessageImpl; +import org.apache.paimon.types.DataField; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.CommitIncrement; +import org.apache.paimon.utils.InternalRowUtils; +import org.apache.paimon.utils.RecordWriter; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +/** End-to-end maintenance tests for source-backed BTree and Bitmap indexes. */ +class PrimaryKeySortedIndexMaintenanceTest { + + @TempDir java.nio.file.Path tempDir; + private IOManager ioManager; + + @BeforeEach + void before() throws IOException { + ioManager = new IOManagerImpl(tempDir.toString()); + } + + @AfterEach + void after() throws Exception { + ioManager.close(); + } + + @Test + void testBuildRestoreAndReplaceSourceGroups() throws Exception { + TestFileStore store = createStore(); + TestKeyValueGenerator generator = new TestKeyValueGenerator(); + List firstBatch = + Arrays.asList( + withKey(generator.nextInsert("20201110", 10, 30L, new int[] {3}, "c"), 0), + withKey(generator.nextInsert("20201110", 10, 10L, new int[] {1}, "a"), 4), + withKey(generator.nextInsert("20201110", 10, null, new int[] {0}, null), 8), + withKey(generator.nextInsert("20201110", 10, 20L, new int[] {2}, "b"), 2), + withKey(generator.nextInsert("20201110", 10, 40L, new int[] {4}, "d"), 6)); + BinaryRow partition = generator.getPartition(firstBatch.get(0)); + + writeAndCommit(store, partition, 0, firstBatch.subList(0, 3), 0); + writeAndCommit(store, partition, 0, firstBatch.subList(3, 5), 1); + assertThat(store.newScan().plan().files()).hasSize(2); + compactAndCommit(store, partition, 0, 2); + + List compactedFiles = + store.newScan().plan().files().stream() + .map(entry -> entry.file()) + .collect(Collectors.toList()); + assertThat(compactedFiles).anyMatch(PrimaryKeyIndexSourcePolicy::shouldRead); + List firstPayloads = sourcePayloads(store, partition, 0); + assertCompleteGroups(firstPayloads, 5); + Set firstPayloadNames = + firstPayloads.stream().map(IndexFileMeta::fileName).collect(Collectors.toSet()); + Set firstSources = sourceNames(firstPayloads); + assertThat(firstSources).isNotEmpty(); + + KeyValue additional = + withKey(generator.nextInsert("20201110", 10, 50L, new int[] {5}, "e"), 5); + writeAndCommit(store, partition, 0, Arrays.asList(additional), 3); + compactAndCommit(store, partition, 0, 4); + + List replacementPayloads = sourcePayloads(store, partition, 0); + assertCompleteGroups(replacementPayloads, 6); + assertThat(replacementPayloads) + .extracting(IndexFileMeta::fileName) + .doesNotContainAnyElementsOf(firstPayloadNames); + assertThat(sourceNames(replacementPayloads)).doesNotContainAnyElementsOf(firstSources); + + AbstractFileStoreWrite restored = store.newWrite(); + restored.withIOManager(ioManager); + AbstractFileStoreWrite.WriterContainer container = + restored.createWriterContainer(partition, 0); + restored.prepareCommit(false, 5); + assertThat(container.primaryKeyIndexMaintainer).isNotNull(); + assertThat(container.primaryKeyIndexMaintainer.buildNotCompleted()).isFalse(); + restored.close(); + } + + @Test + void testRestoreMixedVectorAndScalarPayloads() throws Exception { + TestFileStore store = createMixedStore(); + TestKeyValueGenerator generator = new TestKeyValueGenerator(); + KeyValue first = + withVector( + withKey(generator.nextInsert("20201110", 10, 30L, new int[] {3}, "c"), 0), + new float[] {1, 0}); + KeyValue second = + withVector( + withKey(generator.nextInsert("20201110", 10, 10L, new int[] {1}, "a"), 1), + new float[] {0, 1}); + BinaryRow partition = generator.getPartition(first); + + writeAndCommit(store, partition, 0, Collections.singletonList(first), 0); + writeAndCommit(store, partition, 0, Collections.singletonList(second), 1); + compactAndCommit(store, partition, 0, 2); + + assertThat(sourcePayloads(store, partition, 0)) + .extracting(IndexFileMeta::indexType) + .contains("test-vector-ann", "btree", "bitmap"); + + AbstractFileStoreWrite restored = store.newWrite(); + restored.withIOManager(ioManager); + AbstractFileStoreWrite.WriterContainer container = + restored.createWriterContainer(partition, 0); + restored.prepareCommit(false, 3); + + assertThat(container.primaryKeyIndexMaintainer).isNotNull(); + assertThat(container.primaryKeyIndexMaintainer.buildNotCompleted()).isFalse(); + restored.close(); + } + + private void writeAndCommit( + TestFileStore store, + BinaryRow partition, + int bucket, + List records, + long identifier) + throws Exception { + AbstractFileStoreWrite write = store.newWrite(); + write.withIOManager(ioManager); + AbstractFileStoreWrite.WriterContainer container = + write.createWriterContainer(partition, bucket); + RecordWriter writer = container.writer; + ((MemoryOwner) writer) + .setMemoryPool( + new HeapMemorySegmentPool( + TestFileStore.WRITE_BUFFER_SIZE.getBytes(), + (int) TestFileStore.PAGE_SIZE.getBytes())); + for (KeyValue record : records) { + writer.write(record); + } + commit( + store, + Collections.singletonList(prepareCommit(container, partition, bucket)), + identifier); + container.writer.close(); + container.primaryKeyIndexMaintainer.close(); + write.close(); + } + + private void compactAndCommit( + TestFileStore store, BinaryRow partition, int bucket, long identifier) + throws Exception { + AbstractFileStoreWrite write = store.newWrite(); + write.withIOManager(ioManager); + AbstractFileStoreWrite.WriterContainer container = + write.createWriterContainer(partition, bucket); + ((MemoryOwner) container.writer) + .setMemoryPool( + new HeapMemorySegmentPool( + TestFileStore.WRITE_BUFFER_SIZE.getBytes(), + (int) TestFileStore.PAGE_SIZE.getBytes())); + container.writer.compact(true); + commit( + store, + Collections.singletonList(prepareCommit(container, partition, bucket)), + identifier); + container.writer.close(); + container.primaryKeyIndexMaintainer.close(); + write.close(); + } + + private CommitMessage prepareCommit( + AbstractFileStoreWrite.WriterContainer container, + BinaryRow partition, + int bucket) + throws Exception { + CommitIncrement increment = container.writer.prepareCommit(true); + container.primaryKeyIndexMaintainer.prepareCommit( + increment.newFilesIncrement(), increment.compactIncrement(), true); + return new CommitMessageImpl( + partition, bucket, 1, increment.newFilesIncrement(), increment.compactIncrement()); + } + + private void commit(TestFileStore store, List messages, long identifier) { + ManifestCommittable committable = new ManifestCommittable(identifier, null, messages); + try (FileStoreCommit commit = store.newCommit()) { + commit.withIOManager(ioManager).commit(committable, false); + } + } + + private static List sourcePayloads( + TestFileStore store, BinaryRow partition, int bucket) { + Snapshot snapshot = store.snapshotManager().latestSnapshot(); + return store.newIndexFileHandler().scanSourceIndexes(snapshot, partition, bucket); + } + + private static void assertCompleteGroups(List payloads, long rowCount) { + Map> byType = + payloads.stream().collect(Collectors.groupingBy(IndexFileMeta::indexType)); + assertThat(byType.keySet()).containsExactlyInAnyOrder("btree", "bitmap"); + for (List typePayloads : byType.values()) { + assertThat(typePayloads) + .allMatch(payload -> payload.globalIndexMeta().sourceMeta() != null); + assertThat(typePayloads.stream().mapToLong(IndexFileMeta::rowCount).sum()) + .as("total row count of %s source groups", typePayloads.get(0).indexType()) + .isEqualTo(rowCount); + Map> bySource = + typePayloads.stream() + .collect( + Collectors.groupingBy( + payload -> + PrimaryKeyIndexSourceMeta.fromIndexFile(payload) + .sourceFile())); + for (Map.Entry> sourceGroup : + bySource.entrySet()) { + assertThat(sourceGroup.getValue().stream().mapToLong(IndexFileMeta::rowCount).sum()) + .isEqualTo(sourceGroup.getKey().rowCount()); + } + } + } + + private static Set sourceNames(List payloads) { + Set result = new HashSet<>(); + for (IndexFileMeta payload : payloads) { + result.add(PrimaryKeyIndexSourceMeta.fromIndexFile(payload).sourceFile().fileName()); + } + return result; + } + + private static KeyValue withKey(KeyValue record, int key) { + return new KeyValue() + .replace( + TestKeyValueGenerator.KEY_SERIALIZER + .toBinaryRow(GenericRow.of(key, (long) key)) + .copy(), + record.sequenceNumber(), + record.valueKind(), + record.value()); + } + + private static KeyValue withVector(KeyValue record, float[] vector) { + InternalRow original = record.value(); + GenericRow extended = + new GenericRow(TestKeyValueGenerator.DEFAULT_ROW_TYPE.getFieldCount() + 1); + for (int i = 0; i < TestKeyValueGenerator.DEFAULT_ROW_TYPE.getFieldCount(); i++) { + extended.setField( + i, + InternalRowUtils.get( + original, i, TestKeyValueGenerator.DEFAULT_ROW_TYPE.getTypeAt(i))); + } + extended.setField( + TestKeyValueGenerator.DEFAULT_ROW_TYPE.getFieldCount(), + BinaryVector.fromPrimitiveArray(vector)); + return record.replaceValue(extended); + } + + private TestFileStore createStore() throws Exception { + Map options = new HashMap<>(); + options.put(CoreOptions.BUCKET.key(), "1"); + options.put(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true"); + options.put(CoreOptions.COMPACTION_FORCE_REWRITE_ALL_FILES.key(), "true"); + options.put(CoreOptions.PK_BTREE_INDEX_COLUMNS.key(), "itemId"); + options.put(CoreOptions.PK_BITMAP_INDEX_COLUMNS.key(), "comment"); + options.put( + "fields.itemId.pk-btree.index.options", + "{\"sorted-index.records-per-range\":\"2\"}"); + options.put( + "fields.comment.pk-bitmap.index.options", + "{\"sorted-index.records-per-range\":\"2\"}"); + SchemaManager schemaManager = + new SchemaManager(LocalFileIO.create(), new Path(tempDir.toUri())); + TableSchema schema = + schemaManager.createTable( + new Schema( + TestKeyValueGenerator.DEFAULT_ROW_TYPE.getFields(), + TestKeyValueGenerator.DEFAULT_PART_TYPE.getFieldNames(), + TestKeyValueGenerator.getPrimaryKeys( + TestKeyValueGenerator.GeneratorMode.MULTI_PARTITIONED), + options, + null)); + return new TestFileStore.Builder( + "avro", + tempDir.toString(), + 1, + TestKeyValueGenerator.DEFAULT_PART_TYPE, + TestKeyValueGenerator.KEY_TYPE, + TestKeyValueGenerator.DEFAULT_ROW_TYPE, + TestKeyValueGenerator.TestKeyValueFieldsExtractor.EXTRACTOR, + LookupMergeFunction.wrap( + DeduplicateMergeFunction.factory(), + new CoreOptions(options), + TestKeyValueGenerator.KEY_TYPE, + TestKeyValueGenerator.DEFAULT_ROW_TYPE), + schema) + .build(); + } + + private TestFileStore createMixedStore() throws Exception { + Map options = new HashMap<>(); + options.put(CoreOptions.BUCKET.key(), "1"); + options.put(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true"); + options.put(CoreOptions.COMPACTION_FORCE_REWRITE_ALL_FILES.key(), "true"); + options.put(CoreOptions.PK_BTREE_INDEX_COLUMNS.key(), "itemId"); + options.put(CoreOptions.PK_BITMAP_INDEX_COLUMNS.key(), "comment"); + options.put(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "embedding"); + options.put("fields.embedding.pk-vector.index.type", "test-vector-ann"); + options.put("fields.embedding.pk-vector.distance.metric", "l2"); + + List fields = + new ArrayList<>(TestKeyValueGenerator.DEFAULT_ROW_TYPE.getFields()); + fields.add(new DataField(7, "embedding", DataTypes.VECTOR(2, DataTypes.FLOAT()))); + RowType rowType = new RowType(fields); + SchemaManager schemaManager = + new SchemaManager(LocalFileIO.create(), new Path(tempDir.toUri())); + TableSchema schema = + schemaManager.createTable( + new Schema( + rowType.getFields(), + TestKeyValueGenerator.DEFAULT_PART_TYPE.getFieldNames(), + TestKeyValueGenerator.getPrimaryKeys( + TestKeyValueGenerator.GeneratorMode.MULTI_PARTITIONED), + options, + null)); + return new TestFileStore.Builder( + "avro", + tempDir.toString(), + 1, + TestKeyValueGenerator.DEFAULT_PART_TYPE, + TestKeyValueGenerator.KEY_TYPE, + rowType, + TestKeyValueGenerator.TestKeyValueFieldsExtractor.EXTRACTOR, + LookupMergeFunction.wrap( + DeduplicateMergeFunction.factory(), + new CoreOptions(options), + TestKeyValueGenerator.KEY_TYPE, + rowType), + schema) + .build(); + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainerTest.java b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainerTest.java index 2a6b0dc65ee9..8cd89bb59c7d 100644 --- a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainerTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/BucketedVectorIndexMaintainerTest.java @@ -20,6 +20,7 @@ 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; @@ -69,6 +70,72 @@ void shutdownExecutor() { executor.shutdownNow(); } + @Test + void testRestoreFiltersPayloadsByVectorDefinition() { + DataField vectorField = + new DataField(7, "embedding", DataTypes.VECTOR(2, DataTypes.FLOAT())); + DataFileMeta data = dataFile("data"); + IndexFileMeta vectorPayload = payload("ann", data, 7, "test-vector-ann"); + List restoredPayloads = + Arrays.asList( + payload("btree", data, 3, "btree"), + payload("bitmap", data, 5, "bitmap"), + payload("other-ann", data, 8, "test-vector-ann"), + vectorPayload); + + BucketedVectorIndexMaintainer maintainer = + new BucketedVectorIndexMaintainer( + 7, + new PkVectorAnnSegmentFile(LocalFileIO.create(), pathFactory()), + vectorField, + indexOptions(), + "l2", + "test-vector-ann", + mock(PkVectorDataFileReader.Factory.class), + Collections.singletonList(data), + restoredPayloads, + executor); + + assertThat(maintainer.segments()).containsExactly(vectorPayload); + } + + @Test + void testPreparedCommitCanBeAbortedByCoordinator() throws Exception { + DataField vectorField = + new DataField(7, "embedding", DataTypes.VECTOR(2, DataTypes.FLOAT())); + DataFileMeta data = dataFile("data"); + PkVectorDataFileReader.Factory readerFactory = mock(PkVectorDataFileReader.Factory.class); + PkVectorDataFileReader dataReader = reader(new float[][] {{1, 0}}); + when(readerFactory.create(data)).thenReturn(dataReader); + BucketedVectorIndexMaintainer maintainer = + new BucketedVectorIndexMaintainer( + 7, + new PkVectorAnnSegmentFile(LocalFileIO.create(), pathFactory()), + vectorField, + indexOptions(), + "l2", + "test-vector-ann", + readerFactory, + Collections.emptyList(), + Collections.emptyList(), + executor); + + BucketedVectorIndexMaintainer.VectorIndexCommit commit = + maintainer.prepareCommit( + DataIncrement.emptyIncrement(), + new CompactIncrement( + Collections.emptyList(), + Collections.singletonList(data), + Collections.emptyList()), + true); + assertThat(maintainer.segments()).hasSize(1); + + commit.abort(new IllegalStateException("later definition failed")); + + assertThat(maintainer.segments()).isEmpty(); + assertThat(fileCount()).isZero(); + } + @Test void testNonBlockingPrepareCommitPublishesCompletedAnnLater() throws Exception { LocalFileIO fileIO = LocalFileIO.create(); @@ -659,6 +726,25 @@ private static DataFileMeta dataFile(String fileName) { .upgrade(1); } + private static IndexFileMeta payload( + String fileName, DataFileMeta source, int fieldId, String indexType) { + PrimaryKeyIndexSourceFile sourceFile = + new PrimaryKeyIndexSourceFile(source.fileName(), source.rowCount()); + return new IndexFileMeta( + indexType, + fileName, + 1, + source.rowCount(), + new GlobalIndexMeta( + 0, + source.rowCount() - 1, + fieldId, + null, + new byte[] {1}, + new PrimaryKeyIndexSourceMeta(sourceFile).serialize()), + null); + } + private IndexPathFactory pathFactory() { Path directory = new Path(tempPath.toUri()); return new IndexPathFactory() { diff --git a/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyIndexWriteTest.java b/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyIndexWriteTest.java new file mode 100644 index 000000000000..f864f6152644 --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/operation/PrimaryKeyIndexWriteTest.java @@ -0,0 +1,197 @@ +/* + * 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.operation; + +import org.apache.paimon.CoreOptions; +import org.apache.paimon.KeyValue; +import org.apache.paimon.TestFileStore; +import org.apache.paimon.TestKeyValueGenerator; +import org.apache.paimon.disk.IOManager; +import org.apache.paimon.disk.IOManagerImpl; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.mergetree.compact.DeduplicateMergeFunction; +import org.apache.paimon.postpone.PostponeBucketFileStoreWrite; +import org.apache.paimon.schema.Schema; +import org.apache.paimon.schema.SchemaManager; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.BucketMode; +import org.apache.paimon.types.DataField; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests BTree and Bitmap primary-key index wiring through the file-store writer. */ +class PrimaryKeyIndexWriteTest { + + @TempDir java.nio.file.Path tempDir; + private IOManager ioManager; + + @BeforeEach + void before() throws IOException { + ioManager = new IOManagerImpl(tempDir.toString()); + } + + @AfterEach + void after() throws Exception { + ioManager.close(); + } + + @Test + void testCreatesCoordinatorForBTreeAndBitmapDefinitions() throws Exception { + Map options = new HashMap<>(); + options.put(CoreOptions.BUCKET.key(), "10"); + options.put(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true"); + options.put(CoreOptions.PK_BTREE_INDEX_COLUMNS.key(), "itemId"); + options.put(CoreOptions.PK_BITMAP_INDEX_COLUMNS.key(), "comment"); + TestFileStore store = createStore(options); + KeyValueFileStoreWrite write = (KeyValueFileStoreWrite) store.newWrite(); + write.withIOManager(ioManager); + TestKeyValueGenerator generator = new TestKeyValueGenerator(); + KeyValue record = generator.next(); + + AbstractFileStoreWrite.WriterContainer container = + write.createWriterContainer(generator.getPartition(record), 1); + + assertThat(container.primaryKeyIndexMaintainer).isNotNull(); + assertThat(container.primaryKeyIndexMaintainer.buildNotCompleted()).isFalse(); + write.close(); + } + + @Test + void testCreatesCoordinatorForMultipleScalarDefinitionsAndVector() throws Exception { + Map options = new HashMap<>(); + options.put(CoreOptions.BUCKET.key(), "10"); + options.put(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true"); + options.put(CoreOptions.PK_BTREE_INDEX_COLUMNS.key(), "itemId,score"); + options.put(CoreOptions.PK_BITMAP_INDEX_COLUMNS.key(), "comment,state"); + options.put(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "embedding"); + options.put("fields.embedding.pk-vector.index.type", "hnsw"); + + List fields = + new ArrayList<>(TestKeyValueGenerator.DEFAULT_ROW_TYPE.getFields()); + fields.add(new DataField(7, "score", DataTypes.BIGINT())); + fields.add(new DataField(8, "state", DataTypes.INT())); + fields.add(new DataField(9, "embedding", DataTypes.VECTOR(3, DataTypes.FLOAT()))); + TestFileStore store = createStore(options, new RowType(fields)); + KeyValueFileStoreWrite write = (KeyValueFileStoreWrite) store.newWrite(); + write.withIOManager(ioManager); + TestKeyValueGenerator generator = new TestKeyValueGenerator(); + KeyValue record = generator.next(); + + AbstractFileStoreWrite.WriterContainer container = + write.createWriterContainer(generator.getPartition(record), 1); + + assertThat(readField(container.primaryKeyIndexMaintainer, "vectorMaintainer")).isNotNull(); + assertThat((List) readField(container.primaryKeyIndexMaintainer, "sortedMaintainers")) + .hasSize(4); + write.close(); + } + + @Test + void testPostponeBucketDefersCoordinatorUntilFixedBucketCompaction() throws Exception { + Map options = new HashMap<>(); + options.put(CoreOptions.BUCKET.key(), String.valueOf(BucketMode.POSTPONE_BUCKET)); + options.put(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true"); + options.put(CoreOptions.PK_BTREE_INDEX_COLUMNS.key(), "itemId"); + options.put(CoreOptions.PK_BITMAP_INDEX_COLUMNS.key(), "comment"); + TableSchema schema = createSchema(options, TestKeyValueGenerator.DEFAULT_ROW_TYPE); + TestKeyValueGenerator generator = new TestKeyValueGenerator(); + KeyValue record = generator.next(); + + TestFileStore postponeStore = + createStore( + schema, TestKeyValueGenerator.DEFAULT_ROW_TYPE, BucketMode.POSTPONE_BUCKET); + AbstractFileStoreWrite postponeWrite = postponeStore.newWrite(); + postponeWrite.withIOManager(ioManager); + AbstractFileStoreWrite.WriterContainer postponeContainer = + postponeWrite.createWriterContainer( + generator.getPartition(record), BucketMode.POSTPONE_BUCKET); + + assertThat(postponeWrite).isInstanceOf(PostponeBucketFileStoreWrite.class); + assertThat(postponeContainer.primaryKeyIndexMaintainer).isNull(); + postponeWrite.close(); + + // Postpone compaction keeps the table schema but uses a fixed-bucket runtime view. + TestFileStore compactStore = createStore(schema, TestKeyValueGenerator.DEFAULT_ROW_TYPE, 1); + KeyValueFileStoreWrite compactWrite = (KeyValueFileStoreWrite) compactStore.newWrite(); + compactWrite.withIOManager(ioManager); + AbstractFileStoreWrite.WriterContainer compactContainer = + compactWrite.createWriterContainer(generator.getPartition(record), 0); + + assertThat(compactContainer.primaryKeyIndexMaintainer).isNotNull(); + compactWrite.close(); + } + + private TestFileStore createStore(Map options) throws Exception { + return createStore(options, TestKeyValueGenerator.DEFAULT_ROW_TYPE); + } + + private TestFileStore createStore(Map options, RowType rowType) + throws Exception { + return createStore(createSchema(options, rowType), rowType, 10); + } + + private TableSchema createSchema(Map options, RowType rowType) + throws Exception { + SchemaManager schemaManager = + new SchemaManager(LocalFileIO.create(), new Path(tempDir.toUri())); + return schemaManager.createTable( + new Schema( + rowType.getFields(), + TestKeyValueGenerator.DEFAULT_PART_TYPE.getFieldNames(), + TestKeyValueGenerator.getPrimaryKeys( + TestKeyValueGenerator.GeneratorMode.MULTI_PARTITIONED), + options, + null)); + } + + private TestFileStore createStore(TableSchema schema, RowType rowType, int numBuckets) { + return new TestFileStore.Builder( + "avro", + tempDir.toString(), + numBuckets, + TestKeyValueGenerator.DEFAULT_PART_TYPE, + TestKeyValueGenerator.KEY_TYPE, + rowType, + TestKeyValueGenerator.TestKeyValueFieldsExtractor.EXTRACTOR, + DeduplicateMergeFunction.factory(), + schema) + .build(); + } + + private static Object readField(Object target, String name) throws Exception { + Field field = target.getClass().getDeclaredField(name); + field.setAccessible(true); + return field.get(target); + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/BtreeGlobalIndexTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/BtreeGlobalIndexTableTest.java index 9ae96662002d..5f6a9152fd00 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/BtreeGlobalIndexTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/BtreeGlobalIndexTableTest.java @@ -19,13 +19,18 @@ package org.apache.paimon.table; import org.apache.paimon.CoreOptions; +import org.apache.paimon.Snapshot; import org.apache.paimon.data.BinaryString; import org.apache.paimon.data.GenericRow; import org.apache.paimon.globalindex.DataEvolutionBatchScan; +import org.apache.paimon.globalindex.GlobalIndexCoverage; import org.apache.paimon.globalindex.GlobalIndexResult; import org.apache.paimon.globalindex.GlobalIndexScanner; import org.apache.paimon.globalindex.IndexedSplit; import org.apache.paimon.globalindex.sorted.SortedGlobalIndexBuilder; +import org.apache.paimon.index.GlobalIndexMeta; +import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.manifest.IndexManifestEntry; import org.apache.paimon.partition.PartitionPredicate; import org.apache.paimon.predicate.Predicate; import org.apache.paimon.predicate.PredicateBuilder; @@ -214,6 +219,60 @@ public void testGlobalIndexScannerKeepsUnindexedRowsSeparate() throws Exception } } + @Test + public void testSourceBackedIndexIsExcludedFromGlobalRowIdScan() throws Exception { + write(10L); + FileStoreTable table = + tableWithSearchMode((FileStoreTable) catalog.getTable(identifier()), "full"); + IndexFileMeta sourceBacked = + new IndexFileMeta( + "btree", + "source-backed-index", + 0, + 10, + new GlobalIndexMeta(0, 9, 1, null, null, new byte[] {1}), + null); + + assertThat(GlobalIndexScanner.create(table, Collections.singletonList(sourceBacked))) + .isEmpty(); + + GlobalIndexCoverage coverage = + new GlobalIndexCoverage( + table, + table.snapshotManager().latestSnapshot(), + PartitionPredicate.ALWAYS_TRUE, + Collections.singletonList(sourceBacked)); + assertThat(coverage.unindexedRanges(1)).containsExactly(new Range(0, 9)); + } + + @Test + public void testOrdinaryAndSourceBackedBTreeIndexesCanCoexist() throws Exception { + write(10L); + createIndex("f1"); + FileStoreTable table = (FileStoreTable) catalog.getTable(identifier()); + Snapshot snapshot = table.snapshotManager().latestSnapshot(); + List mixedIndexes = + table.store().newIndexFileHandler().scan(snapshot, "btree").stream() + .map(IndexManifestEntry::indexFile) + .collect(Collectors.toCollection(ArrayList::new)); + mixedIndexes.add( + new IndexFileMeta( + "btree", + "source-backed-index", + 0, + 10, + new GlobalIndexMeta(0, 9, 1, null, null, new byte[] {1}), + null)); + Predicate predicate = + new PredicateBuilder(table.rowType()).equal(1, BinaryString.fromString("a7")); + + try (GlobalIndexScanner scanner = + GlobalIndexScanner.create(table, mixedIndexes).orElseThrow(AssertionError::new)) { + assertThat(scanner.scan(predicate).get().results().toRangeList()) + .containsExactly(new Range(7, 7)); + } + } + @Test public void testBTreeGlobalIndexSearchModeUsesAllPredicateFieldCoverage() throws Exception { write(500L);