diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilderImpl.java b/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilderImpl.java index 41716fc3d339..ab92a1e668d8 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilderImpl.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/BatchVectorSearchBuilderImpl.java @@ -125,7 +125,7 @@ public BatchVectorSearchBuilder withOption(String key, String value) { @Override public VectorScan newVectorScan() { - return new VectorScanImpl(table, partitionFilter, filter, vectorColumn, options); + return new DataEvolutionVectorScan(table, partitionFilter, filter, vectorColumn, options); } @Override 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 new file mode 100644 index 000000000000..b21d1605b55e --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/BucketVectorSearchSplit.java @@ -0,0 +1,110 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.table.source; + +import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.index.IndexFileMetaSerializer; +import org.apache.paimon.io.DataInputViewStreamWrapper; +import org.apache.paimon.io.DataOutputViewStreamWrapper; + +import java.io.IOException; +import java.io.ObjectInputStream; +import java.io.ObjectOutputStream; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Objects; + +import static org.apache.paimon.utils.Preconditions.checkArgument; + +/** All active data files and vector payloads for one snapshot bucket. */ +public class BucketVectorSearchSplit extends VectorSearchSplit { + + private static final long serialVersionUID = 1L; + private static final int VERSION = 1; + + private DataSplit dataSplit; + private transient List payloadFiles; + + public BucketVectorSearchSplit(DataSplit dataSplit, List payloadFiles) { + this.dataSplit = dataSplit; + for (IndexFileMeta payload : payloadFiles) { + checkArgument( + payload.globalIndexMeta() != null + && payload.globalIndexMeta().sourceMeta() != null, + "Primary-key vector payload %s has no source metadata.", + payload.fileName()); + } + this.payloadFiles = Collections.unmodifiableList(new ArrayList<>(payloadFiles)); + } + + public DataSplit dataSplit() { + return dataSplit; + } + + public List payloadFiles() { + return payloadFiles; + } + + private void writeObject(ObjectOutputStream out) throws IOException { + out.defaultWriteObject(); + out.writeInt(VERSION); + out.writeInt(payloadFiles.size()); + IndexFileMetaSerializer serializer = new IndexFileMetaSerializer(); + for (IndexFileMeta payloadFile : payloadFiles) { + serializer.serialize(payloadFile, new DataOutputViewStreamWrapper(out)); + } + } + + private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException { + in.defaultReadObject(); + int version = in.readInt(); + if (version != VERSION) { + throw new IOException("Unsupported BucketVectorSearchSplit version: " + version); + } + int payloadFileCount = in.readInt(); + if (payloadFileCount < 0) { + throw new IOException("Negative primary-key vector payload file count."); + } + List payloadFiles = new ArrayList<>(payloadFileCount); + IndexFileMetaSerializer serializer = new IndexFileMetaSerializer(); + for (int i = 0; i < payloadFileCount; i++) { + payloadFiles.add(serializer.deserialize(new DataInputViewStreamWrapper(in))); + } + this.payloadFiles = Collections.unmodifiableList(payloadFiles); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + BucketVectorSearchSplit that = (BucketVectorSearchSplit) o; + return Objects.equals(dataSplit, that.dataSplit) + && Objects.equals(payloadFiles, that.payloadFiles); + } + + @Override + public int hashCode() { + return Objects.hash(dataSplit, payloadFiles); + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorReadImpl.java b/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionVectorRead.java similarity index 96% rename from paimon-core/src/main/java/org/apache/paimon/table/source/VectorReadImpl.java rename to paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionVectorRead.java index 9d8bd1541ed7..354bdbe15d9b 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorReadImpl.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionVectorRead.java @@ -40,14 +40,14 @@ import static org.apache.paimon.CoreOptions.GLOBAL_INDEX_THREAD_NUM; -/** Implementation for {@link VectorRead}. */ -public class VectorReadImpl extends AbstractVectorRead implements VectorRead { +/** Data-evolution implementation for {@link VectorRead}. */ +public class DataEvolutionVectorRead extends AbstractVectorRead implements VectorRead { private static final long serialVersionUID = 1L; protected final float[] vector; - public VectorReadImpl( + public DataEvolutionVectorRead( FileStoreTable table, @Nullable PartitionPredicate partitionFilter, @Nullable Predicate filter, diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorScanImpl.java b/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionVectorScan.java similarity index 98% rename from paimon-core/src/main/java/org/apache/paimon/table/source/VectorScanImpl.java rename to paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionVectorScan.java index 63b2dadea3a4..79a9fe81d02b 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/source/VectorScanImpl.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/DataEvolutionVectorScan.java @@ -46,8 +46,8 @@ import static org.apache.paimon.predicate.PredicateVisitor.collectFieldIds; import static org.apache.paimon.utils.Preconditions.checkNotNull; -/** Implementation for {@link VectorScan}. */ -public class VectorScanImpl implements VectorScan { +/** Data-evolution implementation for {@link VectorScan}. */ +public class DataEvolutionVectorScan implements VectorScan { private final FileStoreTable table; @Nullable private final PartitionPredicate partitionFilter; @@ -55,7 +55,7 @@ public class VectorScanImpl implements VectorScan { private final DataField vectorColumn; private final Map options; - public VectorScanImpl( + public DataEvolutionVectorScan( FileStoreTable table, @Nullable PartitionPredicate partitionFilter, @Nullable Predicate filter, 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/DataTableBatchScan.java index 634c95e6eaaa..4ba3b19b6b9d 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/DataTableBatchScan.java @@ -19,6 +19,7 @@ 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; @@ -66,6 +67,7 @@ public static void validateScanBucketOption( private final SchemaManager schemaManager; @Nullable private String readProtectionTagName; + @Nullable private GlobalIndexSplitResult globalIndexSplitResult; public DataTableBatchScan( TableSchema schema, @@ -110,6 +112,17 @@ 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); + } if (startingScanner == null) { startingScanner = createStartingScanner(false); } @@ -135,6 +148,14 @@ protected TableScan.Plan planWithoutAuth() { } } + @Override + public DataTableBatchScan withGlobalIndexResult(GlobalIndexResult globalIndexResult) { + if (globalIndexResult instanceof GlobalIndexSplitResult) { + this.globalIndexSplitResult = (GlobalIndexSplitResult) globalIndexResult; + } + return this; + } + @Override public List listPartitionEntries() { if (startingScanner == null) { diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/GlobalIndexSplitResult.java b/paimon-core/src/main/java/org/apache/paimon/table/source/GlobalIndexSplitResult.java new file mode 100644 index 000000000000..913b3f8833f7 --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/GlobalIndexSplitResult.java @@ -0,0 +1,31 @@ +/* + * 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 java.util.List; + +/** Global index result materialized as ready-to-read splits from one table snapshot. */ +public interface GlobalIndexSplitResult extends GlobalIndexResult { + + long snapshotId(); + + List splits(); +} 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 new file mode 100644 index 000000000000..68350ffb83b3 --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorRead.java @@ -0,0 +1,278 @@ +/* + * 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.KeyValueFileStore; +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.deletionvectors.DeletionVector; +import org.apache.paimon.fs.FileIO; +import org.apache.paimon.globalindex.GlobalIndexReadThreadPool; +import org.apache.paimon.globalindex.GlobalIndexResult; +import org.apache.paimon.index.IndexFileHandler; +import org.apache.paimon.index.pkvector.PkVectorAnnSegmentSearcher; +import org.apache.paimon.index.pkvector.PkVectorBucketIndexState; +import org.apache.paimon.index.pkvector.PkVectorDataFileReader; +import org.apache.paimon.index.pkvector.PkVectorSearchResult; +import org.apache.paimon.index.pkvector.PkVectorSourcePolicy; +import org.apache.paimon.index.pkvector.PrimaryKeyVectorBucketSearch; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.io.KeyValueFileReaderFactory; +import org.apache.paimon.options.Options; +import org.apache.paimon.table.DelegatedFileStoreTable; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.types.DataField; +import org.apache.paimon.types.VectorType; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.PriorityQueue; +import java.util.concurrent.ExecutorService; +import java.util.stream.Collectors; + +import static org.apache.paimon.CoreOptions.GLOBAL_INDEX_THREAD_NUM; +import static org.apache.paimon.globalindex.VectorSearchMetric.normalize; +import static org.apache.paimon.utils.Preconditions.checkArgument; + +/** Executes bucket-local primary-key vector search and merges physical candidates globally. */ +public class PrimaryKeyVectorRead implements VectorRead { + + private static final Comparator BEST_FIRST = + (left, right) -> { + int distance = Float.compare(left.distance, right.distance); + if (distance != 0) { + return distance; + } + int partition = compareBytes(left.partition.toBytes(), right.partition.toBytes()); + if (partition != 0) { + return partition; + } + int bucket = Integer.compare(left.bucket, right.bucket); + if (bucket != 0) { + return bucket; + } + int fileName = left.dataFileName.compareTo(right.dataFileName); + return fileName != 0 ? fileName : Long.compare(left.rowPosition, right.rowPosition); + }; + + private final FileIO fileIO; + private final IndexFileHandler indexFileHandler; + private final KeyValueFileReaderFactory.Builder readerFactoryBuilder; + private final DataField vectorField; + private final String indexType; + private final Options indexOptions; + private final ExecutorService executor; + private final Map searchOptions; + private final float[] query; + private final int limit; + private final String metric; + + public PrimaryKeyVectorRead( + FileStoreTable table, + DataField vectorField, + float[] query, + int limit, + Map searchOptions) { + checkArgument( + vectorField.type() instanceof VectorType, "Vector field must use VECTOR type."); + checkArgument( + query.length == ((VectorType) vectorField.type()).getLength(), + "Query dimension %s does not match vector field dimension %s.", + query.length, + ((VectorType) vectorField.type()).getLength()); + checkArgument(limit > 0, "Vector search limit must be positive: %s.", limit); + this.fileIO = table.fileIO(); + this.indexFileHandler = table.store().newIndexFileHandler(); + this.readerFactoryBuilder = keyValueStore(table).newReaderFactoryBuilder(); + this.vectorField = vectorField; + this.indexType = table.coreOptions().primaryKeyVectorIndexType(vectorField.name()); + this.indexOptions = table.coreOptions().primaryKeyVectorIndexOptions(vectorField.name()); + this.executor = + GlobalIndexReadThreadPool.getExecutorService( + table.coreOptions().toConfiguration().get(GLOBAL_INDEX_THREAD_NUM)); + this.searchOptions = Collections.unmodifiableMap(new HashMap<>(searchOptions)); + this.query = query.clone(); + this.limit = limit; + this.metric = + normalize(table.coreOptions().primaryKeyVectorDistanceMetric(vectorField.name())); + } + + private static KeyValueFileStore keyValueStore(FileStoreTable table) { + FileStoreTable unwrapped = table; + while (unwrapped instanceof DelegatedFileStoreTable) { + unwrapped = ((DelegatedFileStoreTable) unwrapped).wrapped(); + } + checkArgument( + unwrapped.store() instanceof KeyValueFileStore, + "Primary-key vector search requires a key-value file store."); + return (KeyValueFileStore) unwrapped.store(); + } + + @Override + public GlobalIndexResult read(VectorScan.Plan plan) { + checkArgument( + plan instanceof PrimaryKeyVectorScan.Plan, + "Primary-key vector read requires a PrimaryKeyVectorScan plan."); + PrimaryKeyVectorScan.Plan primaryKeyPlan = (PrimaryKeyVectorScan.Plan) plan; + try { + List candidates = new ArrayList<>(); + for (VectorSearchSplit searchSplit : primaryKeyPlan.splits()) { + BucketVectorSearchSplit split = (BucketVectorSearchSplit) searchSplit; + checkArgument( + split.dataSplit().snapshotId() == primaryKeyPlan.snapshotId(), + "Vector bucket split snapshot does not match its plan."); + candidates.addAll(search(split)); + } + return new PrimaryKeyVectorResult(primaryKeyPlan, topK(candidates, limit), metric); + } catch (IOException e) { + throw new RuntimeException("Failed to search primary-key vector index.", e); + } + } + + private List search(BucketVectorSearchSplit split) throws IOException { + DataSplit dataSplit = split.dataSplit(); + List activeFiles = + dataSplit.dataFiles().stream() + .filter(PkVectorSourcePolicy::shouldRead) + .collect(Collectors.toList()); + PkVectorBucketIndexState state = + PkVectorBucketIndexState.fromActivePayloads( + vectorField.id(), indexType, split.payloadFiles()); + Map deletionVectors = deletionVectors(dataSplit); + PkVectorDataFileReader.Factory readerFactory = + new PkVectorDataFileReader.Factory( + readerFactoryBuilder, + dataSplit.partition(), + dataSplit.bucket(), + vectorField, + ((VectorType) vectorField.type()).getLength()); + PkVectorAnnSegmentSearcher annSearcher = + new PkVectorAnnSegmentSearcher( + fileIO, + indexFileHandler.pkVectorAnnSegment( + dataSplit.partition(), dataSplit.bucket()), + vectorField, + indexOptions, + metric, + executor); + PrimaryKeyVectorBucketSearch bucketSearch = + new PrimaryKeyVectorBucketSearch(readerFactory, annSearcher, searchOptions, metric); + List candidates = new ArrayList<>(); + for (PkVectorSearchResult result : + bucketSearch.search(state, activeFiles, deletionVectors, query, limit)) { + candidates.add( + new Candidate( + dataSplit.partition(), + dataSplit.bucket(), + result.dataFileName(), + result.rowPosition(), + result.distance())); + } + return candidates; + } + + private Map deletionVectors(DataSplit split) throws IOException { + DeletionVector.Factory factory = + DeletionVector.factory( + fileIO, split.dataFiles(), split.deletionFiles().orElse(null)); + Map result = new HashMap<>(); + for (DataFileMeta file : split.dataFiles()) { + Optional deletionVector = factory.create(file.fileName()); + if (deletionVector.isPresent()) { + result.put(file.fileName(), deletionVector.get()); + } + } + return result; + } + + static List topK(List candidates, int limit) { + checkArgument(limit > 0, "Vector search limit must be positive: %s.", limit); + PriorityQueue nearest = new PriorityQueue<>(limit, BEST_FIRST.reversed()); + for (Candidate candidate : candidates) { + if (nearest.size() < limit) { + nearest.add(candidate); + } else if (BEST_FIRST.compare(candidate, nearest.peek()) < 0) { + nearest.poll(); + nearest.add(candidate); + } + } + List result = new ArrayList<>(nearest); + Collections.sort(result, BEST_FIRST); + return Collections.unmodifiableList(result); + } + + private static int compareBytes(byte[] left, byte[] right) { + int count = Math.min(left.length, right.length); + for (int i = 0; i < count; i++) { + int comparison = Integer.compare(left[i] & 0xFF, right[i] & 0xFF); + if (comparison != 0) { + return comparison; + } + } + return Integer.compare(left.length, right.length); + } + + /** Snapshot-scoped physical row candidate. */ + public static class Candidate { + + private final BinaryRow partition; + private final int bucket; + private final String dataFileName; + private final long rowPosition; + private final float distance; + + Candidate( + BinaryRow partition, + int bucket, + String dataFileName, + long rowPosition, + float distance) { + this.partition = partition.copy(); + this.bucket = bucket; + this.dataFileName = dataFileName; + this.rowPosition = rowPosition; + this.distance = distance; + } + + public BinaryRow partition() { + return partition; + } + + public int bucket() { + return bucket; + } + + public String dataFileName() { + return dataFileName; + } + + public long rowPosition() { + return rowPosition; + } + + public float distance() { + return distance; + } + } +} diff --git a/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorResult.java b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorResult.java new file mode 100644 index 000000000000..3eba24b30cb8 --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorResult.java @@ -0,0 +1,228 @@ +/* + * 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.IndexedSplit; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.utils.Range; +import org.apache.paimon.utils.RoaringNavigableMap64; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.TreeMap; + +import static org.apache.paimon.globalindex.VectorSearchMetric.normalize; +import static org.apache.paimon.utils.Preconditions.checkArgument; + +/** Snapshot-scoped vector result addressed by physical primary-key table row positions. */ +public class PrimaryKeyVectorResult implements GlobalIndexSplitResult { + + private final PrimaryKeyVectorScan.Plan plan; + private final List candidates; + private final String metric; + + PrimaryKeyVectorResult( + PrimaryKeyVectorScan.Plan plan, + List candidates, + String metric) { + this.plan = plan; + this.candidates = Collections.unmodifiableList(new ArrayList<>(candidates)); + this.metric = normalize(metric); + } + + @Override + public long snapshotId() { + return plan.snapshotId(); + } + + @Override + public List splits() { + Map sourceSplits = sourceSplits(); + Map> selectedByFile = new LinkedHashMap<>(); + for (PrimaryKeyVectorRead.Candidate candidate : candidates) { + FileKey key = + new FileKey( + candidate.partition(), candidate.bucket(), candidate.dataFileName()); + BucketVectorSearchSplit bucketSplit = sourceSplits.get(key); + checkArgument( + bucketSplit != null, + "Primary-key vector candidate references unknown data file %s in bucket %s.", + candidate.dataFileName(), + candidate.bucket()); + DataFileMeta dataFile = dataFile(bucketSplit.dataSplit(), candidate.dataFileName()); + checkArgument( + candidate.rowPosition() >= 0 + && candidate.rowPosition() < dataFile.rowCount() + && candidate.rowPosition() <= Integer.MAX_VALUE, + "Primary-key vector row position %s is outside data file %s.", + candidate.rowPosition(), + candidate.dataFileName()); + TreeMap selected = + selectedByFile.computeIfAbsent(key, ignored -> new TreeMap<>()); + checkArgument( + selected.put((int) candidate.rowPosition(), score(candidate.distance())) + == null, + "Primary-key vector candidate contains duplicate row position %s for data file %s.", + candidate.rowPosition(), + candidate.dataFileName()); + } + + List result = new ArrayList<>(selectedByFile.size()); + for (Map.Entry> entry : selectedByFile.entrySet()) { + BucketVectorSearchSplit bucketSplit = sourceSplits.get(entry.getKey()); + DataSplit source = bucketSplit.dataSplit(); + int fileIndex = fileIndex(source, entry.getKey().dataFileName); + DataSplit.Builder builder = + DataSplit.builder() + .withSnapshot(source.snapshotId()) + .withPartition(source.partition()) + .withBucket(source.bucket()) + .withBucketPath(source.bucketPath()) + .withTotalBuckets(source.totalBuckets()) + .withDataFiles( + Collections.singletonList(source.dataFiles().get(fileIndex))) + .isStreaming(false) + .rawConvertible(false); + if (source.deletionFiles().isPresent()) { + builder.withDataDeletionFiles( + Collections.singletonList(source.deletionFiles().get().get(fileIndex))); + } + result.add( + new IndexedSplit( + builder.build(), ranges(entry.getValue()), scores(entry.getValue()))); + } + return Collections.unmodifiableList(result); + } + + @Override + public RoaringNavigableMap64 results() { + throw new UnsupportedOperationException( + "Primary-key vector results use physical file positions, not global row ids."); + } + + private Map sourceSplits() { + Map result = new HashMap<>(); + for (VectorSearchSplit searchSplit : plan.splits()) { + BucketVectorSearchSplit bucketSplit = (BucketVectorSearchSplit) searchSplit; + for (DataFileMeta dataFile : bucketSplit.dataSplit().dataFiles()) { + FileKey key = + new FileKey( + bucketSplit.dataSplit().partition(), + bucketSplit.dataSplit().bucket(), + dataFile.fileName()); + checkArgument( + result.put(key, bucketSplit) == null, + "Data file %s appears more than once in primary-key vector plan.", + dataFile.fileName()); + } + } + return result; + } + + private static List ranges(TreeMap selected) { + List result = new ArrayList<>(); + long from = -1; + long to = -1; + for (int position : selected.keySet()) { + if (from < 0) { + from = position; + } else if (position != to + 1) { + result.add(new Range(from, to)); + from = position; + } + to = position; + } + if (from >= 0) { + result.add(new Range(from, to)); + } + return result; + } + + private static float[] scores(TreeMap selected) { + float[] result = new float[selected.size()]; + int index = 0; + for (float score : selected.values()) { + result[index++] = score; + } + return result; + } + + private static DataFileMeta dataFile(DataSplit split, String dataFileName) { + return split.dataFiles().get(fileIndex(split, dataFileName)); + } + + private static int fileIndex(DataSplit split, String dataFileName) { + for (int i = 0; i < split.dataFiles().size(); i++) { + if (dataFileName.equals(split.dataFiles().get(i).fileName())) { + return i; + } + } + throw new IllegalArgumentException( + "Data file " + dataFileName + " does not exist in vector bucket split."); + } + + private float score(float distance) { + if ("l2".equals(metric)) { + return 1F / (1F + distance); + } else if ("cosine".equals(metric)) { + return 1F - distance; + } else if ("inner_product".equals(metric)) { + return -distance; + } + throw new IllegalArgumentException("Unsupported primary-key vector metric: " + metric); + } + + private static class FileKey { + + private final BinaryRow partition; + private final int bucket; + private final String dataFileName; + + private FileKey(BinaryRow partition, int bucket, String dataFileName) { + this.partition = partition.copy(); + this.bucket = bucket; + this.dataFileName = dataFileName; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + FileKey fileKey = (FileKey) o; + return bucket == fileKey.bucket + && partition.equals(fileKey.partition) + && dataFileName.equals(fileKey.dataFileName); + } + + @Override + public int hashCode() { + return Objects.hash(partition, bucket, dataFileName); + } + } +} 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 new file mode 100644 index 000000000000..bb887996178a --- /dev/null +++ b/paimon-core/src/main/java/org/apache/paimon/table/source/PrimaryKeyVectorScan.java @@ -0,0 +1,232 @@ +/* + * 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.data.BinaryRow; +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.FileStoreTable; +import org.apache.paimon.table.source.snapshot.SnapshotReader; +import org.apache.paimon.table.source.snapshot.TimeTravelUtil; +import org.apache.paimon.utils.Pair; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.apache.paimon.utils.Preconditions.checkArgument; + +/** Plans complete bucket-local vector inputs from one captured primary-key table snapshot. */ +public class PrimaryKeyVectorScan implements VectorScan { + + private final FileStoreTable table; + private final int vectorFieldId; + private final String indexType; + @Nullable private final PartitionPredicate partitionFilter; + + public PrimaryKeyVectorScan( + FileStoreTable table, + int vectorFieldId, + String indexType, + @Nullable PartitionPredicate partitionFilter) { + this.table = table; + this.vectorFieldId = vectorFieldId; + this.indexType = indexType; + this.partitionFilter = partitionFilter; + } + + @Override + 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()); + } + + SnapshotReader snapshotReader = + table.newSnapshotReader().withSnapshot(snapshot).withMode(ScanMode.ALL).keepStats(); + if (partitionFilter != null) { + snapshotReader.withPartitionFilter(partitionFilter); + } + List dataSplits = snapshotReader.read().dataSplits(); + + IndexFileHandler indexFileHandler = table.store().newIndexFileHandler(); + List vectorIndexEntries = + indexFileHandler.scan( + snapshot, + entry -> + entry.indexFile().indexType().equals(indexType) + && entry.indexFile().globalIndexMeta() != null + && entry.indexFile().globalIndexMeta().sourceMeta() != null + && entry.indexFile().globalIndexMeta().indexFieldId() + == vectorFieldId + && (partitionFilter == null + || partitionFilter.test(entry.partition()))); + return plan(snapshot.id(), dataSplits, vectorIndexEntries); + } + + static Plan plan( + long snapshotId, + List dataSplits, + List vectorIndexEntries) { + Map, List> payloads = new LinkedHashMap<>(); + for (IndexManifestEntry entry : vectorIndexEntries) { + checkArgument( + entry.kind() == FileKind.ADD, + "Primary-key vector index file %s is not active.", + entry.indexFile().fileName()); + checkArgument( + entry.indexFile().globalIndexMeta() != null + && entry.indexFile().globalIndexMeta().sourceMeta() != null, + "Primary-key vector index file %s has no source metadata.", + entry.indexFile().fileName()); + Pair key = Pair.of(entry.partition(), entry.bucket()); + payloads.computeIfAbsent(key, ignored -> new ArrayList<>()).add(entry.indexFile()); + } + + Map, BucketAccumulator> buckets = new LinkedHashMap<>(); + for (DataSplit split : dataSplits) { + checkArgument( + split.snapshotId() == snapshotId, + "Data split snapshot %s does not match vector scan snapshot %s.", + split.snapshotId(), + snapshotId); + checkArgument( + !split.isStreaming(), "Primary-key vector search requires a batch split."); + Pair key = Pair.of(split.partition(), split.bucket()); + BucketAccumulator accumulator = buckets.get(key); + if (accumulator == null) { + accumulator = new BucketAccumulator(split); + buckets.put(key, accumulator); + } + accumulator.add(split); + } + + List result = new ArrayList<>(buckets.size()); + for (Map.Entry, BucketAccumulator> entry : buckets.entrySet()) { + result.add( + new BucketVectorSearchSplit( + entry.getValue().build(), + payloads.getOrDefault(entry.getKey(), Collections.emptyList()))); + } + return new Plan(snapshotId, result); + } + + /** Immutable snapshot vector-search plan. */ + public static class Plan implements VectorScan.Plan { + + private final long snapshotId; + private final List splits; + + private Plan(long snapshotId, List splits) { + this.snapshotId = snapshotId; + this.splits = Collections.unmodifiableList(new ArrayList(splits)); + } + + public long snapshotId() { + return snapshotId; + } + + @Override + public List splits() { + return splits; + } + } + + private static class BucketAccumulator { + + private final long snapshotId; + private final BinaryRow partition; + private final int bucket; + private final String bucketPath; + @Nullable private final Integer totalBuckets; + private final List dataFiles = new ArrayList<>(); + private final List deletionFiles = new ArrayList<>(); + private final Set dataFileNames = new HashSet<>(); + private boolean hasDeletionFile; + + private BucketAccumulator(DataSplit split) { + this.snapshotId = split.snapshotId(); + this.partition = split.partition(); + this.bucket = split.bucket(); + this.bucketPath = split.bucketPath(); + this.totalBuckets = split.totalBuckets(); + } + + private void add(DataSplit split) { + checkArgument( + snapshotId == split.snapshotId() + && partition.equals(split.partition()) + && bucket == split.bucket() + && bucketPath.equals(split.bucketPath()), + "Cannot combine data splits from different snapshot buckets."); + checkArgument( + totalBuckets == null + ? split.totalBuckets() == null + : totalBuckets.equals(split.totalBuckets()), + "Bucket split total-bucket metadata is inconsistent."); + + List splitDeletions = split.deletionFiles().orElse(null); + checkArgument( + splitDeletions == null || splitDeletions.size() == split.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); + checkArgument( + dataFileNames.add(file.fileName()), + "Data file %s appears more than once in vector bucket planning.", + file.fileName()); + dataFiles.add(file); + DeletionFile deletion = splitDeletions == null ? null : splitDeletions.get(i); + deletionFiles.add(deletion); + hasDeletionFile |= deletion != null; + } + } + + private DataSplit build() { + DataSplit.Builder builder = + DataSplit.builder() + .withSnapshot(snapshotId) + .withPartition(partition) + .withBucket(bucket) + .withBucketPath(bucketPath) + .withTotalBuckets(totalBuckets) + .withDataFiles(dataFiles) + .isStreaming(false) + .rawConvertible(false); + if (hasDeletionFile) { + builder.withDataDeletionFiles(deletionFiles); + } + return builder.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 9b77c9c6a22c..f2bcb7af65cb 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,6 +33,7 @@ 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}. */ @@ -124,13 +125,34 @@ public VectorSearchBuilder withOption(String key, String value) { @Override public VectorScan newVectorScan() { - return new VectorScanImpl(table, partitionFilter, filter, vectorColumn, options); + 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); + } + return new DataEvolutionVectorScan(table, partitionFilter, filter, vectorColumn, options); } @Override public VectorRead newVectorRead() { checkNotNull(vector, "vector must be set via withVector()"); - return new VectorReadImpl( + 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 DataEvolutionVectorRead( table, partitionFilter, filter, limit, vectorColumn, vector, options); } + + private boolean isPrimaryKeyVectorSearch() { + return vectorColumn != null + && table.coreOptions().primaryKeyVectorIndexColumns().contains(vectorColumn.name()); + } } diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorReadTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorReadTest.java new file mode 100644 index 000000000000..cf7b44773281 --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorReadTest.java @@ -0,0 +1,58 @@ +/* + * 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.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests global candidate merging for primary-key vector search. */ +class PrimaryKeyVectorReadTest { + + @Test + void testMergesGlobalTopKWithDeterministicTies() { + List candidates = + Arrays.asList( + candidate(1, "file-c", 0, 2F), + candidate(1, "file-b", 1, 1F), + candidate(0, "file-a", 2, 1F)); + + List result = PrimaryKeyVectorRead.topK(candidates, 2); + + assertThat(result) + .extracting( + PrimaryKeyVectorRead.Candidate::bucket, + PrimaryKeyVectorRead.Candidate::dataFileName, + PrimaryKeyVectorRead.Candidate::rowPosition) + .containsExactly( + org.assertj.core.groups.Tuple.tuple(0, "file-a", 2L), + org.assertj.core.groups.Tuple.tuple(1, "file-b", 1L)); + } + + private static PrimaryKeyVectorRead.Candidate candidate( + int bucket, String fileName, long position, float distance) { + return new PrimaryKeyVectorRead.Candidate( + BinaryRow.EMPTY_ROW, bucket, fileName, position, distance); + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorResultTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorResultTest.java new file mode 100644 index 000000000000..849fd76ba73d --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorResultTest.java @@ -0,0 +1,96 @@ +/* + * 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.IndexedSplit; +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; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests conversion from primary-key vector candidates to {@link IndexedSplit}s. */ +class PrimaryKeyVectorResultTest { + + @Test + void testGroupsCandidatesByFileAndAlignsScoresWithRanges() { + DataFileMeta dataFile = dataFile("data-1"); + DeletionFile deletionFile = new DeletionFile("dv", 10, 20, 1L); + DataSplit dataSplit = dataSplit(dataFile, deletionFile); + PrimaryKeyVectorScan.Plan plan = + PrimaryKeyVectorScan.plan( + 11, Collections.singletonList(dataSplit), Collections.emptyList()); + List candidates = + Arrays.asList( + new PrimaryKeyVectorRead.Candidate( + BinaryRow.EMPTY_ROW, 0, "data-1", 4, 0.1F), + new PrimaryKeyVectorRead.Candidate(BinaryRow.EMPTY_ROW, 0, "data-1", 1, 1F), + new PrimaryKeyVectorRead.Candidate( + BinaryRow.EMPTY_ROW, 0, "data-1", 2, 4F)); + + List splits = new PrimaryKeyVectorResult(plan, candidates, "l2").splits(); + + assertThat(splits).hasSize(1); + IndexedSplit split = splits.get(0); + assertThat(split.dataSplit().snapshotId()).isEqualTo(11); + assertThat(split.dataSplit().dataFiles()).containsExactly(dataFile); + assertThat(split.dataSplit().deletionFiles().get()).containsExactly(deletionFile); + assertThat(split.rowRanges()).containsExactly(new Range(1, 2), new Range(4, 4)); + assertThat(split.scores()) + .containsExactly(1F / (1F + 1F), 1F / (1F + 4F), 1F / (1F + 0.1F)); + } + + private static DataSplit dataSplit(DataFileMeta dataFile, DeletionFile deletionFile) { + return DataSplit.builder() + .withSnapshot(11) + .withPartition(BinaryRow.EMPTY_ROW) + .withBucket(0) + .withBucketPath("bucket-0") + .withTotalBuckets(1) + .withDataFiles(Collections.singletonList(dataFile)) + .withDataDeletionFiles(Collections.singletonList(deletionFile)) + .build(); + } + + private static DataFileMeta dataFile(String fileName) { + return DataFileMeta.forAppend( + fileName, + 100, + 5, + SimpleStats.EMPTY_STATS, + 0, + 1, + 1, + Collections.emptyList(), + null, + FileSource.COMPACT, + null, + null, + 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 new file mode 100644 index 000000000000..d3db372da388 --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorScanTest.java @@ -0,0 +1,228 @@ +/* + * 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.FileStore; +import org.apache.paimon.Snapshot; +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.index.GlobalIndexMeta; +import org.apache.paimon.index.IndexFileHandler; +import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.index.pkvector.PkVectorSourceFile; +import org.apache.paimon.index.pkvector.PkVectorSourceMeta; +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.stats.SimpleStats; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.source.snapshot.SnapshotReader; +import org.apache.paimon.utils.Filter; + +import org.junit.jupiter.api.Test; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.ObjectInputStream; +import java.io.ObjectOutputStream; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Optional; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Answers.CALLS_REAL_METHODS; +import static org.mockito.Answers.RETURNS_SELF; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** Tests snapshot-consistent planning for bucket-local primary-key vector search. */ +class PrimaryKeyVectorScanTest { + + @Test + @SuppressWarnings({"unchecked", "rawtypes"}) + void testScansOneSnapshotAndFiltersVectorIdentity() { + CoreOptions coreOptions = coreOptions(); + FileStoreTable table = mock(FileStoreTable.class); + Snapshot snapshot = mock(Snapshot.class); + when(snapshot.id()).thenReturn(11L); + when(table.coreOptions()).thenReturn(coreOptions); + when(table.latestSnapshot()).thenReturn(Optional.of(snapshot)); + + SnapshotReader snapshotReader = mock(SnapshotReader.class, RETURNS_SELF); + SnapshotReader.Plan snapshotPlan = mock(SnapshotReader.Plan.class, CALLS_REAL_METHODS); + when(snapshotPlan.splits()) + .thenReturn(Collections.singletonList(dataSplit(dataFile("data-1")))); + when(snapshotReader.read()).thenReturn(snapshotPlan); + when(table.newSnapshotReader()).thenReturn(snapshotReader); + + List entries = + Arrays.asList( + payloadEntry("ivf-pq", 7, "ann-match"), + payloadEntry("ivf-pq", 8, "ann-other-field"), + payloadEntry("hnsw", 7, "ann-other-type")); + IndexFileHandler indexFileHandler = mock(IndexFileHandler.class); + when(indexFileHandler.scan(eq(snapshot), any(Filter.class))) + .thenAnswer( + invocation -> { + Filter filter = invocation.getArgument(1); + List filtered = new ArrayList<>(); + for (IndexManifestEntry entry : entries) { + if (filter.test(entry)) { + filtered.add(entry); + } + } + return filtered; + }); + FileStore store = mock(FileStore.class); + when(store.newIndexFileHandler()).thenReturn(indexFileHandler); + when(table.store()).thenReturn(store); + + PrimaryKeyVectorScan.Plan plan = new PrimaryKeyVectorScan(table, 7, "ivf-pq", null).scan(); + + assertThat(plan.snapshotId()).isEqualTo(11); + assertThat(plan.splits()).hasSize(1); + BucketVectorSearchSplit bucketSplit = (BucketVectorSearchSplit) plan.splits().get(0); + assertThat(bucketSplit.payloadFiles()) + .extracting(IndexFileMeta::fileName) + .containsExactly("ann-match"); + } + + @Test + void testMergesDataSplitsForCompleteBucketCoverage() { + DataFileMeta data1 = dataFile("data-1"); + DataFileMeta data2 = dataFile("data-2"); + DeletionFile deletion = new DeletionFile("dv", 10, 20, 1L); + DataSplit split1 = dataSplit(data1, Collections.singletonList(deletion)); + DataSplit split2 = dataSplit(data2, null); + + PrimaryKeyVectorScan.Plan plan = + PrimaryKeyVectorScan.plan( + 11, + Arrays.asList(split1, split2), + Collections.singletonList(payloadEntry())); + + assertThat(plan.snapshotId()).isEqualTo(11); + assertThat(plan.splits()).hasSize(1); + BucketVectorSearchSplit bucketSplit = (BucketVectorSearchSplit) plan.splits().get(0); + assertThat(bucketSplit.dataSplit().dataFiles()).containsExactly(data1, data2); + assertThat(bucketSplit.dataSplit().deletionFiles()).isPresent(); + assertThat(bucketSplit.dataSplit().deletionFiles().get()).containsExactly(deletion, null); + assertThat(bucketSplit.payloadFiles()).containsExactly(payloadFile()); + } + + @Test + void testBucketSplitSerialization() throws Exception { + IndexFileMeta payload = payloadFile(); + BucketVectorSearchSplit split = + new BucketVectorSearchSplit( + dataSplit(dataFile("data-1")), Collections.singletonList(payload)); + + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (ObjectOutputStream output = new ObjectOutputStream(bytes)) { + output.writeObject(split); + } + BucketVectorSearchSplit restored; + try (ObjectInputStream input = + new ObjectInputStream(new ByteArrayInputStream(bytes.toByteArray()))) { + restored = (BucketVectorSearchSplit) input.readObject(); + } + + assertThat(restored).isEqualTo(split); + assertThat(restored.payloadFiles().get(0).globalIndexMeta().sourceMeta()) + .isEqualTo(payload.globalIndexMeta().sourceMeta()); + } + + private static DataSplit dataSplit(DataFileMeta dataFile) { + return dataSplit(dataFile, null); + } + + private static DataSplit dataSplit( + DataFileMeta dataFile, java.util.List deletionFiles) { + DataSplit.Builder builder = + DataSplit.builder() + .withSnapshot(11) + .withPartition(BinaryRow.EMPTY_ROW) + .withBucket(0) + .withBucketPath("bucket-0") + .withTotalBuckets(1) + .withDataFiles(Collections.singletonList(dataFile)); + if (deletionFiles != null) { + builder.withDataDeletionFiles(deletionFiles); + } + return builder.build(); + } + + private static IndexManifestEntry payloadEntry() { + return new IndexManifestEntry(FileKind.ADD, BinaryRow.EMPTY_ROW, 0, payloadFile()); + } + + private static IndexManifestEntry payloadEntry(String indexType, int fieldId, String fileName) { + return new IndexManifestEntry( + FileKind.ADD, BinaryRow.EMPTY_ROW, 0, payloadFile(indexType, fieldId, fileName)); + } + + private static IndexFileMeta payloadFile() { + return payloadFile("ivf-pq", 7, "ann"); + } + + private static IndexFileMeta payloadFile(String indexType, int fieldId, String fileName) { + byte[] sourceMeta = + new PkVectorSourceMeta( + Collections.singletonList(new PkVectorSourceFile("data-1", 2))) + .serialize(); + return new IndexFileMeta( + indexType, + fileName, + 100, + 2, + new GlobalIndexMeta(0, 1, fieldId, null, null, sourceMeta), + null); + } + + private static CoreOptions coreOptions() { + Options options = new Options(); + options.set(CoreOptions.PK_VECTOR_INDEX_COLUMNS, "embedding"); + options.setString("fields.embedding.pk-vector.index.type", "ivf-pq"); + return new CoreOptions(options); + } + + private static DataFileMeta dataFile(String fileName) { + return DataFileMeta.forAppend( + fileName, + 100, + 2, + SimpleStats.EMPTY_STATS, + 0, + 0, + 1, + Collections.emptyList(), + null, + FileSource.COMPACT, + null, + null, + null, + null); + } +} 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 new file mode 100644 index 000000000000..090488de9e53 --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/PrimaryKeyVectorSearchTest.java @@ -0,0 +1,97 @@ +/* + * 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.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.reader.RecordReader; +import org.apache.paimon.schema.Schema; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.TableTestBase; +import org.apache.paimon.table.sink.BatchTableCommit; +import org.apache.paimon.table.sink.BatchTableWrite; +import org.apache.paimon.table.sink.BatchWriteBuilder; +import org.apache.paimon.types.DataTypes; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** End-to-end tests for bucket-local primary-key vector search. */ +class PrimaryKeyVectorSearchTest extends TableTestBase { + + @Override + protected Schema schemaDefault() { + return Schema.newBuilder() + .column("id", 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_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(); + } + + @Test + void testVectorSearchMaterializesPhysicalRows() throws Exception { + createTableDefault(); + FileStoreTable table = getTableDefault(); + BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder(); + try (BatchTableWrite write = writeBuilder.newWrite(); + BatchTableCommit commit = writeBuilder.newCommit()) { + write.withIOManager(ioManager); + write.write(GenericRow.of(1, BinaryVector.fromPrimitiveArray(new float[] {3, 0}))); + write.write(GenericRow.of(2, BinaryVector.fromPrimitiveArray(new float[] {1, 0}))); + write.write(GenericRow.of(3, BinaryVector.fromPrimitiveArray(new float[] {2, 0}))); + commit.commit(write.prepareCommit()); + } + + GlobalIndexResult result = + table.newVectorSearchBuilder() + .withVectorColumn("embedding") + .withVector(new float[] {0, 0}) + .withLimit(2) + .executeLocal(); + assertThat(result).isInstanceOf(GlobalIndexSplitResult.class); + + ReadBuilder readBuilder = table.newReadBuilder(); + TableScan.Plan plan = readBuilder.newScan().withGlobalIndexResult(result).plan(); + assertThat(plan.splits()).allMatch(IndexedSplit.class::isInstance); + List ids = new ArrayList<>(); + try (RecordReader reader = readBuilder.newRead().createReader(plan)) { + reader.forEachRemaining(row -> ids.add(row.getInt(0))); + } + + assertThat(ids).containsExactly(2, 3); + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchBuilderTest.java b/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchBuilderTest.java index fefe3bce607a..432e5333bf79 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchBuilderTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/source/VectorSearchBuilderTest.java @@ -594,7 +594,7 @@ public void testRawRefineReadTypeContainsOnlyVectorAndRowId() throws Exception { FileStoreTable table = getTableDefault(); Predicate idFilter = new PredicateBuilder(table.rowType()).greaterOrEqual(0, 1); - ExposingVectorRead read = new ExposingVectorRead(table, idFilter); + ExposingDataEvolutionVectorRead read = new ExposingDataEvolutionVectorRead(table, idFilter); assertThat(read.rawReadType(false).getFieldNames()) .containsExactly(VECTOR_FIELD_NAME, SpecialFields.ROW_ID.name()); @@ -611,7 +611,7 @@ public void testRawCandidateSearchScoresOnlyCandidateBitmap() throws Exception { RoaringNavigableMap64 candidates = new RoaringNavigableMap64(); candidates.add(1L); - ExposingVectorRead read = new ExposingVectorRead(table, null); + ExposingDataEvolutionVectorRead read = new ExposingDataEvolutionVectorRead(table, null); ScoredGlobalIndexResult result = read.rawCandidateSearch( Collections.singletonList(new Range(0, 2)), @@ -937,7 +937,7 @@ public void testScanPartialRangeIntersection() throws Exception { // Build ONE btree index covering partial range [3,7] buildAndCommitBTreeIndex(table, new int[] {3, 4, 5, 6, 7}, new Range(3, 7)); - // VectorScanImpl should attach scalar index because [3,7] intersects [0,9] + // DataEvolutionVectorScan should attach scalar index because [3,7] intersects [0,9] Predicate idFilter = new PredicateBuilder(table.rowType()).greaterOrEqual(0, 5); VectorScan.Plan plan = table.newVectorSearchBuilder() @@ -1409,9 +1409,9 @@ private void buildAndCommitIndex(FileStoreTable table, float[][] vectors) throws buildAndCommitIndex(table, VECTOR_FIELD_NAME, vectors); } - private static class ExposingVectorRead extends VectorReadImpl { + private static class ExposingDataEvolutionVectorRead extends DataEvolutionVectorRead { - private ExposingVectorRead(FileStoreTable table, Predicate filter) { + private ExposingDataEvolutionVectorRead(FileStoreTable table, Predicate filter) { super( table, null, @@ -1654,7 +1654,7 @@ public void testVectorSearchWithBTreePreFilter() throws Exception { buildAndCommitBTreeIndex(table, new int[] {0, 1, 2, 3, 4}, range1); buildAndCommitBTreeIndex(table, new int[] {5, 6, 7, 8, 9}, range2); - // --- Test VectorScanImpl: verify splits contain scalar index files --- + // --- Test DataEvolutionVectorScan: verify splits contain scalar index files --- Predicate idFilter = new PredicateBuilder(table.rowType()).greaterOrEqual(0, 5); VectorSearchBuilder searchBuilder = table.newVectorSearchBuilder() @@ -1676,7 +1676,7 @@ public void testVectorSearchWithBTreePreFilter() throws Exception { .count(); assertThat(scalarCount).isGreaterThan(0); - // --- Test VectorReadImpl: pre-filter should narrow results --- + // --- Test DataEvolutionVectorRead: pre-filter should narrow results --- // Query vector near (0,1) with filter id >= 5 // Without filter: rows 5,6,7,8,9 are closest // With filter id >= 5: btree pre-filter restricts to rows 5-9 diff --git a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java index 1b51e9cce103..3c13fb4c0c65 100644 --- a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java +++ b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/read/SparkVectorReadImpl.java @@ -28,9 +28,9 @@ import org.apache.paimon.partition.PartitionPredicate; import org.apache.paimon.predicate.Predicate; import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.source.DataEvolutionVectorRead; import org.apache.paimon.table.source.IndexVectorSearchSplit; import org.apache.paimon.table.source.RawVectorSearchSplit; -import org.apache.paimon.table.source.VectorReadImpl; import org.apache.paimon.table.source.VectorScan; import org.apache.paimon.types.DataField; import org.apache.paimon.utils.InstantiationUtil; @@ -51,10 +51,10 @@ import static org.apache.paimon.CoreOptions.GLOBAL_INDEX_THREAD_NUM; /** - * Spark-aware {@link VectorReadImpl} that distributes grouped vector index evaluation across the - * Spark cluster instead of evaluating them with the local thread pool. + * Spark-aware {@link DataEvolutionVectorRead} that distributes grouped vector index evaluation + * across the Spark cluster instead of evaluating them with the local thread pool. */ -public class SparkVectorReadImpl extends VectorReadImpl { +public class SparkVectorReadImpl extends DataEvolutionVectorRead { private static final long serialVersionUID = 1L;