diff --git a/paimon-core/src/main/java/org/apache/paimon/globalindex/DataEvolutionBatchScan.java b/paimon-core/src/main/java/org/apache/paimon/globalindex/DataEvolutionBatchScan.java index ca83c6b497aa..57be42ff7370 100644 --- a/paimon-core/src/main/java/org/apache/paimon/globalindex/DataEvolutionBatchScan.java +++ b/paimon-core/src/main/java/org/apache/paimon/globalindex/DataEvolutionBatchScan.java @@ -33,8 +33,8 @@ import org.apache.paimon.predicate.RowIdPredicateVisitor; import org.apache.paimon.predicate.TopN; import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.source.AppendBatchTableScan; import org.apache.paimon.table.source.DataSplit; -import org.apache.paimon.table.source.DataTableBatchScan; import org.apache.paimon.table.source.DataTableScan; import org.apache.paimon.table.source.InnerTableScan; import org.apache.paimon.table.source.Split; @@ -65,13 +65,13 @@ public class DataEvolutionBatchScan implements DataTableScan { private static final Logger LOG = LoggerFactory.getLogger(DataEvolutionBatchScan.class); private final FileStoreTable table; - private final DataTableBatchScan batchScan; + private final AppendBatchTableScan batchScan; private Predicate filter; private RowRangeIndex pushedRowRangeIndex; private GlobalIndexResult globalIndexResult; - public DataEvolutionBatchScan(FileStoreTable wrapped, DataTableBatchScan batchScan) { + public DataEvolutionBatchScan(FileStoreTable wrapped, AppendBatchTableScan batchScan) { this.table = wrapped; this.batchScan = batchScan; } diff --git a/paimon-core/src/main/java/org/apache/paimon/index/pk/PrimaryKeyIndexSourceMeta.java b/paimon-core/src/main/java/org/apache/paimon/index/pk/PrimaryKeyIndexSourceMeta.java index 2bc0169a9d74..072829e96c04 100644 --- a/paimon-core/src/main/java/org/apache/paimon/index/pk/PrimaryKeyIndexSourceMeta.java +++ b/paimon-core/src/main/java/org/apache/paimon/index/pk/PrimaryKeyIndexSourceMeta.java @@ -89,7 +89,16 @@ public static PrimaryKeyIndexSourceMeta deserialize(byte[] bytes) { checkArgument(version == VERSION, "Unsupported index source version: %s.", version); int sourceFileCount = input.readInt(); checkArgument(sourceFileCount > 0, "An index must reference source files."); - List sourceFiles = new ArrayList<>(sourceFileCount); + // Each entry needs at least the two-byte writeUTF length and one long. + int maximumSourceFileCount = input.available() / (Short.BYTES + Long.BYTES); + checkArgument( + sourceFileCount <= maximumSourceFileCount, + "Failed to deserialize index source metadata: source file count %s " + + "exceeds the maximum %s allowed by the remaining bytes.", + sourceFileCount, + maximumSourceFileCount); + List sourceFiles = + new ArrayList<>(Math.min(sourceFileCount, 1024)); for (int i = 0; i < sourceFileCount; i++) { sourceFiles.add(new PrimaryKeyIndexSourceFile(input.readUTF(), input.readLong())); } diff --git a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentSearcher.java b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentSearcher.java index e7fd34bd41fc..797afd104caf 100644 --- a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentSearcher.java +++ b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentSearcher.java @@ -128,6 +128,26 @@ public List search( Map deletionVectors, Set activeSourceFiles, Map searchOptions) { + return search( + segment, + sourceMeta, + query, + limit, + deletionVectors, + activeSourceFiles, + Collections.emptyMap(), + searchOptions); + } + + public List search( + IndexFileMeta segment, + PrimaryKeyIndexSourceMeta sourceMeta, + float[] query, + int limit, + Map deletionVectors, + Set activeSourceFiles, + Map> rowRangesByFile, + Map searchOptions) { checkArgument(limit > 0, "Vector search limit must be positive: %s.", limit); GlobalIndexMeta globalIndexMeta = segment.globalIndexMeta(); checkArgument( @@ -161,7 +181,11 @@ public List search( try { VectorSearch search = new VectorSearch(query, limit, vectorField.name(), searchOptions); RoaringNavigableMap64 liveRows = - liveRowPositions(sourceMeta.sourceFiles(), activeSourceFiles, deletionVectors); + liveRowPositions( + sourceMeta.sourceFiles(), + activeSourceFiles, + deletionVectors, + rowRangesByFile); if (liveRows != null) { search.withIncludeRowIds(liveRows); } @@ -193,6 +217,11 @@ public List search( "ANN segment %s returned snapshot-deleted row position %s.", segment.fileName(), filePosition.rowPosition); + List rowRanges = rowRangesByFile.get(filePosition.dataFileName); + checkArgument( + rowRanges == null || contains(rowRanges, filePosition.rowPosition), + "ANN segment %s returned a row outside the pre-filter.", + segment.fileName()); candidates.add( new PkVectorSearchResult( filePosition.dataFileName, @@ -211,7 +240,8 @@ public List search( private static RoaringNavigableMap64 liveRowPositions( List sourceFiles, Set activeSourceFiles, - Map deletionVectors) { + Map deletionVectors, + Map> rowRangesByFile) { boolean allSourcesActive = true; for (PrimaryKeyIndexSourceFile sourceFile : sourceFiles) { if (!activeSourceFiles.contains(sourceFile.fileName())) { @@ -219,7 +249,7 @@ private static RoaringNavigableMap64 liveRowPositions( break; } } - if (allSourcesActive && deletionVectors.isEmpty()) { + if (allSourcesActive && deletionVectors.isEmpty() && rowRangesByFile.isEmpty()) { return null; } RoaringNavigableMap64 live = new RoaringNavigableMap64(); @@ -228,7 +258,18 @@ private static RoaringNavigableMap64 liveRowPositions( for (PrimaryKeyIndexSourceFile sourceFile : sourceFiles) { boolean active = activeSourceFiles.contains(sourceFile.fileName()); if (active && sourceFile.rowCount() > 0) { - live.addRange(new Range(fileOffset, fileOffset + sourceFile.rowCount() - 1)); + List rowRanges = rowRangesByFile.get(sourceFile.fileName()); + if (rowRanges == null) { + live.addRange(new Range(fileOffset, fileOffset + sourceFile.rowCount() - 1)); + } else { + for (Range range : rowRanges) { + checkArgument( + range.from >= 0 && range.to < sourceFile.rowCount(), + "Pre-filter range is outside source file %s.", + sourceFile.fileName()); + live.addRange(range.addOffset(fileOffset)); + } + } } DeletionVector deletionVector = active ? deletionVectors.get(sourceFile.fileName()) : null; @@ -242,6 +283,23 @@ private static RoaringNavigableMap64 liveRowPositions( return live; } + private static boolean contains(List ranges, long position) { + int low = 0; + int high = ranges.size() - 1; + while (low <= high) { + int middle = (low + high) >>> 1; + Range range = ranges.get(middle); + if (position < range.from) { + high = middle - 1; + } else if (position > range.to) { + low = middle + 1; + } else { + return true; + } + } + return false; + } + private static long totalRowCount(List sourceFiles) { long total = 0; for (PrimaryKeyIndexSourceFile sourceFile : sourceFiles) { diff --git a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearch.java b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearch.java index 04d4a327552a..d864a3a3fb3f 100644 --- a/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearch.java +++ b/paimon-core/src/main/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearch.java @@ -24,6 +24,7 @@ import org.apache.paimon.index.pk.PrimaryKeyIndexSourceFile; import org.apache.paimon.index.pk.PrimaryKeyIndexSourceMeta; import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.utils.Range; import javax.annotation.Nullable; @@ -95,6 +96,25 @@ public Result search( int indexedLimit, int exactLimit) throws IOException { + return search( + state, + activeFiles, + deletionVectors, + Collections.emptyMap(), + query, + indexedLimit, + exactLimit); + } + + public Result search( + PkVectorBucketIndexState state, + List activeFiles, + Map deletionVectors, + Map> rowRangesByFile, + float[] query, + int indexedLimit, + int exactLimit) + throws IOException { checkArgument(indexedLimit > 0, "Vector indexed search limit must be positive."); checkArgument(exactLimit > 0, "Vector exact search limit must be positive."); Map filesByName = new HashMap<>(); @@ -121,15 +141,26 @@ public Result search( covered.add(source.fileName()); } checkArgument(annSearcher != null, "ANN search is not configured."); - for (PkVectorSearchResult result : - annSearcher.search( - ann, - sourceMeta, - query, - indexedLimit, - deletionVectors, - activeSourceFiles, - searchOptions)) { + List annResults = + rowRangesByFile.isEmpty() + ? annSearcher.search( + ann, + sourceMeta, + query, + indexedLimit, + deletionVectors, + activeSourceFiles, + searchOptions) + : annSearcher.search( + ann, + sourceMeta, + query, + indexedLimit, + deletionVectors, + activeSourceFiles, + rowRangesByFile, + searchOptions); + for (PkVectorSearchResult result : annResults) { add(indexedNearest, result, indexedLimit); } } @@ -140,7 +171,14 @@ public Result search( continue; } DeletionVector dv = deletionVectors.get(file.fileName()); - LongPredicate excluded = dv == null ? position -> false : dv::isDeleted; + List rowRanges = rowRangesByFile.get(file.fileName()); + if (rowRanges != null && rowRanges.isEmpty()) { + continue; + } + LongPredicate excluded = + position -> + (dv != null && dv.isDeleted(position)) + || (rowRanges != null && !contains(rowRanges, position)); try (PkVectorReader reader = vectorReaderFactory.create(file)) { for (PkVectorSearchResult result : PkVectorExactSearcher.search( @@ -153,6 +191,23 @@ public Result search( return new Result(sorted(indexedNearest), sorted(exactNearest)); } + private static boolean contains(List ranges, long position) { + int low = 0; + int high = ranges.size() - 1; + while (low <= high) { + int middle = (low + high) >>> 1; + Range range = ranges.get(middle); + if (position < range.from) { + high = middle - 1; + } else if (position > range.to) { + low = middle + 1; + } else { + return true; + } + } + return false; + } + private static List sorted(PriorityQueue nearest) { List result = new ArrayList<>(nearest); Collections.sort(result, BEST_FIRST); diff --git a/paimon-core/src/main/java/org/apache/paimon/privilege/PrivilegedFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/privilege/PrivilegedFileStoreTable.java index 35dfd308df0a..e1064e70458a 100644 --- a/paimon-core/src/main/java/org/apache/paimon/privilege/PrivilegedFileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/privilege/PrivilegedFileStoreTable.java @@ -228,6 +228,12 @@ public DataTableScan newScan() { return wrapped.newScan(); } + @Override + public DataTableScan newScan(SnapshotReaderFactory snapshotReaderFactory) { + privilegeChecker.assertCanSelect(identifier); + return wrapped.newScan(snapshotReaderFactory); + } + @Override public StreamDataTableScan newStreamScan() { privilegeChecker.assertCanSelect(identifier); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/AbstractFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/AbstractFileStoreTable.java index 3154e8734129..cf8388fc55ee 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/AbstractFileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/AbstractFileStoreTable.java @@ -24,7 +24,6 @@ import org.apache.paimon.consumer.ConsumerManager; import org.apache.paimon.fs.FileIO; import org.apache.paimon.fs.Path; -import org.apache.paimon.globalindex.DataEvolutionBatchScan; import org.apache.paimon.manifest.IndexManifestEntry; import org.apache.paimon.manifest.ManifestEntry; import org.apache.paimon.manifest.ManifestFileMeta; @@ -45,8 +44,6 @@ import org.apache.paimon.table.sink.RowKindGenerator; import org.apache.paimon.table.sink.TableCommitImpl; import org.apache.paimon.table.sink.WriteSelector; -import org.apache.paimon.table.source.DataTableBatchScan; -import org.apache.paimon.table.source.DataTableScan; import org.apache.paimon.table.source.DataTableStreamScan; import org.apache.paimon.table.source.SplitGenerator; import org.apache.paimon.table.source.StreamDataTableScan; @@ -287,26 +284,6 @@ public SnapshotReader newSnapshotReader() { dvmetaCache); } - @Override - public DataTableScan newScan() { - DataTableBatchScan scan = - new DataTableBatchScan( - tableSchema, - schemaManager(), - coreOptions(), - newSnapshotReader(), - catalogEnvironment.tableQueryAuth(coreOptions())); - Integer scanBucket = coreOptions().scanBucket(); - if (scanBucket != null) { - DataTableBatchScan.validateScanBucketOption(tableSchema, coreOptions(), scanBucket); - scan.withBucket(scanBucket); - } - if (coreOptions().dataEvolutionEnabled()) { - return new DataEvolutionBatchScan(this, scan); - } - return scan; - } - @Override public StreamDataTableScan newStreamScan() { DataTableStreamScan scan = @@ -321,7 +298,6 @@ public StreamDataTableScan newStreamScan() { !tableSchema.primaryKeys().isEmpty()); Integer scanBucket = coreOptions().scanBucket(); if (scanBucket != null) { - DataTableBatchScan.validateScanBucketOption(tableSchema, coreOptions(), scanBucket); scan.withBucket(scanBucket); } return scan; diff --git a/paimon-core/src/main/java/org/apache/paimon/table/AppendOnlyFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/AppendOnlyFileStoreTable.java index d65c84fd5e65..4d35147c17a4 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/AppendOnlyFileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/AppendOnlyFileStoreTable.java @@ -23,6 +23,7 @@ import org.apache.paimon.data.InternalRow; import org.apache.paimon.fs.FileIO; import org.apache.paimon.fs.Path; +import org.apache.paimon.globalindex.DataEvolutionBatchScan; import org.apache.paimon.operation.AppendOnlyFileStoreScan; import org.apache.paimon.operation.BaseAppendFileStoreWrite; import org.apache.paimon.operation.FileStoreScan; @@ -30,10 +31,12 @@ import org.apache.paimon.schema.TableSchema; import org.apache.paimon.table.query.LocalTableQuery; import org.apache.paimon.table.sink.TableWriteImpl; +import org.apache.paimon.table.source.AppendBatchTableScan; import org.apache.paimon.table.source.AppendOnlySplitGenerator; import org.apache.paimon.table.source.AppendTableRead; import org.apache.paimon.table.source.DataEvolutionSplitGenerator; import org.apache.paimon.table.source.DataEvolutionTableRead; +import org.apache.paimon.table.source.DataTableScan; import org.apache.paimon.table.source.InnerTableRead; import org.apache.paimon.table.source.SplitGenerator; import org.apache.paimon.table.source.splitread.AppendTableRawFileSplitReadProvider; @@ -133,6 +136,24 @@ public InnerTableRead newRead() { : new AppendTableRead(providerFactories, schema()); } + @Override + public DataTableScan newScan() { + return newScan(FileStoreTable::newSnapshotReader); + } + + @Override + public DataTableScan newScan(SnapshotReaderFactory snapshotReaderFactory) { + CoreOptions options = coreOptions(); + AppendBatchTableScan scan = + new AppendBatchTableScan( + schema(), + schemaManager(), + options, + snapshotReaderFactory.create(this), + catalogEnvironment.tableQueryAuth(options)); + return options.dataEvolutionEnabled() ? new DataEvolutionBatchScan(this, scan) : scan; + } + @Override public TableWriteImpl newWrite(String commitUser) { return newWrite(commitUser, null); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/ChainGroupReadTable.java b/paimon-core/src/main/java/org/apache/paimon/table/ChainGroupReadTable.java index 1428d26e9d5f..535b4a575b52 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/ChainGroupReadTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/ChainGroupReadTable.java @@ -56,6 +56,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.function.Function; import java.util.stream.Collectors; import static org.apache.paimon.utils.Preconditions.checkArgument; @@ -102,12 +103,21 @@ public DataTableScan newScan() { return new ChainTableBatchScan(((AbstractFileStoreTable) wrapped).tableSchema, this); } - private DataTableScan newSnapshotScan() { - return wrapped.newScan(); + @Override + public DataTableScan newScan(SnapshotReaderFactory snapshotReaderFactory) { + super.validateSchema(); + return new ChainTableBatchScan( + ((AbstractFileStoreTable) wrapped).tableSchema, + this, + table -> table.newScan(snapshotReaderFactory)); + } + + private DataTableScan newSnapshotScan(Function scanCreator) { + return scanCreator.apply(wrapped); } - private DataTableScan newDeltaScan() { - return other().newScan(); + private DataTableScan newDeltaScan(Function scanCreator) { + return scanCreator.apply(other()); } @Override @@ -207,11 +217,18 @@ public static class ChainTableBatchScan extends FallbackReadScan { public ChainTableBatchScan( TableSchema tableSchema, ChainGroupReadTable chainGroupReadTable) { + this(tableSchema, chainGroupReadTable, FileStoreTable::newScan); + } + + private ChainTableBatchScan( + TableSchema tableSchema, + ChainGroupReadTable chainGroupReadTable, + Function scanCreator) { super( chainGroupReadTable.wrapped, chainGroupReadTable.other(), tableSchema, - FileStoreTable::newScan); + scanCreator); this.options = CoreOptions.fromMap(tableSchema.options()); this.chainGroupReadTable = chainGroupReadTable; this.partitionConverter = @@ -507,8 +524,8 @@ private DataTableScan newChainPartitionListingScan( boolean snapshot, PartitionPredicate scanPartitionPredicate) { DataTableScan scan = snapshot - ? chainGroupReadTable.newSnapshotScan() - : chainGroupReadTable.newDeltaScan(); + ? chainGroupReadTable.newSnapshotScan(scanCreator) + : chainGroupReadTable.newDeltaScan(scanCreator); if (scanPartitionPredicate != null) { scan.withPartitionFilter(scanPartitionPredicate); } @@ -546,8 +563,8 @@ private Set preloadTargetSnapshotSplits(List splits) { private DataTableScan newFilteredScan(boolean snapshot) { DataTableScan scan = snapshot - ? chainGroupReadTable.newSnapshotScan() - : chainGroupReadTable.newDeltaScan(); + ? chainGroupReadTable.newSnapshotScan(scanCreator) + : chainGroupReadTable.newDeltaScan(scanCreator); if (dataPredicate != null) { scan.withFilter(dataPredicate); } diff --git a/paimon-core/src/main/java/org/apache/paimon/table/DelegatedFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/DelegatedFileStoreTable.java index 7ba4bc20a9d7..f9bcbeee01c0 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/DelegatedFileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/DelegatedFileStoreTable.java @@ -315,6 +315,11 @@ public DataTableScan newScan() { return wrapped.newScan(); } + @Override + public DataTableScan newScan(SnapshotReaderFactory snapshotReaderFactory) { + return wrapped.newScan(snapshotReaderFactory); + } + @Override public StreamDataTableScan newStreamScan() { return wrapped.newStreamScan(); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/FallbackReadFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/FallbackReadFileStoreTable.java index b4521cfc216c..bc84bca88afe 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/FallbackReadFileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/FallbackReadFileStoreTable.java @@ -212,10 +212,15 @@ protected Map rewriteOtherOptions(Map options) { @Override public DataTableScan newScan() { - return newScan(FileStoreTable::newScan); + return newFallbackScan(FileStoreTable::newScan); } - public DataTableScan newScan(Function scanCreator) { + @Override + public DataTableScan newScan(SnapshotReaderFactory snapshotReaderFactory) { + return newFallbackScan(table -> table.newScan(snapshotReaderFactory)); + } + + public DataTableScan newFallbackScan(Function scanCreator) { validateSchema(); FileStoreTable first = wrappedFirst ? wrapped : other; FileStoreTable second = wrappedFirst ? other : wrapped; diff --git a/paimon-core/src/main/java/org/apache/paimon/table/FileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/FileStoreTable.java index b07465a25828..5a3f87d0edf1 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/FileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/FileStoreTable.java @@ -32,6 +32,8 @@ import org.apache.paimon.table.sink.RowKeyExtractor; import org.apache.paimon.table.sink.TableCommitImpl; import org.apache.paimon.table.sink.TableWriteImpl; +import org.apache.paimon.table.source.DataTableScan; +import org.apache.paimon.table.source.snapshot.SnapshotReader; import org.apache.paimon.tag.TagAutoManager; import org.apache.paimon.types.RowType; import org.apache.paimon.utils.BranchManager; @@ -55,6 +57,13 @@ */ public interface FileStoreTable extends DataTable { + /** Factory to create a snapshot reader for the actual table scanned. */ + @FunctionalInterface + interface SnapshotReaderFactory { + + SnapshotReader create(FileStoreTable table); + } + void setManifestCache(SegmentsCache manifestCache); @Nullable @@ -103,6 +112,9 @@ default Optional comment() { FileStore store(); + /** Creates a scan with a customized snapshot reader. */ + DataTableScan newScan(SnapshotReaderFactory snapshotReaderFactory); + CatalogEnvironment catalogEnvironment(); @Override diff --git a/paimon-core/src/main/java/org/apache/paimon/table/PrimaryKeyFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/PrimaryKeyFileStoreTable.java index a2fee49bfb88..3030b3504e1d 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/PrimaryKeyFileStoreTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/PrimaryKeyFileStoreTable.java @@ -33,9 +33,11 @@ import org.apache.paimon.schema.TableSchema; import org.apache.paimon.table.query.LocalTableQuery; import org.apache.paimon.table.sink.TableWriteImpl; +import org.apache.paimon.table.source.DataTableScan; import org.apache.paimon.table.source.InnerTableRead; import org.apache.paimon.table.source.KeyValueTableRead; import org.apache.paimon.table.source.MergeTreeSplitGenerator; +import org.apache.paimon.table.source.PrimaryKeyBatchScan; import org.apache.paimon.table.source.SplitGenerator; import org.apache.paimon.types.RowType; import org.apache.paimon.utils.RowKindFilter; @@ -150,6 +152,20 @@ public InnerTableRead newRead() { () -> store().newRead(), () -> store().newBatchRawFileRead(), schema()); } + @Override + public DataTableScan newScan() { + return newScan(FileStoreTable::newSnapshotReader); + } + + @Override + public DataTableScan newScan(SnapshotReaderFactory snapshotReaderFactory) { + return new PrimaryKeyBatchScan( + this, + snapshotReaderFactory.create(this), + catalogEnvironment.tableQueryAuth(coreOptions()), + null); + } + @Override public TableWriteImpl newWrite(String commitUser) { return newWrite(commitUser, null); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/DataTableBatchScan.java b/paimon-core/src/main/java/org/apache/paimon/table/source/AbstractBatchTableScan.java similarity index 79% rename from paimon-core/src/main/java/org/apache/paimon/table/source/DataTableBatchScan.java rename to paimon-core/src/main/java/org/apache/paimon/table/source/AbstractBatchTableScan.java index 4ba3b19b6b9d..ef9d289b6b4e 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/DataTableBatchScan.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/AbstractBatchTableScan.java @@ -19,7 +19,6 @@ package org.apache.paimon.table.source; import org.apache.paimon.CoreOptions; -import org.apache.paimon.globalindex.GlobalIndexResult; import org.apache.paimon.manifest.PartitionEntry; import org.apache.paimon.predicate.Predicate; import org.apache.paimon.predicate.SortValue; @@ -48,16 +47,10 @@ import static org.apache.paimon.table.source.PushDownUtils.minmaxAvailable; -/** {@link TableScan} implementation for batch planning. */ -public class DataTableBatchScan extends AbstractDataTableScan { +/** Base {@link TableScan} implementation for batch planning. */ +public abstract class AbstractBatchTableScan extends AbstractDataTableScan { - private static final Logger LOG = LoggerFactory.getLogger(DataTableBatchScan.class); - - /** Validates {@link CoreOptions#SCAN_BUCKET} for primary-key fixed-bucket tables. */ - public static void validateScanBucketOption( - TableSchema schema, CoreOptions coreOptions, int bucket) { - AbstractDataTableScan.validateScanBucketOption(schema, coreOptions, bucket); - } + private static final Logger LOG = LoggerFactory.getLogger(AbstractBatchTableScan.class); private StartingScanner startingScanner; private boolean hasNext; @@ -67,9 +60,8 @@ public static void validateScanBucketOption( private final SchemaManager schemaManager; @Nullable private String readProtectionTagName; - @Nullable private GlobalIndexSplitResult globalIndexSplitResult; - public DataTableBatchScan( + protected AbstractBatchTableScan( TableSchema schema, SchemaManager schemaManager, CoreOptions options, @@ -89,6 +81,10 @@ public DataTableBatchScan( if (options.bucket() == BucketMode.POSTPONE_BUCKET) { snapshotReader.onlyReadRealBuckets(); } + Integer scanBucket = options.scanBucket(); + if (scanBucket != null) { + snapshotReader.withBucket(scanBucket); + } } @Override @@ -111,49 +107,44 @@ public InnerTableScan withTopN(TopN topN) { } @Override - protected TableScan.Plan planWithoutAuth() { - if (globalIndexSplitResult != null) { - if (!hasNext) { - throw new EndOfScanException(); - } - hasNext = false; - if (globalIndexSplitResult.snapshotId() > 0) { - maybeCreateReadProtectionTag(globalIndexSplitResult.snapshotId()); - } - List splits = new ArrayList<>(globalIndexSplitResult.splits()); - return new PlanImpl(null, globalIndexSplitResult.snapshotId(), splits); + protected final TableScan.Plan planWithoutAuth() { + if (!hasNext) { + throw new EndOfScanException(); } + hasNext = false; + + Plan preProcessedPlan = preProcessPlan(); + if (preProcessedPlan != null) { + return preProcessedPlan; + } + if (startingScanner == null) { startingScanner = createStartingScanner(false); } - if (hasNext) { - hasNext = false; - StartingScanner.Result result; - Optional pushed = applyPushDownLimit(); - if (pushed.isPresent()) { - result = pushed.get(); - } else { - pushed = applyPushDownTopN(); - result = pushed.orElseGet(() -> startingScanner.scan(snapshotReader)); - } - - if (result instanceof ScannedResult) { - maybeCreateReadProtectionTag(((ScannedResult) result).currentSnapshotId()); - } - - return DataFilePlan.fromResult(result); + StartingScanner.Result result; + Optional pushed = applyPushDownLimit(); + if (pushed.isPresent()) { + result = pushed.get(); } else { - throw new EndOfScanException(); + pushed = applyPushDownTopN(); + result = pushed.orElseGet(() -> startingScanner.scan(snapshotReader)); } - } - @Override - public DataTableBatchScan withGlobalIndexResult(GlobalIndexResult globalIndexResult) { - if (globalIndexResult instanceof GlobalIndexSplitResult) { - this.globalIndexSplitResult = (GlobalIndexSplitResult) globalIndexResult; + if (result instanceof ScannedResult) { + maybeCreateReadProtectionTag(((ScannedResult) result).currentSnapshotId()); } - return this; + + return postProcessPlan(DataFilePlan.fromResult(result)); + } + + @Nullable + protected Plan preProcessPlan() { + return null; + } + + protected Plan postProcessPlan(Plan plan) { + return plan; } @Override @@ -263,7 +254,7 @@ public String readProtectionTagName() { return readProtectionTagName; } - private void maybeCreateReadProtectionTag(long snapshotId) { + protected final void maybeCreateReadProtectionTag(long snapshotId) { Duration timeRetained = options().scanPlanAutoTagTimeRetained(); if (timeRetained == null) { return; diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/AbstractDataTableScan.java b/paimon-core/src/main/java/org/apache/paimon/table/source/AbstractDataTableScan.java index 7bf5e6860351..6695e5cefe84 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/AbstractDataTableScan.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/AbstractDataTableScan.java @@ -31,7 +31,6 @@ import org.apache.paimon.partition.PartitionPredicate; import org.apache.paimon.predicate.Predicate; import org.apache.paimon.schema.TableSchema; -import org.apache.paimon.table.BucketMode; import org.apache.paimon.table.source.snapshot.CompactedStartingScanner; import org.apache.paimon.table.source.snapshot.ContinuousCompactorStartingScanner; import org.apache.paimon.table.source.snapshot.ContinuousFromSnapshotFullStartingScanner; @@ -92,7 +91,7 @@ abstract class AbstractDataTableScan implements DataTableScan { // Last applied auth predicate; guards redundant re-application across plan()s. @Nullable private Predicate appliedAuthPredicate; // Whether the auth predicate has a non-partition part (enforced only at read time). Used by - // DataTableBatchScan to disable limit push down; not pushed through withFilter. + // AbstractBatchTableScan to disable limit push down; not pushed through withFilter. protected boolean authHasNonPartitionFilter; protected AbstractDataTableScan( @@ -146,7 +145,8 @@ private void applyAuthFilter(@Nullable Predicate authPredicate) { // changed/removed auth leaves no stale pruning. The full filter is enforced at read time. snapshotReader.manifestsReader().withAuthPartitionFilter(authPartitionFilter); // A non-partition auth part is enforced only at read time, so limit push down is unsafe - // (DataTableBatchScan reads this). Kept off SnapshotReader since it is not a pushed filter. + // (AbstractBatchTableScan reads this). Kept off SnapshotReader since it is not a pushed + // filter. this.authHasNonPartitionFilter = hasNonPartitionPart; } @@ -162,43 +162,6 @@ public AbstractDataTableScan withBucket(int bucket) { return this; } - /** Validates {@link CoreOptions#SCAN_BUCKET} for primary-key fixed-bucket tables. */ - static void validateScanBucketOption(TableSchema schema, CoreOptions coreOptions, int bucket) { - checkArgument( - !schema.primaryKeys().isEmpty(), - "Bucket scan is only supported for primary key tables."); - checkArgument( - bucketModeFromOption(coreOptions.bucket()) == BucketMode.HASH_FIXED, - "Bucket scan is only supported for fixed-bucket tables, but got bucket mode %s.", - bucketModeFromOption(coreOptions.bucket())); - validateFixedBucketRange(coreOptions, bucket); - } - - private static void validateFixedBucketRange(CoreOptions coreOptions, int bucket) { - checkArgument(bucket >= 0, "Bucket id must be non-negative, but is %s.", bucket); - int numBuckets = coreOptions.bucket(); - checkArgument( - numBuckets > 0, - "Bucket scan is only supported for tables with bucket > 0, but got bucket %s.", - numBuckets); - checkArgument( - bucket < numBuckets, - "Bucket id %s must be less than table bucket number %s.", - bucket, - numBuckets); - } - - private static BucketMode bucketModeFromOption(int bucketOption) { - switch (bucketOption) { - case -2: - return BucketMode.POSTPONE_MODE; - case -1: - return BucketMode.HASH_DYNAMIC; - default: - return BucketMode.HASH_FIXED; - } - } - @Override public AbstractDataTableScan withBucketFilter(Filter bucketFilter) { snapshotReader.withBucketFilter(bucketFilter); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/AppendBatchTableScan.java b/paimon-core/src/main/java/org/apache/paimon/table/source/AppendBatchTableScan.java new file mode 100644 index 000000000000..6570bb66980a --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/AppendBatchTableScan.java @@ -0,0 +1,37 @@ +/* + * 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.table.source; + +import org.apache.paimon.CoreOptions; +import org.apache.paimon.schema.SchemaManager; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.source.snapshot.SnapshotReader; + +/** Batch scan for append tables. */ +public class AppendBatchTableScan extends AbstractBatchTableScan { + + public AppendBatchTableScan( + TableSchema schema, + SchemaManager schemaManager, + CoreOptions options, + SnapshotReader snapshotReader, + TableQueryAuth queryAuth) { + super(schema, schemaManager, options, snapshotReader, queryAuth); + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/BucketVectorSearchSplit.java b/paimon-core/src/main/java/org/apache/paimon/table/source/BucketVectorSearchSplit.java index b21d1605b55e..31c2f7b0cba1 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/BucketVectorSearchSplit.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/BucketVectorSearchSplit.java @@ -22,13 +22,16 @@ import org.apache.paimon.index.IndexFileMetaSerializer; import org.apache.paimon.io.DataInputViewStreamWrapper; import org.apache.paimon.io.DataOutputViewStreamWrapper; +import org.apache.paimon.utils.Range; import java.io.IOException; import java.io.ObjectInputStream; import java.io.ObjectOutputStream; import java.util.ArrayList; import java.util.Collections; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; import java.util.Objects; import static org.apache.paimon.utils.Preconditions.checkArgument; @@ -41,8 +44,16 @@ public class BucketVectorSearchSplit extends VectorSearchSplit { private DataSplit dataSplit; private transient List payloadFiles; + private Map> rowRangesByFile; public BucketVectorSearchSplit(DataSplit dataSplit, List payloadFiles) { + this(dataSplit, payloadFiles, Collections.emptyMap()); + } + + public BucketVectorSearchSplit( + DataSplit dataSplit, + List payloadFiles, + Map> rowRangesByFile) { this.dataSplit = dataSplit; for (IndexFileMeta payload : payloadFiles) { checkArgument( @@ -52,6 +63,13 @@ public BucketVectorSearchSplit(DataSplit dataSplit, List payloadF payload.fileName()); } this.payloadFiles = Collections.unmodifiableList(new ArrayList<>(payloadFiles)); + Map> ranges = new LinkedHashMap<>(); + for (Map.Entry> entry : rowRangesByFile.entrySet()) { + ranges.put( + entry.getKey(), + Collections.unmodifiableList(new ArrayList<>(entry.getValue()))); + } + this.rowRangesByFile = Collections.unmodifiableMap(ranges); } public DataSplit dataSplit() { @@ -62,6 +80,10 @@ public List payloadFiles() { return payloadFiles; } + public Map> rowRangesByFile() { + return rowRangesByFile; + } + private void writeObject(ObjectOutputStream out) throws IOException { out.defaultWriteObject(); out.writeInt(VERSION); @@ -100,11 +122,12 @@ public boolean equals(Object o) { } BucketVectorSearchSplit that = (BucketVectorSearchSplit) o; return Objects.equals(dataSplit, that.dataSplit) - && Objects.equals(payloadFiles, that.payloadFiles); + && Objects.equals(payloadFiles, that.payloadFiles) + && Objects.equals(rowRangesByFile, that.rowRangesByFile); } @Override public int hashCode() { - return Objects.hash(dataSplit, payloadFiles); + return Objects.hash(dataSplit, payloadFiles, rowRangesByFile); } } diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyBatchScan.java b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyBatchScan.java new file mode 100644 index 000000000000..ac25c7692093 --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyBatchScan.java @@ -0,0 +1,171 @@ +/* + * 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.table.source; + +import org.apache.paimon.Snapshot; +import org.apache.paimon.globalindex.GlobalIndexResult; +import org.apache.paimon.index.GlobalIndexMeta; +import org.apache.paimon.index.IndexFileHandler; +import org.apache.paimon.index.pk.PrimaryKeyIndexDefinition; +import org.apache.paimon.index.pk.PrimaryKeyIndexDefinitions; +import org.apache.paimon.manifest.FileKind; +import org.apache.paimon.manifest.IndexManifestEntry; +import org.apache.paimon.predicate.Predicate; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.source.snapshot.SnapshotReader; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +/** Batch scan for primary-key tables and indexes. */ +public class PrimaryKeyBatchScan extends AbstractBatchTableScan { + + private final FileStoreTable table; + private final @Nullable PrimaryKeySortedIndexScan.ReaderFactory readerFactory; + + @Nullable private Predicate filter; + @Nullable private GlobalIndexSplitResult globalIndexSplitResult; + + public PrimaryKeyBatchScan( + FileStoreTable table, + SnapshotReader snapshotReader, + TableQueryAuth queryAuth, + @Nullable PrimaryKeySortedIndexScan.ReaderFactory readerFactory) { + super( + table.schema(), + table.schemaManager(), + table.coreOptions(), + snapshotReader, + queryAuth); + this.table = table; + this.readerFactory = readerFactory; + } + + @Override + public PrimaryKeyBatchScan withFilter(Predicate predicate) { + this.filter = predicate; + super.withFilter(predicate); + return this; + } + + @Override + public PrimaryKeyBatchScan withGlobalIndexResult(GlobalIndexResult globalIndexResult) { + if (globalIndexResult instanceof GlobalIndexSplitResult) { + this.globalIndexSplitResult = (GlobalIndexSplitResult) globalIndexResult; + } + return this; + } + + @Override + @Nullable + protected Plan preProcessPlan() { + if (globalIndexSplitResult == null) { + return null; + } + if (globalIndexSplitResult.snapshotId() > 0) { + maybeCreateReadProtectionTag(globalIndexSplitResult.snapshotId()); + } + List splits = new ArrayList<>(globalIndexSplitResult.splits()); + return new PlanImpl(null, globalIndexSplitResult.snapshotId(), splits); + } + + @Override + protected Plan postProcessPlan(Plan dataPlan) { + if (!(dataPlan instanceof SnapshotReader.Plan)) { + return dataPlan; + } + SnapshotReader.Plan snapshotPlan = (SnapshotReader.Plan) dataPlan; + if (filter == null + || !options().globalIndexEnabled() + || !snapshotReader.hasNonPartitionFilter() + || table.schema().primaryKeys().isEmpty() + || !options().deletionVectorsEnabled() + || options().deletionVectorsMergeOnRead() + || options().bucket() <= 0 + || snapshotPlan.snapshotId() == null + || snapshotPlan.splits().isEmpty()) { + return dataPlan; + } + + List dataSplits = new ArrayList<>(); + for (Split split : snapshotPlan.splits()) { + if (!(split instanceof DataSplit) || ((DataSplit) split).isStreaming()) { + return dataPlan; + } + dataSplits.add((DataSplit) split); + } + + long snapshotId = snapshotPlan.snapshotId(); + Snapshot snapshot = snapshotReader.snapshotManager().snapshot(snapshotId); + if (snapshot == null) { + return dataPlan; + } + TableSchema snapshotSchema = table.schemaManager().schema(snapshot.schemaId()); + List definitions = + PrimaryKeyIndexDefinitions.create(snapshotSchema).definitions(); + Set scalarFields = new HashSet<>(); + for (PrimaryKeyIndexDefinition definition : definitions) { + if (definition.family() == PrimaryKeyIndexDefinition.Family.BTREE + || definition.family() == PrimaryKeyIndexDefinition.Family.BITMAP) { + scalarFields.add(definition.fieldId()); + } + } + if (scalarFields.isEmpty()) { + return dataPlan; + } + + IndexFileHandler indexFileHandler = snapshotReader.indexFileHandler(); + if (indexFileHandler == null) { + return dataPlan; + } + List indexEntries = + indexFileHandler.scan( + snapshot, + entry -> { + GlobalIndexMeta meta = entry.indexFile().globalIndexMeta(); + return entry.kind() == FileKind.ADD + && meta != null + && meta.sourceMeta() != null + && scalarFields.contains(meta.indexFieldId()); + }); + PrimaryKeySortedIndexScan.Plan indexPlan = + PrimaryKeySortedIndexScan.plan(snapshotId, dataSplits, definitions, indexEntries); + PrimaryKeySortedIndexScan.ReaderFactory factory = + readerFactory == null + ? PrimaryKeySortedIndexScan.readerFactory( + snapshotReader.snapshotManager().fileIO(), + snapshotReader.pathFactory(), + snapshotSchema.logicalRowType(), + options().toConfiguration()) + : readerFactory; + PrimaryKeySortedIndexScan.EvaluatedPlan evaluated = + PrimaryKeySortedIndexScan.evaluate( + indexPlan, snapshotSchema.logicalRowType(), filter, definitions, factory); + PrimaryKeySortedIndexResult result = new PrimaryKeySortedIndexResult(evaluated); + return new PlanImpl( + snapshotPlan.watermark(), + snapshotPlan.snapshotId(), + new ArrayList<>(result.splits())); + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeySortedIndexResult.java b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeySortedIndexResult.java new file mode 100644 index 000000000000..a29bab53d9fc --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeySortedIndexResult.java @@ -0,0 +1,143 @@ +/* + * 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.table.source; + +import org.apache.paimon.globalindex.GlobalIndexResult; +import org.apache.paimon.globalindex.IndexedSplit; +import org.apache.paimon.utils.Range; +import org.apache.paimon.utils.RoaringNavigableMap64; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Optional; +import java.util.Set; + +/** Snapshot-scoped scalar-index result addressed by physical data-file row positions. */ +public class PrimaryKeySortedIndexResult implements GlobalIndexSplitResult { + + private static final Logger LOG = LoggerFactory.getLogger(PrimaryKeySortedIndexResult.class); + private static final int MAX_INDEXED_RANGES_PER_FILE = 4096; + + private final long snapshotId; + private final List splits; + + PrimaryKeySortedIndexResult(PrimaryKeySortedIndexScan.EvaluatedPlan plan) { + this.snapshotId = plan.snapshotId(); + List converted = new ArrayList<>(); + Set preservedNonRawSplits = Collections.newSetFromMap(new IdentityHashMap<>()); + for (PrimaryKeySortedIndexScan.EvaluatedFile evaluatedFile : plan.files()) { + PrimaryKeySortedIndexScan.FilePlan file = evaluatedFile.file(); + DataSplit sourceSplit = file.sourceSplit(); + if (!sourceSplit.rawConvertible()) { + if (preservedNonRawSplits.add(sourceSplit)) { + converted.add(sourceSplit); + } + continue; + } + + Optional result = evaluatedFile.result(); + if (!result.isPresent()) { + converted.add(toSingleFileSplit(file)); + continue; + } + + RoaringNavigableMap64 positions = result.get().results(); + if (positions.isEmpty()) { + continue; + } + List ranges = ranges(positions, file.dataFile().rowCount()); + if (ranges == null) { + LOG.warn( + "Primary-key sorted index returned an invalid row position for data file " + + "{}; falling back to a raw scan for this file.", + file.dataFile().fileName()); + converted.add(toSingleFileSplit(file)); + } else { + converted.add(new IndexedSplit(toSingleFileSplit(file), ranges, null)); + } + } + this.splits = Collections.unmodifiableList(converted); + } + + @Override + public long snapshotId() { + return snapshotId; + } + + @Override + public List splits() { + return splits; + } + + @Override + public RoaringNavigableMap64 results() { + throw new UnsupportedOperationException( + "Primary-key sorted-index results use physical file positions, not global row ids."); + } + + private static DataSplit toSingleFileSplit(PrimaryKeySortedIndexScan.FilePlan file) { + DataSplit source = file.sourceSplit(); + DataSplit.Builder builder = + DataSplit.builder() + .withSnapshot(source.snapshotId()) + .withPartition(source.partition()) + .withBucket(source.bucket()) + .withBucketPath(source.bucketPath()) + .withTotalBuckets(source.totalBuckets()) + .withDataFiles(Collections.singletonList(file.dataFile())) + .isStreaming(false) + .rawConvertible(false); + if (source.deletionFiles().isPresent()) { + builder.withDataDeletionFiles( + Collections.singletonList(source.deletionFiles().get().get(file.fileIndex()))); + } + return builder.build(); + } + + private static List ranges(RoaringNavigableMap64 positions, long rowCount) { + List ranges = new ArrayList<>(); + long from = -1; + long to = -1; + for (long position : positions) { + if (position < 0 || position >= rowCount || position > Integer.MAX_VALUE) { + return null; + } + if (from < 0) { + from = position; + } else if (position != to + 1) { + if (ranges.size() >= MAX_INDEXED_RANGES_PER_FILE) { + return null; + } + ranges.add(new Range(from, to)); + from = position; + } + to = position; + } + if (ranges.size() >= MAX_INDEXED_RANGES_PER_FILE) { + return null; + } + ranges.add(new Range(from, to)); + return ranges; + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeySortedIndexScan.java b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeySortedIndexScan.java new file mode 100644 index 000000000000..9e8c842fd2f8 --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeySortedIndexScan.java @@ -0,0 +1,358 @@ +/* + * 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.table.source; + +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.fs.FileIO; +import org.apache.paimon.globalindex.GlobalIndexEvaluator; +import org.apache.paimon.globalindex.GlobalIndexIOMeta; +import org.apache.paimon.globalindex.GlobalIndexReadThreadPool; +import org.apache.paimon.globalindex.GlobalIndexReader; +import org.apache.paimon.globalindex.GlobalIndexResult; +import org.apache.paimon.globalindex.GlobalIndexer; +import org.apache.paimon.globalindex.io.GlobalIndexFileReader; +import org.apache.paimon.index.GlobalIndexMeta; +import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.index.IndexPathFactory; +import org.apache.paimon.index.pk.PrimaryKeyIndexDefinition; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceFile; +import org.apache.paimon.index.pksorted.PkSortedBucketIndexState; +import org.apache.paimon.index.pksorted.PkSortedIndexGroup; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.manifest.FileKind; +import org.apache.paimon.manifest.IndexManifestEntry; +import org.apache.paimon.options.Options; +import org.apache.paimon.predicate.Predicate; +import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.FileStorePathFactory; +import org.apache.paimon.utils.IndexFilePathFactories; +import org.apache.paimon.utils.Pair; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ExecutorService; + +import static org.apache.paimon.CoreOptions.GLOBAL_INDEX_THREAD_NUM; +import static org.apache.paimon.utils.Preconditions.checkArgument; +import static org.apache.paimon.utils.Preconditions.checkNotNull; + +/** Plans source-backed scalar index groups in file-local row-position space. */ +public final class PrimaryKeySortedIndexScan { + + private static final Logger LOG = LoggerFactory.getLogger(PrimaryKeySortedIndexScan.class); + + private PrimaryKeySortedIndexScan() {} + + /** Factory to create a sorted-index reader for a source data file. */ + @FunctionalInterface + public interface ReaderFactory { + + GlobalIndexReader create( + FilePlan file, PrimaryKeyIndexDefinition definition, List payloads); + } + + static ReaderFactory readerFactory( + FileIO fileIO, FileStorePathFactory pathFactory, RowType rowType, Options options) { + IndexFilePathFactories pathFactories = new IndexFilePathFactories(pathFactory); + ExecutorService executor = + GlobalIndexReadThreadPool.getExecutorService(options.get(GLOBAL_INDEX_THREAD_NUM)); + GlobalIndexFileReader fileReader = meta -> fileIO.newInputStream(meta.filePath()); + return (file, definition, payloads) -> { + IndexPathFactory indexPathFactory = + pathFactories.get(file.sourceSplit().partition(), file.sourceSplit().bucket()); + List ioMetas = new ArrayList<>(payloads.size()); + for (IndexFileMeta payload : payloads) { + GlobalIndexMeta meta = checkNotNull(payload.globalIndexMeta()); + ioMetas.add( + new GlobalIndexIOMeta( + indexPathFactory.toPath(payload), + payload.fileSize(), + meta.indexMeta())); + } + GlobalIndexer indexer = + GlobalIndexer.create( + definition.indexType(), + rowType.getField(definition.fieldId()), + definition.options()); + return indexer.createReader(fileReader, ioMetas, executor); + }; + } + + static Plan plan( + long snapshotId, + List dataSplits, + List definitions, + List indexEntries) { + Map, List> payloadsByBucket = new LinkedHashMap<>(); + for (IndexManifestEntry entry : indexEntries) { + IndexFileMeta payload = entry.indexFile(); + GlobalIndexMeta meta = payload.globalIndexMeta(); + if (entry.kind() != FileKind.ADD || meta == null || meta.sourceMeta() == null) { + continue; + } + Pair bucket = Pair.of(entry.partition(), entry.bucket()); + payloadsByBucket.computeIfAbsent(bucket, ignored -> new ArrayList<>()).add(payload); + } + + List scalarDefinitions = new ArrayList<>(); + for (PrimaryKeyIndexDefinition definition : definitions) { + if (definition.family() == PrimaryKeyIndexDefinition.Family.BTREE + || definition.family() == PrimaryKeyIndexDefinition.Family.BITMAP) { + scalarDefinitions.add(definition); + } + } + + Map, List> sourcesByBucket = + new LinkedHashMap<>(); + for (DataSplit split : dataSplits) { + checkArgument( + split.snapshotId() == snapshotId, + "Data split snapshot %s does not match sorted-index scan snapshot %s.", + split.snapshotId(), + snapshotId); + checkArgument( + !split.isStreaming(), "Primary-key sorted-index scan requires batch splits."); + List deletions = split.deletionFiles().orElse(null); + checkArgument( + deletions == null || deletions.size() == split.dataFiles().size(), + "Deletion files must align with data files in a sorted-index split."); + List sources = + sourcesByBucket.computeIfAbsent( + Pair.of(split.partition(), split.bucket()), + ignored -> new ArrayList<>()); + for (DataFileMeta dataFile : split.dataFiles()) { + sources.add( + new PrimaryKeyIndexSourceFile(dataFile.fileName(), dataFile.rowCount())); + } + } + + Map, Map>> + groupsByBucket = new LinkedHashMap<>(); + for (Map.Entry, List> bucketEntry : + sourcesByBucket.entrySet()) { + Pair bucket = bucketEntry.getKey(); + List bucketPayloads = + payloadsByBucket.getOrDefault(bucket, Collections.emptyList()); + Map> groupsBySource = new LinkedHashMap<>(); + for (PrimaryKeyIndexDefinition definition : scalarDefinitions) { + List definitionPayloads = new ArrayList<>(); + for (IndexFileMeta payload : bucketPayloads) { + GlobalIndexMeta meta = payload.globalIndexMeta(); + if (meta != null + && definition.indexType().equals(payload.indexType()) + && definition.fieldId() == meta.indexFieldId()) { + definitionPayloads.add(payload); + } + } + try { + PkSortedBucketIndexState state = + PkSortedBucketIndexState.fromActivePayloads( + definition.fieldId(), + definition.indexType(), + bucketEntry.getValue(), + definitionPayloads); + for (PkSortedIndexGroup group : state.groups()) { + groupsBySource + .computeIfAbsent( + group.sourceFile().fileName(), + ignored -> new LinkedHashMap<>()) + .put(definition.fieldId(), group); + } + } catch (RuntimeException e) { + rethrowIfInterrupted(e); + LOG.warn( + "Failed to plan primary-key sorted index for partition {}, bucket {} " + + "and field {}; falling back to a raw scan for this field.", + bucket.getKey(), + bucket.getValue(), + definition.fieldId(), + e); + } + } + groupsByBucket.put(bucket, groupsBySource); + } + + List files = new ArrayList<>(); + for (DataSplit split : dataSplits) { + Map> groupsBySource = + groupsByBucket.getOrDefault( + Pair.of(split.partition(), split.bucket()), Collections.emptyMap()); + for (int fileIndex = 0; fileIndex < split.dataFiles().size(); fileIndex++) { + DataFileMeta dataFile = split.dataFiles().get(fileIndex); + Map groups = + groupsBySource.getOrDefault(dataFile.fileName(), Collections.emptyMap()); + files.add(new FilePlan(split, fileIndex, groups)); + } + } + return new Plan(snapshotId, files); + } + + static EvaluatedPlan evaluate( + Plan plan, + RowType rowType, + Predicate predicate, + List definitions, + ReaderFactory readerFactory) { + Map definitionsByField = new LinkedHashMap<>(); + for (PrimaryKeyIndexDefinition definition : definitions) { + if (definition.family() == PrimaryKeyIndexDefinition.Family.BTREE + || definition.family() == PrimaryKeyIndexDefinition.Family.BITMAP) { + definitionsByField.put(definition.fieldId(), definition); + } + } + + List files = new ArrayList<>(); + for (FilePlan file : plan.files()) { + GlobalIndexEvaluator evaluator = + new GlobalIndexEvaluator( + rowType, + fieldId -> { + PrimaryKeyIndexDefinition definition = + definitionsByField.get(fieldId); + Optional group = file.group(fieldId); + if (definition == null || !group.isPresent()) { + return Collections.emptyList(); + } + return Collections.singletonList( + readerFactory.create( + file, definition, group.get().payloads())); + }); + Optional result; + try { + result = evaluator.evaluate(predicate); + } catch (RuntimeException e) { + rethrowIfInterrupted(e); + LOG.warn( + "Failed to evaluate primary-key sorted index for data file {}; " + + "falling back to a raw scan for this file.", + file.dataFile().fileName(), + e); + result = Optional.empty(); + } finally { + evaluator.close(); + } + files.add(new EvaluatedFile(file, result)); + } + return new EvaluatedPlan(plan.snapshotId(), files); + } + + private static void rethrowIfInterrupted(RuntimeException exception) { + if (Thread.currentThread().isInterrupted()) { + throw exception; + } + } + + /** Immutable groups for all source files in one captured snapshot. */ + public static final class Plan { + + private final long snapshotId; + private final List files; + + private Plan(long snapshotId, List files) { + this.snapshotId = snapshotId; + this.files = Collections.unmodifiableList(new ArrayList<>(files)); + } + + public long snapshotId() { + return snapshotId; + } + + public List files() { + return files; + } + } + + /** One active data file and its complete field-local payload groups. */ + public static final class FilePlan { + + private final DataSplit sourceSplit; + private final int fileIndex; + private final Map groups; + + private FilePlan( + DataSplit sourceSplit, int fileIndex, Map groups) { + this.sourceSplit = sourceSplit; + this.fileIndex = fileIndex; + this.groups = Collections.unmodifiableMap(new LinkedHashMap<>(groups)); + } + + public DataFileMeta dataFile() { + return sourceSplit.dataFiles().get(fileIndex); + } + + public Optional group(int fieldId) { + return Optional.ofNullable(groups.get(fieldId)); + } + + DataSplit sourceSplit() { + return sourceSplit; + } + + int fileIndex() { + return fileIndex; + } + } + + /** Predicate results for all source files in one captured snapshot. */ + public static final class EvaluatedPlan { + + private final long snapshotId; + private final List files; + + private EvaluatedPlan(long snapshotId, List files) { + this.snapshotId = snapshotId; + this.files = Collections.unmodifiableList(new ArrayList<>(files)); + } + + public long snapshotId() { + return snapshotId; + } + + public List files() { + return files; + } + } + + /** Optional file-local index result; empty means that the file requires a raw scan. */ + public static final class EvaluatedFile { + + private final FilePlan file; + private final Optional result; + + private EvaluatedFile(FilePlan file, Optional result) { + this.file = file; + this.result = result; + } + + public FilePlan file() { + return file; + } + + public Optional result() { + return result; + } + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java index 6daf3591488d..315050bb5903 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java @@ -37,6 +37,8 @@ import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.io.KeyValueFileReaderFactory; import org.apache.paimon.options.Options; +import org.apache.paimon.predicate.Predicate; +import org.apache.paimon.predicate.PredicateVisitor; import org.apache.paimon.reader.RecordReader; import org.apache.paimon.reader.ScoreRecordIterator; import org.apache.paimon.reader.ScoreRecordReader; @@ -45,6 +47,10 @@ import org.apache.paimon.types.DataField; import org.apache.paimon.types.RowType; import org.apache.paimon.types.VectorType; +import org.apache.paimon.utils.Range; +import org.apache.paimon.utils.RoaringNavigableMap64; + +import javax.annotation.Nullable; import java.io.IOException; import java.io.Serializable; @@ -52,11 +58,13 @@ import java.util.Collections; import java.util.Comparator; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.PriorityQueue; +import java.util.Set; import java.util.concurrent.ExecutorService; import java.util.stream.Collectors; @@ -97,6 +105,7 @@ public class PrimaryKeyVectorRead implements VectorRead, Serializable { private final String metric; private final int refineFactor; private final int indexedLimit; + @Nullable private final Predicate filter; public PrimaryKeyVectorRead( FileStoreTable table, @@ -104,6 +113,16 @@ public PrimaryKeyVectorRead( float[] query, int limit, Map searchOptions) { + this(table, vectorField, query, limit, searchOptions, null); + } + + public PrimaryKeyVectorRead( + FileStoreTable table, + DataField vectorField, + float[] query, + int limit, + Map searchOptions, + @Nullable Predicate filter) { checkArgument( vectorField.type() instanceof VectorType, "Vector field must use VECTOR type."); checkArgument( @@ -125,6 +144,7 @@ public PrimaryKeyVectorRead( VectorSearchRefineOptions.resolve( this.searchOptions, table.options(), vectorField.name(), indexType); this.indexedLimit = VectorSearchRefineOptions.searchLimit(limit, refineFactor); + this.filter = filter; } private static KeyValueFileStore keyValueStore(FileStoreTable table) { @@ -243,12 +263,102 @@ private SearchResult search(BucketVectorSearchSplit split, SearchContext context table.coreOptions().globalIndexSearchMode()); PrimaryKeyVectorBucketSearch.Result result = bucketSearch.search( - state, activeFiles, deletionVectors, query, indexedLimit, limit); + state, + activeFiles, + deletionVectors, + rowRangesByFile(split), + query, + indexedLimit, + limit); return new SearchResult( candidates(dataSplit, result.indexedCandidates()), candidates(dataSplit, result.exactCandidates())); } + private Map> rowRangesByFile(BucketVectorSearchSplit split) + throws IOException { + Map> result = new LinkedHashMap<>(split.rowRangesByFile()); + if (filter == null) { + return result; + } + + DataSplit dataSplit = split.dataSplit(); + for (int i = 0; i < dataSplit.dataFiles().size(); i++) { + DataFileMeta dataFile = dataSplit.dataFiles().get(i); + result.put( + dataFile.fileName(), + residualRowRanges(dataSplit, i, result.get(dataFile.fileName()))); + } + return result; + } + + private List residualRowRanges( + DataSplit source, int fileIndex, @Nullable List candidateRanges) + throws IOException { + DataFileMeta dataFile = source.dataFiles().get(fileIndex); + if (dataFile.rowCount() == 0) { + return Collections.emptyList(); + } + if (candidateRanges != null && candidateRanges.isEmpty()) { + return Collections.emptyList(); + } + + DataSplit.Builder builder = + DataSplit.builder() + .withSnapshot(source.snapshotId()) + .withPartition(source.partition()) + .withBucket(source.bucket()) + .withBucketPath(source.bucketPath()) + .withTotalBuckets(source.totalBuckets()) + .withDataFiles(Collections.singletonList(dataFile)) + .isStreaming(false) + .rawConvertible(false); + if (source.deletionFiles().isPresent()) { + builder.withDataDeletionFiles( + Collections.singletonList(source.deletionFiles().get().get(fileIndex))); + } + IndexedSplit allRows = + new IndexedSplit( + builder.build(), + candidateRanges == null + ? Collections.singletonList(new Range(0, dataFile.rowCount() - 1)) + : candidateRanges, + null); + ReadBuilder readBuilder = + table.newReadBuilder().withReadType(filterReadType()).withFilter(filter); + RoaringNavigableMap64 positions = new RoaringNavigableMap64(); + try (RecordReader reader = + readBuilder.newRead().executeFilter().createReader(allRows)) { + RecordReader.RecordIterator batch; + while ((batch = reader.readBatch()) != null) { + checkArgument( + batch instanceof ScoreRecordIterator, + "Residual primary-key vector filter requires physical row positions."); + ScoreRecordIterator positionsBatch = + (ScoreRecordIterator) batch; + try { + while (positionsBatch.next() != null) { + positions.add(positionsBatch.returnedRowId()); + } + } finally { + positionsBatch.releaseBatch(); + } + } + } + return positions.toRangeList(); + } + + private RowType filterReadType() { + Set filterFields = PredicateVisitor.collectFieldNames(filter); + List readFields = new ArrayList<>(); + for (String field : table.rowType().getFieldNames()) { + if (filterFields.contains(field)) { + readFields.add(field); + } + } + return table.rowType().project(readFields); + } + private static List candidates( DataSplit split, List searchResults) { List candidates = new ArrayList<>(searchResults.size()); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorScan.java b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorScan.java index 221d4afaafc2..e4bcb15d3c5f 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorScan.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorScan.java @@ -20,17 +20,18 @@ import org.apache.paimon.Snapshot; import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.globalindex.IndexedSplit; import org.apache.paimon.index.IndexFileHandler; import org.apache.paimon.index.IndexFileMeta; import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.manifest.FileKind; import org.apache.paimon.manifest.IndexManifestEntry; import org.apache.paimon.partition.PartitionPredicate; -import org.apache.paimon.table.BucketMode; +import org.apache.paimon.predicate.Predicate; import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.source.snapshot.SnapshotReader; -import org.apache.paimon.table.source.snapshot.TimeTravelUtil; import org.apache.paimon.utils.Pair; +import org.apache.paimon.utils.Range; import javax.annotation.Nullable; @@ -51,16 +52,27 @@ public class PrimaryKeyVectorScan implements VectorScan { private final int vectorFieldId; private final String indexType; @Nullable private final PartitionPredicate partitionFilter; + @Nullable private final Predicate filter; public PrimaryKeyVectorScan( FileStoreTable table, int vectorFieldId, String indexType, @Nullable PartitionPredicate partitionFilter) { + this(table, vectorFieldId, indexType, partitionFilter, null); + } + + public PrimaryKeyVectorScan( + FileStoreTable table, + int vectorFieldId, + String indexType, + @Nullable PartitionPredicate partitionFilter, + @Nullable Predicate filter) { this.table = table; this.vectorFieldId = vectorFieldId; this.indexType = indexType; this.partitionFilter = partitionFilter; + this.filter = filter; } @Override @@ -68,22 +80,39 @@ public Plan scan() { checkArgument( table.coreOptions().primaryKeyVectorIndexEnabled(), "Primary-key vector search requires a configured primary-key vector index."); - @Nullable Snapshot snapshot = TimeTravelUtil.tryTravelOrLatest(table); - if (snapshot == null) { - return new Plan(0, Collections.emptyList()); + checkArgument( + filter == null + || (table.coreOptions().deletionVectorsEnabled() + && !table.coreOptions().deletionVectorsMergeOnRead()), + "Primary-key vector pre-filter requires deletion vectors without merge-on-read."); + SnapshotReader snapshotReader = table.newSnapshotReader().keepStats(); + DataTableScan dataScan = table.newScan(ignored -> snapshotReader); + checkArgument( + dataScan instanceof PrimaryKeyBatchScan, + "Primary-key vector search requires a primary-key batch scan."); + PrimaryKeyBatchScan batchScan = (PrimaryKeyBatchScan) dataScan; + if (partitionFilter != null) { + batchScan.withPartitionFilter(partitionFilter); } - - SnapshotReader snapshotReader = - table.newSnapshotReader().withSnapshot(snapshot).withMode(ScanMode.ALL).keepStats(); - if (table.coreOptions().bucket() == BucketMode.POSTPONE_BUCKET) { - snapshotReader.onlyReadRealBuckets(); + if (filter != null) { + batchScan.withFilter(filter); } - if (partitionFilter != null) { - snapshotReader.withPartitionFilter(partitionFilter); + TableScan.Plan tablePlan = batchScan.planWithoutAuth(); + if (!(tablePlan instanceof SnapshotReader.Plan)) { + checkArgument( + tablePlan.splits().isEmpty(), + "Primary-key vector search requires a snapshot plan."); + return new Plan(0, Collections.emptyList()); + } + SnapshotReader.Plan snapshotPlan = (SnapshotReader.Plan) tablePlan; + if (snapshotPlan.snapshotId() == null) { + return new Plan(0, Collections.emptyList()); } - List dataSplits = snapshotReader.read().dataSplits(); + Snapshot snapshot = snapshotReader.snapshotManager().snapshot(snapshotPlan.snapshotId()); + checkArgument(snapshot != null, "Primary-key vector snapshot does not exist."); - IndexFileHandler indexFileHandler = table.store().newIndexFileHandler(); + IndexFileHandler indexFileHandler = snapshotReader.indexFileHandler(); + checkArgument(indexFileHandler != null, "Primary-key vector index handler is unavailable."); List vectorIndexEntries = indexFileHandler.scan( snapshot, @@ -95,12 +124,12 @@ public Plan scan() { == vectorFieldId && (partitionFilter == null || partitionFilter.test(entry.partition()))); - return plan(snapshot.id(), dataSplits, vectorIndexEntries); + return plan(snapshot.id(), snapshotPlan.splits(), vectorIndexEntries); } static Plan plan( long snapshotId, - List dataSplits, + List dataSplits, List vectorIndexEntries) { Map, List> payloads = new LinkedHashMap<>(); for (IndexManifestEntry entry : vectorIndexEntries) { @@ -118,18 +147,19 @@ static Plan plan( } Map, BucketAccumulator> buckets = new LinkedHashMap<>(); - for (DataSplit split : dataSplits) { + for (Split split : dataSplits) { + DataSplit dataSplit = unwrapDataSplit(split); checkArgument( - split.snapshotId() == snapshotId, + dataSplit.snapshotId() == snapshotId, "Data split snapshot %s does not match vector scan snapshot %s.", - split.snapshotId(), + dataSplit.snapshotId(), snapshotId); checkArgument( - !split.isStreaming(), "Primary-key vector search requires a batch split."); - Pair key = Pair.of(split.partition(), split.bucket()); + !dataSplit.isStreaming(), "Primary-key vector search requires a batch split."); + Pair key = Pair.of(dataSplit.partition(), dataSplit.bucket()); BucketAccumulator accumulator = buckets.get(key); if (accumulator == null) { - accumulator = new BucketAccumulator(split); + accumulator = new BucketAccumulator(dataSplit); buckets.put(key, accumulator); } accumulator.add(split); @@ -140,11 +170,23 @@ static Plan plan( result.add( new BucketVectorSearchSplit( entry.getValue().build(), - payloads.getOrDefault(entry.getKey(), Collections.emptyList()))); + payloads.getOrDefault(entry.getKey(), Collections.emptyList()), + entry.getValue().rowRangesByFile())); } return new Plan(snapshotId, result); } + private static DataSplit unwrapDataSplit(Split split) { + if (split instanceof IndexedSplit) { + return ((IndexedSplit) split).dataSplit(); + } + checkArgument( + split instanceof DataSplit, + "Unsupported primary-key vector source split: %s.", + split.getClass().getName()); + return (DataSplit) split; + } + /** Immutable snapshot vector-search plan. */ public static class Plan implements VectorScan.Plan { @@ -176,6 +218,7 @@ private static class BucketAccumulator { private final List dataFiles = new ArrayList<>(); private final List deletionFiles = new ArrayList<>(); private final Set dataFileNames = new HashSet<>(); + private final Map> rowRangesByFile = new LinkedHashMap<>(); private boolean hasDeletionFile; private BucketAccumulator(DataSplit split) { @@ -186,25 +229,37 @@ private BucketAccumulator(DataSplit split) { this.totalBuckets = split.totalBuckets(); } - private void add(DataSplit split) { + private void add(Split split) { + DataSplit dataSplit; + List rowRanges = null; + if (split instanceof IndexedSplit) { + IndexedSplit indexedSplit = (IndexedSplit) split; + dataSplit = indexedSplit.dataSplit(); + checkArgument( + dataSplit.dataFiles().size() == 1, + "Primary-key vector pre-filter split must contain one data file."); + rowRanges = indexedSplit.rowRanges(); + } else { + dataSplit = unwrapDataSplit(split); + } checkArgument( - snapshotId == split.snapshotId() - && partition.equals(split.partition()) - && bucket == split.bucket() - && bucketPath.equals(split.bucketPath()), + snapshotId == dataSplit.snapshotId() + && partition.equals(dataSplit.partition()) + && bucket == dataSplit.bucket() + && bucketPath.equals(dataSplit.bucketPath()), "Cannot combine data splits from different snapshot buckets."); checkArgument( totalBuckets == null - ? split.totalBuckets() == null - : totalBuckets.equals(split.totalBuckets()), + ? dataSplit.totalBuckets() == null + : totalBuckets.equals(dataSplit.totalBuckets()), "Bucket split total-bucket metadata is inconsistent."); - List splitDeletions = split.deletionFiles().orElse(null); + List splitDeletions = dataSplit.deletionFiles().orElse(null); checkArgument( - splitDeletions == null || splitDeletions.size() == split.dataFiles().size(), + splitDeletions == null || splitDeletions.size() == dataSplit.dataFiles().size(), "Deletion files must align with data files in a bucket split."); - for (int i = 0; i < split.dataFiles().size(); i++) { - DataFileMeta file = split.dataFiles().get(i); + for (int i = 0; i < dataSplit.dataFiles().size(); i++) { + DataFileMeta file = dataSplit.dataFiles().get(i); checkArgument( dataFileNames.add(file.fileName()), "Data file %s appears more than once in vector bucket planning.", @@ -214,6 +269,14 @@ private void add(DataSplit split) { deletionFiles.add(deletion); hasDeletionFile |= deletion != null; } + if (rowRanges != null) { + rowRangesByFile.put( + dataSplit.dataFiles().get(0).fileName(), new ArrayList<>(rowRanges)); + } + } + + private Map> rowRangesByFile() { + return rowRangesByFile; } private DataSplit build() { diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilderImpl.java b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilderImpl.java index e706e40ba6ca..e37f4e9d438a 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilderImpl.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/VectorSearchBuilderImpl.java @@ -33,7 +33,6 @@ import java.util.Optional; import static org.apache.paimon.partition.PartitionPredicate.splitPartitionPredicatesAndDataPredicates; -import static org.apache.paimon.utils.Preconditions.checkArgument; import static org.apache.paimon.utils.Preconditions.checkNotNull; /** Implementation for {@link VectorSearchBuilder}. */ @@ -126,14 +125,12 @@ public VectorSearchBuilder withOption(String key, String value) { @Override public VectorScan newVectorScan() { if (isPrimaryKeyVectorSearch()) { - checkArgument( - filter == null, - "Primary-key vector search does not support non-partition filters."); return new PrimaryKeyVectorScan( table, vectorColumn.id(), table.coreOptions().primaryKeyVectorIndexType(vectorColumn.name()), - partitionFilter); + partitionFilter, + filter); } return new DataEvolutionVectorScan(table, partitionFilter, filter, vectorColumn, options); } @@ -142,10 +139,7 @@ public VectorScan newVectorScan() { public VectorRead newVectorRead() { checkNotNull(vector, "vector must be set via withVector()"); if (isPrimaryKeyVectorSearch()) { - checkArgument( - filter == null, - "Primary-key vector search does not support non-partition filters."); - return new PrimaryKeyVectorRead(table, vectorColumn, vector, limit, options); + return new PrimaryKeyVectorRead(table, vectorColumn, vector, limit, options, filter); } return new DataEvolutionVectorRead( table, partitionFilter, filter, limit, vectorColumn, vector, options); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/snapshot/SnapshotReader.java b/paimon-core/src/main/java/org/apache/paimon/table/source/snapshot/SnapshotReader.java index 90e680e6d13f..7bc7150556c9 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/snapshot/SnapshotReader.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/snapshot/SnapshotReader.java @@ -21,6 +21,7 @@ import org.apache.paimon.Snapshot; import org.apache.paimon.consumer.ConsumerManager; import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.index.IndexFileHandler; import org.apache.paimon.manifest.BucketEntry; import org.apache.paimon.manifest.ManifestEntry; import org.apache.paimon.manifest.ManifestFileMeta; @@ -69,6 +70,9 @@ public interface SnapshotReader { FileStorePathFactory pathFactory(); + @Nullable + IndexFileHandler indexFileHandler(); + SnapshotReader withSnapshot(long snapshotId); SnapshotReader withSnapshot(Snapshot snapshot); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/snapshot/SnapshotReaderImpl.java b/paimon-core/src/main/java/org/apache/paimon/table/source/snapshot/SnapshotReaderImpl.java index 90301d3a0f65..9826965039b2 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/snapshot/SnapshotReaderImpl.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/snapshot/SnapshotReaderImpl.java @@ -177,6 +177,11 @@ public FileStorePathFactory pathFactory() { return pathFactory; } + @Override + public IndexFileHandler indexFileHandler() { + return indexFileHandler; + } + @Override public SnapshotReader withSnapshot(long snapshotId) { scan.withSnapshot(snapshotId); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/system/AuditLogTable.java b/paimon-core/src/main/java/org/apache/paimon/table/system/AuditLogTable.java index 71ca70d8fe27..0725bc1db477 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/system/AuditLogTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/system/AuditLogTable.java @@ -27,6 +27,7 @@ import org.apache.paimon.disk.IOManager; import org.apache.paimon.fs.FileIO; import org.apache.paimon.fs.Path; +import org.apache.paimon.index.IndexFileHandler; import org.apache.paimon.manifest.BucketEntry; import org.apache.paimon.manifest.IndexManifestEntry; import org.apache.paimon.manifest.ManifestEntry; @@ -184,12 +185,20 @@ public List primaryKeys() { @Override public SnapshotReader newSnapshotReader() { - return new AuditLogDataReader(wrapped.newSnapshotReader()); + return newSnapshotReader(wrapped); + } + + private SnapshotReader newSnapshotReader(FileStoreTable table) { + return new AuditLogDataReader(table.newSnapshotReader()); } @Override public DataTableScan newScan() { - return new AuditLogBatchScan(wrapped.newScan()); + return new AuditLogBatchScan(wrapped.newScan(this::newScanSnapshotReader)); + } + + private SnapshotReader newScanSnapshotReader(FileStoreTable table) { + return new AuditLogDataReader(table.newSnapshotReader(), false); } @Override @@ -280,9 +289,15 @@ private Optional convert(Predicate predicate) { private class AuditLogDataReader implements SnapshotReader { private final SnapshotReader wrapped; + private final boolean convertFilter; private AuditLogDataReader(SnapshotReader wrapped) { + this(wrapped, true); + } + + private AuditLogDataReader(SnapshotReader wrapped, boolean convertFilter) { this.wrapped = wrapped; + this.convertFilter = convertFilter; } @Override @@ -325,6 +340,12 @@ public FileStorePathFactory pathFactory() { return wrapped.pathFactory(); } + @Override + @Nullable + public IndexFileHandler indexFileHandler() { + return null; + } + public SnapshotReader withSnapshot(long snapshotId) { wrapped.withSnapshot(snapshotId); return this; @@ -336,12 +357,20 @@ public SnapshotReader withSnapshot(Snapshot snapshot) { } public SnapshotReader withFilter(Predicate predicate) { - convert(predicate).ifPresent(wrapped::withFilter); + if (convertFilter) { + convert(predicate).ifPresent(wrapped::withFilter); + } else { + wrapped.withFilter(predicate); + } return this; } @Override public SnapshotReader withFilter(Predicate predicate, Predicate pushdownPredicate) { + if (!convertFilter) { + wrapped.withFilter(predicate, pushdownPredicate); + return this; + } Optional converted = convert(predicate); Optional convertedPushdown = convert(pushdownPredicate); if (converted.isPresent()) { diff --git a/paimon-core/src/main/java/org/apache/paimon/table/system/ReadOptimizedTable.java b/paimon-core/src/main/java/org/apache/paimon/table/system/ReadOptimizedTable.java index 0a96533a995e..11bb7a353c8a 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/system/ReadOptimizedTable.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/system/ReadOptimizedTable.java @@ -28,11 +28,9 @@ import org.apache.paimon.manifest.ManifestFileMeta; import org.apache.paimon.schema.SchemaManager; import org.apache.paimon.table.DataTable; -import org.apache.paimon.table.FallbackReadFileStoreTable; import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.ReadonlyTable; import org.apache.paimon.table.Table; -import org.apache.paimon.table.source.DataTableBatchScan; import org.apache.paimon.table.source.DataTableScan; import org.apache.paimon.table.source.DataTableStreamScan; import org.apache.paimon.table.source.InnerTableRead; @@ -128,7 +126,7 @@ public SnapshotReader newSnapshotReader() { private SnapshotReader newSnapshotReader(FileStoreTable wrapped) { if (!wrapped.schema().primaryKeys().isEmpty()) { return wrapped.newSnapshotReader() - .withLevel(coreOptions().numLevels() - 1) + .withLevel(wrapped.coreOptions().numLevels() - 1) .enableValueFilter(); } else { return wrapped.newSnapshotReader(); @@ -137,20 +135,7 @@ private SnapshotReader newSnapshotReader(FileStoreTable wrapped) { @Override public DataTableScan newScan() { - if (wrapped instanceof FallbackReadFileStoreTable) { - return ((FallbackReadFileStoreTable) wrapped).newScan(this::newScan); - } - return newScan(wrapped); - } - - private DataTableScan newScan(FileStoreTable wrapped) { - CoreOptions options = wrapped.coreOptions(); - return new DataTableBatchScan( - wrapped.schema(), - schemaManager(), - options, - newSnapshotReader(wrapped), - wrapped.catalogEnvironment().tableQueryAuth(options)); + return wrapped.newScan(this::newSnapshotReader); } @Override diff --git a/paimon-core/src/test/java/org/apache/paimon/globalindex/DataEvolutionBatchScanTest.java b/paimon-core/src/test/java/org/apache/paimon/globalindex/DataEvolutionBatchScanTest.java index b406b9793f35..78771719d068 100644 --- a/paimon-core/src/test/java/org/apache/paimon/globalindex/DataEvolutionBatchScanTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/globalindex/DataEvolutionBatchScanTest.java @@ -22,8 +22,8 @@ import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.predicate.Predicate; import org.apache.paimon.predicate.PredicateBuilder; +import org.apache.paimon.table.source.AppendBatchTableScan; import org.apache.paimon.table.source.DataSplit; -import org.apache.paimon.table.source.DataTableBatchScan; import org.apache.paimon.table.source.DataTableScan; import org.apache.paimon.table.source.Split; import org.apache.paimon.table.source.snapshot.SnapshotReader; @@ -60,7 +60,7 @@ public void testWithFilterKeepsMixedOrWhenRowRangeExtractionFails() { PredicateBuilder builder = new PredicateBuilder(rowTypeWithRowId()); Predicate predicate = PredicateBuilder.or(builder.equal(2, 1L), builder.greaterThan(0, 5)); - DataTableBatchScan batchScan = mock(DataTableBatchScan.class); + AppendBatchTableScan batchScan = mock(AppendBatchTableScan.class); SnapshotReader snapshotReader = mockSnapshotReader(batchScan); new DataEvolutionBatchScan(null, batchScan).withFilter(predicate); @@ -74,7 +74,7 @@ public void testWithFilterRemovesRowIdAfterRowRangeExtractionSucceeds() { Predicate nonRowIdPredicate = builder.greaterThan(0, 5); Predicate predicate = PredicateBuilder.and(builder.equal(2, 1L), nonRowIdPredicate); - DataTableBatchScan batchScan = mock(DataTableBatchScan.class); + AppendBatchTableScan batchScan = mock(AppendBatchTableScan.class); SnapshotReader snapshotReader = mockSnapshotReader(batchScan); new DataEvolutionBatchScan(null, batchScan).withFilter(predicate); @@ -91,7 +91,7 @@ public void testWithFilterDropsNestedMixedOrFromStatsResidual() { Predicate predicate = PredicateBuilder.and(builder.between(2, 0L, 10L), nonRowIdPredicate, mixedOr); - DataTableBatchScan batchScan = mock(DataTableBatchScan.class); + AppendBatchTableScan batchScan = mock(AppendBatchTableScan.class); SnapshotReader snapshotReader = mockSnapshotReader(batchScan); new DataEvolutionBatchScan(null, batchScan).withFilter(predicate); @@ -102,7 +102,7 @@ public void testWithFilterDropsNestedMixedOrFromStatsResidual() { @Test public void testWithShardKeepsDataEvolutionWrapper() { - DataTableBatchScan batchScan = mock(DataTableBatchScan.class); + AppendBatchTableScan batchScan = mock(AppendBatchTableScan.class); when(batchScan.withShard(0, 2)).thenReturn(batchScan); DataEvolutionBatchScan scan = new DataEvolutionBatchScan(null, batchScan); @@ -217,7 +217,7 @@ private static RowType rowTypeWithRowId() { new DataField(2, ROW_ID.name(), DataTypes.BIGINT())); } - private static SnapshotReader mockSnapshotReader(DataTableBatchScan batchScan) { + private static SnapshotReader mockSnapshotReader(AppendBatchTableScan batchScan) { SnapshotReader snapshotReader = mock(SnapshotReader.class); when(batchScan.snapshotReader()).thenReturn(snapshotReader); return snapshotReader; diff --git a/paimon-core/src/test/java/org/apache/paimon/index/pk/PrimaryKeyIndexSourceMetaTest.java b/paimon-core/src/test/java/org/apache/paimon/index/pk/PrimaryKeyIndexSourceMetaTest.java index 9cdcbfd2a5b3..e3bdab3eba38 100644 --- a/paimon-core/src/test/java/org/apache/paimon/index/pk/PrimaryKeyIndexSourceMetaTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/index/pk/PrimaryKeyIndexSourceMetaTest.java @@ -77,6 +77,18 @@ void testRejectsUnsupportedVersion() throws Exception { .hasMessageContaining("Unsupported index source version: 2"); } + @Test + void testRejectsSourceCountBeforeAllocation() throws Exception { + DataOutputSerializer output = new DataOutputSerializer(8); + output.writeInt(1); + output.writeInt(Integer.MAX_VALUE); + + assertThatThrownBy(() -> PrimaryKeyIndexSourceMeta.deserialize(output.getCopyOfBuffer())) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("source file count") + .hasMessageContaining("exceeds the maximum"); + } + @Test void testRejectsTruncatedAndTrailingMetadata() throws Exception { DataOutputSerializer truncated = new DataOutputSerializer(64); diff --git a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentFileTest.java b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentFileTest.java index 76c9fa832cb5..a316af28a7fc 100644 --- a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentFileTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PkVectorAnnSegmentFileTest.java @@ -31,6 +31,7 @@ import org.apache.paimon.stats.SimpleStats; import org.apache.paimon.types.DataField; import org.apache.paimon.types.DataTypes; +import org.apache.paimon.utils.Range; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; @@ -39,6 +40,7 @@ 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.UUID; @@ -103,6 +105,9 @@ void testBuildsAndSearchesMultiSourceSegment() throws Exception { Map deletionVectors = new HashMap<>(); deletionVectors.put("data-2", data2Deletes); + Map> rowRangesByFile = new HashMap<>(); + rowRangesByFile.put("data-1", Collections.singletonList(new Range(1, 1))); + rowRangesByFile.put("data-2", Collections.singletonList(new Range(1, 1))); ExecutorService executor = Executors.newSingleThreadExecutor(); List candidates; @@ -116,6 +121,8 @@ fileIO, annFile, vectorField(), indexOptions(), "l2", executor) new float[] {0, 0}, 3, deletionVectors, + new HashSet<>(Arrays.asList("data-1", "data-2")), + rowRangesByFile, Collections.emptyMap()); } finally { executor.shutdownNow(); @@ -125,7 +132,6 @@ fileIO, annFile, vectorField(), indexOptions(), "l2", executor) .extracting(PkVectorSearchResult::dataFileName, PkVectorSearchResult::rowPosition) .containsExactly( org.assertj.core.groups.Tuple.tuple("data-2", 1L), - org.assertj.core.groups.Tuple.tuple("data-1", 0L), org.assertj.core.groups.Tuple.tuple("data-1", 1L)); } diff --git a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearchTest.java b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearchTest.java index 3437d94f240d..c4c53d6fe4c7 100644 --- a/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearchTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/index/pkvector/PrimaryKeyVectorBucketSearchTest.java @@ -28,6 +28,7 @@ import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.manifest.FileSource; import org.apache.paimon.stats.SimpleStats; +import org.apache.paimon.utils.Range; import org.junit.jupiter.api.Test; @@ -206,8 +207,11 @@ void testExactFallbackMergesFilesAndAppliesDeletionVectors() throws Exception { data1Deletes.delete(0); Map deletionVectors = new HashMap<>(); deletionVectors.put("data-1", data1Deletes); + Map> rowRangesByFile = new HashMap<>(); + rowRangesByFile.put("data-1", Collections.singletonList(new Range(1, 1))); + rowRangesByFile.put("data-2", Collections.singletonList(new Range(1, 1))); - List results = + PrimaryKeyVectorBucketSearch.Result results = new PrimaryKeyVectorBucketSearch( readerFactory, null, @@ -219,17 +223,17 @@ void testExactFallbackMergesFilesAndAppliesDeletionVectors() throws Exception { 7, "test-vector-ann", Collections.emptyList()), Arrays.asList(data1, data2), deletionVectors, + rowRangesByFile, new float[] {0, 0}, + 2, 2); - assertThat(results) + assertThat(results.exactCandidates()) .extracting( PkVectorSearchResult::dataFileName, PkVectorSearchResult::rowPosition, PkVectorSearchResult::distance) - .containsExactly( - org.assertj.core.groups.Tuple.tuple("data-2", 0L, 1F), - org.assertj.core.groups.Tuple.tuple("data-1", 1L, 4F)); + .containsExactly(org.assertj.core.groups.Tuple.tuple("data-1", 1L, 4F)); } private static PkVectorDataFileReader reader(float[][] vectors) throws IOException { diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexBatchScanTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexBatchScanTest.java new file mode 100644 index 000000000000..35de4f3145c0 --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexBatchScanTest.java @@ -0,0 +1,243 @@ +/* + * 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.table.source; + +import org.apache.paimon.CoreOptions; +import org.apache.paimon.Snapshot; +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.globalindex.GlobalIndexReader; +import org.apache.paimon.globalindex.GlobalIndexResult; +import org.apache.paimon.globalindex.IndexedSplit; +import org.apache.paimon.globalindex.btree.BTreeGlobalIndexerFactory; +import org.apache.paimon.index.GlobalIndexMeta; +import org.apache.paimon.index.IndexFileHandler; +import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceFile; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceMeta; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.manifest.FileKind; +import org.apache.paimon.manifest.FileSource; +import org.apache.paimon.manifest.IndexManifestEntry; +import org.apache.paimon.predicate.Predicate; +import org.apache.paimon.predicate.PredicateBuilder; +import org.apache.paimon.schema.SchemaManager; +import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.stats.SimpleStats; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.source.snapshot.SnapshotReader; +import org.apache.paimon.types.DataField; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.utils.Filter; +import org.apache.paimon.utils.Range; +import org.apache.paimon.utils.RoaringNavigableMap64; +import org.apache.paimon.utils.SnapshotManager; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.RETURNS_SELF; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** Tests automatic source-backed BTree/Bitmap evaluation in ordinary batch planning. */ +class PrimaryKeySortedIndexBatchScanTest { + + @Test + void testOrdinaryBatchScanUsesSnapshotScopedSortedIndex() { + ScanFixture fixture = fixture(reader(2)); + + TableScan.Plan result = fixture.scan.plan(); + + assertThat(result.splits()).singleElement().isInstanceOf(IndexedSplit.class); + IndexedSplit indexedSplit = (IndexedSplit) result.splits().get(0); + assertThat(indexedSplit.dataSplit().dataFiles()).containsExactly(fixture.dataFile); + assertThat(indexedSplit.rowRanges()).containsExactly(new Range(2, 2)); + } + + @Test + void testOrdinaryBatchScanFailsWhenApplyingSortedIndexFails() { + ScanFixture fixture = fixture(reader(2)); + when(fixture.scan + .snapshotReader + .indexFileHandler() + .scan( + any(Snapshot.class), + org.mockito.ArgumentMatchers.>any())) + .thenThrow(new RuntimeException("corrupt index")); + + assertThatThrownBy(fixture.scan::plan) + .isInstanceOf(RuntimeException.class) + .hasMessage("corrupt index"); + } + + @Test + void testDisabledGlobalIndexUsesOrdinaryDataPlan() { + ScanFixture fixture = fixture(reader(2), false); + + TableScan.Plan result = fixture.scan.plan(); + + assertThat(result.splits()).singleElement().isInstanceOf(DataSplit.class); + } + + private static TableSchema tableSchema(boolean globalIndexEnabled) { + Map options = new HashMap<>(); + options.put(CoreOptions.BUCKET.key(), "2"); + options.put(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true"); + options.put(CoreOptions.DELETION_VECTORS_MERGE_ON_READ.key(), "false"); + options.put(CoreOptions.GLOBAL_INDEX_ENABLED.key(), Boolean.toString(globalIndexEnabled)); + options.put(CoreOptions.PK_BTREE_INDEX_COLUMNS.key(), "f7"); + return new TableSchema( + 5, + Arrays.asList( + new DataField(1, "id", DataTypes.INT().notNull()), + new DataField(7, "f7", DataTypes.INT())), + 7, + Collections.emptyList(), + Collections.singletonList("id"), + options, + null); + } + + private static ScanFixture fixture(GlobalIndexReader reader) { + return fixture(reader, true); + } + + private static ScanFixture fixture(GlobalIndexReader reader, boolean globalIndexEnabled) { + TableSchema schema = tableSchema(globalIndexEnabled); + CoreOptions options = new CoreOptions(schema.options()); + DataFileMeta dataFile = dataFile("data-1", 4); + DataSplit dataSplit = dataSplit(11, dataFile); + Snapshot snapshot = mock(Snapshot.class); + when(snapshot.id()).thenReturn(11L); + when(snapshot.schemaId()).thenReturn(schema.id()); + SnapshotManager snapshotManager = mock(SnapshotManager.class); + when(snapshotManager.latestSnapshot()).thenReturn(snapshot); + when(snapshotManager.snapshot(11)).thenReturn(snapshot); + SnapshotReader snapshotReader = mock(SnapshotReader.class, RETURNS_SELF); + when(snapshotReader.snapshotManager()).thenReturn(snapshotManager); + when(snapshotReader.hasNonPartitionFilter()).thenReturn(true); + when(snapshotReader.read()) + .thenReturn(new PlanImpl(null, 11L, Collections.singletonList(dataSplit))); + IndexFileHandler indexFileHandler = mock(IndexFileHandler.class); + when(snapshotReader.indexFileHandler()).thenReturn(indexFileHandler); + when(indexFileHandler.scan( + eq(snapshot), + org.mockito.ArgumentMatchers.>any())) + .thenReturn( + Collections.singletonList( + new IndexManifestEntry( + FileKind.ADD, + BinaryRow.EMPTY_ROW, + 0, + payload("btree-0", "data-1", 4)))); + SchemaManager schemaManager = mock(SchemaManager.class); + when(schemaManager.schema(schema.id())).thenReturn(schema); + Predicate predicate = new PredicateBuilder(schema.logicalRowType()).equal(1, 42); + FileStoreTable table = mock(FileStoreTable.class); + when(table.schema()).thenReturn(schema); + when(table.schemaManager()).thenReturn(schemaManager); + when(table.coreOptions()).thenReturn(options); + PrimaryKeyBatchScan scan = + new PrimaryKeyBatchScan( + table, + snapshotReader, + mock(TableQueryAuth.class), + (ignoredFile, ignoredDefinition, ignoredPayloads) -> reader); + scan.withFilter(predicate); + return new ScanFixture(dataFile, scan); + } + + private static GlobalIndexReader reader(long rowPosition) { + RoaringNavigableMap64 positions = new RoaringNavigableMap64(); + positions.add(rowPosition); + GlobalIndexReader reader = mock(GlobalIndexReader.class); + when(reader.visitEqual(any(), eq(42))) + .thenReturn( + CompletableFuture.completedFuture( + Optional.of(GlobalIndexResult.create(positions)))); + return reader; + } + + private static DataSplit dataSplit(long snapshotId, DataFileMeta dataFile) { + return DataSplit.builder() + .withSnapshot(snapshotId) + .withPartition(BinaryRow.EMPTY_ROW) + .withBucket(0) + .withBucketPath("bucket-0") + .withTotalBuckets(2) + .withDataFiles(Collections.singletonList(dataFile)) + .withDataDeletionFiles(Collections.singletonList(null)) + .isStreaming(false) + .rawConvertible(true) + .build(); + } + + private static DataFileMeta dataFile(String fileName, long rowCount) { + return DataFileMeta.forAppend( + fileName, + 100, + rowCount, + SimpleStats.EMPTY_STATS, + 0, + 0, + 1, + Collections.emptyList(), + null, + FileSource.COMPACT, + null, + null, + null, + null); + } + + private static IndexFileMeta payload(String fileName, String sourceName, long sourceRowCount) { + byte[] sourceMeta = + new PrimaryKeyIndexSourceMeta( + new PrimaryKeyIndexSourceFile(sourceName, sourceRowCount)) + .serialize(); + return new IndexFileMeta( + BTreeGlobalIndexerFactory.IDENTIFIER, + fileName, + 100, + sourceRowCount, + new GlobalIndexMeta(0, sourceRowCount - 1, 7, null, new byte[] {1}, sourceMeta), + null); + } + + private static class ScanFixture { + + private final DataFileMeta dataFile; + private final PrimaryKeyBatchScan scan; + + private ScanFixture(DataFileMeta dataFile, PrimaryKeyBatchScan scan) { + this.dataFile = dataFile; + this.scan = scan; + } + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexReadTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexReadTest.java new file mode 100644 index 000000000000..0cca7fb5e2f3 --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexReadTest.java @@ -0,0 +1,149 @@ +/* + * 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.table.source; + +import org.apache.paimon.CoreOptions; +import org.apache.paimon.Snapshot; +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.data.BinaryString; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.data.InternalRow; +import org.apache.paimon.globalindex.IndexedSplit; +import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceMeta; +import org.apache.paimon.predicate.Predicate; +import org.apache.paimon.predicate.PredicateBuilder; +import org.apache.paimon.reader.RecordReader; +import org.apache.paimon.schema.Schema; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.TableTestBase; +import org.apache.paimon.types.DataTypes; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** End-to-end reads for source-backed scalar indexes, residual filters, and deletion vectors. */ +class PrimaryKeySortedIndexReadTest extends TableTestBase { + + @Override + protected Schema schemaDefault() { + return Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("score", DataTypes.INT()) + .column("tag", DataTypes.STRING()) + .primaryKey("id") + .option(CoreOptions.BUCKET.key(), "1") + .option(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true") + .option(CoreOptions.PK_BTREE_INDEX_COLUMNS.key(), "score") + .option( + "fields.score.pk-btree.index.options", + "{\"sorted-index.records-per-range\":\"2\"}") + .build(); + } + + @Test + void testDeletionVectorAndResidualPredicateRemainActive() throws Exception { + createTableDefault(); + FileStoreTable table = getTableDefault(); + BinaryString keep = BinaryString.fromString("keep"); + BinaryString drop = BinaryString.fromString("drop"); + write(table, ioManager, GenericRow.of(1, 10, keep), GenericRow.of(2, 10, drop)); + write(table, ioManager, GenericRow.of(3, 10, keep), GenericRow.of(4, 20, keep)); + compact(table, BinaryRow.EMPTY_ROW, 0, ioManager, true); + + Snapshot compactedSnapshot = table.store().snapshotManager().latestSnapshot(); + assertThat( + table.store() + .newIndexFileHandler() + .scanSourceIndexes(compactedSnapshot, BinaryRow.EMPTY_ROW, 0)) + .isNotEmpty(); + write( + table, + ioManager, + GenericRow.ofKind(org.apache.paimon.types.RowKind.DELETE, 3, 10, keep)); + + Snapshot snapshot = table.store().snapshotManager().latestSnapshot(); + List payloads = + table.store() + .newIndexFileHandler() + .scanSourceIndexes(snapshot, BinaryRow.EMPTY_ROW, 0); + assertThat(payloads).isNotEmpty(); + + PredicateBuilder builder = new PredicateBuilder(table.rowType()); + Predicate predicate = PredicateBuilder.and(builder.equal(1, 10), builder.equal(2, keep)); + ReadBuilder readBuilder = table.newReadBuilder().withFilter(predicate); + TableScan.Plan plan = readBuilder.newScan().plan(); + + assertThat(plan.splits()) + .extracting( + split -> + ((split instanceof IndexedSplit) + ? ((IndexedSplit) split).dataSplit() + : (DataSplit) split) + .dataFiles() + .get(0) + .fileName()) + .allMatch( + source -> + payloads.stream() + .map(PrimaryKeyIndexSourceMeta::fromIndexFile) + .anyMatch( + meta -> + meta.sourceFile() + .fileName() + .equals(source))); + + assertThat(plan.splits()).anyMatch(IndexedSplit.class::isInstance); + assertThat(plan.splits()) + .filteredOn(IndexedSplit.class::isInstance) + .map(IndexedSplit.class::cast) + .anyMatch( + split -> + split.dataSplit().deletionFiles().isPresent() + && split.dataSplit().deletionFiles().get().get(0) != null); + + List ids = new ArrayList<>(); + try (RecordReader reader = + readBuilder.newRead().executeFilter().createReader(plan)) { + reader.forEachRemaining(row -> ids.add(row.getInt(0))); + } + + assertThat(ids).containsExactly(1); + + FileStoreTable historicTable = + table.copy( + Collections.singletonMap( + CoreOptions.SCAN_SNAPSHOT_ID.key(), + Long.toString(compactedSnapshot.id()))); + ReadBuilder historicReadBuilder = historicTable.newReadBuilder().withFilter(predicate); + TableScan.Plan historicPlan = historicReadBuilder.newScan().plan(); + assertThat(historicPlan.splits()).anyMatch(IndexedSplit.class::isInstance); + List historicIds = new ArrayList<>(); + try (RecordReader reader = + historicReadBuilder.newRead().executeFilter().createReader(historicPlan)) { + reader.forEachRemaining(row -> historicIds.add(row.getInt(0))); + } + assertThat(historicIds).containsExactlyInAnyOrder(1, 3); + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexResultTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexResultTest.java new file mode 100644 index 000000000000..3b7089f776e9 --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexResultTest.java @@ -0,0 +1,287 @@ +/* + * 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.table.source; + +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.globalindex.GlobalIndexReader; +import org.apache.paimon.globalindex.GlobalIndexResult; +import org.apache.paimon.globalindex.IndexedSplit; +import org.apache.paimon.globalindex.btree.BTreeGlobalIndexerFactory; +import org.apache.paimon.index.GlobalIndexMeta; +import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.index.pk.PrimaryKeyIndexDefinition; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceFile; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceMeta; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.manifest.FileKind; +import org.apache.paimon.manifest.FileSource; +import org.apache.paimon.manifest.IndexManifestEntry; +import org.apache.paimon.options.Options; +import org.apache.paimon.predicate.Predicate; +import org.apache.paimon.predicate.PredicateBuilder; +import org.apache.paimon.stats.SimpleStats; +import org.apache.paimon.types.DataField; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.Range; +import org.apache.paimon.utils.RoaringNavigableMap64; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** Tests conversion from file-local sorted-index results to physical-position splits. */ +class PrimaryKeySortedIndexResultTest { + + @Test + void testIndexedEmptyRawAndInvalidFiles() { + DataFileMeta indexed = dataFile("indexed", 5); + DataFileMeta empty = dataFile("empty", 5); + DataFileMeta raw = dataFile("raw", 5); + DataFileMeta invalid = dataFile("invalid", 5); + List deletionFiles = + Arrays.asList( + new DeletionFile("dv-indexed", 0, 1, 1L), + new DeletionFile("dv-empty", 1, 1, 1L), + new DeletionFile("dv-raw", 2, 1, 1L), + new DeletionFile("dv-invalid", 3, 1, 1L)); + DataSplit split = dataSplit(11, Arrays.asList(indexed, empty, raw, invalid), deletionFiles); + PrimaryKeyIndexDefinition definition = + new PrimaryKeyIndexDefinition( + "f7", + 7, + BTreeGlobalIndexerFactory.IDENTIFIER, + new Options(), + PrimaryKeyIndexDefinition.Family.BTREE); + PrimaryKeySortedIndexScan.Plan plan = + PrimaryKeySortedIndexScan.plan( + 11, + Collections.singletonList(split), + Collections.singletonList(definition), + Arrays.asList( + payloadEntry(payload("btree-indexed", "indexed", 5)), + payloadEntry(payload("btree-empty", "empty", 5)), + payloadEntry(payload("btree-invalid", "invalid", 5)))); + RowType rowType = RowType.of(new DataField(7, "f7", DataTypes.INT())); + Predicate predicate = new PredicateBuilder(rowType).equal(0, 42); + PrimaryKeySortedIndexScan.EvaluatedPlan evaluated = + PrimaryKeySortedIndexScan.evaluate( + plan, + rowType, + predicate, + Collections.singletonList(definition), + (file, ignoredDefinition, ignoredPayloads) -> { + if (file.dataFile().fileName().equals("indexed")) { + return reader(1, 2, 4); + } else if (file.dataFile().fileName().equals("empty")) { + return reader(); + } + return reader(5); + }); + + PrimaryKeySortedIndexResult result = new PrimaryKeySortedIndexResult(evaluated); + + assertThat(result.snapshotId()).isEqualTo(11); + assertThat(result.splits()).hasSize(3); + + assertThat(result.splits().get(0)).isInstanceOf(IndexedSplit.class); + IndexedSplit indexedSplit = (IndexedSplit) result.splits().get(0); + assertThat(indexedSplit.dataSplit().dataFiles()).containsExactly(indexed); + assertThat(indexedSplit.dataSplit().deletionFiles().get()) + .containsExactly(deletionFiles.get(0)); + assertThat(indexedSplit.rowRanges()).containsExactly(new Range(1, 2), new Range(4, 4)); + + assertRawSplit(result.splits().get(1), raw, deletionFiles.get(2)); + assertRawSplit(result.splits().get(2), invalid, deletionFiles.get(3)); + } + + @Test + void testNonRawConvertibleSplitPreservesMergeBoundary() { + DataFileMeta first = dataFile("first", 5); + DataFileMeta second = dataFile("second", 5); + List deletionFiles = + Arrays.asList( + new DeletionFile("dv-first", 0, 1, 1L), + new DeletionFile("dv-second", 1, 1, 1L)); + DataSplit split = dataSplit(11, Arrays.asList(first, second), deletionFiles, false); + PrimaryKeyIndexDefinition definition = + new PrimaryKeyIndexDefinition( + "f7", + 7, + BTreeGlobalIndexerFactory.IDENTIFIER, + new Options(), + PrimaryKeyIndexDefinition.Family.BTREE); + PrimaryKeySortedIndexScan.Plan plan = + PrimaryKeySortedIndexScan.plan( + 11, + Collections.singletonList(split), + Collections.singletonList(definition), + Arrays.asList( + payloadEntry(payload("btree-first", "first", 5)), + payloadEntry(payload("btree-second", "second", 5)))); + RowType rowType = RowType.of(new DataField(7, "f7", DataTypes.INT())); + Predicate predicate = new PredicateBuilder(rowType).equal(0, 42); + PrimaryKeySortedIndexScan.EvaluatedPlan evaluated = + PrimaryKeySortedIndexScan.evaluate( + plan, + rowType, + predicate, + Collections.singletonList(definition), + (file, ignoredDefinition, ignoredPayloads) -> + file.dataFile().fileName().equals("first") ? reader(1) : reader(2)); + + PrimaryKeySortedIndexResult result = new PrimaryKeySortedIndexResult(evaluated); + + assertThat(result.splits()).singleElement().isSameAs(split); + } + + @Test + void testFragmentedIndexResultFallsBackToRawSplit() { + DataFileMeta file = dataFile("fragmented", 8193); + DeletionFile deletionFile = new DeletionFile("dv-fragmented", 0, 1, 1L); + DataSplit split = + dataSplit( + 11, + Collections.singletonList(file), + Collections.singletonList(deletionFile)); + PrimaryKeyIndexDefinition definition = + new PrimaryKeyIndexDefinition( + "f7", + 7, + BTreeGlobalIndexerFactory.IDENTIFIER, + new Options(), + PrimaryKeyIndexDefinition.Family.BTREE); + PrimaryKeySortedIndexScan.Plan plan = + PrimaryKeySortedIndexScan.plan( + 11, + Collections.singletonList(split), + Collections.singletonList(definition), + Collections.singletonList( + payloadEntry(payload("btree-fragmented", "fragmented", 8193)))); + RowType rowType = RowType.of(new DataField(7, "f7", DataTypes.INT())); + Predicate predicate = new PredicateBuilder(rowType).equal(0, 42); + long[] positions = new long[4097]; + for (int i = 0; i < positions.length; i++) { + positions[i] = i * 2L; + } + PrimaryKeySortedIndexScan.EvaluatedPlan evaluated = + PrimaryKeySortedIndexScan.evaluate( + plan, + rowType, + predicate, + Collections.singletonList(definition), + (ignoredFile, ignoredDefinition, ignoredPayloads) -> reader(positions)); + + PrimaryKeySortedIndexResult result = new PrimaryKeySortedIndexResult(evaluated); + + assertThat(result.splits()).singleElement().isInstanceOf(DataSplit.class); + assertRawSplit(result.splits().get(0), file, deletionFile); + } + + private static void assertRawSplit( + Split split, DataFileMeta expectedFile, DeletionFile expectedDeletionFile) { + assertThat(split).isInstanceOf(DataSplit.class); + DataSplit rawSplit = (DataSplit) split; + assertThat(rawSplit.snapshotId()).isEqualTo(11); + assertThat(rawSplit.dataFiles()).containsExactly(expectedFile); + assertThat(rawSplit.deletionFiles().get()).containsExactly(expectedDeletionFile); + assertThat(rawSplit.rawConvertible()).isFalse(); + } + + private static GlobalIndexReader reader(long... rowPositions) { + RoaringNavigableMap64 positions = new RoaringNavigableMap64(); + for (long rowPosition : rowPositions) { + positions.add(rowPosition); + } + GlobalIndexReader reader = mock(GlobalIndexReader.class); + when(reader.visitEqual(any(), any())) + .thenReturn( + CompletableFuture.completedFuture( + Optional.of(GlobalIndexResult.create(positions)))); + return reader; + } + + private static DataSplit dataSplit( + long snapshotId, List dataFiles, List deletionFiles) { + return dataSplit(snapshotId, dataFiles, deletionFiles, true); + } + + private static DataSplit dataSplit( + long snapshotId, + List dataFiles, + List deletionFiles, + boolean rawConvertible) { + return DataSplit.builder() + .withSnapshot(snapshotId) + .withPartition(BinaryRow.EMPTY_ROW) + .withBucket(0) + .withBucketPath("bucket-0") + .withTotalBuckets(2) + .withDataFiles(dataFiles) + .withDataDeletionFiles(deletionFiles) + .isStreaming(false) + .rawConvertible(rawConvertible) + .build(); + } + + private static DataFileMeta dataFile(String fileName, long rowCount) { + return DataFileMeta.forAppend( + fileName, + 100, + rowCount, + SimpleStats.EMPTY_STATS, + 0, + 0, + 1, + Collections.emptyList(), + null, + FileSource.COMPACT, + null, + null, + null, + null); + } + + private static IndexManifestEntry payloadEntry(IndexFileMeta payload) { + return new IndexManifestEntry(FileKind.ADD, BinaryRow.EMPTY_ROW, 0, payload); + } + + private static IndexFileMeta payload(String fileName, String sourceName, long sourceRowCount) { + byte[] sourceMeta = + new PrimaryKeyIndexSourceMeta( + new PrimaryKeyIndexSourceFile(sourceName, sourceRowCount)) + .serialize(); + return new IndexFileMeta( + BTreeGlobalIndexerFactory.IDENTIFIER, + fileName, + 100, + sourceRowCount, + new GlobalIndexMeta(0, sourceRowCount - 1, 7, null, new byte[] {1}, sourceMeta), + null); + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexScanTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexScanTest.java new file mode 100644 index 000000000000..c0343638cf27 --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeySortedIndexScanTest.java @@ -0,0 +1,368 @@ +/* + * 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.table.source; + +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.globalindex.GlobalIndexReader; +import org.apache.paimon.globalindex.GlobalIndexResult; +import org.apache.paimon.globalindex.bitmap.BitmapGlobalIndexerFactory; +import org.apache.paimon.globalindex.btree.BTreeGlobalIndexerFactory; +import org.apache.paimon.index.GlobalIndexMeta; +import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.index.pk.PrimaryKeyIndexDefinition; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceFile; +import org.apache.paimon.index.pk.PrimaryKeyIndexSourceMeta; +import org.apache.paimon.index.pksorted.PkSortedBucketIndexState; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.manifest.FileKind; +import org.apache.paimon.manifest.FileSource; +import org.apache.paimon.manifest.IndexManifestEntry; +import org.apache.paimon.options.Options; +import org.apache.paimon.predicate.Predicate; +import org.apache.paimon.predicate.PredicateBuilder; +import org.apache.paimon.stats.SimpleStats; +import org.apache.paimon.types.DataField; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.RoaringNavigableMap64; + +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.when; + +/** Tests source-backed BTree and Bitmap planning in file-local row-position space. */ +class PrimaryKeySortedIndexScanTest { + + @Test + void testPayloadStateIsBuiltOncePerBucketAndDefinition() { + DataSplit split = + dataSplit( + 11, 0, dataFile("data-1", 4), dataFile("data-2", 4), dataFile("data-3", 4)); + List definitions = + Arrays.asList( + definition( + 7, + BTreeGlobalIndexerFactory.IDENTIFIER, + PrimaryKeyIndexDefinition.Family.BTREE), + definition( + 8, + BitmapGlobalIndexerFactory.IDENTIFIER, + PrimaryKeyIndexDefinition.Family.BITMAP)); + List entries = new java.util.ArrayList<>(); + for (int i = 1; i <= 3; i++) { + String source = "data-" + i; + entries.add(payloadEntry(0, payload("btree-" + i, source, 4, "btree", 7, 4))); + entries.add(payloadEntry(0, payload("bitmap-" + i, source, 4, "bitmap", 8, 4))); + } + + try (MockedStatic states = + mockStatic( + PkSortedBucketIndexState.class, org.mockito.Mockito.CALLS_REAL_METHODS)) { + PrimaryKeySortedIndexScan.Plan plan = + PrimaryKeySortedIndexScan.plan( + 11, Collections.singletonList(split), definitions, entries); + + assertThat(plan.files()).hasSize(3); + assertThat(plan.files()) + .allSatisfy( + file -> { + assertThat(file.group(7)).isPresent(); + assertThat(file.group(8)).isPresent(); + }); + states.verify( + times(1), + () -> + PkSortedBucketIndexState.fromActivePayloads( + eq(7), eq("btree"), anyList(), anyList())); + states.verify( + times(1), + () -> + PkSortedBucketIndexState.fromActivePayloads( + eq(8), eq("bitmap"), anyList(), anyList())); + } + } + + @Test + void testSnapshotScopedGroupPlanning() { + DataSplit split = dataSplit(11, 0, dataFile("data-1", 4)); + List definitions = + Arrays.asList( + definition( + 7, + BTreeGlobalIndexerFactory.IDENTIFIER, + PrimaryKeyIndexDefinition.Family.BTREE), + definition( + 8, + BitmapGlobalIndexerFactory.IDENTIFIER, + PrimaryKeyIndexDefinition.Family.BITMAP)); + List entries = + Arrays.asList( + payloadEntry(0, payload("btree-0", "data-1", 4, "btree", 7, 2)), + payloadEntry(0, payload("btree-1", "data-1", 4, "btree", 7, 2)), + payloadEntry(1, payload("wrong-bucket", "data-1", 4, "btree", 7, 4)), + payloadEntry(0, payload("wrong-field", "data-1", 4, "btree", 9, 4)), + payloadEntry(0, payload("wrong-type", "data-1", 4, "bitmap", 7, 4)), + payloadEntry(0, payload("wrong-source", "data-2", 4, "bitmap", 8, 4)), + payloadEntry(0, ordinaryPayload("ordinary", "btree", 7, 4))); + + PrimaryKeySortedIndexScan.Plan plan = + PrimaryKeySortedIndexScan.plan( + 11, Collections.singletonList(split), definitions, entries); + + assertThat(plan.snapshotId()).isEqualTo(11); + assertThat(plan.files()).hasSize(1); + PrimaryKeySortedIndexScan.FilePlan file = plan.files().get(0); + assertThat(file.dataFile().fileName()).isEqualTo("data-1"); + assertThat(file.group(7)).isPresent(); + assertThat(file.group(7).get().payloads()) + .extracting(IndexFileMeta::fileName) + .containsExactly("btree-0", "btree-1"); + assertThat(file.group(8)).isEmpty(); + } + + @Test + void testRotatedPayloadsAreUnionedBeforeEvaluation() { + DataSplit split = dataSplit(11, 0, dataFile("data-1", 4)); + PrimaryKeyIndexDefinition definition = + definition( + 7, + BTreeGlobalIndexerFactory.IDENTIFIER, + PrimaryKeyIndexDefinition.Family.BTREE); + PrimaryKeySortedIndexScan.Plan plan = + PrimaryKeySortedIndexScan.plan( + 11, + Collections.singletonList(split), + Collections.singletonList(definition), + Arrays.asList( + payloadEntry(0, payload("btree-0", "data-1", 4, "btree", 7, 2)), + payloadEntry(0, payload("btree-1", "data-1", 4, "btree", 7, 2)))); + RowType rowType = RowType.of(new DataField(7, "f7", DataTypes.INT())); + Predicate predicate = new PredicateBuilder(rowType).equal(0, 42); + RoaringNavigableMap64 positions = new RoaringNavigableMap64(); + positions.add(3); + GlobalIndexReader reader = mock(GlobalIndexReader.class); + when(reader.visitEqual(any(), eq(42))) + .thenReturn( + CompletableFuture.completedFuture( + Optional.of(GlobalIndexResult.create(positions)))); + AtomicInteger readersCreated = new AtomicInteger(); + + PrimaryKeySortedIndexScan.EvaluatedPlan evaluated = + PrimaryKeySortedIndexScan.evaluate( + plan, + rowType, + predicate, + Collections.singletonList(definition), + (ignoredFile, ignoredDefinition, payloads) -> { + readersCreated.incrementAndGet(); + assertThat(payloads) + .extracting(IndexFileMeta::fileName) + .containsExactly("btree-0", "btree-1"); + return reader; + }); + + assertThat(readersCreated).hasValue(1); + assertThat(evaluated.files()).hasSize(1); + assertThat(evaluated.files().get(0).result()).isPresent(); + assertThat(evaluated.files().get(0).result().get().results()).containsExactly(3L); + } + + @Test + void testPerFileBooleanFallbackSemantics() { + DataSplit split = dataSplit(11, 0, dataFile("data-1", 4)); + PrimaryKeyIndexDefinition definition = + definition( + 7, + BTreeGlobalIndexerFactory.IDENTIFIER, + PrimaryKeyIndexDefinition.Family.BTREE); + PrimaryKeySortedIndexScan.Plan plan = + PrimaryKeySortedIndexScan.plan( + 11, + Collections.singletonList(split), + Collections.singletonList(definition), + Collections.singletonList( + payloadEntry(0, payload("btree-0", "data-1", 4, "btree", 7, 4)))); + RowType rowType = + RowType.of( + new DataField(7, "f7", DataTypes.INT()), + new DataField(8, "f8", DataTypes.INT())); + PredicateBuilder builder = new PredicateBuilder(rowType); + GlobalIndexReader reader = readerWithPositions(2); + + PrimaryKeySortedIndexScan.EvaluatedPlan andResult = + PrimaryKeySortedIndexScan.evaluate( + plan, + rowType, + PredicateBuilder.and(builder.equal(0, 42), builder.equal(1, 99)), + Collections.singletonList(definition), + (ignoredFile, ignoredDefinition, ignoredPayloads) -> reader); + PrimaryKeySortedIndexScan.EvaluatedPlan orResult = + PrimaryKeySortedIndexScan.evaluate( + plan, + rowType, + PredicateBuilder.or(builder.equal(0, 42), builder.equal(1, 99)), + Collections.singletonList(definition), + (ignoredFile, ignoredDefinition, ignoredPayloads) -> reader); + + assertThat(andResult.files().get(0).result()).isPresent(); + assertThat(andResult.files().get(0).result().get().results()).containsExactly(2L); + assertThat(orResult.files().get(0).result()).isEmpty(); + } + + @Test + void testReaderFailureFallsBackOnlyCurrentFile() { + DataSplit split = dataSplit(11, 0, dataFile("data-1", 4), dataFile("data-2", 4)); + PrimaryKeyIndexDefinition definition = + definition( + 7, + BTreeGlobalIndexerFactory.IDENTIFIER, + PrimaryKeyIndexDefinition.Family.BTREE); + PrimaryKeySortedIndexScan.Plan plan = + PrimaryKeySortedIndexScan.plan( + 11, + Collections.singletonList(split), + Collections.singletonList(definition), + Arrays.asList( + payloadEntry(0, payload("btree-1", "data-1", 4, "btree", 7, 4)), + payloadEntry(0, payload("btree-2", "data-2", 4, "btree", 7, 4)))); + RowType rowType = RowType.of(new DataField(7, "f7", DataTypes.INT())); + Predicate predicate = new PredicateBuilder(rowType).equal(0, 42); + GlobalIndexReader failedReader = mock(GlobalIndexReader.class); + when(failedReader.visitEqual(any(), eq(42))) + .thenThrow(new RuntimeException("corrupt index")); + GlobalIndexReader successfulReader = readerWithPositions(1); + + PrimaryKeySortedIndexScan.EvaluatedPlan evaluated = + PrimaryKeySortedIndexScan.evaluate( + plan, + rowType, + predicate, + Collections.singletonList(definition), + (file, ignoredDefinition, ignoredPayloads) -> + file.dataFile().fileName().equals("data-1") + ? failedReader + : successfulReader); + + assertThat(evaluated.files()).hasSize(2); + assertThat(evaluated.files().get(0).result()).isEmpty(); + assertThat(evaluated.files().get(1).result()).isPresent(); + assertThat(evaluated.files().get(1).result().get().results()).containsExactly(1L); + } + + private static PrimaryKeyIndexDefinition definition( + int fieldId, String indexType, PrimaryKeyIndexDefinition.Family family) { + return new PrimaryKeyIndexDefinition( + "f" + fieldId, fieldId, indexType, new Options(), family); + } + + private static GlobalIndexReader readerWithPositions(long... rowPositions) { + RoaringNavigableMap64 positions = new RoaringNavigableMap64(); + for (long rowPosition : rowPositions) { + positions.add(rowPosition); + } + GlobalIndexReader reader = mock(GlobalIndexReader.class); + when(reader.visitEqual(any(), any())) + .thenReturn( + CompletableFuture.completedFuture( + Optional.of(GlobalIndexResult.create(positions)))); + return reader; + } + + private static DataSplit dataSplit(long snapshotId, int bucket, DataFileMeta... files) { + return DataSplit.builder() + .withSnapshot(snapshotId) + .withPartition(BinaryRow.EMPTY_ROW) + .withBucket(bucket) + .withBucketPath("bucket-" + bucket) + .withTotalBuckets(2) + .withDataFiles(Arrays.asList(files)) + .isStreaming(false) + .rawConvertible(false) + .build(); + } + + private static DataFileMeta dataFile(String fileName, long rowCount) { + return DataFileMeta.forAppend( + fileName, + 100, + rowCount, + SimpleStats.EMPTY_STATS, + 0, + 0, + 1, + Collections.emptyList(), + null, + FileSource.COMPACT, + null, + null, + null, + null); + } + + private static IndexManifestEntry payloadEntry(int bucket, IndexFileMeta payload) { + return new IndexManifestEntry(FileKind.ADD, BinaryRow.EMPTY_ROW, bucket, payload); + } + + private static IndexFileMeta payload( + String fileName, + String sourceName, + long sourceRowCount, + String indexType, + int fieldId, + long payloadRowCount) { + byte[] sourceMeta = + new PrimaryKeyIndexSourceMeta( + new PrimaryKeyIndexSourceFile(sourceName, sourceRowCount)) + .serialize(); + return new IndexFileMeta( + indexType, + fileName, + 100, + payloadRowCount, + new GlobalIndexMeta( + 0, sourceRowCount - 1, fieldId, null, new byte[] {1}, sourceMeta), + null); + } + + private static IndexFileMeta ordinaryPayload( + String fileName, String indexType, int fieldId, long rowCount) { + return new IndexFileMeta( + indexType, + fileName, + 100, + rowCount, + new GlobalIndexMeta(0, rowCount - 1, fieldId, null, new byte[] {1}, null), + null); + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorScanTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorScanTest.java index 5ff206752f46..369c90d71134 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorScanTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorScanTest.java @@ -19,7 +19,6 @@ package org.apache.paimon.table.source; import org.apache.paimon.CoreOptions; -import org.apache.paimon.FileStore; import org.apache.paimon.Snapshot; import org.apache.paimon.data.BinaryRow; import org.apache.paimon.index.GlobalIndexMeta; @@ -32,11 +31,15 @@ import org.apache.paimon.manifest.FileSource; import org.apache.paimon.manifest.IndexManifestEntry; import org.apache.paimon.options.Options; +import org.apache.paimon.schema.SchemaManager; +import org.apache.paimon.schema.TableSchema; import org.apache.paimon.stats.SimpleStats; import org.apache.paimon.table.BucketMode; import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.source.snapshot.SnapshotReader; import org.apache.paimon.utils.Filter; +import org.apache.paimon.utils.Range; +import org.apache.paimon.utils.SnapshotManager; import org.junit.jupiter.api.Test; @@ -78,6 +81,7 @@ void testPostponeTableOnlyScansRealBuckets() { SnapshotReader reader = mock(SnapshotReader.class, RETURNS_SELF); SnapshotReader.Plan snapshotPlan = mock(SnapshotReader.Plan.class, CALLS_REAL_METHODS); + when(snapshotPlan.snapshotId()).thenReturn(11L); when(snapshotPlan.splits()).thenReturn(Collections.emptyList()); when(reader.read()).thenReturn(snapshotPlan); when(table.newSnapshotReader()).thenReturn(reader); @@ -85,9 +89,8 @@ void testPostponeTableOnlyScansRealBuckets() { IndexFileHandler indexFileHandler = mock(IndexFileHandler.class); when(indexFileHandler.scan(eq(snapshot), any(Filter.class))) .thenReturn(Collections.emptyList()); - FileStore store = mock(FileStore.class); - when(store.newIndexFileHandler()).thenReturn(indexFileHandler); - when(table.store()).thenReturn(store); + when(reader.indexFileHandler()).thenReturn(indexFileHandler); + configureBatchScan(table, reader, snapshot); new PrimaryKeyVectorScan(table, 7, "ivf-pq", null).scan(); @@ -106,6 +109,7 @@ void testScansOneSnapshotAndFiltersVectorIdentity() { SnapshotReader snapshotReader = mock(SnapshotReader.class, RETURNS_SELF); SnapshotReader.Plan snapshotPlan = mock(SnapshotReader.Plan.class, CALLS_REAL_METHODS); + when(snapshotPlan.snapshotId()).thenReturn(11L); when(snapshotPlan.splits()) .thenReturn(Collections.singletonList(dataSplit(dataFile("data-1")))); when(snapshotReader.read()).thenReturn(snapshotPlan); @@ -129,9 +133,8 @@ void testScansOneSnapshotAndFiltersVectorIdentity() { } return filtered; }); - FileStore store = mock(FileStore.class); - when(store.newIndexFileHandler()).thenReturn(indexFileHandler); - when(table.store()).thenReturn(store); + when(snapshotReader.indexFileHandler()).thenReturn(indexFileHandler); + configureBatchScan(table, snapshotReader, snapshot); PrimaryKeyVectorScan.Plan plan = new PrimaryKeyVectorScan(table, 7, "ivf-pq", null).scan(); @@ -171,7 +174,10 @@ void testBucketSplitSerialization() throws Exception { IndexFileMeta payload = payloadFile(); BucketVectorSearchSplit split = new BucketVectorSearchSplit( - dataSplit(dataFile("data-1")), Collections.singletonList(payload)); + dataSplit(dataFile("data-1")), + Collections.singletonList(payload), + Collections.singletonMap( + "data-1", Collections.singletonList(new Range(1, 1)))); ByteArrayOutputStream bytes = new ByteArrayOutputStream(); try (ObjectOutputStream output = new ObjectOutputStream(bytes)) { @@ -243,6 +249,26 @@ private static CoreOptions coreOptions() { return new CoreOptions(options); } + private static void configureBatchScan( + FileStoreTable table, SnapshotReader snapshotReader, Snapshot snapshot) { + TableSchema schema = mock(TableSchema.class); + when(schema.primaryKeys()).thenReturn(Collections.singletonList("id")); + when(table.schema()).thenReturn(schema); + when(table.schemaManager()).thenReturn(mock(SchemaManager.class)); + SnapshotManager snapshotManager = mock(SnapshotManager.class); + when(snapshotManager.latestSnapshot()).thenReturn(snapshot); + when(snapshotManager.snapshot(snapshot.id())).thenReturn(snapshot); + when(snapshotReader.snapshotManager()).thenReturn(snapshotManager); + when(table.newScan(any(FileStoreTable.SnapshotReaderFactory.class))) + .thenAnswer( + invocation -> { + FileStoreTable.SnapshotReaderFactory factory = + invocation.getArgument(0); + return new PrimaryKeyBatchScan( + table, factory.create(table), mock(TableQueryAuth.class), null); + }); + } + private static DataFileMeta dataFile(String fileName) { return DataFileMeta.forAppend( fileName, diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorSearchTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorSearchTest.java index dd34a2663887..d1fb7c25563b 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorSearchTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorSearchTest.java @@ -19,12 +19,14 @@ package org.apache.paimon.table.source; import org.apache.paimon.CoreOptions; +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.globalindex.GlobalIndexResult; import org.apache.paimon.globalindex.IndexedSplit; import org.apache.paimon.globalindex.testvector.TestVectorGlobalIndexerFactory; +import org.apache.paimon.predicate.PredicateBuilder; import org.apache.paimon.reader.RecordReader; import org.apache.paimon.schema.Schema; import org.apache.paimon.table.FileStoreTable; @@ -151,6 +153,84 @@ void testVectorSearchMaterializesPhysicalRows() throws Exception { assertThat(ids).containsExactly(2, 3); } + @Test + void testEmptyVectorSearch() throws Exception { + createTableDefault(); + FileStoreTable table = getTableDefault(); + + GlobalIndexResult result = + table.newVectorSearchBuilder() + .withVectorColumn("embedding") + .withVector(new float[] {0, 0}) + .withLimit(1) + .executeLocal(); + + assertThat(((GlobalIndexSplitResult) result).splits()).isEmpty(); + } + + @Test + void testVectorSearchUsesSortedIndexPreFilter() throws Exception { + Schema schema = + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("score", DataTypes.INT()) + .column("embedding", DataTypes.VECTOR(2, DataTypes.FLOAT())) + .primaryKey("id") + .option(CoreOptions.BUCKET.key(), "1") + .option(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true") + .option(CoreOptions.PK_BTREE_INDEX_COLUMNS.key(), "score") + .option(CoreOptions.PK_VECTOR_INDEX_COLUMNS.key(), "embedding") + .option( + "fields.embedding.pk-vector.index.type", + TestVectorGlobalIndexerFactory.IDENTIFIER) + .option("fields.embedding.pk-vector.distance.metric", "l2") + .option("test.vector.dimension", "2") + .option("test.vector.metric", "l2") + .build(); + catalog.createTable(identifier(), schema, false); + FileStoreTable table = getTableDefault(); + write( + table, + ioManager, + GenericRow.of(1, 0, BinaryVector.fromPrimitiveArray(new float[] {0, 0})), + GenericRow.of(2, 1, BinaryVector.fromPrimitiveArray(new float[] {10, 0})), + GenericRow.of(3, 1, BinaryVector.fromPrimitiveArray(new float[] {20, 0}))); + compact(table, BinaryRow.EMPTY_ROW, 0, ioManager, true); + + GlobalIndexResult result = + table.newVectorSearchBuilder() + .withVectorColumn("embedding") + .withVector(new float[] {0, 0}) + .withFilter(new PredicateBuilder(table.rowType()).equal(1, 1)) + .withLimit(1) + .executeLocal(); + + assertThat(readIds(table, result)).containsExactly(2); + + GlobalIndexResult residualResult = + table.newVectorSearchBuilder() + .withVectorColumn("embedding") + .withVector(new float[] {0, 0}) + .withFilter(new PredicateBuilder(table.rowType()).equal(0, 2)) + .withLimit(1) + .executeLocal(); + + assertThat(readIds(table, residualResult)).containsExactly(2); + + PredicateBuilder predicateBuilder = new PredicateBuilder(table.rowType()); + GlobalIndexResult indexedResidualResult = + table.newVectorSearchBuilder() + .withVectorColumn("embedding") + .withVector(new float[] {0, 0}) + .withFilter( + PredicateBuilder.and( + predicateBuilder.equal(1, 1), predicateBuilder.equal(0, 3))) + .withLimit(1) + .executeLocal(); + + assertThat(readIds(table, indexedResidualResult)).containsExactly(3); + } + @Test void testFirstRowVectorSearch() throws Exception { catalog.createTable(identifier(), vectorSchema("first-row", false), false); diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/ScanBucketTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/ScanBucketTest.java index 5a56db2e6483..257b6c3d2432 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/source/ScanBucketTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/ScanBucketTest.java @@ -40,7 +40,6 @@ import java.util.stream.Collectors; import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Tests for {@link CoreOptions#SCAN_BUCKET}. */ public class ScanBucketTest { @@ -59,43 +58,33 @@ public void testWithBucketAllowsAppendOnlyFixedBucketTable() throws Exception { } @Test - public void testScanBucketOptionRejectsOutOfRangeBucketId() throws Exception { + public void testScanBucketOptionAllowsOutOfRangeBucketId() throws Exception { FileStoreTable table = createTableWithScanBucket("4", true, "5"); - assertThatThrownBy(() -> table.newScan().plan()) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("Bucket id 5 must be less than table bucket number 4"); + assertThat(table.newScan().plan().splits()).isEmpty(); } @Test - public void testScanBucketOptionRejectsDynamicBucketTable() throws Exception { + public void testScanBucketOptionAllowsDynamicBucketTable() throws Exception { FileStoreTable table = createTableWithScanBucket("-1", true, "0"); - assertThatThrownBy(() -> table.newScan().plan()) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("fixed-bucket tables"); + assertThat(table.newScan().plan().splits()).isEmpty(); } @Test - public void testScanBucketOptionRejectsPostponeBucketTable() throws Exception { + public void testScanBucketOptionAllowsPostponeBucketTable() throws Exception { FileStoreTable table = createTableWithScanBucket("-2", true, "0"); - assertThatThrownBy(() -> table.newScan().plan()) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("fixed-bucket tables"); + assertThat(table.newScan().plan().splits()).isEmpty(); } @Test - public void testScanBucketOptionRejectsBucketUnawareTable() throws Exception { + public void testScanBucketOptionAllowsBucketUnawareTable() throws Exception { FileStoreTable table = createBucketUnawareTableWithScanBucket("0"); - assertThatThrownBy(() -> table.newScan().plan()) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("primary key tables"); + assertThat(table.newScan().plan().splits()).isEmpty(); } @Test - public void testScanBucketOptionRejectsTableWithoutPrimaryKey() throws Exception { + public void testScanBucketOptionAllowsTableWithoutPrimaryKey() throws Exception { FileStoreTable table = createAppendOnlyTableWithScanBucket("4", "0"); - assertThatThrownBy(() -> table.newScan().plan()) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("primary key tables"); + assertThat(table.newScan().plan().splits()).isEmpty(); } @Test @@ -136,11 +125,9 @@ public void testScanBucketOptionViaReadBuilder() throws Exception { } @Test - public void testScanBucketOptionRejectsDirectTableScanForDynamicBucketTable() throws Exception { + public void testScanBucketOptionAllowsDirectTableScanForDynamicBucketTable() throws Exception { FileStoreTable table = createTableWithScanBucket("-1", true, "0"); - assertThatThrownBy(() -> table.newScan().plan()) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("fixed-bucket tables"); + assertThat(table.newScan().plan().splits()).isEmpty(); } private static List extractBuckets(List splits) { diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/TableScanTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/TableScanTest.java index cd417058a8d4..75863dd6427b 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/source/TableScanTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/TableScanTest.java @@ -56,6 +56,22 @@ /** Tests for {@link TableScan}. */ public class TableScanTest extends ScannerTestBase { + @Test + public void testBatchScanTypes() throws Exception { + assertThat(AbstractBatchTableScan.class.getDeclaredMethods()) + .noneMatch(method -> method.getName().equals("create")); + assertThat(AppendBatchTableScan.class.getDeclaredMethods()) + .noneMatch(method -> method.getName().equals("create")); + assertThat(PrimaryKeyBatchScan.class.getDeclaredMethods()) + .noneMatch(method -> method.getName().equals("create")); + assertThat(PrimaryKeyBatchScan.class.getDeclaredConstructors()).hasSize(1); + + assertThat(table.newScan()).isInstanceOf(PrimaryKeyBatchScan.class); + + createAppendOnlyTable(); + assertThat(table.newScan()).isInstanceOf(AppendBatchTableScan.class); + } + @Test public void testPlan() throws Exception { SnapshotManager snapshotManager = table.snapshotManager(); diff --git a/paimon-core/src/test/java/org/apache/paimon/table/system/AuditLogTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/system/AuditLogTableTest.java index cbba947bcb9b..93c7933b0b72 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/system/AuditLogTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/system/AuditLogTableTest.java @@ -19,27 +19,39 @@ package org.apache.paimon.table.system; import org.apache.paimon.CoreOptions; +import org.apache.paimon.Snapshot; import org.apache.paimon.catalog.Identifier; +import org.apache.paimon.data.BinaryRow; import org.apache.paimon.data.BinaryString; import org.apache.paimon.data.GenericRow; import org.apache.paimon.data.InternalRow; import org.apache.paimon.fs.FileIO; import org.apache.paimon.fs.Path; import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.globalindex.IndexedSplit; +import org.apache.paimon.options.Options; +import org.apache.paimon.predicate.PredicateBuilder; +import org.apache.paimon.reader.RecordReader; import org.apache.paimon.schema.Schema; import org.apache.paimon.schema.SchemaChange; import org.apache.paimon.schema.SchemaManager; import org.apache.paimon.schema.SchemaUtils; import org.apache.paimon.schema.TableSchema; +import org.apache.paimon.table.CatalogEnvironment; +import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplit; import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.FileStoreTableFactory; import org.apache.paimon.table.TableTestBase; +import org.apache.paimon.table.source.ChainSplit; +import org.apache.paimon.table.source.ReadBuilder; +import org.apache.paimon.table.source.TableScan; import org.apache.paimon.types.DataTypes; import org.apache.paimon.types.RowKind; import org.junit.jupiter.api.Test; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import static org.apache.paimon.catalog.Identifier.SYSTEM_TABLE_SPLITTER; @@ -58,6 +70,138 @@ public void testReadAuditLogFromLatest() throws Exception { assertThat(result).containsExactlyInAnyOrderElementsOf(expectRow); } + @Test + public void testSnapshotReaderDisablesIndexFileHandler() throws Exception { + AuditLogTable auditLogTable = createAuditLogTable("audit_table_index_handler", false); + + assertThat(auditLogTable.newSnapshotReader().indexFileHandler()).isNull(); + } + + @Test + public void testReadAuditLogWithPrimaryKeySortedIndex() throws Exception { + String tableName = "audit_table_with_sorted_index"; + Path tablePath = new Path(String.format("%s/%s.db/%s", warehouse, database, tableName)); + FileIO fileIO = LocalFileIO.create(); + TableSchema tableSchema = + SchemaUtils.forceCommit( + new SchemaManager(fileIO, tablePath), + Schema.newBuilder() + .column("pk", DataTypes.INT()) + .column("score", DataTypes.INT()) + .primaryKey("pk") + .option(CoreOptions.CHANGELOG_PRODUCER.key(), "input") + .option(CoreOptions.BUCKET.key(), "1") + .option(CoreOptions.DELETION_VECTORS_ENABLED.key(), "true") + .option(CoreOptions.PK_BTREE_INDEX_COLUMNS.key(), "score") + .build()); + FileStoreTable dataTable = FileStoreTableFactory.create(fileIO, tablePath, tableSchema); + write(dataTable, ioManager, GenericRow.of(1, 10)); + write(dataTable, ioManager, GenericRow.of(2, 20)); + compact(dataTable, BinaryRow.EMPTY_ROW, 0, ioManager, true); + Snapshot snapshot = dataTable.snapshotManager().latestSnapshot(); + assertThat( + dataTable + .store() + .newIndexFileHandler() + .scanSourceIndexes(snapshot, BinaryRow.EMPTY_ROW, 0)) + .isNotEmpty(); + + AuditLogTable auditLogTable = + (AuditLogTable) + catalog.getTable( + identifier( + tableName + + SYSTEM_TABLE_SPLITTER + + AuditLogTable.AUDIT_LOG)); + PredicateBuilder predicateBuilder = new PredicateBuilder(auditLogTable.rowType()); + ReadBuilder readBuilder = + auditLogTable.newReadBuilder().withFilter(predicateBuilder.equal(2, 10)); + TableScan.Plan plan = readBuilder.newScan().plan(); + assertThat(plan.splits()).isNotEmpty().noneMatch(IndexedSplit.class::isInstance); + try (RecordReader reader = readBuilder.newRead().createReader(plan)) { + reader.forEachRemaining(ignored -> {}); + } + } + + @Test + public void testChainTableAuditLogPreservesChainScan() throws Exception { + Path tablePath = new Path(String.format("%s/%s.db/chain_audit_table", warehouse, database)); + FileIO fileIO = LocalFileIO.create(); + SchemaManager schemaManager = new SchemaManager(fileIO, tablePath); + schemaManager.createTable( + Schema.newBuilder() + .column("dt", DataTypes.STRING()) + .column("pk", DataTypes.STRING()) + .column("v", DataTypes.STRING()) + .partitionKeys("dt") + .primaryKey("dt", "pk") + .option(CoreOptions.BUCKET.key(), "1") + .option(CoreOptions.CHANGELOG_PRODUCER.key(), "input") + .option(CoreOptions.MERGE_ENGINE.key(), "deduplicate") + .option(CoreOptions.SEQUENCE_FIELD.key(), "v") + .build()); + + FileStoreTable initialTable = + FileStoreTableFactory.create( + fileIO, + tablePath, + schemaManager.latest().get(), + CatalogEnvironment.empty()); + initialTable.createBranch("snapshot"); + initialTable.createBranch("delta"); + List chainOptions = + Arrays.asList( + SchemaChange.setOption("chain-table.enabled", "true"), + SchemaChange.setOption("scan.fallback-snapshot-branch", "snapshot"), + SchemaChange.setOption("scan.fallback-delta-branch", "delta"), + SchemaChange.setOption("partition.timestamp-pattern", "$dt"), + SchemaChange.setOption("partition.timestamp-formatter", "yyyyMMdd")); + schemaManager.commitChanges(chainOptions); + new SchemaManager(fileIO, tablePath, "snapshot").commitChanges(chainOptions); + new SchemaManager(fileIO, tablePath, "delta").commitChanges(chainOptions); + + FileStoreTable snapshotTable = branchTable(fileIO, tablePath, "snapshot"); + FileStoreTable deltaTable = branchTable(fileIO, tablePath, "delta"); + write( + snapshotTable, + ioManager, + GenericRow.of( + BinaryString.fromString("20240101"), + BinaryString.fromString("k"), + BinaryString.fromString("snapshot"))); + write( + deltaTable, + ioManager, + GenericRow.of( + BinaryString.fromString("20240102"), + BinaryString.fromString("k"), + BinaryString.fromString("delta"))); + + FileStoreTable chainTable = + FileStoreTableFactory.create( + fileIO, + tablePath, + schemaManager.latest().get(), + CatalogEnvironment.empty()); + AuditLogTable auditLogTable = new AuditLogTable(chainTable); + PredicateBuilder predicateBuilder = new PredicateBuilder(auditLogTable.rowType()); + TableScan.Plan plan = + auditLogTable + .newReadBuilder() + .withFilter(predicateBuilder.equal(1, BinaryString.fromString("20240102"))) + .newScan() + .plan(); + + assertThat(plan.splits()) + .singleElement() + .satisfies( + split -> { + assertThat(split).isInstanceOf(FallbackSplit.class); + assertThat(((FallbackSplit) split).wrapped()) + .isInstanceOf(ChainSplit.class); + }); + } + @Test public void testReadSequenceNumberWithTableOption() throws Exception { AuditLogTable auditLogTable = createAuditLogTable("audit_table_with_seq", true); @@ -129,6 +273,17 @@ private AuditLogTable createAuditLogTable(String tableName, boolean enableSequen return (AuditLogTable) catalog.getTable(auditLogTableId); } + private FileStoreTable branchTable(FileIO fileIO, Path tablePath, String branch) { + TableSchema branchSchema = + new SchemaManager(fileIO, tablePath, branch) + .latest() + .orElseThrow(AssertionError::new); + Options dynamicOptions = new Options(); + dynamicOptions.set(CoreOptions.BRANCH, branch); + return FileStoreTableFactory.create( + fileIO, tablePath, branchSchema, dynamicOptions, CatalogEnvironment.empty()); + } + private void writeTestData(FileStoreTable table) throws Exception { write(table, GenericRow.ofKind(RowKind.INSERT, 1, 1, 1)); write(table, GenericRow.ofKind(RowKind.DELETE, 1, 1, 1)); diff --git a/paimon-core/src/test/java/org/apache/paimon/tag/BatchReadTagCreatorTest.java b/paimon-core/src/test/java/org/apache/paimon/tag/BatchReadTagCreatorTest.java index 4f225e2f07ea..15172ef4d2dc 100644 --- a/paimon-core/src/test/java/org/apache/paimon/tag/BatchReadTagCreatorTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/tag/BatchReadTagCreatorTest.java @@ -23,8 +23,9 @@ import org.apache.paimon.data.GenericRow; import org.apache.paimon.options.Options; import org.apache.paimon.table.FileStoreTable; -import org.apache.paimon.table.source.DataTableBatchScan; +import org.apache.paimon.table.source.AbstractBatchTableScan; import org.apache.paimon.table.source.InnerTableScan; +import org.apache.paimon.table.source.PrimaryKeyBatchScan; import org.apache.paimon.table.source.TableScan; import org.apache.paimon.utils.SnapshotManager; import org.apache.paimon.utils.TagManager; @@ -109,9 +110,9 @@ public void testScanCreatesProtectionTag() throws Exception { TableScan.Plan plan = scan.plan(); assertThat(plan.splits()).isNotEmpty(); - assertThat(scan).isInstanceOf(DataTableBatchScan.class); + assertThat(scan).isInstanceOf(PrimaryKeyBatchScan.class); - DataTableBatchScan batchScan = (DataTableBatchScan) scan; + AbstractBatchTableScan batchScan = (AbstractBatchTableScan) scan; String tagName = batchScan.readProtectionTagName(); assertThat(tagName).isNotNull(); assertThat(tagName).startsWith(BatchReadTagCreator.BATCH_READ_TAG_PREFIX); @@ -132,8 +133,8 @@ public void testScanDoesNotCreateTagWhenDisabled() throws Exception { InnerTableScan scan = table.newScan(); scan.plan(); - assertThat(scan).isInstanceOf(DataTableBatchScan.class); - DataTableBatchScan batchScan = (DataTableBatchScan) scan; + assertThat(scan).isInstanceOf(PrimaryKeyBatchScan.class); + AbstractBatchTableScan batchScan = (AbstractBatchTableScan) scan; assertThat(batchScan.readProtectionTagName()).isNull(); } } diff --git a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/ScanBucketITCase.java b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/ScanBucketITCase.java index 4e4ae6a52587..29fddf715c5c 100644 --- a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/ScanBucketITCase.java +++ b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/ScanBucketITCase.java @@ -40,9 +40,7 @@ import java.util.List; import static java.util.Collections.singletonList; -import static org.apache.paimon.testutils.assertj.PaimonAssertions.anyCauseMatches; import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatThrownBy; /** ITCase for {@link CoreOptions#SCAN_BUCKET}. */ public class ScanBucketITCase extends CatalogITCaseBase { @@ -92,52 +90,6 @@ public void testScanBucketFilter() throws Exception { .containsExactly(Row.of((long) expected.size())); } - @Test - public void testScanBucketRejectsDynamicBucketTable() { - sql( - "CREATE TABLE dynamic_t (id INT, v INT, PRIMARY KEY (id) NOT ENFORCED) " - + "WITH ('bucket' = '-1')"); - - assertThatThrownBy( - () -> - batchSql( - "SELECT * FROM dynamic_t /*+ OPTIONS('scan.bucket' = '0') */")) - .satisfies( - anyCauseMatches( - IllegalArgumentException.class, - "Bucket scan is only supported for fixed-bucket tables")); - } - - @Test - public void testScanBucketRejectsBucketUnawareTable() { - sql("CREATE TABLE append_t (id INT, v INT) WITH ('bucket' = '-1')"); - - assertThatThrownBy( - () -> - batchSql( - "SELECT * FROM append_t /*+ OPTIONS('scan.bucket' = '0') */")) - .satisfies( - anyCauseMatches( - IllegalArgumentException.class, - "Bucket scan is only supported for primary key tables")); - } - - @Test - public void testScanBucketRejectsPostponeBucketTable() { - sql( - "CREATE TABLE postpone_t (id INT, v INT, PRIMARY KEY (id) NOT ENFORCED) " - + "WITH ('bucket' = '-2')"); - - assertThatThrownBy( - () -> - batchSql( - "SELECT * FROM postpone_t /*+ OPTIONS('scan.bucket' = '0') */")) - .satisfies( - anyCauseMatches( - IllegalArgumentException.class, - "Bucket scan is only supported for fixed-bucket tables")); - } - private void writeRows(FileStoreTable table, int... idAndValues) throws Exception { BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder(); BatchTableWrite write = writeBuilder.newWrite(); diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala index 924950746764..df600a6a75e8 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala @@ -27,7 +27,7 @@ import org.apache.paimon.spark.read.{BaseScan, BatchReadTagCleanupListener, Paim import org.apache.paimon.spark.sources.PaimonMicroBatchStream import org.apache.paimon.spark.util.OptionUtils import org.apache.paimon.table.{DataTable, FileStoreTable, InnerTable} -import org.apache.paimon.table.source.{DataTableBatchScan, InnerTableScan, Split} +import org.apache.paimon.table.source.{InnerTableScan, Split} import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.SQLConfHelper diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/read/BatchReadTagConcurrentExpireTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/read/BatchReadTagConcurrentExpireTest.scala index 45d1118d8517..d3af071ad060 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/read/BatchReadTagConcurrentExpireTest.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/read/BatchReadTagConcurrentExpireTest.scala @@ -20,7 +20,7 @@ package org.apache.paimon.spark.read import org.apache.paimon.options.ExpireConfig import org.apache.paimon.spark.PaimonSparkTestBase -import org.apache.paimon.table.source.DataTableBatchScan +import org.apache.paimon.table.source.AbstractBatchTableScan import java.time.Duration import java.util @@ -110,7 +110,7 @@ class BatchReadTagConcurrentExpireTest extends PaimonSparkTestBase { assert(!splits.isEmpty) // Verify protection tag was created - val batchScan = scan.asInstanceOf[DataTableBatchScan] + val batchScan = scan.asInstanceOf[AbstractBatchTableScan] val tagName = batchScan.readProtectionTagName assert(tagName != null, "Protection tag should be created during scan planning") assert(table.tagManager().tagExists(tagName))