From 2fafc3dc4605285444d0dc2148c5318c3b9a2534 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Thu, 24 Jul 2025 19:04:27 +0200 Subject: [PATCH 01/34] save point -- in the middle of just mess --- fdb-extensions/fdb-extensions.gradle | 1 + .../async/hnsw/AbstractChangeSet.java | 94 + .../foundationdb/async/hnsw/AbstractNode.java | 73 + .../async/hnsw/AbstractStorageAdapter.java | 213 ++ .../async/hnsw/ByNodeStorageAdapter.java | 324 +++ .../foundationdb/async/hnsw/DataNode.java | 57 + .../apple/foundationdb/async/hnsw/HNSW.java | 2435 +++++++++++++++++ .../async/hnsw/IntermediateNode.java | 57 + .../apple/foundationdb/async/hnsw/Metric.java | 129 + .../foundationdb/async/hnsw/Neighbor.java | 39 + .../async/hnsw/NeighborWithVector.java | 46 + .../apple/foundationdb/async/hnsw/Node.java | 70 + .../foundationdb/async/hnsw/NodeHelpers.java | 80 + .../foundationdb/async/hnsw/NodeKind.java | 60 + .../async/hnsw/OnReadListener.java | 54 + .../async/hnsw/OnWriteListener.java | 63 + .../async/hnsw/StorageAdapter.java | 165 ++ .../apple/foundationdb/async/hnsw/Vector.java | 129 + .../foundationdb/async/hnsw/package-info.java | 24 + .../indexes/VectorIndexSimpleTest.java | 45 + .../indexes/VectorIndexTestBase.java | 223 ++ .../evolution/test_field_type_change.proto | 2 +- .../evolution/test_header_as_group.proto | 2 +- .../evolution/test_swap_union_fields.proto | 2 +- .../src/test/proto/expression_tests.proto | 2 +- .../src/test/proto/test_no_record_types.proto | 2 +- .../src/test/proto/test_records_8.proto | 2 +- .../test/proto/test_records_chained_2.proto | 2 +- .../test_records_duplicate_union_fields.proto | 2 +- ...rds_duplicate_union_fields_reordered.proto | 2 +- .../src/test/proto/test_records_oneof.proto | 2 +- .../test/proto/test_records_transform.proto | 2 +- .../src/test/proto/test_records_vector.proto | 38 + gradle/libs.versions.toml | 2 + 34 files changed, 4432 insertions(+), 11 deletions(-) create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractChangeSet.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Neighbor.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NeighborWithVector.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeHelpers.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKind.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/package-info.java create mode 100644 fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java create mode 100644 fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java create mode 100644 fdb-record-layer-core/src/test/proto/test_records_vector.proto diff --git a/fdb-extensions/fdb-extensions.gradle b/fdb-extensions/fdb-extensions.gradle index bf281a2314..45b2e09302 100644 --- a/fdb-extensions/fdb-extensions.gradle +++ b/fdb-extensions/fdb-extensions.gradle @@ -26,6 +26,7 @@ dependencies { } api(libs.fdbJava) implementation(libs.guava) + implementation(libs.half4j) implementation(libs.slf4j.api) compileOnly(libs.jsr305) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractChangeSet.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractChangeSet.java new file mode 100644 index 0000000000..1b11d58856 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractChangeSet.java @@ -0,0 +1,94 @@ +/* + * AbstractChangeSet.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.Transaction; +import com.apple.foundationdb.async.hnsw.Node.ChangeSet; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +/** + * Abstract base implementations for all {@link ChangeSet}s. + * @param slot type class + * @param node type class (self type) + */ +public abstract class AbstractChangeSet> implements ChangeSet { + @Nullable + private final ChangeSet previousChangeSet; + + @Nonnull + private final N node; + + private final int level; + + AbstractChangeSet(@Nullable final ChangeSet previousChangeSet, @Nonnull final N node, final int level) { + this.previousChangeSet = previousChangeSet; + this.node = node; + this.level = level; + } + + @Override + public void apply(@Nonnull final Transaction transaction) { + if (previousChangeSet != null) { + previousChangeSet.apply(transaction); + } + } + + /** + * Previous change set in the chain of change sets. Can be {@code null} if there is no previous change set. + * @return the previous change set in the chain of change sets + */ + @Nullable + public ChangeSet getPreviousChangeSet() { + return previousChangeSet; + } + + /** + * The node this change set applies to. + * @return the node this change set applies to + */ + @Nonnull + public N getNode() { + return node; + } + + /** + * The level we should use when maintaining the node slot index. If {@code level < 0}, do not maintain the node slot + * index. + * @return the level used when maintaing the node slot index + */ + public int getLevel() { + return level; + } + + /** + * Returns whether this change set needs to also update the node slot index. There are scenarios where we + * do not need to update such an index in general. For instance, the user may not want to use such an index. + * In addition to that, there are change set implementations that should not update the index even if such and index + * is maintained in general. For instance, the moved-in slots were already persisted in the database before the + * move-in operation. We should not update the node slot index in such a case. + * @return {@code true} if we need to update the node slot index, {@code false} otherwise + */ + public boolean isUpdateNodeSlotIndex() { + return level >= 0; + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java new file mode 100644 index 0000000000..87e03b393b --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java @@ -0,0 +1,73 @@ +/* + * AbstractNode.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.tuple.Tuple; +import com.christianheina.langx.half4j.Half; +import com.google.common.collect.ImmutableList; + +import javax.annotation.Nonnull; +import java.util.List; + +/** + * TODO. + * @param node type class. + */ +abstract class AbstractNode implements Node { + @Nonnull + private final Tuple primaryKey; + + @Nonnull + private final Vector vector; + + @Nonnull + private final List neighbors; + + protected AbstractNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, + @Nonnull final List neighbors) { + this.primaryKey = primaryKey; + this.vector = vector; + this.neighbors = ImmutableList.copyOf(neighbors); + } + + @Nonnull + @Override + public Tuple getPrimaryKey() { + return primaryKey; + } + + @Nonnull + public Vector getVector() { + return vector; + } + + @Nonnull + @Override + public List getNeighbors() { + return neighbors; + } + + @Nonnull + @Override + public N getNeighbor(final int index) { + return neighbors.get(index); + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java new file mode 100644 index 0000000000..57c8f3d730 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java @@ -0,0 +1,213 @@ +/* + * AbstractStorageAdapter.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.ReadTransaction; +import com.apple.foundationdb.Transaction; +import com.apple.foundationdb.subspace.Subspace; +import com.apple.foundationdb.tuple.Tuple; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.math.BigInteger; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; + +/** + * Implementations and attributes common to all concrete implementations of {@link StorageAdapter}. + */ +abstract class AbstractStorageAdapter implements StorageAdapter { + public static final byte SUBSPACE_PREFIX_ENTRY_NODE = 0x01; + public static final byte SUBSPACE_PREFIX_DATA = 0x02; + + @Nonnull + private final HNSW.Config config; + @Nonnull + private final Subspace subspace; + @Nonnull + private final OnWriteListener onWriteListener; + @Nonnull + private final OnReadListener onReadListener; + + private final Subspace entryNodeSubspace; + private final Subspace dataSubspace; + + protected AbstractStorageAdapter(@Nonnull final HNSW.Config config, @Nonnull final Subspace subspace, + @Nonnull final Subspace nodeSlotIndexSubspace, + @Nonnull final Function hilbertValueFunction, + @Nonnull final OnWriteListener onWriteListener, + @Nonnull final OnReadListener onReadListener) { + this.config = config; + this.subspace = subspace; + this.onWriteListener = onWriteListener; + this.onReadListener = onReadListener; + + this.entryNodeSubspace = subspace.subspace(Tuple.from(SUBSPACE_PREFIX_ENTRY_NODE)); + this.dataSubspace = subspace.subspace(Tuple.from(SUBSPACE_PREFIX_DATA)); + } + + @Override + @Nonnull + public HNSW.Config getConfig() { + return config; + } + + @Override + @Nonnull + public Subspace getSubspace() { + return subspace; + } + + @Nullable + @Override + public Subspace getSecondarySubspace() { + return null; + } + + @Override + @Nonnull + public Subspace getEntryNodeSubspace() { + return entryNodeSubspace; + } + + @Override + @Nonnull + public Subspace getDataSubspace() { + return dataSubspace; + } + + @Override + @Nonnull + public OnWriteListener getOnWriteListener() { + return onWriteListener; + } + + @Override + @Nonnull + public OnReadListener getOnReadListener() { + return onReadListener; + } + + @Override + public void writeNodes(@Nonnull final Transaction transaction, @Nonnull final List nodes) { + for (final Node node : nodes) { + writeNode(transaction, node); + } + } + + protected void writeNode(@Nonnull final Transaction transaction, @Nonnull final Node node) { + final Node.ChangeSet changeSet = node.getChangeSet(); + if (changeSet == null) { + return; + } + + changeSet.apply(transaction); + getOnWriteListener().onNodeWritten(node); + } + + @Nonnull + public byte[] packWithSubspace(final byte[] key) { + return getSubspace().pack(key); + } + + @Nonnull + public byte[] packWithSubspace(final Tuple tuple) { + return getSubspace().pack(tuple); + } + + @Nonnull + @Override + public CompletableFuture scanNodeIndexAndFetchNode(@Nonnull final ReadTransaction transaction, + final int level, + @Nonnull final BigInteger hilbertValue, + @Nonnull final Tuple key, + final boolean isInsertUpdate) { + Objects.requireNonNull(nodeSlotIndexAdapter); + return nodeSlotIndexAdapter.scanIndexForNodeId(transaction, level, hilbertValue, key, isInsertUpdate) + .thenCompose(nodeId -> nodeId == null + ? CompletableFuture.completedFuture(null) + : fetchNode(transaction, nodeId)); + } + + @Override + public void insertIntoNodeIndexIfNecessary(@Nonnull final Transaction transaction, final int level, + @Nonnull final NodeSlot nodeSlot) { + if (!getConfig().isUseNodeSlotIndex() || !(nodeSlot instanceof ChildSlot)) { + return; + } + + Objects.requireNonNull(nodeSlotIndexAdapter); + nodeSlotIndexAdapter.writeChildSlot(transaction, level, (ChildSlot)nodeSlot); + } + + @Override + public void deleteFromNodeIndexIfNecessary(@Nonnull final Transaction transaction, final int level, + @Nonnull final NodeSlot nodeSlot) { + if (!getConfig().isUseNodeSlotIndex() || !(nodeSlot instanceof ChildSlot)) { + return; + } + + Objects.requireNonNull(nodeSlotIndexAdapter); + nodeSlotIndexAdapter.clearChildSlot(transaction, level, (ChildSlot)nodeSlot); + } + + @Nonnull + @Override + public CompletableFuture fetchNode(@Nonnull final ReadTransaction transaction, @Nonnull final byte[] nodeId) { + return getOnWriteListener().onAsyncReadForWrite(fetchNodeInternal(transaction, nodeId).thenApply(this::checkNode)); + } + + @Nonnull + protected abstract CompletableFuture fetchNodeInternal(@Nonnull ReadTransaction transaction, @Nonnull byte[] nodeId); + + /** + * Method to perform basic invariant check(s) on a newly-fetched node. + * + * @param node the node to check + * @param the type param for the node in order for this method to not be lossy on the type of the node that + * was passed in + * + * @return the node that was passed in + */ + @Nullable + private N checkNode(@Nullable final N node) { + if (node != null && (node.size() < getConfig().getMinM() || node.size() > getConfig().getMaxM())) { + if (!node.isRoot()) { + throw new IllegalStateException("packing of non-root is out of valid range"); + } + } + return node; + } + + @Nonnull + abstract > AbstractChangeSet + newInsertChangeSet(@Nonnull N node, int level, @Nonnull List insertedSlots); + + @Nonnull + abstract > AbstractChangeSet + newUpdateChangeSet(@Nonnull N node, int level, @Nonnull S originalSlot, @Nonnull S updatedSlot); + + @Nonnull + abstract > AbstractChangeSet + newDeleteChangeSet(@Nonnull N node, int level, @Nonnull List deletedSlots); +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java new file mode 100644 index 0000000000..84b8899d67 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java @@ -0,0 +1,324 @@ +/* + * ByNodeStorageAdapter.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.ReadTransaction; +import com.apple.foundationdb.Transaction; +import com.apple.foundationdb.subspace.Subspace; +import com.apple.foundationdb.tuple.Tuple; +import com.christianheina.langx.half4j.Half; +import com.google.common.base.Verify; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Lists; +import com.google.common.collect.Streams; + +import javax.annotation.Nonnull; +import java.math.BigInteger; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; + +/** + * Storage adapter that represents each node as a single key/value pair in the database. That paradigm + * greatly simplifies how nodes, node slots, and operations on these data structures are managed. Also, a + * node that is serialized into a key/value pair using this storage adapter leaves a significantly smaller + * footprint when compared to the multitude of key/value pairs that are serialized/deserialized using + * {@link BySlotStorageAdapter}. + *
+ * These advantages are offset by the key disadvantage of having to read/deserialize and the serialize/write + * the entire node to realize/persist a minute change to one of its slots. That, in turn, may cause a higher + * likelihood of conflicting with another transaction. + *
+ * Each node is serialized as follows (each {@code (thing)} is a tuple containing {@code thing}): + * 
+ * {@code
+ *    key: nodeId: byte[16]
+ *    value: Tuple(nodeKind: Long, slotList: (slot1, ..., slotn])
+ *           slot:
+ *               for leaf nodes:
+ *                   (hilbertValue: BigInteger, itemKey: (point: Tuple(d1, ..., dk),
+ *                    keySuffix: (...)),
+ *                    value: (...))
+ *               for intermediate nodes:
+ *                   (smallestHV: BigInteger, smallestKey: (...),
+ *                    largestHV: BigInteger, largestKey: (...),
+ *                    childId: byte[16],
+ *                    mbr: (minD1, minD2, ..., minDk, maxD1, maxD2, ..., maxDk)))
+ * }
+ *
+ */ +class ByNodeStorageAdapter extends AbstractStorageAdapter implements StorageAdapter { + public ByNodeStorageAdapter(@Nonnull final HNSW.Config config, @Nonnull final Subspace subspace, + @Nonnull final Subspace nodeSlotIndexSubspace, + @Nonnull final Function hilbertValueFunction, + @Nonnull final OnWriteListener onWriteListener, + @Nonnull final OnReadListener onReadListener) { + super(config, subspace, nodeSlotIndexSubspace, hilbertValueFunction, onWriteListener, onReadListener); + } + + @Override + public CompletableFuture> fetchEntryNode(@Nonnull final Transaction transaction) { + final byte[] key = getEntryNodeSubspace().pack(); + + return transaction.get(key) + .thenApply(valueBytes -> { + if (valueBytes == null) { + throw new IllegalStateException("cannot fetch entry point"); + } + final Node node = fromTuple(Tuple.fromBytes(valueBytes)); + final OnReadListener onReadListener = getOnReadListener(); + onReadListener.onNodeRead(node); + onReadListener.onKeyValueRead(node, key, valueBytes); + return node; + }); + } + + @Override + public void writeLeafNodeSlot(@Nonnull final Transaction transaction, @Nonnull final DataNode node, + @Nonnull final ItemSlot itemSlot) { + persistNode(transaction, node); + } + + @Override + public void clearLeafNodeSlot(@Nonnull final Transaction transaction, @Nonnull final DataNode node, + @Nonnull final ItemSlot itemSlot) { + persistNode(transaction, node); + } + + private void persistNode(@Nonnull final Transaction transaction, @Nonnull final Node node) { + final byte[] packedKey = packWithSubspace(node.getId()); + + if (node.isEmpty()) { + // this can only happen when we just deleted the last slot; delete the entire node + transaction.clear(packedKey); + getOnWriteListener().onKeyCleared(node, packedKey); + } else { + // updateNodeIndexIfNecessary(transaction, level, node); + final byte[] packedValue = toTuple(node).pack(); + transaction.set(packedKey, packedValue); + getOnWriteListener().onKeyValueWritten(node, packedKey, packedValue); + } + } + + @Nonnull + private Tuple toTuple(@Nonnull final Node node) { + final HNSW.Config config = getConfig(); + final List slotTuples = Lists.newArrayListWithExpectedSize(node.size()); + for (final NodeSlot nodeSlot : node.getSlots()) { + final Tuple slotTuple = Tuple.fromStream( + Streams.concat(nodeSlot.getSlotKey(config.isStoreHilbertValues()).getItems().stream(), + nodeSlot.getSlotValue().getItems().stream())); + slotTuples.add(slotTuple); + } + return Tuple.from(node.getKind().getSerialized(), slotTuples); + } + + @Nonnull + @Override + public CompletableFuture fetchNodeInternal(@Nonnull final ReadTransaction transaction, + @Nonnull final byte[] nodeId) { + final byte[] key = packWithSubspace(nodeId); + return transaction.get(key) + .thenApply(valueBytes -> { + if (valueBytes == null) { + return null; + } + final Node node = fromTuple(nodeId, Tuple.fromBytes(valueBytes)); + final OnReadListener onReadListener = getOnReadListener(); + onReadListener.onNodeRead(node); + onReadListener.onKeyValueRead(node, key, valueBytes); + return node; + }); + } + + @Nonnull + private Node fromTuple(@Nonnull final Tuple tuple) { + final NodeKind nodeKind = NodeKind.fromSerializedNodeKind((byte)tuple.getLong(0)); + final Tuple neighborsTuple = tuple.getNestedTuple(1); + + List neighborsWithVectors = null; + Half[] neighborVectorHalfs = null; + List neighbors = null; + + for (final Object neighborObject : neighborsTuple) { + final Tuple neighborTuple = (Tuple)neighborObject; + switch (nodeKind) { + case DATA: + final Tuple neighborPrimaryKey = neighborTuple.getNestedTuple(0); + final Tuple neighborVectorTuple = neighborTuple.getNestedTuple(1); + if (neighborsWithVectors == null) { + neighborsWithVectors = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); + neighborVectorHalfs = new Half[neighborVectorTuple.size()]; + } + + for (int i = 0; i < neighborVectorTuple.size(); i ++) { + neighborVectorHalfs[i] = Half.shortBitsToHalf(shortFromBytes(neighborVectorTuple.getBytes(i))); + } + neighborsWithVectors.add(new NeighborWithVector(neighborPrimaryKey, new Vector.HalfVector(neighborVectorHalfs))); + break; + + case INTERMEDIATE: + if (neighbors == null) { + neighbors = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); + } + neighbors.add(new Neighbor(neighborTuple)); + break; + + default: + throw new IllegalStateException("unknown node kind"); + } + } + + Verify.verify((nodeKind == NodeKind.DATA && neighborsWithVectors != null) || + (nodeKind == NodeKind.INTERMEDIATE && neighbors != null)); + + return nodeKind == NodeKind.DATA + ? new DataNode(nodeId, itemSlots) + : new IntermediateNode(nodeId, childSlots); + } + + @Nonnull + @Override + public > AbstractChangeSet + newInsertChangeSet(@Nonnull final N node, final int level, @Nonnull final List insertedSlots) { + return new InsertChangeSet<>(node, level, insertedSlots); + } + + @Nonnull + @Override + public > AbstractChangeSet + newUpdateChangeSet(@Nonnull final N node, final int level, + @Nonnull final S originalSlot, @Nonnull final S updatedSlot) { + return new UpdateChangeSet<>(node, level, originalSlot, updatedSlot); + } + + @Nonnull + @Override + public > AbstractChangeSet + newDeleteChangeSet(@Nonnull final N node, final int level, @Nonnull final List deletedSlots) { + return new DeleteChangeSet<>(node, level, deletedSlots); + } + + private class InsertChangeSet> extends AbstractChangeSet { + @Nonnull + private final List insertedSlots; + + public InsertChangeSet(@Nonnull final N node, final int level, @Nonnull final List insertedSlots) { + super(node.getChangeSet(), node, level); + this.insertedSlots = ImmutableList.copyOf(insertedSlots); + } + + @Override + public void apply(@Nonnull final Transaction transaction) { + super.apply(transaction); + + // + // If this change set is the first, we persist the node, don't persist the node otherwise. This is a + // performance optimization to avoid writing and rewriting the node for each change set in the chain + // of change sets. + // + if (getPreviousChangeSet() == null) { + persistNode(transaction, getNode()); + } + if (isUpdateNodeSlotIndex()) { + for (final S insertedSlot : insertedSlots) { + insertIntoNodeIndexIfNecessary(transaction, getLevel(), insertedSlot); + } + } + } + } + + private class UpdateChangeSet> extends AbstractChangeSet { + @Nonnull + private final S originalSlot; + @Nonnull + private final S updatedSlot; + + public UpdateChangeSet(@Nonnull final N node, final int level, @Nonnull final S originalSlot, + @Nonnull final S updatedSlot) { + super(node.getChangeSet(), node, level); + this.originalSlot = originalSlot; + this.updatedSlot = updatedSlot; + } + + @Override + public void apply(@Nonnull final Transaction transaction) { + super.apply(transaction); + + // + // If this change set is the first, we persist the node, don't persist the node otherwise. This is a + // performance optimization to avoid writing and rewriting the node for each change set in the chain + // of change sets. + // + if (getPreviousChangeSet() == null) { + persistNode(transaction, getNode()); + } + if (isUpdateNodeSlotIndex()) { + deleteFromNodeIndexIfNecessary(transaction, getLevel(), originalSlot); + insertIntoNodeIndexIfNecessary(transaction, getLevel(), updatedSlot); + } + } + } + + private class DeleteChangeSet> extends AbstractChangeSet { + @Nonnull + private final List deletedSlots; + + public DeleteChangeSet(@Nonnull final N node, final int level, @Nonnull final List deletedSlots) { + super(node.getChangeSet(), node, level); + this.deletedSlots = ImmutableList.copyOf(deletedSlots); + } + + @Override + public void apply(@Nonnull final Transaction transaction) { + super.apply(transaction); + + // + // If this change set is the first, we persist the node, don't persist the node otherwise. This is a + // performance optimization to avoid writing and rewriting the node for each change set in the chain + // of change sets. + // + if (getPreviousChangeSet() == null) { + persistNode(transaction, getNode()); + } + if (isUpdateNodeSlotIndex()) { + for (final S deletedSlot : deletedSlots) { + deleteFromNodeIndexIfNecessary(transaction, getLevel(), deletedSlot); + } + } + } + } + + private short shortFromBytes(byte[] bytes) { + Verify.verify(bytes.length == 2); + int high = bytes[0] & 0xFF; // Convert to unsigned int + int low = bytes[1] & 0xFF; + + return (short) ((high << 8) | low); + } + + private byte[] bytesFromShort(short value) { + byte[] result = new byte[2]; + result[0] = (byte) ((value >> 8) & 0xFF); // high byte first + result[1] = (byte) (value & 0xFF); // low byte second + return result; + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java new file mode 100644 index 0000000000..d9bded06d4 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java @@ -0,0 +1,57 @@ +/* + * DataNode.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.tuple.Tuple; +import com.christianheina.langx.half4j.Half; +import com.google.common.collect.Lists; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.util.List; + +/** + * A leaf node of the R-tree. A leaf node holds the actual data in {@link ItemSlot}s. + */ +class DataNode extends AbstractNode { + public DataNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, + @Nonnull final List neighbors) { + super(primaryKey, vector, neighbors); + } + + @Nonnull + @Override + public DataNode asDataNode() { + return this; + } + + @Nonnull + @Override + public IntermediateNode asIntermediateNode() { + throw new IllegalStateException("this is not a data node"); + } + + @Nonnull + @Override + public NodeKind getKind() { + return NodeKind.DATA; + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java new file mode 100644 index 0000000000..eb7c8aa7ef --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -0,0 +1,2435 @@ +/* + * HNSW.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.Database; +import com.apple.foundationdb.ReadTransaction; +import com.apple.foundationdb.Transaction; +import com.apple.foundationdb.TransactionContext; +import com.apple.foundationdb.annotation.API; +import com.apple.foundationdb.annotation.SpotBugsSuppressWarnings; +import com.apple.foundationdb.async.AsyncIterator; +import com.apple.foundationdb.async.AsyncUtil; +import com.apple.foundationdb.subspace.Subspace; +import com.apple.foundationdb.tuple.Tuple; +import com.apple.foundationdb.tuple.TupleHelpers; +import com.google.common.base.Preconditions; +import com.google.common.base.Verify; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; +import com.google.common.collect.Lists; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.math.BigInteger; +import java.util.ArrayDeque; +import java.util.Arrays; +import java.util.Collections; +import java.util.Deque; +import java.util.Iterator; +import java.util.List; +import java.util.NoSuchElementException; +import java.util.Objects; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Executor; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiPredicate; +import java.util.function.Function; +import java.util.function.Predicate; +import java.util.function.Supplier; +import java.util.stream.Collectors; + +/** + * TODO. + */ +@API(API.Status.EXPERIMENTAL) +public class HNSW { + private static final Logger logger = LoggerFactory.getLogger(HNSW.class); + + /** + * root id. The root id is always only zeros. + */ + static final byte[] rootId = new byte[16]; + + public static final int MAX_CONCURRENT_READS = 16; + + /** + * Indicator if we should maintain a secondary node index consisting of hilbet value and key to speed up + * update/deletes. + */ + public static final boolean DEFAULT_USE_NODE_SLOT_INDEX = false; + + /** + * The minimum number of slots a node has (if not the root node). {@code M} should be chosen in a way that the + * minimum is half of the maximum. That in turn guarantees that overflow/underflow handling can be performed without + * causing further underflow/overflow. + */ + public static final int DEFAULT_MIN_M = 16; + /** + * The maximum number of slots a node has. This value is derived from {@link #DEFAULT_MIN_M}. + */ + public static final int DEFAULT_MAX_M = 2 * DEFAULT_MIN_M; + + /** + * The magic split number. We split {@code S} to {@code S + 1} nodes while inserting data and fuse + * {@code S + 1} to {@code S} nodes while deleting data. Academically, 2-to-3 splits and 3-to-2 fuses + * seem to yield the best results. Please be aware of the following constraints: + *
    + *
  1. When splitting {@code S} to {@code S + 1} nodes, we re-distribute the children of {@code S} nodes + * into {@code S + 1} nodes which may cause an underflow if {@code S} and {@code M} are not set carefully with + * respect to each other. Example: {@code MIN_M = 25}, {@code MAX_M = 32}, {@code S = 2}, two nodes at + * already at maximum capacity containing a combined total of 64 children when a new child is inserted. + * We split the two nodes into three as indicated by {@code S = 2}. We have 65 children but there is no way + * of distributing them among three nodes such that none of them underflows. This constraint can be + * formulated as {@code S * MAX_M / (S + 1) >= MIN_M}.
    2. + *
  2. When fusing {@code S + 1} to {@code S} nodes, we re-distribute the children of {@code S + 1} nodes + * into {@code S + 1} nodes which may cause an overflow if {@code S} and {@code M} are not set carefully with + * respect to each other. Example: {@code MIN_M = 25}, {@code MAX_M = 32}, {@code S = 2}, three nodes at + * already at minimum capacity containing a combined total of 75 children when a child is deleted. + * We fuse the three nodes into two as indicated by {@code S = 2}. We have 75 children but there is no way + * of distributing them among two nodes such that none of them overflows. This constraint can be formulated as + * {@code (S + 1) * MIN_M / S <= MAX_M}.
    3. + *
+ * Both constraints are in fact the same constraint and can be written as {@code MAX_M / MIN_M >= (S + 1) / S}. + */ + public static final int DEFAULT_S = 2; + + /** + * Default storage layout. Can be either {@code BY_SLOT} or {@code BY_NODE}. {@code BY_SLOT} encodes all information + * pertaining to a {@link NodeSlot} as one key/value pair in the database; {@code BY_NODE} encodes all information + * pertaining to a {@link Node} as one key/value pair in the database. While {@code BY_SLOT} avoids conflicts as + * most inserts/updates only need to update one slot, it is by far less compact as some information is stored + * in a normalized fashion and therefore repeated multiple times (i.e. node identifiers, etc.). {@code BY_NODE} + * inlines slot information into the node leading to a more size-efficient layout of the data. That advantage is + * offset by a higher likelihood of conflicts. + */ + @Nonnull + public static final Storage DEFAULT_STORAGE = Storage.BY_NODE; + + /** + * Indicator if Hilbert values should be stored or not with the data (in leaf nodes). A Hilbert value can always + * be recomputed from the point. + */ + public static final boolean DEFAULT_STORE_HILBERT_VALUES = true; + + @Nonnull + public static final Config DEFAULT_CONFIG = new Config(); + + @Nonnull + private final StorageAdapter storageAdapter; + @Nonnull + private final Executor executor; + @Nonnull + private final Config config; + @Nonnull + private final Function hilbertValueFunction; + @Nonnull + private final Supplier nodeIdSupplier; + @Nonnull + private final OnWriteListener onWriteListener; + @Nonnull + private final OnReadListener onReadListener; + + /** + * Different kinds of storage layouts. + */ + public enum Storage { + /** + * Every node slot is serialized as a key/value pair in FDB. + */ + BY_SLOT(BySlotStorageAdapter::new), + /** + * Every node with all its slots is serialized as one key/value pair. + */ + BY_NODE(ByNodeStorageAdapter::new); + + @Nonnull + private final StorageAdapterCreator storageAdapterCreator; + + Storage(@Nonnull final StorageAdapterCreator storageAdapterCreator) { + this.storageAdapterCreator = storageAdapterCreator; + } + + @Nonnull + private StorageAdapter newStorageAdapter(@Nonnull final Config config, @Nonnull final Subspace subspace, + @Nonnull final Subspace nodeSlotIndexSubspace, + @Nonnull final Function hilbertValueFunction, + @Nonnull final OnWriteListener onWriteListener, + @Nonnull final OnReadListener onReadListener) { + return storageAdapterCreator.create(config, subspace, nodeSlotIndexSubspace, + hilbertValueFunction, onWriteListener, onReadListener); + } + } + + /** + * Functional interface to create a {@link StorageAdapter}. + */ + private interface StorageAdapterCreator { + StorageAdapter create(@Nonnull Config config, @Nonnull Subspace subspace, @Nonnull Subspace nodeSlotIndexSubspace, + @Nonnull Function hilbertValueFunction, + @Nonnull OnWriteListener onWriteListener, + @Nonnull OnReadListener onReadListener); + } + + /** + * Configuration settings for a {@link HNSW}. + */ + public static class Config { + private final boolean useNodeSlotIndex; + private final int minM; + private final int maxM; + private final int splitS; + @Nonnull + private final Storage storage; + + private final boolean storeHilbertValues; + + protected Config() { + this.useNodeSlotIndex = DEFAULT_USE_NODE_SLOT_INDEX; + this.minM = DEFAULT_MIN_M; + this.maxM = DEFAULT_MAX_M; + this.splitS = DEFAULT_S; + this.storage = DEFAULT_STORAGE; + this.storeHilbertValues = DEFAULT_STORE_HILBERT_VALUES; + } + + protected Config(final boolean useNodeSlotIndex, final int minM, final int maxM, final int splitS, + @Nonnull final Storage storage, final boolean storeHilbertValues) { + this.useNodeSlotIndex = useNodeSlotIndex; + this.minM = minM; + this.maxM = maxM; + this.splitS = splitS; + this.storage = storage; + this.storeHilbertValues = storeHilbertValues; + } + + public boolean isUseNodeSlotIndex() { + return useNodeSlotIndex; + } + + public int getMinM() { + return minM; + } + + public int getMaxM() { + return maxM; + } + + public int getSplitS() { + return splitS; + } + + @Nonnull + public Storage getStorage() { + return storage; + } + + public boolean isStoreHilbertValues() { + return storeHilbertValues; + } + + public ConfigBuilder toBuilder() { + return new ConfigBuilder(useNodeSlotIndex, minM, maxM, splitS, storage, storeHilbertValues); + } + + @Override + public String toString() { + return storage + ", M=" + minM + "-" + maxM + ", S=" + splitS + + (useNodeSlotIndex ? ", slotIndex" : "") + + (storeHilbertValues ? ", storeHV" : ""); + } + } + + /** + * Builder for {@link Config}. + * + * @see #newConfigBuilder + */ + @CanIgnoreReturnValue + public static class ConfigBuilder { + private boolean useNodeSlotIndex = DEFAULT_USE_NODE_SLOT_INDEX; + private int minM = DEFAULT_MIN_M; + private int maxM = DEFAULT_MAX_M; + private int splitS = DEFAULT_S; + @Nonnull + private Storage storage = DEFAULT_STORAGE; + private boolean storeHilbertValues = DEFAULT_STORE_HILBERT_VALUES; + + public ConfigBuilder() { + } + + public ConfigBuilder(final boolean useNodeSlotIndex, final int minM, final int maxM, final int splitS, + @Nonnull final Storage storage, final boolean storeHilbertValues) { + this.useNodeSlotIndex = useNodeSlotIndex; + this.minM = minM; + this.maxM = maxM; + this.splitS = splitS; + this.storage = storage; + this.storeHilbertValues = storeHilbertValues; + } + + public int getMinM() { + return minM; + } + + public ConfigBuilder setMinM(final int minM) { + this.minM = minM; + return this; + } + + public int getMaxM() { + return maxM; + } + + public ConfigBuilder setMaxM(final int maxM) { + this.maxM = maxM; + return this; + } + + public int getSplitS() { + return splitS; + } + + public ConfigBuilder setSplitS(final int splitS) { + this.splitS = splitS; + return this; + } + + @Nonnull + public Storage getStorage() { + return storage; + } + + public ConfigBuilder setStorage(@Nonnull final Storage storage) { + this.storage = storage; + return this; + } + + public boolean isStoreHilbertValues() { + return storeHilbertValues; + } + + public ConfigBuilder setStoreHilbertValues(final boolean storeHilbertValues) { + this.storeHilbertValues = storeHilbertValues; + return this; + } + + public boolean isUseNodeSlotIndex() { + return useNodeSlotIndex; + } + + public ConfigBuilder setUseNodeSlotIndex(final boolean useNodeSlotIndex) { + this.useNodeSlotIndex = useNodeSlotIndex; + return this; + } + + public Config build() { + return new Config(isUseNodeSlotIndex(), getMinM(), getMaxM(), getSplitS(), getStorage(), isStoreHilbertValues()); + } + } + + /** + * Start building a {@link Config}. + * @return a new {@code Config} that can be altered and then built for use with a {@link HNSW} + * @see ConfigBuilder#build + */ + public static ConfigBuilder newConfigBuilder() { + return new ConfigBuilder(); + } + + /** + * Initialize a new R-tree with the default configuration. + * @param subspace the subspace where the r-tree is stored + * @param secondarySubspace the subspace where the node index (if used is stored) + * @param executor an executor to use when running asynchronous tasks + * @param hilbertValueFunction function to compute the Hilbert value from a {@link Point} + */ + public HNSW(@Nonnull final Subspace subspace, @Nonnull final Subspace secondarySubspace, + @Nonnull final Executor executor, @Nonnull final Function hilbertValueFunction) { + this(subspace, secondarySubspace, executor, DEFAULT_CONFIG, hilbertValueFunction, NodeHelpers::newRandomNodeId, + OnWriteListener.NOOP, OnReadListener.NOOP); + } + + /** + * Initialize a new R-tree. + * @param subspace the subspace where the r-tree is stored + * @param nodeSlotIndexSubspace the subspace where the node index (if used is stored) + * @param executor an executor to use when running asynchronous tasks + * @param config configuration to use + * @param hilbertValueFunction function to compute the Hilbert value for a {@link Point} + * @param nodeIdSupplier supplier to be invoked when new nodes are created + * @param onWriteListener an on-write listener to be called after writes take place + * @param onReadListener an on-read listener to be called after reads take place + */ + public HNSW(@Nonnull final Subspace subspace, @Nonnull final Subspace nodeSlotIndexSubspace, + @Nonnull final Executor executor, @Nonnull final Config config, + @Nonnull final Function hilbertValueFunction, + @Nonnull final Supplier nodeIdSupplier, + @Nonnull final OnWriteListener onWriteListener, + @Nonnull final OnReadListener onReadListener) { + this.storageAdapter = config.getStorage() + .newStorageAdapter(config, subspace, nodeSlotIndexSubspace, hilbertValueFunction, onWriteListener, + onReadListener); + this.executor = executor; + this.config = config; + this.hilbertValueFunction = hilbertValueFunction; + this.nodeIdSupplier = nodeIdSupplier; + this.onWriteListener = onWriteListener; + this.onReadListener = onReadListener; + } + + /** + * Get the {@link StorageAdapter} used to manage this r-tree. + * @return r-tree subspace + */ + @Nonnull + StorageAdapter getStorageAdapter() { + return storageAdapter; + } + + /** + * Get the executer used by this r-tree. + * @return executor used when running asynchronous tasks + */ + @Nonnull + public Executor getExecutor() { + return executor; + } + + /** + * Get this r-tree's configuration. + * @return r-tree configuration + */ + @Nonnull + public Config getConfig() { + return config; + } + + /** + * Get the on-write listener. + * @return the on-write listener + */ + @Nonnull + public OnWriteListener getOnWriteListener() { + return onWriteListener; + } + + /** + * Get the on-read listener. + * @return the on-read listener + */ + @Nonnull + public OnReadListener getOnReadListener() { + return onReadListener; + } + + // + // Read Path + // + + /** + * Perform a scan over the tree within the transaction passed in using a predicate that is also passed in to + * eliminate subtrees from the scan. This predicate may be stateful which allows for dynamic adjustments of the + * queried area while the scan is active. + *
+ * A scan of the tree offers all items that pass the {@code mbrPredicate} test in Hilbert Value order using an + * {@link AsyncIterator}. The predicate that is passed in is applied to intermediate nodes as well as leaf nodes, + * but not to elements contained by a leaf node. The caller should filter out items in a downstream operation. + * A scan of the tree will not prefetch the next node before the items of the current node have been consumed. This + * guarantees that the semantics of the mbr predicate can be adapted in response to the items being consumed. + * (this allows for efficient scans for {@code ORDER BY x, y LIMIT n} queries). + * @param readTransaction the transaction to use + * @param mbrPredicate a predicate on an mbr {@link Rectangle} + * @param suffixKeyPredicate a predicate on the suffix key + * @return an {@link AsyncIterator} of {@link ItemSlot}s. + */ + @Nonnull + public AsyncIterator scan(@Nonnull final ReadTransaction readTransaction, + @Nonnull final Predicate mbrPredicate, + @Nonnull final BiPredicate suffixKeyPredicate) { + return scan(readTransaction, null, null, mbrPredicate, suffixKeyPredicate); + } + + /** + * Perform a scan over the tree within the transaction passed in using a predicate that is also passed in to + * eliminate subtrees from the scan. This predicate may be stateful which allows for dynamic adjustments of the + * queried area while the scan is active. + *
+ * A scan of the tree offers all items that pass the {@code mbrPredicate} test in Hilbert Value order using an + * {@link AsyncIterator}. The predicate that is passed in is applied to intermediate nodes as well as leaf nodes, + * but not to elements contained in a leaf node. The caller should filter out items in a downstream operation. + * A scan of the tree will not prefetch the next node before the items of the current node have been consumed. This + * guarantees that the semantics of the mbr predicate can be adapted in response to the items being consumed. + * (this allows for efficient scans for {@code ORDER BY x, y LIMIT n} queries). + * @param readTransaction the transaction to use + * @param lastHilbertValue the last Hilbert value that was returned by a previous call to this method + * @param lastKey the last key that was returned by a previous call to this method + * @param mbrPredicate a predicate on an mbr {@link Rectangle} + * @param suffixKeyPredicate a predicate on the suffix key + * @return an {@link AsyncIterator} of {@link ItemSlot}s. + */ + @Nonnull + public AsyncIterator scan(@Nonnull final ReadTransaction readTransaction, + @Nullable final BigInteger lastHilbertValue, + @Nullable final Tuple lastKey, + @Nonnull final Predicate mbrPredicate, + @Nonnull final BiPredicate suffixKeyPredicate) { + Preconditions.checkArgument((lastHilbertValue == null && lastKey == null) || + (lastHilbertValue != null && lastKey != null)); + AsyncIterator leafIterator = + new LeafIterator(readTransaction, rootId, lastHilbertValue, lastKey, mbrPredicate, suffixKeyPredicate); + return new ItemSlotIterator(leafIterator); + } + + /** + * Returns the left-most path from a given node id to a leaf node containing items as a {@link TraversalState}. + * The term left-most used here is defined by comparing {@code (largestHilbertValue, largestKey)} when + * comparing nodes (the left one being the smaller, the right one being the greater). + * @param readTransaction the transaction to use + * @param nodeId node id to start from. This may be the actual root of the tree or some other node within the tree. + * @param lastHilbertValue hilbert value serving as a watermark to return only items that are larger than the + * {@code (lastHilbertValue, lastKey)} pair + * @param lastKey key serving as a watermark to return only items that are larger than the + * {@code (lastHilbertValue, lastKey)} pair + * @param mbrPredicate a predicate on an mbr {@link Rectangle}. This predicate is evaluated on the way down to the + * leaf node. + * @param suffixPredicate predicate to be invoked on a range of suffixes + * @return a {@link TraversalState} of the left-most path from {@code nodeId} to a {@link DataNode} whose + * {@link Node}s all pass the mbr predicate test. + */ + @Nonnull + private CompletableFuture fetchLeftmostPathToLeaf(@Nonnull final ReadTransaction readTransaction, + @Nonnull final byte[] nodeId, + @Nullable final BigInteger lastHilbertValue, + @Nullable final Tuple lastKey, + @Nonnull final Predicate mbrPredicate, + @Nonnull final BiPredicate suffixPredicate) { + final AtomicReference currentId = new AtomicReference<>(nodeId); + final List> toBeProcessed = Lists.newArrayList(); + final AtomicReference leafNode = new AtomicReference<>(null); + return AsyncUtil.whileTrue(() -> onReadListener.onAsyncRead(storageAdapter.fetchNode(readTransaction, currentId.get())) + .thenApply(node -> { + if (node == null) { + if (Arrays.equals(currentId.get(), rootId)) { + Verify.verify(leafNode.get() == null); + return false; + } + throw new IllegalStateException("unable to fetch node for scan"); + } + if (node.getKind() == NodeKind.INTERMEDIATE) { + final Iterable childSlots = ((IntermediateNode)node).getSlots(); + Deque toBeProcessedThisLevel = new ArrayDeque<>(); + for (final Iterator iterator = childSlots.iterator(); iterator.hasNext(); ) { + final ChildSlot childSlot = iterator.next(); + if (lastHilbertValue != null && + lastKey != null) { + final int hilbertValueAndKeyCompare = + childSlot.compareLargestHilbertValueAndKey(lastHilbertValue, lastKey); + if (hilbertValueAndKeyCompare < 0) { + // + // The (lastHilbertValue, lastKey) pair is larger than the + // (largestHilbertValue, largestKey) pair of the current child. Advance to the next + // child. + // + continue; + } + } + + if (!mbrPredicate.test(childSlot.getMbr())) { + onReadListener.onChildNodeDiscard(childSlot); + continue; + } + + if (childSlot.suffixPredicateCanBeApplied()) { + if (!suffixPredicate.test(childSlot.getSmallestKeySuffix(), + childSlot.getLargestKeySuffix())) { + onReadListener.onChildNodeDiscard(childSlot); + continue; + } + } + + toBeProcessedThisLevel.addLast(childSlot); + iterator.forEachRemaining(toBeProcessedThisLevel::addLast); + } + toBeProcessed.add(toBeProcessedThisLevel); + + final ChildSlot nextChildSlot = resolveNextIdForFetch(toBeProcessed, mbrPredicate, + suffixPredicate, onReadListener); + if (nextChildSlot == null) { + return false; + } + + currentId.set(Objects.requireNonNull(nextChildSlot.getChildId())); + return true; + } else { + leafNode.set((DataNode)node); + return false; + } + }), executor).thenApply(vignore -> leafNode.get() == null + ? TraversalState.end() + : TraversalState.of(toBeProcessed, leafNode.get())); + } + + /** + * Returns the next left-most path from a given {@link TraversalState} to a leaf node containing items as + * a {@link TraversalState}. The term left-most used here is defined by comparing + * {@code (largestHilbertValue, largestKey)} when comparing nodes (the left one being the smaller, the right one + * being the greater). + * @param readTransaction the transaction to use + * @param traversalState traversal state to start from. The initial traversal state is always obtained by initially + * calling {@link #fetchLeftmostPathToLeaf(ReadTransaction, byte[], BigInteger, Tuple, Predicate, BiPredicate)}. + * @param mbrPredicate a predicate on an mbr {@link Rectangle}. This predicate is evaluated for each node that + * is processed. + * @return a {@link TraversalState} of the left-most path from {@code nodeId} to a {@link DataNode} whose + * {@link Node}s all pass the mbr predicate test. + */ + @Nonnull + private CompletableFuture fetchNextPathToLeaf(@Nonnull final ReadTransaction readTransaction, + @Nonnull final TraversalState traversalState, + @Nullable final BigInteger lastHilbertValue, + @Nullable final Tuple lastKey, + @Nonnull final Predicate mbrPredicate, + @Nonnull final BiPredicate suffixPredicate) { + + final List> toBeProcessed = traversalState.getToBeProcessed(); + final AtomicReference leafNode = new AtomicReference<>(null); + + return AsyncUtil.whileTrue(() -> { + final ChildSlot nextChildSlot = resolveNextIdForFetch(toBeProcessed, mbrPredicate, suffixPredicate, + onReadListener); + if (nextChildSlot == null) { + return AsyncUtil.READY_FALSE; + } + + // fetch the left-most path rooted at the current child to its left-most leaf and concatenate the paths + return fetchLeftmostPathToLeaf(readTransaction, nextChildSlot.getChildId(), lastHilbertValue, + lastKey, mbrPredicate, suffixPredicate) + .thenApply(nestedTraversalState -> { + if (nestedTraversalState.isEnd()) { + // no more data in this subtree + return true; + } + // combine the traversal states + leafNode.set(nestedTraversalState.getCurrentLeafNode()); + toBeProcessed.addAll(nestedTraversalState.getToBeProcessed()); + return false; + }); + }, executor).thenApply(v -> leafNode.get() == null + ? TraversalState.end() + : TraversalState.of(toBeProcessed, leafNode.get())); + } + + /** + * Return the next {@link ChildSlot} that needs to be processed given a list of deques that need to be processed + * as part of the current scan. + * @param toBeProcessed list of deques + * @param mbrPredicate a predicate on an mbr {@link Rectangle} + * @param suffixPredicate a predicate that is tested if applicable on the key suffix + * @return The next child slot that needs to be processed or {@code null} if there is no next child slot. + * As a side effect of calling this method the child slot is removed from {@code toBeProcessed}. + */ + @Nullable + @SuppressWarnings("PMD.AvoidBranchingStatementAsLastInLoop") + private static ChildSlot resolveNextIdForFetch(@Nonnull final List> toBeProcessed, + @Nonnull final Predicate mbrPredicate, + @Nonnull final BiPredicate suffixPredicate, + @Nonnull final OnReadListener onReadListener) { + for (int level = toBeProcessed.size() - 1; level >= 0; level--) { + final Deque toBeProcessedThisLevel = toBeProcessed.get(level); + + while (!toBeProcessedThisLevel.isEmpty()) { + final ChildSlot childSlot = toBeProcessedThisLevel.pollFirst(); + if (!mbrPredicate.test(childSlot.getMbr())) { + onReadListener.onChildNodeDiscard(childSlot); + continue; + } + if (childSlot.suffixPredicateCanBeApplied()) { + if (!suffixPredicate.test(childSlot.getSmallestKeySuffix(), + childSlot.getLargestKeySuffix())) { + onReadListener.onChildNodeDiscard(childSlot); + continue; + } + } + toBeProcessed.subList(level + 1, toBeProcessed.size()).clear(); + return childSlot; + } + } + return null; + } + + // + // Insert/Update path + // + + + /** + * Method to insert an object/item into the R-tree. The item is treated unique per its point in space as well as its + * additional key that is also passed in. The Hilbert value of the point is passed in as to allow the caller to + * compute Hilbert values themselves. Note that there is a bijective mapping between point and Hilbert + * value which allows us to recompute point from Hilbert value as well as Hilbert value from point. We currently + * treat point and Hilbert value independent, however, they are redundant and not independent at all. The implication + * is that we do not have to store both point and Hilbert value (but we currently do). + * @param tc transaction context + * @param point the point to be used in space + * @param keySuffix the additional key to be stored with the item + * @param value the additional value to be stored with the item + * @return a completable future that completes when the insert is completed + */ + @Nonnull + public CompletableFuture insertOrUpdate(@Nonnull final TransactionContext tc, + @Nonnull final Point point, + @Nonnull final Tuple keySuffix, + @Nonnull final Tuple value) { + final BigInteger hilbertValue = hilbertValueFunction.apply(point); + final Tuple itemKey = Tuple.from(point.getCoordinates(), keySuffix); + + // + // Get to the leaf node we need to start the insert from and then call the appropriate method to perform + // the actual insert/update. + // + return tc.runAsync(transaction -> fetchPathForModification(transaction, hilbertValue, itemKey, true) + .thenCompose(leafNode -> { + if (leafNode == null) { + leafNode = new DataNode(rootId, Lists.newArrayList()); + } + return insertOrUpdateSlot(transaction, leafNode, point, hilbertValue, itemKey, value); + })); + } + + /** + * Inserts a new slot into the {@link DataNode} passed in or updates an existing slot of the {@link DataNode} passed + * in. + * @param transaction transaction + * @param targetNode leaf node that is the target of this insert or update + * @param point the point to be used in space + * @param hilbertValue the hilbert value of the point + * @param key the additional key to be stored with the item + * @param value the additional value to be stored with the item + * @return a completable future that completes when the insert/update is completed + */ + @Nonnull + private CompletableFuture insertOrUpdateSlot(@Nonnull final Transaction transaction, + @Nonnull final DataNode targetNode, + @Nonnull final Point point, + @Nonnull final BigInteger hilbertValue, + @Nonnull final Tuple key, + @Nonnull final Tuple value) { + Verify.verify(targetNode.size() <= config.getMaxM()); + + final AtomicInteger level = new AtomicInteger(0); + final ItemSlot newSlot = new ItemSlot(hilbertValue, point, key, value); + final AtomicInteger insertSlotIndex = new AtomicInteger(findInsertUpdateItemSlotIndex(targetNode, hilbertValue, key)); + if (insertSlotIndex.get() < 0) { + // just update the slot with the potentially new value + storageAdapter.writeLeafNodeSlot(transaction, targetNode, newSlot); + return AsyncUtil.DONE; + } + + // + // This is an insert. + // + + final AtomicReference currentNode = new AtomicReference<>(targetNode); + final AtomicReference parentSlot = new AtomicReference<>(newSlot); + + // + // Inch our way upwards in the tree to perform the necessary adjustments. What needs to be done next + // is informed by the result of the current operation: + // 1. A split happened; we need to insert a new slot into the parent node -- prime current node and + // current slot and continue. + // 2. The slot was inserted but mbrs, largest Hilbert Values and largest Keys need to be adjusted upwards. + // 3. We are done as no further adjustments are necessary. + // + return AsyncUtil.whileTrue(() -> { + final NodeSlot currentNewSlot = parentSlot.get(); + + if (currentNewSlot != null) { + return insertSlotIntoTargetNode(transaction, level.get(), hilbertValue, key, currentNode.get(), currentNewSlot, insertSlotIndex.get()) + .thenApply(nodeOrAdjust -> { + if (currentNode.get().isRoot()) { + return false; + } + currentNode.set(currentNode.get().getParentNode()); + parentSlot.set(nodeOrAdjust.getSlotInParent()); + insertSlotIndex.set(nodeOrAdjust.getSplitNode() == null ? -1 : nodeOrAdjust.getSplitNode().getSlotIndexInParent()); + level.incrementAndGet(); + return nodeOrAdjust.getSplitNode() != null || nodeOrAdjust.parentNeedsAdjustment(); + }); + } else { + // adjustment only + return updateSlotsAndAdjustNode(transaction, level.get(), hilbertValue, key, currentNode.get(), true) + .thenApply(nodeOrAdjust -> { + Verify.verify(nodeOrAdjust.getSlotInParent() == null); + if (currentNode.get().isRoot()) { + return false; + } + currentNode.set(currentNode.get().getParentNode()); + level.incrementAndGet(); + return nodeOrAdjust.parentNeedsAdjustment(); + }); + } + }, executor); + } + + /** + * Insert a new slot into the target node passed in. + * @param transaction transaction + * @param level the current level of target node, {@code 0} indicating the leaf level + * @param hilbertValue the Hilbert Value of the record that is being inserted + * @param key the key of the record that is being inserted + * @param targetNode target node + * @param newSlot new slot + * @param slotIndexInTargetNode The index of the new slot that we should use when inserting the new slot. While + * this information can be computed from the other arguments passed in, the caller already knows this + * information; we can avoid searching for the proper spot on our own. + * @return a completable future that when completed indicates what needs to be done next (see {@link NodeOrAdjust}). + */ + @Nonnull + private CompletableFuture insertSlotIntoTargetNode(@Nonnull final Transaction transaction, + final int level, + @Nonnull final BigInteger hilbertValue, + @Nonnull final Tuple key, + @Nonnull final Node targetNode, + @Nonnull final NodeSlot newSlot, + final int slotIndexInTargetNode) { + if (targetNode.size() < config.getMaxM()) { + // enough space left in target + + if (logger.isTraceEnabled()) { + logger.trace("regular insert without splitting; node={}; size={}", targetNode, targetNode.size()); + } + targetNode.insertSlot(storageAdapter, level - 1, slotIndexInTargetNode, newSlot); + + if (targetNode.getKind() == NodeKind.INTERMEDIATE) { + // + // If this is an insert for an intermediate node, the child node referred to by newSlot + // is a split node from a lower level meaning a split has happened on a lower level and the + // participating siblings of that split have potentially changed. + // + storageAdapter.writeNodes(transaction, Collections.singletonList(targetNode)); + } else { + // if this is an insert for a leaf node we can just write the slot + Verify.verify(targetNode.getKind() == NodeKind.LEAF); + storageAdapter.writeLeafNodeSlot(transaction, (DataNode)targetNode, (ItemSlot)newSlot); + } + + // node has left some space -- indicate that we are done splitting at the current node + if (!targetNode.isRoot()) { + return fetchParentNodeIfNecessary(transaction, targetNode, level, hilbertValue, key, true) + .thenApply(ignored -> adjustSlotInParent(targetNode, level) + ? NodeOrAdjust.ADJUST + : NodeOrAdjust.NONE); + } + + // no split and no adjustment + return CompletableFuture.completedFuture(NodeOrAdjust.NONE); + } else { + // + // If this is the root we need to grow the tree taller by splitting the root to get a new root + // with two children each containing half of the slots previously contained by the old root node. + // + if (targetNode.isRoot()) { + if (logger.isTraceEnabled()) { + logger.trace("splitting root node; size={}", targetNode.size()); + } + // temporarily overfill the old root node + targetNode.insertSlot(storageAdapter, level - 1, slotIndexInTargetNode, newSlot); + + splitRootNode(transaction, level, targetNode); + return CompletableFuture.completedFuture(NodeOrAdjust.NONE); + } + + // + // Node is full -- borrow some space from the siblings if possible. The paper does overflow handling and + // node splitting separately -- we do it in one path. + // + final CompletableFuture> siblings = + fetchParentNodeIfNecessary(transaction, targetNode, level, hilbertValue, key, true) + .thenCompose(ignored -> + fetchSiblings(transaction, targetNode)); + + return siblings.thenApply(siblingNodes -> { + int numSlots = + Math.toIntExact(siblingNodes + .stream() + .mapToLong(Node::size) + .sum()); + + // First determine if we actually need to split; create the split node if we do; for the remainder of + // this method splitNode != null <=> we are splitting; otherwise we handle overflow. + final Node splitNode; + final List newSiblingNodes; + if (numSlots == siblingNodes.size() * config.getMaxM()) { + if (logger.isTraceEnabled()) { + logger.trace("splitting node; node={}, siblings={}", + targetNode, + siblingNodes.stream().map(Node::toString) + .collect(Collectors.joining(","))); + } + splitNode = targetNode.newOfSameKind(nodeIdSupplier.get()); + // link this split node to become the last node of the siblings + splitNode.linkToParent(Objects.requireNonNull(targetNode.getParentNode()), + siblingNodes.get(siblingNodes.size() - 1).getSlotIndexInParent() + 1); + newSiblingNodes = Lists.newArrayList(siblingNodes); + newSiblingNodes.add(splitNode); + } else { + if (logger.isTraceEnabled()) { + logger.trace("handling overflow; node={}, numSlots={}, siblings={}", + targetNode, + numSlots, + siblingNodes.stream().map(Node::toString) + .collect(Collectors.joining(","))); + } + splitNode = null; + newSiblingNodes = siblingNodes; + } + + // temporarily overfill targetNode + numSlots++; + targetNode.insertSlot(storageAdapter, level - 1, slotIndexInTargetNode, newSlot); + + // sibling nodes are in hilbert value order + final Iterator slotIterator = + siblingNodes + .stream() + .flatMap(Node::slotsStream) + .iterator(); + + // + // Distribute all slots (including the new one which is now at its correct position among its brethren) + // across all siblings (which includes the targetNode and (if we are splitting) the splitNode). + // At the end of this modification all siblings have and (almost) equal count of slots that is + // guaranteed to be between minM and maxM. + // + + final int base = numSlots / newSiblingNodes.size(); + int rest = numSlots % newSiblingNodes.size(); + + List> newNodeSlotLists = Lists.newArrayList(); + List currentNodeSlots = Lists.newArrayList(); + while (slotIterator.hasNext()) { + final NodeSlot slot = slotIterator.next(); + currentNodeSlots.add(slot); + if (currentNodeSlots.size() == base + (rest > 0 ? 1 : 0)) { + if (rest > 0) { + // one fewer to distribute + rest--; + } + + newNodeSlotLists.add(currentNodeSlots); + currentNodeSlots = Lists.newArrayList(); + } + } + + Verify.verify(newSiblingNodes.size() == newNodeSlotLists.size()); + + final Iterator newSiblingNodesIterator = newSiblingNodes.iterator(); + final Iterator> newNodeSlotsIterator = newNodeSlotLists.iterator(); + + // assign slots to nodes + while (newSiblingNodesIterator.hasNext()) { + final Node newSiblingNode = newSiblingNodesIterator.next(); + Verify.verify(newNodeSlotsIterator.hasNext()); + final List newNodeSlots = newNodeSlotsIterator.next(); + newSiblingNode.moveOutAllSlots(storageAdapter); + newSiblingNode.moveInSlots(storageAdapter, newNodeSlots); + } + + // update nodes + storageAdapter.writeNodes(transaction, newSiblingNodes); + + // + // Adjust the parent's slot information in memory only; we'll write it in the next iteration when + // we go one level up. + // + for (final Node siblingNode : siblingNodes) { + adjustSlotInParent(siblingNode, level); + } + + if (splitNode == null) { + // didn't split -- just continue adjusting + return NodeOrAdjust.ADJUST; + } + + // + // Manufacture a new slot for the splitNode; the caller will then use that slot to insert it into the + // parent. + // + final NodeSlot firstSlotOfSplitNode = splitNode.getSlot(0); + final NodeSlot lastSlotOfSplitNode = splitNode.getSlot(splitNode.size() - 1); + return new NodeOrAdjust( + new ChildSlot(firstSlotOfSplitNode.getSmallestHilbertValue(), firstSlotOfSplitNode.getSmallestKey(), + lastSlotOfSplitNode.getLargestHilbertValue(), lastSlotOfSplitNode.getLargestKey(), + splitNode.getId(), NodeHelpers.computeMbr(splitNode.getSlots())), + splitNode, true); + }); + } + } + + /** + * Split the root node. This method first creates two nodes {@code left} and {@code right}. The root node, + * whose ID is always a string of {@code 0x00}, contains some number {@code n} of slots. {@code n / 2} slots of those + * {@code n} slots are moved to {@code left}, the rest to {@code right}. The root node is then updated to have two + * children: {@code left} and {@code right}. All three nodes are then updated in the database. + * @param transaction transaction to use + * @param level the level counting starting at {@code 0} indicating the leaf level increasing upwards + * @param oldRootNode the old root node + */ + private void splitRootNode(@Nonnull final Transaction transaction, + final int level, + @Nonnull final Node oldRootNode) { + final Node leftNode = oldRootNode.newOfSameKind(nodeIdSupplier.get()); + final Node rightNode = oldRootNode.newOfSameKind(nodeIdSupplier.get()); + final int leftSize = oldRootNode.size() / 2; + final List leftSlots = ImmutableList.copyOf(oldRootNode.getSlots(0, leftSize)); + leftNode.moveInSlots(storageAdapter, leftSlots); + final int rightSize = oldRootNode.size() - leftSize; + final List rightSlots = ImmutableList.copyOf(oldRootNode.getSlots(leftSize, leftSize + rightSize)); + rightNode.moveInSlots(storageAdapter, rightSlots); + + final NodeSlot firstSlotOfLeftNode = leftSlots.get(0); + final NodeSlot lastSlotOfLeftNode = leftSlots.get(leftSlots.size() - 1); + final NodeSlot firstSlotOfRightNode = rightSlots.get(0); + final NodeSlot lastSlotOfRightNode = rightSlots.get(rightSlots.size() - 1); + + final ChildSlot leftChildSlot = new ChildSlot(firstSlotOfLeftNode.getSmallestHilbertValue(), firstSlotOfLeftNode.getSmallestKey(), + lastSlotOfLeftNode.getLargestHilbertValue(), lastSlotOfLeftNode.getLargestKey(), + leftNode.getId(), NodeHelpers.computeMbr(leftNode.getSlots())); + final ChildSlot rightChildSlot = new ChildSlot(firstSlotOfRightNode.getSmallestHilbertValue(), firstSlotOfRightNode.getSmallestKey(), + lastSlotOfRightNode.getLargestHilbertValue(), lastSlotOfRightNode.getLargestKey(), + rightNode.getId(), NodeHelpers.computeMbr(rightNode.getSlots())); + + oldRootNode.moveOutAllSlots(storageAdapter); + final IntermediateNode newRootNode = new IntermediateNode(rootId) + .insertSlot(storageAdapter, level, 0, leftChildSlot) + .insertSlot(storageAdapter, level, 1, rightChildSlot); + + storageAdapter.writeNodes(transaction, Lists.newArrayList(oldRootNode, newRootNode, leftNode, rightNode)); + } + + // Delete Path + + /** + * Method to delete from the R-tree. The item is treated unique per its point in space as well as its + * additional key that is passed in. + * @param tc transaction context + * @param point the point + * @param keySuffix the additional key to be stored with the item + * @return a completable future that completes when the delete operation is completed + */ + @Nonnull + public CompletableFuture delete(@Nonnull final TransactionContext tc, + @Nonnull final Point point, + @Nonnull final Tuple keySuffix) { + final BigInteger hilbertValue = hilbertValueFunction.apply(point); + final Tuple itemKey = Tuple.from(point.getCoordinates(), keySuffix); + + // + // Get to the leaf node we need to start the delete operation from and then call the appropriate method to + // perform the actual delete. + // + return tc.runAsync(transaction -> fetchPathForModification(transaction, hilbertValue, itemKey, false) + .thenCompose(leafNode -> { + if (leafNode == null) { + return AsyncUtil.DONE; + } + return deleteSlotIfExists(transaction, leafNode, hilbertValue, itemKey); + })); + } + + /** + * Deletes a slot from the {@link DataNode} passed or exits if the slot could not be found in the target node. + * in. + * @param transaction transaction + * @param targetNode leaf node that is the target of this delete operation + * @param hilbertValue the hilbert value of the point + * @param key the additional key to be stored with the item + * @return a completable future that completes when the delete is completed + */ + @Nonnull + private CompletableFuture deleteSlotIfExists(@Nonnull final Transaction transaction, + @Nonnull final DataNode targetNode, + @Nonnull final BigInteger hilbertValue, + @Nonnull final Tuple key) { + Verify.verify(targetNode.size() <= config.getMaxM()); + + final AtomicInteger level = new AtomicInteger(0); + final AtomicInteger deleteSlotIndex = new AtomicInteger(findDeleteItemSlotIndex(targetNode, hilbertValue, key)); + if (deleteSlotIndex.get() < 0) { + // + // The slot was not found meaning that the item was not found and that means we don't have to do anything + // here. + // + return AsyncUtil.DONE; + } + + // + // We found the slot and therefore the item. + // + + final NodeSlot deleteSlot = targetNode.getSlot(deleteSlotIndex.get()); + final AtomicReference currentNode = new AtomicReference<>(targetNode); + final AtomicReference parentSlot = new AtomicReference<>(deleteSlot); + + // + // Inch our way upwards in the tree to perform the necessary adjustments. What needs to be done next + // is informed by the result of the current operation: + // 1. A fuse happened; we need to delete an existing slot from the parent node -- prime current node and + // current slot and continue. + // 2. The slot was deleted but mbrs, largest Hilbert Values and largest Keys need to be adjusted upwards. + // 3. We are done as no further adjustments are necessary. + // + return AsyncUtil.whileTrue(() -> { + final NodeSlot currentDeleteSlot = parentSlot.get(); + + if (currentDeleteSlot != null) { + return deleteSlotFromTargetNode(transaction, level.get(), hilbertValue, key, currentNode.get(), currentDeleteSlot, deleteSlotIndex.get()) + .thenApply(nodeOrAdjust -> { + if (currentNode.get().isRoot()) { + return false; + } + currentNode.set(currentNode.get().getParentNode()); + parentSlot.set(nodeOrAdjust.getSlotInParent()); + deleteSlotIndex.set(nodeOrAdjust.getTombstoneNode() == null ? -1 : nodeOrAdjust.getTombstoneNode().getSlotIndexInParent()); + level.incrementAndGet(); + return nodeOrAdjust.getTombstoneNode() != null || nodeOrAdjust.parentNeedsAdjustment(); + }); + } else { + // adjustment only + return updateSlotsAndAdjustNode(transaction, level.get(), hilbertValue, key, currentNode.get(), false) + .thenApply(nodeOrAdjust -> { + Verify.verify(nodeOrAdjust.getSlotInParent() == null); + if (currentNode.get().isRoot()) { + return false; + } + currentNode.set(currentNode.get().getParentNode()); + level.incrementAndGet(); + return nodeOrAdjust.parentNeedsAdjustment(); + }); + } + }, executor); + } + + /** + * Delete and existing slot from the target node passed in. + * @param transaction transaction + * @param level the current level of target node, {@code 0} indicating the leaf level + * @param hilbertValue the Hilbert Value of the record that is being deleted + * @param key the key of the record that is being deleted + * @param targetNode target node + * @param deleteSlot existing slot that is to be deleted + * @param slotIndexInTargetNode The index of the new slot that we should use when inserting the new slot. While + * this information can be computed from the other arguments passed in, the caller already knows this + * information; we can avoid searching for the proper spot on our own. + * @return a completable future that when completed indicates what needs to be done next (see {@link NodeOrAdjust}). + */ + @Nonnull + private CompletableFuture deleteSlotFromTargetNode(@Nonnull final Transaction transaction, + final int level, + final BigInteger hilbertValue, + final Tuple key, + @Nonnull final Node targetNode, + @Nonnull final NodeSlot deleteSlot, + final int slotIndexInTargetNode) { + // + // We need to keep the number of slots per node between minM <= size() <= maxM unless this is the root node. + // + if (targetNode.isRoot() || targetNode.size() > config.getMinM()) { + if (logger.isTraceEnabled()) { + logger.trace("regular delete; node={}; size={}", targetNode, targetNode.size()); + } + targetNode.deleteSlot(storageAdapter, level - 1, slotIndexInTargetNode); + + if (targetNode.getKind() == NodeKind.INTERMEDIATE) { + // + // If this node is the root and the root node is an intermediate node, then it should at least have two + // children. + // + Verify.verify(!targetNode.isRoot() || targetNode.size() >= 2); + // + // If this is a delete operation within an intermediate node, the slot being deleted results from a + // fuse operation meaning a fuse has occurred on a lower level and the participating siblings of that split have + // potentially changed. + // + storageAdapter.writeNodes(transaction, Collections.singletonList(targetNode)); + } else { + Verify.verify(targetNode.getKind() == NodeKind.LEAF); + storageAdapter.clearLeafNodeSlot(transaction, (DataNode)targetNode, (ItemSlot)deleteSlot); + } + + // node is not under-flowing -- indicate that we are done fusing at the current node + if (!targetNode.isRoot()) { + return fetchParentNodeIfNecessary(transaction, targetNode, level, hilbertValue, key, false) + .thenApply(ignored -> adjustSlotInParent(targetNode, level) + ? NodeOrAdjust.ADJUST + : NodeOrAdjust.NONE); + } + + // no fuse and no adjustment + return CompletableFuture.completedFuture(NodeOrAdjust.NONE); // no fuse and no adjustment + } else { + // + // Node is under min-capacity -- borrow some children/items from the siblings if possible. + // + final CompletableFuture> siblings = + fetchParentNodeIfNecessary(transaction, targetNode, level, hilbertValue, key, false) + .thenCompose(ignored -> fetchSiblings(transaction, targetNode)); + + return siblings.thenApply(siblingNodes -> { + int numSlots = + Math.toIntExact(siblingNodes + .stream() + .mapToLong(Node::size) + .sum()); + + final Node tombstoneNode; + final List newSiblingNodes; + if (numSlots == siblingNodes.size() * config.getMinM()) { + if (logger.isTraceEnabled()) { + logger.trace("fusing nodes; node={}, siblings={}", + targetNode, + siblingNodes.stream().map(Node::toString).collect(Collectors.joining(","))); + } + tombstoneNode = siblingNodes.get(siblingNodes.size() - 1); + newSiblingNodes = siblingNodes.subList(0, siblingNodes.size() - 1); + } else { + if (logger.isTraceEnabled()) { + logger.trace("handling underflow; node={}, numSlots={}, siblings={}", + targetNode, + numSlots, + siblingNodes.stream().map(Node::toString).collect(Collectors.joining(","))); + } + tombstoneNode = null; + newSiblingNodes = siblingNodes; + } + + // temporarily underfill targetNode + numSlots--; + targetNode.deleteSlot(storageAdapter, level - 1, slotIndexInTargetNode); + + // sibling nodes are in hilbert value order + final Iterator slotIterator = + siblingNodes + .stream() + .flatMap(Node::slotsStream) + .iterator(); + + // + // Distribute all slots (excluding the one we want to delete) across all siblings (which also excludes + // the targetNode and (if we are fusing) the tombstoneNode). + // At the end of this modification all siblings have and (almost) equal count of slots that is + // guaranteed to be between minM and maxM. + // + + final int base = numSlots / newSiblingNodes.size(); + int rest = numSlots % newSiblingNodes.size(); + + List> newNodeSlotLists = Lists.newArrayList(); + List currentNodeSlots = Lists.newArrayList(); + while (slotIterator.hasNext()) { + final NodeSlot slot = slotIterator.next(); + currentNodeSlots.add(slot); + if (currentNodeSlots.size() == base + (rest > 0 ? 1 : 0)) { + if (rest > 0) { + // one fewer to distribute + rest--; + } + + newNodeSlotLists.add(currentNodeSlots); + currentNodeSlots = Lists.newArrayList(); + } + } + + Verify.verify(newSiblingNodes.size() == newNodeSlotLists.size()); + + if (tombstoneNode != null) { + // remove the slots for the tombstone node and update + tombstoneNode.moveOutAllSlots(storageAdapter); + storageAdapter.writeNodes(transaction, Collections.singletonList(tombstoneNode)); + } + + final Iterator newSiblingNodesIterator = newSiblingNodes.iterator(); + final Iterator> newNodeSlotsIterator = newNodeSlotLists.iterator(); + + // assign the slots to the appropriate nodes + while (newSiblingNodesIterator.hasNext()) { + final Node newSiblingNode = newSiblingNodesIterator.next(); + Verify.verify(newNodeSlotsIterator.hasNext()); + final List newNodeSlots = newNodeSlotsIterator.next(); + newSiblingNode.moveOutAllSlots(storageAdapter); + newSiblingNode.moveInSlots(storageAdapter, newNodeSlots); + } + + final IntermediateNode parentNode = Objects.requireNonNull(targetNode.getParentNode()); + if (parentNode.isRoot() && parentNode.size() == 2 && tombstoneNode != null) { + // + // The parent node (root) would only have one child after this delete. + // We shrink the tree by removing the root and making the last remaining sibling the root. + // + final Node toBePromotedNode = Iterables.getOnlyElement(newSiblingNodes); + promoteNodeToRoot(transaction, level, parentNode, toBePromotedNode); + return NodeOrAdjust.NONE; + } + + storageAdapter.writeNodes(transaction, newSiblingNodes); + + for (final Node newSiblingNode : newSiblingNodes) { + adjustSlotInParent(newSiblingNode, level); + } + + if (tombstoneNode == null) { + // + // We only handled underfill (and didn't need to fuse) but still need to continue adjusting + // mbrs, largest Hilbert values, and largest keys upward the tree. + // + return NodeOrAdjust.ADJUST; + } + + // + // We need to signal that the current operation ended in a fuse, and we need to delete the slot for + // the tombstoneNode one level higher. + // + return new NodeOrAdjust(parentNode.getSlot(tombstoneNode.getSlotIndexInParent()), + tombstoneNode, true); + }); + } + } + + /** + * Promote the given node to become the new root node. The node that is passed only changes its node id but retains + * all of it slots. This operation is the opposite of {@link #splitRootNode(Transaction, int, Node)} which can be + * invoked by the insert code path. + * @param transaction transaction + * @param level the level counting starting at {@code 0} indicating the leaf level increasing upwards + * @param oldRootNode the old root node + * @param toBePromotedNode node to be promoted. + */ + private void promoteNodeToRoot(final @Nonnull Transaction transaction, final int level, final IntermediateNode oldRootNode, + final Node toBePromotedNode) { + oldRootNode.deleteAllSlots(storageAdapter, level); + + // hold on to the slots of the to-be-promoted node -- copy them as moveOutAllSlots() will mutate the slot list + final List newRootSlots = ImmutableList.copyOf(toBePromotedNode.getSlots()); + toBePromotedNode.moveOutAllSlots(storageAdapter); + final Node newRootNode = toBePromotedNode.newOfSameKind(rootId).moveInSlots(storageAdapter, newRootSlots); + + // We need to update the node and the new root node in order to clear out the existing slots of the pre-promoted + // node. + storageAdapter.writeNodes(transaction, ImmutableList.of(oldRootNode, newRootNode, toBePromotedNode)); + } + + // + // Helper methods that may be called from more than one code path. + // + + /** + * Updates (persists) the slots for a target node and then computes the necessary adjustments in its parent + * node (without persisting those). + * @param transaction the transaction to use + * @param level the current level of target node, {@code 0} indicating the leaf level + * @param targetNode the target node + * @return A future containing either {@link NodeOrAdjust#NONE} if no further adjustments need to be persisted or + * {@link NodeOrAdjust#ADJUST} if the slots of the parent node of the target node need to be adjusted as + * well. + */ + @Nonnull + private CompletableFuture updateSlotsAndAdjustNode(@Nonnull final Transaction transaction, + final int level, + @Nonnull final BigInteger hilbertValue, + @Nonnull final Tuple key, + @Nonnull final Node targetNode, + final boolean isInsertUpdate) { + storageAdapter.writeNodes(transaction, Collections.singletonList(targetNode)); + + if (targetNode.isRoot()) { + return CompletableFuture.completedFuture(NodeOrAdjust.NONE); + } + + return fetchParentNodeIfNecessary(transaction, targetNode, level, hilbertValue, key, isInsertUpdate) + .thenApply(ignored -> adjustSlotInParent(targetNode, level) + ? NodeOrAdjust.ADJUST + : NodeOrAdjust.NONE); + } + + /** + * Updates the target node's mbr, largest Hilbert value as well its largest key in the target node's parent slot. + * @param targetNode target node + * @return {@code true} if any attributes of the target slot were modified, {@code false} otherwise. This will + * inform the caller if modifications need to be persisted and/or if the parent node itseld=f needs to be + * adjusted as well. + */ + private boolean adjustSlotInParent(@Nonnull final Node targetNode, final int level) { + Preconditions.checkArgument(!targetNode.isRoot()); + boolean slotHasChanged; + final IntermediateNode parentNode = Objects.requireNonNull(targetNode.getParentNode()); + final int slotIndexInParent = targetNode.getSlotIndexInParent(); + final ChildSlot childSlot = parentNode.getSlot(slotIndexInParent); + final Rectangle newMbr = NodeHelpers.computeMbr(targetNode.getSlots()); + slotHasChanged = !childSlot.getMbr().equals(newMbr); + final NodeSlot firstSlotOfTargetNode = targetNode.getSlot(0); + slotHasChanged |= !childSlot.getSmallestHilbertValue().equals(firstSlotOfTargetNode.getSmallestHilbertValue()); + slotHasChanged |= !childSlot.getSmallestKey().equals(firstSlotOfTargetNode.getSmallestKey()); + final NodeSlot lastSlotOfTargetNode = targetNode.getSlot(targetNode.size() - 1); + slotHasChanged |= !childSlot.getLargestHilbertValue().equals(lastSlotOfTargetNode.getLargestHilbertValue()); + slotHasChanged |= !childSlot.getLargestKey().equals(lastSlotOfTargetNode.getLargestKey()); + + if (slotHasChanged) { + parentNode.updateSlot(storageAdapter, level, slotIndexInParent, + new ChildSlot(firstSlotOfTargetNode.getSmallestHilbertValue(), firstSlotOfTargetNode.getSmallestKey(), + lastSlotOfTargetNode.getLargestHilbertValue(), lastSlotOfTargetNode.getLargestKey(), childSlot.getChildId(), + newMbr)); + } + return slotHasChanged; + } + + @Nonnull + private CompletableFuture fetchPathForModification(@Nonnull final Transaction transaction, + @Nonnull final BigInteger hilbertValue, + @Nonnull final Tuple key, + final boolean isInsertUpdate) { + if (config.isUseNodeSlotIndex()) { + return scanIndexAndFetchLeafNode(transaction, hilbertValue, key, isInsertUpdate); + } else { + return fetchUpdatePathToLeaf(transaction, hilbertValue, key, isInsertUpdate); + } + } + + @Nonnull + private CompletableFuture scanIndexAndFetchLeafNode(@Nonnull final ReadTransaction transaction, + @Nonnull final BigInteger hilbertValue, + @Nonnull final Tuple key, + final boolean isInsertUpdate) { + return storageAdapter.scanNodeIndexAndFetchNode(transaction, 0, hilbertValue, key, isInsertUpdate) + .thenApply(node -> { + Verify.verify(node == null || + (node.getKind() == NodeKind.LEAF && node instanceof DataNode)); + return (DataNode)node; + }); + } + + @Nonnull + private CompletableFuture scanIndexAndFetchIntermediateNode(@Nonnull final ReadTransaction transaction, + final int level, + @Nonnull final BigInteger hilbertValue, + @Nonnull final Tuple key, + final boolean isInsertUpdate) { + Verify.verify(level > 0); + return storageAdapter.scanNodeIndexAndFetchNode(transaction, level, hilbertValue, key, isInsertUpdate) + .thenApply(node -> { + // + // Note that there is no non-error scenario where node can be null here; either the node is + // not in the node slot index but is the root node which has already been resolved and fetched OR + // this node is a legitimate parent node of a node we know must exist as level > 0. If node were + // null here, it would mean that there is a node that is not the root but its parent is not in + // the R-tree. + // + Verify.verify(node.getKind() == NodeKind.INTERMEDIATE && node instanceof IntermediateNode); + return (IntermediateNode)node; + }); + } + + @Nonnull + private CompletableFuture fetchParentNodeIfNecessary(@Nonnull final ReadTransaction transaction, + @Nonnull final Node node, + final int level, + @Nonnull final BigInteger hilbertValue, + @Nonnull final Tuple key, + final boolean isInsertUpdate) { + Verify.verify(!node.isRoot()); + final IntermediateNode linkedParentNode = node.getParentNode(); + if (linkedParentNode != null) { + return CompletableFuture.completedFuture(linkedParentNode); + } + + Verify.verify(getConfig().isUseNodeSlotIndex()); + return scanIndexAndFetchIntermediateNode(transaction, level + 1, hilbertValue, key, isInsertUpdate) + .thenApply(parentNode -> { + final int slotIndexInParent = findChildSlotIndex(parentNode, node.getId()); + Verify.verify(slotIndexInParent >= 0); + node.linkToParent(parentNode, slotIndexInParent); + return parentNode; + }); + } + + /** + * Method to fetch the update path of a given {@code (hilbertValue, key)} pair. The update path is a {@link DataNode} + * and all its parent nodes to the root node. The caller can invoke {@link Node#getParentNode()} to navigate to + * all nodes in the update path starting from the {@link DataNode} that is returned. The {@link DataNode} that is + * returned may or may not already contain a slot for the {@code (hilbertValue, key)} pair passed in. This logic is + * invoked for insert, updates, as well as delete operations. If it is used for insert and the item is not yet + * part of the leaf node, the leaf node that is returned can be understood as the correct place to insert the item + * in question. + * @param transaction the transaction to use + * @param hilbertValue the Hilbert value to look for + * @param key the key to look for + * @param isInsertUpdate is this call part of and index/update operation or a delete operation + * @return A completable future containing a {@link DataNode} and by extension (through {@link Node#getParentNode()}) + * all intermediate nodes up to the root node that may get affected by an insert, update, or delete + * of the specified item. + */ + @Nonnull + private CompletableFuture fetchUpdatePathToLeaf(@Nonnull final Transaction transaction, + @Nonnull final BigInteger hilbertValue, + @Nonnull final Tuple key, + final boolean isInsertUpdate) { + final AtomicReference parentNode = new AtomicReference<>(null); + final AtomicInteger slotInParent = new AtomicInteger(-1); + final AtomicReference currentId = new AtomicReference<>(rootId); + final AtomicReference leafNode = new AtomicReference<>(null); + return AsyncUtil.whileTrue(() -> storageAdapter.fetchNode(transaction, currentId.get()) + .thenApply(node -> { + if (node == null) { + if (Arrays.equals(currentId.get(), rootId)) { + Verify.verify(leafNode.get() == null); + return false; + } + throw new IllegalStateException("unable to fetch node for insert or update"); + } + if (parentNode.get() != null) { + node.linkToParent(parentNode.get(), slotInParent.get()); + } + if (node.getKind() == NodeKind.INTERMEDIATE) { + final IntermediateNode intermediateNode = (IntermediateNode)node; + final int slotIndex = findChildSlotIndex(intermediateNode, hilbertValue, key, isInsertUpdate); + if (slotIndex < 0) { + Verify.verify(!isInsertUpdate); + // + // This is for a delete operation and we were unable to find a child that covers + // the Hilbert Value/key to be deleted + return false; + } + + parentNode.set(intermediateNode); + slotInParent.set(slotIndex); + final ChildSlot childSlot = intermediateNode.getSlot(slotIndex); + currentId.set(childSlot.getChildId()); + return true; + } else { + leafNode.set((DataNode)node); + return false; + } + }), executor) + .thenApply(ignored -> { + final DataNode node = leafNode.get(); + if (logger.isTraceEnabled()) { + logger.trace("update path; path={}", NodeHelpers.nodeIdPath(node)); + } + return node; + }); + } + + /** + * Method to fetch the siblings of a given node. The node passed in must not be the root node and must be linked up + * to its parent. The parent already has information obout the children ids. This method (through the slot + * information of the node passed in) can then determine adjacent nodes. + * @param transaction the transaction to use + * @param node the node to fetch siblings for + * @return a completable future containing a list of {@link Node}s that contain the {@link Config#getSplitS()} + * number of siblings (where the node passed in is counted as a sibling) if that many siblings exist. In + * the case (i.e. for a small root node) where there are not enough siblings we return the maximum possible + * number of siblings. The returned sibling nodes are returned in Hilbert value order and contain the node + * passed in at the correct position in the returned list. The siblings will also attempt to hug the nodes + * passed in as good as possible meaning that we attempt to return the node passed in as middle-most element + * of the returned list. + */ + @Nonnull + private CompletableFuture> fetchSiblings(@Nonnull final Transaction transaction, + @Nonnull final Node node) { + // this deque is only modified by once upon creation + final ArrayDeque toBeProcessed = new ArrayDeque<>(); + final List> working = Lists.newArrayList(); + final int numSiblings = config.getSplitS(); + final Node[] siblings = new Node[numSiblings]; + + // + // Do some acrobatics to find the best start/end positions for the siblings. Take into account how many + // are warranted, if the node that was passed occupies a slot in its parent node that is touching the end or the + // beginning of the parent's slots, and the total number of slots in the parent of the node that was + // passed in. + // + final IntermediateNode parentNode = Objects.requireNonNull(node.getParentNode()); + int slotIndexInParent = node.getSlotIndexInParent(); + int start = slotIndexInParent - numSiblings / 2; + int end = start + numSiblings; + if (start < 0) { + start = 0; + end = numSiblings; + } else if (end > parentNode.size()) { + end = parentNode.size(); + start = end - numSiblings; + } + + // because lambdas + final int minSibling = start; + + for (int i = start; i < end; i++) { + toBeProcessed.addLast(parentNode.getSlot(i).getChildId()); + } + + // Fetch all sibling nodes (in parallel if possible). + return AsyncUtil.whileTrue(() -> { + working.removeIf(CompletableFuture::isDone); + + while (working.size() <= MAX_CONCURRENT_READS) { + final int index = numSiblings - toBeProcessed.size(); + final byte[] currentId = toBeProcessed.pollFirst(); + if (currentId == null) { + break; + } + + final int slotIndex = minSibling + index; + if (slotIndex != slotIndexInParent) { + working.add(storageAdapter.fetchNode(transaction, currentId) + .thenAccept(siblingNode -> { + Objects.requireNonNull(siblingNode); + siblingNode.linkToParent(parentNode, slotIndex); + siblings[index] = siblingNode; + })); + } else { + // put node in the list of siblings -- even though node is strictly speaking not a sibling of itself + siblings[index] = node; + } + } + + if (working.isEmpty()) { + return AsyncUtil.READY_FALSE; + } + return AsyncUtil.whenAny(working).thenApply(v -> true); + }, executor).thenApply(vignore -> Lists.newArrayList(siblings)); + } + + /** + * Method to compute the depth of this R-tree. + * @param transactionContext transaction context to be used + * @return the depth of the R-tree + */ + public int depth(@Nonnull final TransactionContext transactionContext) { + // + // find the number of levels in this tree + // + Node node = + transactionContext.run(tr -> fetchUpdatePathToLeaf(tr, BigInteger.ONE, new Tuple(), true).join()); + if (node == null) { + logger.trace("R-tree is empty."); + return 0; + } + + int numLevels = 1; + while (node.getParentNode() != null) { + numLevels ++; + node = node.getParentNode(); + } + Verify.verify(node.isRoot(), "end of update path should be the root"); + logger.trace("numLevels = {}", numLevels); + return numLevels; + } + + /** + * Method to validate the Hilbert R-tree. + * @param db the database to use + */ + public void validate(@Nonnull final Database db) { + validate(db, Integer.MAX_VALUE); + } + + /** + * Method to validate the Hilbert R-tree. + * @param db the database to use + * @param maxNumNodesToBeValidated a maximum number of nodes this call should attempt to validate + */ + public void validate(@Nonnull final Database db, + final int maxNumNodesToBeValidated) { + + ArrayDeque toBeProcessed = new ArrayDeque<>(); + toBeProcessed.addLast(new ValidationTraversalState(depth(db) - 1, null, rootId)); + + while (!toBeProcessed.isEmpty()) { + db.run(tr -> validate(tr, maxNumNodesToBeValidated, toBeProcessed).join()); + } + } + + /** + * Method to validate the Hilbert R-tree. + * @param transaction the transaction to use + * @param maxNumNodesToBeValidated a maximum number of nodes this call should attempt to validate + * @param toBeProcessed a deque with node information that still needs to be processed + * @return a completable future that completes successfully with the current deque of to-be-processed nodes if the + * portion of the tree that was validated is in fact valid, completes with failure otherwise + */ + @Nonnull + private CompletableFuture> validate(@Nonnull final Transaction transaction, + final int maxNumNodesToBeValidated, + @Nonnull final ArrayDeque toBeProcessed) { + final AtomicInteger numNodesEnqueued = new AtomicInteger(0); + final List>> working = Lists.newArrayList(); + + // Fetch the entire tree. + return AsyncUtil.whileTrue(() -> { + final Iterator>> workingIterator = working.iterator(); + while (workingIterator.hasNext()) { + final CompletableFuture> nextFuture = workingIterator.next(); + if (nextFuture.isDone()) { + toBeProcessed.addAll(nextFuture.join()); + workingIterator.remove(); + } + } + + while (working.size() <= MAX_CONCURRENT_READS && numNodesEnqueued.get() < maxNumNodesToBeValidated) { + final ValidationTraversalState currentValidationTraversalState = toBeProcessed.pollFirst(); + if (currentValidationTraversalState == null) { + break; + } + + final IntermediateNode parentNode = currentValidationTraversalState.getParentNode(); + final int level = currentValidationTraversalState.getLevel(); + final ChildSlot childSlotInParentNode; + final int slotIndexInParent; + if (parentNode != null) { + int slotIndex; + ChildSlot childSlot = null; + for (slotIndex = 0; slotIndex < parentNode.size(); slotIndex++) { + childSlot = parentNode.getSlot(slotIndex); + if (Arrays.equals(childSlot.getChildId(), currentValidationTraversalState.getChildId())) { + break; + } + } + + if (slotIndex == parentNode.size()) { + throw new IllegalStateException("child slot not found in parent for child node"); + } else { + childSlotInParentNode = childSlot; + slotIndexInParent = slotIndex; + } + } else { + childSlotInParentNode = null; + slotIndexInParent = -1; + } + + final CompletableFuture fetchedNodeFuture = + onReadListener.onAsyncRead(storageAdapter.fetchNode(transaction, currentValidationTraversalState.getChildId()) + .thenApply(node -> { + if (parentNode != null) { + Objects.requireNonNull(node); + node.linkToParent(parentNode, slotIndexInParent); + } + return node; + }) + .thenCompose(childNode -> { + if (parentNode != null && getConfig().isUseNodeSlotIndex()) { + final var childSlot = parentNode.getSlot(slotIndexInParent); + return storageAdapter.scanNodeIndexAndFetchNode(transaction, level, + childSlot.getLargestHilbertValue(), childSlot.getLargestKey(), false) + .thenApply(nodeFromIndex -> { + Objects.requireNonNull(nodeFromIndex); + if (!Arrays.equals(nodeFromIndex.getId(), childNode.getId())) { + logger.warn("corrupt node slot index at level {}, parentNode = {}", level, parentNode); + throw new IllegalStateException("corrupt node index"); + } + return childNode; + }); + } + return CompletableFuture.completedFuture(childNode); + })); + working.add(fetchedNodeFuture.thenApply(childNode -> { + if (childNode == null) { + // Starting at root node but root node was not fetched since the R-tree has no entries. + return ImmutableList.of(); + } + childNode.validate(); + childNode.validateParentNode(parentNode, childSlotInParentNode); + + // add all children to the to be processed queue + if (childNode.getKind() == NodeKind.INTERMEDIATE) { + return ((IntermediateNode)childNode).getSlots() + .stream() + .map(childSlot -> new ValidationTraversalState(level - 1, + (IntermediateNode)childNode, childSlot.getChildId())) + .collect(ImmutableList.toImmutableList()); + } else { + return ImmutableList.of(); + } + })); + numNodesEnqueued.addAndGet(1); + } + + if (working.isEmpty()) { + return AsyncUtil.READY_FALSE; + } + return AsyncUtil.whenAny(working).thenApply(v -> true); + }, executor).thenApply(vignore -> toBeProcessed); + } + + /** + * Method to find the appropriate child slot index for a given Hilbert value and key. This method is used + * to find the proper slot indexes for the insert/update path and for the delete path. Note that if + * {@code (largestHilbertValue, largestKey)} of the last child is less than {@code (hilbertValue, key)}, we insert + * through the last child as we treat the (non-existing) next item as {@code (infinity, infinity)}. + * @param intermediateNode the intermediate node to search + * @param hilbertValue hilbert value + * @param key key + * @param isInsertUpdate indicator if the caller + * @return the 0-based slot index that corresponds to the given {@code (hilbertValue, key)} pair {@code p} if a slot + * covers that pair. If such a slot cannot be found while a new record is inserted, slot {@code 0} is + * returned if that slot is compared larger than {@code p}, the last slot ({@code size - 1}) if that slot is + * compared smaller than {@code p}. If, on the contrary, a record is deleted and a slot covering {@code p} + * cannot be found, this method returns {@code -1}. + */ + private static int findChildSlotIndex(@Nonnull final IntermediateNode intermediateNode, + @Nonnull final BigInteger hilbertValue, + @Nonnull final Tuple key, + final boolean isInsertUpdate) { + Verify.verify(!intermediateNode.isEmpty()); + + if (!isInsertUpdate) { + // make sure that the node covers the Hilbert Value/key we would like to delete + final ChildSlot firstChildSlot = intermediateNode.getSlot(0); + + final int compare = NodeSlot.compareHilbertValueKeyPair(firstChildSlot.getSmallestHilbertValue(), firstChildSlot.getSmallestKey(), + hilbertValue, key); + if (compare > 0) { + // child smallest HV/key > target HV/key + return -1; + } + } + + for (int slotIndex = 0; slotIndex < intermediateNode.size(); slotIndex++) { + final ChildSlot childSlot = intermediateNode.getSlot(slotIndex); + + // + // Choose subtree with the minimum Hilbert value that is greater than the target + // Hilbert value. If there is no such subtree, i.e. the target Hilbert value is the + // largest Hilbert value, we choose the largest one in the current node. + // + final int compare = NodeSlot.compareHilbertValueKeyPair(childSlot.getLargestHilbertValue(), childSlot.getLargestKey(), hilbertValue, key); + if (compare >= 0) { + // child largest HV/key > target HV/key + return slotIndex; + } + } + + // + // This is an intermediate node; we insert through the last child, but return -1 if this is for a delete + // operation. + return isInsertUpdate ? intermediateNode.size() - 1 : - 1; + } + + /** + * Method to find the appropriate child slot index for a given child it. + * @param parentNode the intermediate node to search + * @param childId the child id to search for + * @return if found the 0-based slot index that corresponds to slot using holding the given {@code childId}; + * {@code -1} otherwise + */ + private static int findChildSlotIndex(@Nonnull final IntermediateNode parentNode, @Nonnull final byte[] childId) { + for (int slotIndex = 0; slotIndex < parentNode.size(); slotIndex++) { + final ChildSlot childSlot = parentNode.getSlot(slotIndex); + + if (Arrays.equals(childSlot.getChildId(), childId)) { + return slotIndex; + } + } + return -1; + } + + /** + * Method to find the appropriate item slot index for a given Hilbert value and key. This method is used + * to find the proper item slot index for the insert/update path. + * @param leafNode the leaf node to search + * @param hilbertValue hilbert value + * @param key key + * @return {@code -1} if the item specified by {@code (hilbertValue, key)} already exists in {@code leafNode}; + * the 0-based slot index that represents the insertion point index of the given {@code (hilbertValue, key)} + * pair, otherwise + */ + private static int findInsertUpdateItemSlotIndex(@Nonnull final DataNode leafNode, + @Nonnull final BigInteger hilbertValue, + @Nonnull final Tuple key) { + for (int slotIndex = 0; slotIndex < leafNode.size(); slotIndex++) { + final ItemSlot slot = leafNode.getSlot(slotIndex); + + final int compare = NodeSlot.compareHilbertValueKeyPair(slot.getHilbertValue(), slot.getKey(), hilbertValue, key); + if (compare == 0) { + return -1; + } + + if (compare > 0) { + return slotIndex; + } + } + + return leafNode.size(); + } + + /** + * Method to find the appropriate item slot index for a given Hilbert value and key. This method is used + * to find the proper item slot index for the delete path. + * @param leafNode the leaf node to search + * @param hilbertValue hilbert value + * @param key key + * @return {@code -1} if the item specified by {@code (hilbertValue, key)} does not exist in {@code leafNode}; + * the 0-based slot index that corresponds to the slot for the given {@code (hilbertValue, key)} + * pair, otherwise + */ + private static int findDeleteItemSlotIndex(@Nonnull final DataNode leafNode, + @Nonnull final BigInteger hilbertValue, + @Nonnull final Tuple key) { + for (int slotIndex = 0; slotIndex < leafNode.size(); slotIndex++) { + final ItemSlot slot = leafNode.getSlot(slotIndex); + + final int compare = NodeSlot.compareHilbertValueKeyPair(slot.getHilbertValue(), slot.getKey(), hilbertValue, key); + if (compare == 0) { + return slotIndex; + } + + if (compare > 0) { + return -1; + } + } + + return -1; + } + + /** + * Traversal state of a scan over the tree. A scan consists of an initial walk to the left-most applicable leaf node + * potentially containing items relevant to the scan. The caller then consumes that leaf node and advances to the + * next leaf node that is relevant to the scan. The notion of next emerges using the order defined by the + * composite {@code (hilbertValue, key)} for items in leaf nodes and {@code (largestHilbertValue, largestKey)} in + * intermediate nodes. The traversal state captures the node ids that still have to be processed on each discovered + * level in order to fulfill the requirements of the scan operation. + */ + private static class TraversalState { + @Nullable + private final List> toBeProcessed; + + @Nullable + private final DataNode currentLeafNode; + + private TraversalState(@Nullable final List> toBeProcessed, @Nullable final DataNode currentLeafNode) { + this.toBeProcessed = toBeProcessed; + this.currentLeafNode = currentLeafNode; + } + + @Nonnull + public List> getToBeProcessed() { + return Objects.requireNonNull(toBeProcessed); + } + + @Nonnull + public DataNode getCurrentLeafNode() { + return Objects.requireNonNull(currentLeafNode); + } + + public boolean isEnd() { + return currentLeafNode == null; + } + + public static TraversalState of(@Nonnull final List> toBeProcessed, @Nonnull final DataNode currentLeafNode) { + return new TraversalState(toBeProcessed, currentLeafNode); + } + + public static TraversalState end() { + return new TraversalState(null, null); + } + } + + /** + * An {@link AsyncIterator} over the leaf nodes that represent the result of a scan over the tree. This iterator + * interfaces with the scan logic + * (see {@link #fetchLeftmostPathToLeaf(ReadTransaction, byte[], BigInteger, Tuple, Predicate, BiPredicate)} and + * {@link #fetchNextPathToLeaf(ReadTransaction, TraversalState, BigInteger, Tuple, Predicate, BiPredicate)}) and wraps + * intermediate {@link TraversalState}s created by these methods. + */ + private class LeafIterator implements AsyncIterator { + @Nonnull + private final ReadTransaction readTransaction; + @Nonnull + private final byte[] rootId; + @Nullable + private final BigInteger lastHilbertValue; + @Nullable + private final Tuple lastKey; + @Nonnull + private final Predicate mbrPredicate; + @Nonnull + private final BiPredicate suffixKeyPredicate; + + @Nullable + private TraversalState currentState; + @Nullable + private CompletableFuture nextStateFuture; + + @SpotBugsSuppressWarnings("EI_EXPOSE_REP2") + public LeafIterator(@Nonnull final ReadTransaction readTransaction, @Nonnull final byte[] rootId, + @Nullable final BigInteger lastHilbertValue, @Nullable final Tuple lastKey, + @Nonnull final Predicate mbrPredicate, @Nonnull final BiPredicate suffixKeyPredicate) { + Preconditions.checkArgument((lastHilbertValue == null && lastKey == null) || + (lastHilbertValue != null && lastKey != null)); + this.readTransaction = readTransaction; + this.rootId = rootId; + this.lastHilbertValue = lastHilbertValue; + this.lastKey = lastKey; + this.mbrPredicate = mbrPredicate; + this.suffixKeyPredicate = suffixKeyPredicate; + this.currentState = null; + this.nextStateFuture = null; + } + + @Override + public CompletableFuture onHasNext() { + if (nextStateFuture == null) { + if (currentState == null) { + nextStateFuture = fetchLeftmostPathToLeaf(readTransaction, rootId, lastHilbertValue, lastKey, + mbrPredicate, suffixKeyPredicate); + } else { + nextStateFuture = fetchNextPathToLeaf(readTransaction, currentState, lastHilbertValue, lastKey, + mbrPredicate, suffixKeyPredicate); + } + } + return nextStateFuture.thenApply(traversalState -> !traversalState.isEnd()); + } + + @Override + public boolean hasNext() { + return onHasNext().join(); + } + + @Override + public DataNode next() { + if (hasNext()) { + // underlying has already completed + currentState = Objects.requireNonNull(nextStateFuture).join(); + nextStateFuture = null; + return currentState.getCurrentLeafNode(); + } + throw new NoSuchElementException("called next() on exhausted iterator"); + } + + @Override + public void cancel() { + if (nextStateFuture != null) { + nextStateFuture.cancel(false); + } + } + } + + /** + * Iterator for iterating the items contained in the leaf nodes produced by an underlying {@link LeafIterator}. + * This iterator is the async equivalent of + * {@code Streams.stream(leafIterator).flatMap(leafNode -> leafNode.getItems().stream()).toIterator()}. + */ + public static class ItemSlotIterator implements AsyncIterator { + @Nonnull + private final AsyncIterator leafIterator; + @Nullable + private DataNode currentLeafNode; + @Nullable + private Iterator currenLeafItemsIterator; + + private ItemSlotIterator(@Nonnull final AsyncIterator leafIterator) { + this.leafIterator = leafIterator; + this.currentLeafNode = null; + this.currenLeafItemsIterator = null; + } + + @Override + public CompletableFuture onHasNext() { + if (currenLeafItemsIterator != null && currenLeafItemsIterator.hasNext()) { + return CompletableFuture.completedFuture(true); + } + // we know that each leaf has items (or if it doesn't it is the root; we are done if there are no items + return leafIterator.onHasNext() + .thenApply(hasNext -> { + if (hasNext) { + this.currentLeafNode = leafIterator.next(); + this.currenLeafItemsIterator = currentLeafNode.getSlots().iterator(); + return currenLeafItemsIterator.hasNext(); + } + return false; + }); + } + + @Override + public boolean hasNext() { + return onHasNext().join(); + } + + @Override + public ItemSlot next() { + if (hasNext()) { + return Objects.requireNonNull(currenLeafItemsIterator).next(); + } + throw new NoSuchElementException("called next() on exhausted iterator"); + } + + @Override + public void cancel() { + leafIterator.cancel(); + } + } + + /** + * Class to signal the caller of insert/update/delete code paths what the next action in that path should be. + * The indicated action is either another insert/delete on a higher level in the tree, further adjustments of + * secondary attributes on a higher level in the tree, or an indication that the insert/update/delete path is done + * with all necessary modifications. + */ + private static class NodeOrAdjust { + public static final NodeOrAdjust NONE = new NodeOrAdjust(null, null, false); + public static final NodeOrAdjust ADJUST = new NodeOrAdjust(null, null, true); + + @Nullable + private final ChildSlot slotInParent; + @Nullable + private final Node node; + + private final boolean parentNeedsAdjustment; + + private NodeOrAdjust(@Nullable final ChildSlot slotInParent, @Nullable final Node node, final boolean parentNeedsAdjustment) { + Verify.verify((slotInParent == null && node == null) || + (slotInParent != null && node != null)); + this.slotInParent = slotInParent; + this.node = node; + this.parentNeedsAdjustment = parentNeedsAdjustment; + } + + @Nullable + public ChildSlot getSlotInParent() { + return slotInParent; + } + + @Nullable + public Node getSplitNode() { + return node; + } + + @Nullable + public Node getTombstoneNode() { + return node; + } + + public boolean parentNeedsAdjustment() { + return parentNeedsAdjustment; + } + } + + /** + * Helper class for the traversal of nodes during tree validation. + */ + private static class ValidationTraversalState { + final int level; + @Nullable + private final IntermediateNode parentNode; + @Nonnull + private final byte[] childId; + + public ValidationTraversalState(final int level, @Nullable final IntermediateNode parentNode, @Nonnull final byte[] childId) { + this.level = level; + this.parentNode = parentNode; + this.childId = childId; + } + + public int getLevel() { + return level; + } + + @Nullable + public IntermediateNode getParentNode() { + return parentNode; + } + + @Nonnull + public byte[] getChildId() { + return childId; + } + } + + /** + * Class to capture an N-dimensional point. It wraps a {@link Tuple} mostly due to proximity with its serialization + * format and provides helpers for Euclidean operations. Note that the coordinates used here do not need to be + * numbers. + */ + public static class Point { + @Nonnull + private final Tuple coordinates; + + public Point(@Nonnull final Tuple coordinates) { + Preconditions.checkArgument(!coordinates.isEmpty()); + this.coordinates = coordinates; + } + + @Nonnull + public Tuple getCoordinates() { + return coordinates; + } + + public int getNumDimensions() { + return coordinates.size(); + } + + @Nullable + public Object getCoordinate(final int dimension) { + return coordinates.get(dimension); + } + + @Nullable + public Number getCoordinateAsNumber(final int dimension) { + return (Number)getCoordinate(dimension); + } + + @Override + public boolean equals(final Object o) { + if (this == o) { + return true; + } + if (!(o instanceof Point)) { + return false; + } + final Point point = (Point)o; + return TupleHelpers.equals(coordinates, point.coordinates); + } + + @Override + public int hashCode() { + return coordinates.hashCode(); + } + + @Nonnull + @Override + public String toString() { + return coordinates.toString(); + } + } + + /** + * Class to capture an N-dimensional rectangle/cube/hypercube. It wraps a {@link Tuple} mostly due to proximity + * with its serialization format and provides helpers for Euclidean operations. Note that the coordinates used here + * do not need to be numbers. + */ + public static class Rectangle { + /** + * A tuple that holds the coordinates of this N-dimensional rectangle. The layout is defined as + * {@code (low1, low2, ..., lowN, high1, high2, ..., highN}. Note that we don't use nested {@link Tuple}s for + * space-saving reasons (when the tuple is serialized). + */ + @Nonnull + private final Tuple ranges; + + public Rectangle(final Tuple ranges) { + Preconditions.checkArgument(!ranges.isEmpty() && ranges.size() % 2 == 0); + this.ranges = ranges; + } + + public int getNumDimensions() { + return ranges.size() >> 1; + } + + @Nonnull + public Tuple getRanges() { + return ranges; + } + + @Nonnull + public Object getLow(final int dimension) { + return ranges.get(dimension); + } + + @Nonnull + public Object getHigh(final int dimension) { + return ranges.get((ranges.size() >> 1) + dimension); + } + + @Nonnull + public BigInteger area() { + BigInteger currentArea = BigInteger.ONE; + for (int d = 0; d < getNumDimensions(); d++) { + currentArea = currentArea.multiply(BigInteger.valueOf(((Number)getHigh(d)).longValue() - ((Number)getLow(d)).longValue())); + } + return currentArea; + } + + @Nonnull + public Rectangle unionWith(@Nonnull final Point point) { + Preconditions.checkArgument(getNumDimensions() == point.getNumDimensions()); + boolean isModified = false; + Object[] ranges = new Object[getNumDimensions() << 1]; + + for (int d = 0; d < getNumDimensions(); d++) { + final Object coordinate = point.getCoordinate(d); + final Tuple coordinateTuple = Tuple.from(coordinate); + final Object low = getLow(d); + final Tuple lowTuple = Tuple.from(low); + if (TupleHelpers.compare(coordinateTuple, lowTuple) < 0) { + ranges[d] = coordinate; + isModified = true; + } else { + ranges[d] = low; + } + + final Object high = getHigh(d); + final Tuple highTuple = Tuple.from(high); + if (TupleHelpers.compare(coordinateTuple, highTuple) > 0) { + ranges[getNumDimensions() + d] = coordinate; + isModified = true; + } else { + ranges[getNumDimensions() + d] = high; + } + } + + if (!isModified) { + return this; + } + + return new Rectangle(Tuple.from(ranges)); + } + + @Nonnull + public Rectangle unionWith(@Nonnull final Rectangle other) { + Preconditions.checkArgument(getNumDimensions() == other.getNumDimensions()); + boolean isModified = false; + Object[] ranges = new Object[getNumDimensions() << 1]; + + for (int d = 0; d < getNumDimensions(); d++) { + final Object otherLow = other.getLow(d); + final Tuple otherLowTuple = Tuple.from(otherLow); + final Object otherHigh = other.getHigh(d); + final Tuple otherHighTuple = Tuple.from(otherHigh); + + final Object low = getLow(d); + final Tuple lowTuple = Tuple.from(low); + if (TupleHelpers.compare(otherLowTuple, lowTuple) < 0) { + ranges[d] = otherLow; + isModified = true; + } else { + ranges[d] = low; + } + final Object high = getHigh(d); + final Tuple highTuple = Tuple.from(high); + if (TupleHelpers.compare(otherHighTuple, highTuple) > 0) { + ranges[getNumDimensions() + d] = otherHigh; + isModified = true; + } else { + ranges[getNumDimensions() + d] = high; + } + } + + if (!isModified) { + return this; + } + + return new Rectangle(Tuple.from(ranges)); + } + + public boolean isOverlapping(@Nonnull final Rectangle other) { + Preconditions.checkArgument(getNumDimensions() == other.getNumDimensions()); + + for (int d = 0; d < getNumDimensions(); d++) { + final Tuple otherLowTuple = Tuple.from(other.getLow(d)); + final Tuple otherHighTuple = Tuple.from(other.getHigh(d)); + + final Tuple lowTuple = Tuple.from(getLow(d)); + final Tuple highTuple = Tuple.from(getHigh(d)); + + if (TupleHelpers.compare(highTuple, otherLowTuple) < 0 || + TupleHelpers.compare(lowTuple, otherHighTuple) > 0) { + return false; + } + } + return true; + } + + public boolean contains(@Nonnull final Point point) { + Preconditions.checkArgument(getNumDimensions() == point.getNumDimensions()); + + for (int d = 0; d < getNumDimensions(); d++) { + final Tuple otherTuple = Tuple.from(point.getCoordinate(d)); + + final Tuple lowTuple = Tuple.from(getLow(d)); + final Tuple highTuple = Tuple.from(getHigh(d)); + + if (TupleHelpers.compare(highTuple, otherTuple) < 0 || + TupleHelpers.compare(lowTuple, otherTuple) > 0) { + return false; + } + } + return true; + } + + @Override + public boolean equals(final Object o) { + if (this == o) { + return true; + } + if (!(o instanceof Rectangle)) { + return false; + } + final Rectangle rectangle = (Rectangle)o; + return TupleHelpers.equals(ranges, rectangle.ranges); + } + + @Override + public int hashCode() { + return ranges.hashCode(); + } + + @Nonnull + public String toPlotString() { + final StringBuilder builder = new StringBuilder(); + for (int d = 0; d < getNumDimensions(); d++) { + builder.append(((Number)getLow(d)).longValue()); + if (d + 1 < getNumDimensions()) { + builder.append(","); + } + } + + builder.append(","); + + for (int d = 0; d < getNumDimensions(); d++) { + builder.append(((Number)getHigh(d)).longValue()); + if (d + 1 < getNumDimensions()) { + builder.append(","); + } + } + return builder.toString(); + } + + @Nonnull + @Override + public String toString() { + return ranges.toString(); + } + + @Nonnull + public static Rectangle fromPoint(@Nonnull final Point point) { + final Object[] mbrRanges = new Object[point.getNumDimensions() * 2]; + for (int d = 0; d < point.getNumDimensions(); d++) { + final Object coordinate = point.getCoordinate(d); + mbrRanges[d] = coordinate; + mbrRanges[point.getNumDimensions() + d] = coordinate; + } + return new Rectangle(Tuple.from(mbrRanges)); + } + } + +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java new file mode 100644 index 0000000000..040653e4e0 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java @@ -0,0 +1,57 @@ +/* + * IntermediateNode.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.tuple.Tuple; +import com.christianheina.langx.half4j.Half; + +import javax.annotation.Nonnull; +import java.util.List; + +/** + * An intermediate node of the R-tree. An intermediate node holds the holds information about its children nodes that + * be intermediate nodes or leaf nodes. The secondary attributes such as {@code largestHilbertValue}, + * {@code largestKey} can be derived (and recomputed) if the children of this node are available to be introspected. + */ +class IntermediateNode extends AbstractNode { + public IntermediateNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, + @Nonnull final List neighbors) { + super(primaryKey, vector, neighbors); + } + + @Nonnull + @Override + public DataNode asDataNode() { + throw new IllegalStateException("this is not a data node"); + } + + @Nonnull + @Override + public IntermediateNode asIntermediateNode() { + return this; + } + + @Nonnull + @Override + public NodeKind getKind() { + return NodeKind.INTERMEDIATE; + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java new file mode 100644 index 0000000000..afa0e14fc5 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java @@ -0,0 +1,129 @@ +/* + * Metric.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +public interface Metric { + double distance(Double[] vector1, Double[] vector2); + + default double comparativeDistance(Double[] vector1, Double[] vector2) { + return distance(vector1, vector2); + } + + /** + * A helper method to validate that vectors can be compared. + * @param vector1 The first vector. + * @param vector2 The second vector. + */ + private static void validate(Double[] vector1, Double[] vector2) { + if (vector1 == null || vector2 == null) { + throw new IllegalArgumentException("Vectors cannot be null"); + } + if (vector1.length != vector2.length) { + throw new IllegalArgumentException( + "Vectors must have the same dimensionality. Got " + vector1.length + " and " + vector2.length + ); + } + if (vector1.length == 0) { + throw new IllegalArgumentException("Vectors cannot be empty."); + } + } + + class ManhattanMetric implements Metric { + @Override + public double distance(final Double[] vector1, final Double[] vector2) { + Metric.validate(vector1, vector2); + + double sumOfAbsDiffs = 0.0; + for (int i = 0; i < vector1.length; i++) { + sumOfAbsDiffs += Math.abs(vector1[i] - vector2[i]); + } + return sumOfAbsDiffs; + } + } + + class EuclideanMetric implements Metric { + @Override + public double distance(final Double[] vector1, final Double[] vector2) { + Metric.validate(vector1, vector2); + + return Math.sqrt(EuclideanSquareMetric.distanceInternal(vector1, vector2)); + } + } + + class EuclideanSquareMetric implements Metric { + @Override + public double distance(final Double[] vector1, final Double[] vector2) { + Metric.validate(vector1, vector2); + return distanceInternal(vector1, vector2); + } + + private static double distanceInternal(final Double[] vector1, final Double[] vector2) { + double sumOfSquares = 0.0d; + for (int i = 0; i < vector1.length; i++) { + double diff = vector1[i] - vector2[i]; + sumOfSquares += diff * diff; + } + return sumOfSquares; + } + } + + class CosineMetric implements Metric { + @Override + public double distance(final Double[] vector1, final Double[] vector2) { + Metric.validate(vector1, vector2); + + double dotProduct = 0.0; + double normA = 0.0; + double normB = 0.0; + + for (int i = 0; i < vector1.length; i++) { + dotProduct += vector1[i] * vector2[i]; + normA += vector1[i] * vector1[i]; + normB += vector2[i] * vector2[i]; + } + + // Handle the case of zero-vectors to avoid division by zero + if (normA == 0.0 || normB == 0.0) { + return Double.POSITIVE_INFINITY; + } + + return 1.0d - dotProduct / (Math.sqrt(normA) * Math.sqrt(normB)); + } + } + + class DotProductMetric implements Metric { + @Override + public double distance(final Double[] vector1, final Double[] vector2) { + throw new UnsupportedOperationException("dot product metric is not a true metric and can only be used for ranking"); + } + + @Override + public double comparativeDistance(final Double[] vector1, final Double[] vector2) { + Metric.validate(vector1, vector2); + + double product = 0.0d; + for (int i = 0; i < vector1.length; i++) { + product += vector1[i] * vector2[i]; + } + return -product; + } + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Neighbor.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Neighbor.java new file mode 100644 index 0000000000..9f3107ddcb --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Neighbor.java @@ -0,0 +1,39 @@ +/* + * Neighbor.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.tuple.Tuple; + +import javax.annotation.Nonnull; + +public class Neighbor { + @Nonnull + private final Tuple primaryKey; + + public Neighbor(@Nonnull final Tuple primaryKey) { + this.primaryKey = primaryKey; + } + + @Nonnull + public Tuple getPrimaryKey() { + return primaryKey; + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NeighborWithVector.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NeighborWithVector.java new file mode 100644 index 0000000000..a531cb3e7c --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NeighborWithVector.java @@ -0,0 +1,46 @@ +/* + * Neighbor.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.tuple.Tuple; +import com.christianheina.langx.half4j.Half; + +import javax.annotation.Nonnull; + +public class NeighborWithVector extends Neighbor { + @Nonnull + private final Vector vector; + + public NeighborWithVector(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector) { + super(primaryKey); + this.vector = vector; + } + + @Nonnull + public Vector getVector() { + return vector; + } + + @Nonnull + public Vector getDoubleVector() { + return vector.toDoubleVector(); + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java new file mode 100644 index 0000000000..cdaa5ad0d1 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java @@ -0,0 +1,70 @@ +/* + * Node.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.tuple.Tuple; +import com.christianheina.langx.half4j.Half; +import com.google.errorprone.annotations.CanIgnoreReturnValue; + +import javax.annotation.Nonnull; + +/** + * TODO. + * @param neighbor type + */ +public interface Node { + @Nonnull + Tuple getPrimaryKey(); + + @Nonnull + Vector getVector(); + + @Nonnull + Iterable getNeighbors(); + + @Nonnull + N getNeighbor(int index); + + @CanIgnoreReturnValue + @Nonnull + Node insert(@Nonnull StorageAdapter storageAdapter, int level, int slotIndex, @Nonnull NodeSlot slot); + + @CanIgnoreReturnValue + @Nonnull + Node update(@Nonnull StorageAdapter storageAdapter, int level, int slotIndex, @Nonnull NodeSlot updatedSlot); + + @CanIgnoreReturnValue + @Nonnull + Node delete(@Nonnull StorageAdapter storageAdapter, int level, int slotIndex); + + /** + * Return the kind of the node, i.e. {@link NodeKind#DATA} or {@link NodeKind#INTERMEDIATE}. + * @return the kind of this node as a {@link NodeKind} + */ + @Nonnull + NodeKind getKind(); + + @Nonnull + DataNode asDataNode(); + + @Nonnull + IntermediateNode asIntermediateNode(); +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeHelpers.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeHelpers.java new file mode 100644 index 0000000000..965f5742cc --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeHelpers.java @@ -0,0 +1,80 @@ +/* + * NodeHelpers.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.google.common.collect.Lists; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicLong; + +/** + * Some helper methods for {@link Node}s. + */ +public class NodeHelpers { + private static final char[] hexArray = "0123456789ABCDEF".toCharArray(); + + private NodeHelpers() { + // nothing + } + + /** + * Helper method to format bytes as hex strings for logging and debugging. + * @param bytes an array of bytes + * @return a {@link String} containing the hexadecimal representation of the byte array passed in + */ + @Nonnull + static String bytesToHex(byte[] bytes) { + char[] hexChars = new char[bytes.length * 2]; + for (int j = 0; j < bytes.length; j++) { + int v = bytes[j] & 0xFF; + hexChars[j * 2] = hexArray[v >>> 4]; + hexChars[j * 2 + 1] = hexArray[v & 0x0F]; + } + return "0x" + new String(hexChars).replaceFirst("^0+(?!$)", ""); + } + + /** + * Helper method to format the node ids of an insert/update path as a string. + * @param node a node that is usually linked up to its parents to form an insert/update path + * @return a {@link String} containing the string presentation of the insert/update path starting at {@code node} + */ + @Nonnull + static String nodeIdPath(@Nullable Node node) { + final List nodeIds = Lists.newArrayList(); + do { + if (node != null) { + nodeIds.add(bytesToHex(node.getId())); + node = node.getParentNode(); + } else { + nodeIds.add(""); + } + } while (node != null); + Collections.reverse(nodeIds); + return String.join(", ", nodeIds); + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKind.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKind.java new file mode 100644 index 0000000000..98f0c1adfd --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKind.java @@ -0,0 +1,60 @@ +/* + * NodeKind.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.google.common.base.Verify; + +import javax.annotation.Nonnull; + +/** + * Enum to capture the kind of node. + */ +public enum NodeKind { + DATA((byte)0x00), + INTERMEDIATE((byte)0x01); + + private final byte serialized; + + NodeKind(final byte serialized) { + this.serialized = serialized; + } + + public byte getSerialized() { + return serialized; + } + + @Nonnull + static NodeKind fromSerializedNodeKind(byte serializedNodeKind) { + final NodeKind nodeKind; + switch (serializedNodeKind) { + case 0x00: + nodeKind = NodeKind.DATA; + break; + case 0x01: + nodeKind = NodeKind.INTERMEDIATE; + break; + default: + throw new IllegalArgumentException("unknown node kind"); + } + Verify.verify(nodeKind.getSerialized() == serializedNodeKind); + return nodeKind; + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java new file mode 100644 index 0000000000..98e74bf8f6 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java @@ -0,0 +1,54 @@ +/* + * OnReadListener.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import javax.annotation.Nonnull; +import java.util.concurrent.CompletableFuture; + +/** + * Function interface for a call back whenever we read the slots for a node. + */ +public interface OnReadListener { + OnReadListener NOOP = new OnReadListener() { + }; + + default void onSlotIndexEntryRead(@Nonnull final byte[] key) { + // nothing + } + + default CompletableFuture onAsyncRead(@Nonnull CompletableFuture future) { + return future; + } + + default void onNodeRead(@Nonnull Node node) { + // nothing + } + + default void onKeyValueRead(@Nonnull Node node, + @Nonnull byte[] key, + @Nonnull byte[] value) { + // nothing + } + + default void onChildNodeDiscard(@Nonnull final ChildSlot childSlot) { + // nothing + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java new file mode 100644 index 0000000000..a2b7ad3697 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java @@ -0,0 +1,63 @@ +/* + * OnWriteListener.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import javax.annotation.Nonnull; +import java.util.concurrent.CompletableFuture; + +/** + * Function interface for a call back whenever we read the slots for a node. + */ +public interface OnWriteListener { + OnWriteListener NOOP = new OnWriteListener() { + }; + + default void onSlotIndexEntryWritten(@Nonnull final byte[] key) { + // nothing + } + + default void onSlotIndexEntryCleared(@Nonnull final byte[] key) { + // nothing + } + + default CompletableFuture onAsyncReadForWrite(@Nonnull CompletableFuture future) { + return future; + } + + default void onNodeWritten(@Nonnull Node node) { + // nothing + } + + default void onKeyValueWritten(@Nonnull Node node, + @Nonnull byte[] key, + @Nonnull byte[] value) { + // nothing + } + + default void onNodeCleared(@Nonnull Node node) { + // nothing + } + + default void onKeyCleared(@Nonnull Node node, + @Nonnull byte[] key) { + // nothing + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java new file mode 100644 index 0000000000..6e6ecde9d7 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java @@ -0,0 +1,165 @@ +/* + * StorageAdapter.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.ReadTransaction; +import com.apple.foundationdb.Transaction; +import com.apple.foundationdb.subspace.Subspace; +import com.apple.foundationdb.tuple.Tuple; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.math.BigInteger; +import java.util.List; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; + +/** + * Storage adapter used for serialization and deserialization of nodes. + */ +interface StorageAdapter { + + /** + * Get the {@link HNSW.Config} associated with this storage adapter. + * @return the configuration used by this storage adapter + */ + @Nonnull + HNSW.Config getConfig(); + + /** + * Get the subspace used to store this r-tree. + * + * @return r-tree subspace + */ + @Nonnull + Subspace getSubspace(); + + /** + * Get the subspace used to store a node slot index if in warranted by the {@link HNSW.Config}. + * + * @return secondary subspace or {@code null} if we do not maintain a node slot index + */ + @Nullable + Subspace getSecondarySubspace(); + + @Nonnull + Subspace getEntryNodeSubspace(); + + @Nonnull + Subspace getDataSubspace(); + + /** + * Get the on-write listener. + * + * @return the on-write listener. + */ + @Nonnull + OnWriteListener getOnWriteListener(); + + /** + * Get the on-read listener. + * + * @return the on-read listener. + */ + @Nonnull + OnReadListener getOnReadListener(); + + CompletableFuture> fetchEntryNode(@Nonnull Transaction transaction); + + /** + * Insert a new entry into the node index if configuration indicates we should maintain such an index. + * + * @param transaction the transaction to use + * @param level the level counting starting at {@code 0} indicating the leaf level increasing upwards + * @param nodeSlot the {@link NodeSlot} to be inserted + */ + void insertIntoNodeIndexIfNecessary(@Nonnull Transaction transaction, int level, @Nonnull NodeSlot nodeSlot); + + /** + * Deletes an entry from the node index if configuration indicates we should maintain such an index. + * + * @param transaction the transaction to use + * @param level the level counting starting at {@code 0} indicating the leaf level increasing upwards + * @param nodeSlot the {@link NodeSlot} to be deleted + */ + void deleteFromNodeIndexIfNecessary(@Nonnull Transaction transaction, int level, @Nonnull NodeSlot nodeSlot); + + /** + * Persist a node slot. + * + * @param transaction the transaction to use + * @param node node whose slot to persist + * @param itemSlot the node slot to persist + */ + void writeLeafNodeSlot(@Nonnull Transaction transaction, @Nonnull DataNode node, @Nonnull ItemSlot itemSlot); + + /** + * Clear out a leaf node slot. + * + * @param transaction the transaction to use + * @param node node whose slot is cleared out + * @param itemSlot the node slot to clear out + */ + void clearLeafNodeSlot(@Nonnull Transaction transaction, @Nonnull DataNode node, @Nonnull ItemSlot itemSlot); + + /** + * Method to (re-)persist a list of nodes passed in. + * + * @param transaction the transaction to use + * @param nodes a list of nodes to be (re-persisted) + */ + void writeNodes(@Nonnull Transaction transaction, @Nonnull List nodes); + + /** + * Scan the node slot index for the given Hilbert Value/key pair and return the appropriate {@link Node}. + * Note that this method requires a node slot index to be maintained. + * + * @param transaction the transaction to use + * @param level the level we should search counting upwards starting from level {@code 0} for the leaf node + * level. + * @param hilbertValue the Hilbert Value of the {@code (Hilbert Value, key)} pair to search for + * @param key the key of the {@code (Hilbert Value, key)} pair to search for + * @param isInsertUpdate a use case indicator determining if this search is going to be used for an + * update operation or a delete operation + * + * @return a future that when completed holds the appropriate {@link Node} or {@code null} if such a + * {@link Node} could not be found. + */ + @Nonnull + CompletableFuture scanNodeIndexAndFetchNode(@Nonnull ReadTransaction transaction, int level, + @Nonnull BigInteger hilbertValue, @Nonnull Tuple key, + boolean isInsertUpdate); + + /** + * Method to fetch the data needed to construct a {@link Node}. Note that a node on disk is represented by its + * slots. Each slot is represented by a key/value pair in FDB. Each key (common for both leaf and intermediate + * nodes) starts with an 8-byte node id (which is usually a serialized {@link UUID}) followed by one byte which + * indicates the {@link NodeKind} of node the slot belongs to. + * + * @param transaction the transaction to use + * @param nodeId the node id we should use + * + * @return A completable future containing the {@link Node} that was fetched from the database once completed. + * The node may be an object of {@link DataNode} or of {@link IntermediateNode}. + */ + @Nonnull + CompletableFuture fetchNode(@Nonnull ReadTransaction transaction, @Nonnull byte[] nodeId); +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java new file mode 100644 index 0000000000..d8482e2ac4 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java @@ -0,0 +1,129 @@ +/* + * NodeHelpers.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.christianheina.langx.half4j.Half; +import com.google.common.base.Suppliers; + +import javax.annotation.Nonnull; +import java.util.function.Supplier; + +/** + * TODO. + * @param representation type + */ +public abstract class Vector { + @Nonnull + protected R[] data; + + public Vector(@Nonnull final R[] data) { + this.data = data; + } + + public int size() { + return data.length; + } + + @Nonnull + public R[] getData() { + return data; + } + + @Nonnull + public abstract Vector toHalfVector(); + + @Nonnull + public abstract DoubleVector toDoubleVector(); + + public static class HalfVector extends Vector { + @Nonnull + private final Supplier toDoubleVectorSupplier; + + public HalfVector(@Nonnull final Half[] data) { + super(data); + this.toDoubleVectorSupplier = Suppliers.memoize(this::computeDoubleVector); + } + + @Nonnull + @Override + public Vector toHalfVector() { + return this; + } + + @Nonnull + @Override + public DoubleVector toDoubleVector() { + return toDoubleVectorSupplier.get(); + } + + @Nonnull + public DoubleVector computeDoubleVector() { + Double[] result = new Double[data.length]; + for (int i = 0; i < data.length; i ++) { + result[i] = data[i].doubleValue(); + } + return new DoubleVector(result); + } + } + + public static class DoubleVector extends Vector { + @Nonnull + private final Supplier toHalfVectorSupplier; + + public DoubleVector(@Nonnull final Double[] data) { + super(data); + this.toHalfVectorSupplier = Suppliers.memoize(this::computeHalfVector); + } + + @Nonnull + @Override + public HalfVector toHalfVector() { + return toHalfVectorSupplier.get(); + } + + @Nonnull + public HalfVector computeHalfVector() { + Half[] result = new Half[data.length]; + for (int i = 0; i < data.length; i ++) { + result[i] = Half.valueOf(data[i]); + } + return new HalfVector(result); + } + + @Nonnull + @Override + public DoubleVector toDoubleVector() { + return this; + } + } + + static double distance(@Nonnull Metric metric, + @Nonnull final Vector vector1, + @Nonnull final Vector vector2) { + return metric.distance(vector1.toDoubleVector().getData(), vector2.toDoubleVector().getData()); + } + + static double comparativeDistance(@Nonnull Metric metric, + @Nonnull final Vector vector1, + @Nonnull final Vector vector2) { + return metric.comparativeDistance(vector1.toDoubleVector().getData(), vector2.toDoubleVector().getData()); + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/package-info.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/package-info.java new file mode 100644 index 0000000000..5565b7f9f6 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/package-info.java @@ -0,0 +1,24 @@ +/* + * package-info.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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. + */ + +/** + * Classes and interfaces related to the Hilbert R-tree implementation. + */ +package com.apple.foundationdb.async.hnsw; diff --git a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java new file mode 100644 index 0000000000..5b069a802e --- /dev/null +++ b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java @@ -0,0 +1,45 @@ +/* + * MultidimensionalIndexTestBase.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.record.provider.foundationdb.indexes; + +import com.apple.test.Tags; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Tests for multidimensional type indexes. + */ +@Tag(Tags.RequiresFDB) +public class VectorIndexSimpleTest extends VectorIndexTestBase { + private static final Logger logger = LoggerFactory.getLogger(VectorIndexSimpleTest.class); + + @Test + void basicReadTest() throws Exception { + super.basicReadTest(false); + } + + @Test + void basicConcurrentReadTest() throws Exception { + super.basicConcurrentReadTest(false); + } +} diff --git a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java new file mode 100644 index 0000000000..9a968e161d --- /dev/null +++ b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java @@ -0,0 +1,223 @@ +/* + * MultidimensionalIndexTestBase.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.record.provider.foundationdb.indexes; + +import com.apple.foundationdb.async.AsyncUtil; +import com.apple.foundationdb.record.RecordMetaData; +import com.apple.foundationdb.record.RecordMetaDataBuilder; +import com.apple.foundationdb.record.provider.foundationdb.FDBRecordContext; +import com.apple.foundationdb.record.provider.foundationdb.FDBStoredRecord; +import com.apple.foundationdb.record.provider.foundationdb.query.FDBRecordStoreQueryTestBase; +import com.apple.foundationdb.record.vector.TestRecordsVectorsProto; +import com.apple.foundationdb.tuple.Tuple; +import com.apple.test.Tags; +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.protobuf.Message; +import org.assertj.core.util.Streams; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Tag; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nonnull; +import java.text.SimpleDateFormat; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Random; +import java.util.concurrent.CompletableFuture; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +import static com.apple.foundationdb.record.metadata.Key.Expressions.field; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +/** + * Tests for multidimensional type indexes. + */ +@Tag(Tags.RequiresFDB) +public abstract class VectorIndexTestBase extends FDBRecordStoreQueryTestBase { + private static final Logger logger = LoggerFactory.getLogger(VectorIndexTestBase.class); + + private static final SimpleDateFormat timeFormat = new SimpleDateFormat("MM/dd/yyyy HH:mm:ss"); + + protected void openRecordStore(FDBRecordContext context) throws Exception { + openRecordStore(context, NO_HOOK); + } + + protected void openRecordStore(final FDBRecordContext context, final RecordMetaDataHook hook) throws Exception { + RecordMetaDataBuilder metaDataBuilder = RecordMetaData.newBuilder().setRecords(TestRecordsVectorsProto.getDescriptor()); + metaDataBuilder.getRecordType("NaiveVectorRecord").setPrimaryKey(field("rec_no")); + hook.apply(metaDataBuilder); + createOrOpenRecordStore(context, metaDataBuilder.getRecordMetaData()); + } + + static Function getRecordGenerator(@Nonnull final Random random) { + return recNo -> { + final Iterable vector = + () -> IntStream.range(0, 4000).mapToDouble(index -> random.nextDouble()).iterator(); + //logRecord(recNo, vector); + return TestRecordsVectorsProto.NaiveVectorRecord.newBuilder() + .setRecNo(recNo) + .addAllVectorData(vector) + .build(); + }; + } + + public void loadRecords(final boolean useAsync, @Nonnull final RecordMetaDataHook hook, final long seed, + final int numSamples) { + final Random random = new Random(seed); + final var recordGenerator = getRecordGenerator(random); + if (useAsync) { + Assertions.assertDoesNotThrow(() -> batchAsync(hook, numSamples, 100, recNo -> recordStore.saveRecordAsync(recordGenerator.apply(recNo)))); + } else { + Assertions.assertDoesNotThrow(() -> batch(hook, numSamples, 100, recNo -> recordStore.saveRecord(recordGenerator.apply(recNo)))); + } + } + + public void deleteRecords(final boolean useAsync, @Nonnull final RecordMetaDataHook hook, final long seed, final int numRecords, + final int numDeletes) throws Exception { + Preconditions.checkArgument(numDeletes <= numRecords); + final Random random = new Random(seed); + final List recNos = IntStream.range(0, numRecords) + .boxed() + .collect(Collectors.toList()); + Collections.shuffle(recNos, random); + final List recNosToBeDeleted = recNos.subList(0, numDeletes); + if (useAsync) { + batchAsync(hook, recNosToBeDeleted.size(), 500, recNo -> recordStore.deleteRecordAsync(Tuple.from(recNo))); + } else { + batch(hook, recNosToBeDeleted.size(), 500, recNo -> recordStore.deleteRecord(Tuple.from(recNo))); + } + } + + private long batch(final RecordMetaDataHook hook, final int numRecords, final int batchSize, Consumer recordConsumer) throws Exception { + long numRecordsCommitted = 0; + while (numRecordsCommitted < numRecords) { + try (FDBRecordContext context = openContext()) { + openRecordStore(context, hook); + int recNoInBatch; + + for (recNoInBatch = 0; numRecordsCommitted + recNoInBatch < numRecords && recNoInBatch < batchSize; recNoInBatch++) { + recordConsumer.accept(numRecordsCommitted + recNoInBatch); + } + commit(context); + numRecordsCommitted += recNoInBatch; + logger.info("committed batch, numRecordsCommitted = {}", numRecordsCommitted); + } + } + return numRecordsCommitted; + } + + private long batchAsync(final RecordMetaDataHook hook, final int numRecords, final int batchSize, Function> recordConsumer) throws Exception { + long numRecordsCommitted = 0; + while (numRecordsCommitted < numRecords) { + try (FDBRecordContext context = openContext()) { + openRecordStore(context, hook); + int recNoInBatch; + final var futures = new ArrayList>(); + + for (recNoInBatch = 0; numRecordsCommitted + recNoInBatch < numRecords && recNoInBatch < batchSize; recNoInBatch++) { + futures.add(recordConsumer.apply(numRecordsCommitted + recNoInBatch)); + } + + // wait and then commit + AsyncUtil.whenAll(futures).get(); + commit(context); + numRecordsCommitted += recNoInBatch; + logger.info("committed batch, numRecordsCommitted = {}", numRecordsCommitted); + } + } + return numRecordsCommitted; + } + + private static void logRecord(final long recNo, @Nonnull final Iterable vector) { + if (logger.isInfoEnabled()) { + logger.info("recNo: {}; vectorData: [{})", + recNo, Streams.stream(vector).map(String::valueOf).collect(Collectors.joining(","))); + } + } + + void basicReadTest(final boolean useAsync) throws Exception { + loadRecords(useAsync, NO_HOOK, 0, 5000); + try (FDBRecordContext context = openContext()) { + openRecordStore(context); + for (long l = 0; l < 5000; l ++) { + FDBStoredRecord rec = recordStore.loadRecord(Tuple.from(l)); + //Thread.sleep(10); + assertNotNull(rec); + TestRecordsVectorsProto.NaiveVectorRecord.Builder recordBuilder = + TestRecordsVectorsProto.NaiveVectorRecord.newBuilder(); + recordBuilder.mergeFrom(rec.getRecord()); + final var record = recordBuilder.build(); + //logRecord(record.getRecNo(), record.getVectorDataList()); + } + commit(context); + } + } + + void basicConcurrentReadTest(final boolean useAsync) throws Exception { + loadRecords(useAsync, NO_HOOK, 0, 5000); + for (int i = 0; i < 1; i ++) { + try (FDBRecordContext context = openContext()) { + openRecordStore(context); + for (long l = 0; l < 5000; l += 20) { + final var batch = fetchBatchConcurrently(l, 20); + //batch.forEach(record -> logRecord(record.getRecNo(), record.getVectorDataList())); + } + commit(context); + } + } + } + + private List fetchBatchConcurrently(long startRecNo, int batchSize) throws Exception { + final ImmutableList.Builder>> futureBatchBuilder = ImmutableList.builder(); + for (int i = 0; i < batchSize; i ++) { + final long task = startRecNo + i; + //System.out.println("task " + task + " scheduled"); + final CompletableFuture> recordFuture = recordStore.loadRecordAsync(Tuple.from(startRecNo + i)).thenApply(r -> { +// try { +// Thread.sleep(10); +// } catch (InterruptedException e) { +// throw new RuntimeException(e); +// } + //System.out.println("task " + task + " done, thread: " + Thread.currentThread()); + return r; + }); + futureBatchBuilder.add(recordFuture); + } + final var batchFuture = AsyncUtil.getAll(futureBatchBuilder.build()); + final var batch = batchFuture.get(); + + return batch.stream() + .map(rec -> { + assertNotNull(rec); + TestRecordsVectorsProto.NaiveVectorRecord.Builder recordBuilder = + TestRecordsVectorsProto.NaiveVectorRecord.newBuilder(); + recordBuilder.mergeFrom(rec.getRecord()); + return recordBuilder.build(); + }) + .collect(ImmutableList.toImmutableList()); + } +} diff --git a/fdb-record-layer-core/src/test/proto/evolution/test_field_type_change.proto b/fdb-record-layer-core/src/test/proto/evolution/test_field_type_change.proto index eb5805405e..a7194761b2 100644 --- a/fdb-record-layer-core/src/test/proto/evolution/test_field_type_change.proto +++ b/fdb-record-layer-core/src/test/proto/evolution/test_field_type_change.proto @@ -24,7 +24,7 @@ package com.apple.foundationdb.record.evolution.fieldtypechange; option java_package = "com.apple.foundationdb.record.evolution"; option java_outer_classname = "TestFieldTypeChangeProto"; -import "record_metadata_options.proto"; +//import "record_metadata_options.proto"; import "test_records_1.proto"; // This needs to match the "MySimpleRecord" definition test_records_1.proto diff --git a/fdb-record-layer-core/src/test/proto/evolution/test_header_as_group.proto b/fdb-record-layer-core/src/test/proto/evolution/test_header_as_group.proto index 22478fc812..54ef33d2e7 100644 --- a/fdb-record-layer-core/src/test/proto/evolution/test_header_as_group.proto +++ b/fdb-record-layer-core/src/test/proto/evolution/test_header_as_group.proto @@ -24,7 +24,7 @@ package com.apple.foundationdb.record.evolution.headergroup; option java_package = "com.apple.foundationdb.record.evolution"; option java_outer_classname = "TestHeaderAsGroupProto"; -import "record_metadata_options.proto"; +//import "record_metadata_options.proto"; // This is taken from test_records_with_header.proto but the header has become a group message MyRecord { diff --git a/fdb-record-layer-core/src/test/proto/evolution/test_swap_union_fields.proto b/fdb-record-layer-core/src/test/proto/evolution/test_swap_union_fields.proto index 3b2a6a29e6..fab35adf8b 100644 --- a/fdb-record-layer-core/src/test/proto/evolution/test_swap_union_fields.proto +++ b/fdb-record-layer-core/src/test/proto/evolution/test_swap_union_fields.proto @@ -24,7 +24,7 @@ package com.apple.foundationdb.record.evolution.swap; option java_package = "com.apple.foundationdb.record.evolution"; option java_outer_classname = "TestSwapUnionFieldsProto"; -import "record_metadata_options.proto"; +//import "record_metadata_options.proto"; import "test_records_1.proto"; message RecordTypeUnion { diff --git a/fdb-record-layer-core/src/test/proto/expression_tests.proto b/fdb-record-layer-core/src/test/proto/expression_tests.proto index 570a0d920b..4b9f728d69 100644 --- a/fdb-record-layer-core/src/test/proto/expression_tests.proto +++ b/fdb-record-layer-core/src/test/proto/expression_tests.proto @@ -23,7 +23,7 @@ package com.apple.foundationdb.record.metadata; option java_outer_classname = "ExpressionTestsProto"; -import "record_metadata_options.proto"; +//import "record_metadata_options.proto"; import "tuple_fields.proto"; message TestScalarFieldAccess { diff --git a/fdb-record-layer-core/src/test/proto/test_no_record_types.proto b/fdb-record-layer-core/src/test/proto/test_no_record_types.proto index 92fcd8a13b..e8f0a4c64a 100644 --- a/fdb-record-layer-core/src/test/proto/test_no_record_types.proto +++ b/fdb-record-layer-core/src/test/proto/test_no_record_types.proto @@ -25,7 +25,7 @@ package com.apple.foundationdb.record.testnorecords; option java_package = "com.apple.foundationdb.record"; option java_outer_classname = "TestNoRecordTypesProto"; -import "record_metadata_options.proto"; +//import "record_metadata_options.proto"; message RecordTypeUnion { } diff --git a/fdb-record-layer-core/src/test/proto/test_records_8.proto b/fdb-record-layer-core/src/test/proto/test_records_8.proto index 81e1d2cf3a..d67700e899 100644 --- a/fdb-record-layer-core/src/test/proto/test_records_8.proto +++ b/fdb-record-layer-core/src/test/proto/test_records_8.proto @@ -24,7 +24,7 @@ package com.apple.foundationdb.record.test8; option java_package = "com.apple.foundationdb.record"; option java_outer_classname = "TestRecords8Proto"; -import "record_metadata_options.proto"; +//import "record_metadata_options.proto"; message StringRecordId { required string rec_id = 1; diff --git a/fdb-record-layer-core/src/test/proto/test_records_chained_2.proto b/fdb-record-layer-core/src/test/proto/test_records_chained_2.proto index c0165dd764..ec6f877a25 100644 --- a/fdb-record-layer-core/src/test/proto/test_records_chained_2.proto +++ b/fdb-record-layer-core/src/test/proto/test_records_chained_2.proto @@ -25,7 +25,7 @@ option java_package = "com.apple.foundationdb.record"; option java_outer_classname = "TestRecordsChained2Proto"; import "record_metadata_options.proto"; -import "test_records_1.proto"; +//import "test_records_1.proto"; import "test_records_2.proto"; message MyChainedRecord2 { diff --git a/fdb-record-layer-core/src/test/proto/test_records_duplicate_union_fields.proto b/fdb-record-layer-core/src/test/proto/test_records_duplicate_union_fields.proto index 2c5f997caa..b9dec24353 100644 --- a/fdb-record-layer-core/src/test/proto/test_records_duplicate_union_fields.proto +++ b/fdb-record-layer-core/src/test/proto/test_records_duplicate_union_fields.proto @@ -24,7 +24,7 @@ package com.apple.foundationdb.record.test.duplicateunionfields; option java_package = "com.apple.foundationdb.record"; option java_outer_classname = "TestRecordsDuplicateUnionFields"; -import "record_metadata_options.proto"; +//import "record_metadata_options.proto"; import "test_records_1.proto"; message RecordTypeUnion { diff --git a/fdb-record-layer-core/src/test/proto/test_records_duplicate_union_fields_reordered.proto b/fdb-record-layer-core/src/test/proto/test_records_duplicate_union_fields_reordered.proto index b10dbf4d76..f9857dc3e0 100644 --- a/fdb-record-layer-core/src/test/proto/test_records_duplicate_union_fields_reordered.proto +++ b/fdb-record-layer-core/src/test/proto/test_records_duplicate_union_fields_reordered.proto @@ -24,7 +24,7 @@ package com.apple.foundationdb.record.test.duplicateunionfields.reordered; option java_package = "com.apple.foundationdb.record"; option java_outer_classname = "TestRecordsDuplicateUnionFieldsReordered"; -import "record_metadata_options.proto"; +//import "record_metadata_options.proto"; import "test_records_1.proto"; message RecordTypeUnion { diff --git a/fdb-record-layer-core/src/test/proto/test_records_oneof.proto b/fdb-record-layer-core/src/test/proto/test_records_oneof.proto index 53b6a43c4a..ed25d8ccb0 100644 --- a/fdb-record-layer-core/src/test/proto/test_records_oneof.proto +++ b/fdb-record-layer-core/src/test/proto/test_records_oneof.proto @@ -24,7 +24,7 @@ package com.apple.foundationdb.record.testOneOf; option java_package = "com.apple.foundationdb.record"; option java_outer_classname = "TestRecordsOneOfProto"; -import "record_metadata_options.proto"; +//import "record_metadata_options.proto"; message MySimpleRecord { optional int64 rec_no = 1; diff --git a/fdb-record-layer-core/src/test/proto/test_records_transform.proto b/fdb-record-layer-core/src/test/proto/test_records_transform.proto index c20c657542..0c7c7e5732 100644 --- a/fdb-record-layer-core/src/test/proto/test_records_transform.proto +++ b/fdb-record-layer-core/src/test/proto/test_records_transform.proto @@ -24,7 +24,7 @@ package com.apple.foundationdb.record.transform; option java_package = "com.apple.foundationdb.record"; option java_outer_classname = "TestRecordsTransformProto"; -import "record_metadata_options.proto"; +//import "record_metadata_options.proto"; message DefaultTransformMessage { message MessageAa { diff --git a/fdb-record-layer-core/src/test/proto/test_records_vector.proto b/fdb-record-layer-core/src/test/proto/test_records_vector.proto new file mode 100644 index 0000000000..bafe614416 --- /dev/null +++ b/fdb-record-layer-core/src/test/proto/test_records_vector.proto @@ -0,0 +1,38 @@ +/* + * test_records_multidimensional.proto + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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. + */ +syntax = "proto2"; + +package com.apple.foundationdb.record.test.vector; + +option java_package = "com.apple.foundationdb.record.vector"; +option java_outer_classname = "TestRecordsVectorsProto"; + +import "record_metadata_options.proto"; + +option (schema).store_record_versions = true; + +message NaiveVectorRecord { + optional int64 rec_no = 1 [(field).primary_key = true]; + repeated double vector_data = 2; +} + +message RecordTypeUnion { + optional NaiveVectorRecord _NaiveVectorRecord = 1; +} diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index a9abaf9a20..0eabf0d87d 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -37,6 +37,7 @@ generatedAnnotation = "1.3.2" grpc = "1.64.1" grpc-commonProtos = "2.37.0" guava = "33.3.1-jre" +half4j = "0.0.2" h2 = "1.3.148" icu = "69.1" lucene = "8.11.1" @@ -95,6 +96,7 @@ grpc-services = { module = "io.grpc:grpc-services", version.ref = "grpc" } grpc-stub = { module = "io.grpc:grpc-stub", version.ref = "grpc" } grpc-util = { module = "io.grpc:grpc-util", version.ref = "grpc" } guava = { module = "com.google.guava:guava", version.ref = "guava" } +half4j = { module = "com.christianheina.langx:half4j", version.ref = "half4j"} icu = { module = "com.ibm.icu:icu4j", version.ref = "icu" } javaPoet = { module = "com.squareup:javapoet", version.ref = "javaPoet" } jsr305 = { module = "com.google.code.findbugs:jsr305", version.ref = "jsr305" } From 0b3fee0217374eba5d2565261a018c71fec51112 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Thu, 24 Jul 2025 22:59:27 +0200 Subject: [PATCH 02/34] save point -- in the middle of just mess --- .../async/hnsw/ByNodeStorageAdapter.java | 46 ++++--- .../foundationdb/async/hnsw/DataNode.java | 6 + .../apple/foundationdb/async/hnsw/HNSW.java | 127 +++++++++--------- .../async/hnsw/IntermediateNode.java | 6 + .../apple/foundationdb/async/hnsw/Node.java | 3 + .../async/hnsw/NodeWithLayer.java | 43 ++++++ .../async/hnsw/StorageAdapter.java | 2 +- 7 files changed, 154 insertions(+), 79 deletions(-) create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeWithLayer.java diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java index 84b8899d67..63a4df1905 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java @@ -75,19 +75,22 @@ public ByNodeStorageAdapter(@Nonnull final HNSW.Config config, @Nonnull final Su } @Override - public CompletableFuture> fetchEntryNode(@Nonnull final Transaction transaction) { + public CompletableFuture> fetchEntryNode(@Nonnull final ReadTransaction readTransaction) { final byte[] key = getEntryNodeSubspace().pack(); - return transaction.get(key) + return readTransaction.get(key) .thenApply(valueBytes -> { if (valueBytes == null) { throw new IllegalStateException("cannot fetch entry point"); } - final Node node = fromTuple(Tuple.fromBytes(valueBytes)); + + final Tuple entryTuple = Tuple.fromBytes(valueBytes); + final int lMax = (int)entryTuple.getLong(0); + final Node node = nodeFromTuple(entryTuple.getNestedTuple(1)); final OnReadListener onReadListener = getOnReadListener(); onReadListener.onNodeRead(node); onReadListener.onKeyValueRead(node, key, valueBytes); - return node; + return node.withLayer(lMax); }); } @@ -141,7 +144,7 @@ public CompletableFuture fetchNodeInternal(@Nonnull final ReadTransaction if (valueBytes == null) { return null; } - final Node node = fromTuple(nodeId, Tuple.fromBytes(valueBytes)); + final Node node = nodeFromTuple(nodeId, Tuple.fromBytes(valueBytes)); final OnReadListener onReadListener = getOnReadListener(); onReadListener.onNodeRead(node); onReadListener.onKeyValueRead(node, key, valueBytes); @@ -150,10 +153,17 @@ public CompletableFuture fetchNodeInternal(@Nonnull final ReadTransaction } @Nonnull - private Node fromTuple(@Nonnull final Tuple tuple) { + private Node nodeFromTuple(@Nonnull final Tuple tuple) { final NodeKind nodeKind = NodeKind.fromSerializedNodeKind((byte)tuple.getLong(0)); - final Tuple neighborsTuple = tuple.getNestedTuple(1); + final Tuple primaryKey = tuple.getNestedTuple(1); + final Tuple vectorTuple = tuple.getNestedTuple(2); + final Tuple neighborsTuple = tuple.getNestedTuple(3); + final Half[] vectorHalfs = new Half[vectorTuple.size()]; + for (int i = 0; i < vectorTuple.size(); i ++) { + vectorHalfs[i] = Half.shortBitsToHalf(shortFromBytes(vectorTuple.getBytes(i))); + } + final Vector.HalfVector vector = new Vector.HalfVector(vectorHalfs); List neighborsWithVectors = null; Half[] neighborVectorHalfs = null; List neighbors = null; @@ -162,6 +172,13 @@ private Node fromTuple(@Nonnull final Tuple tuple) { final Tuple neighborTuple = (Tuple)neighborObject; switch (nodeKind) { case DATA: + if (neighbors == null) { + neighbors = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); + } + neighbors.add(new Neighbor(neighborTuple)); + break; + + case INTERMEDIATE: final Tuple neighborPrimaryKey = neighborTuple.getNestedTuple(0); final Tuple neighborVectorTuple = neighborTuple.getNestedTuple(1); if (neighborsWithVectors == null) { @@ -175,24 +192,17 @@ private Node fromTuple(@Nonnull final Tuple tuple) { neighborsWithVectors.add(new NeighborWithVector(neighborPrimaryKey, new Vector.HalfVector(neighborVectorHalfs))); break; - case INTERMEDIATE: - if (neighbors == null) { - neighbors = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); - } - neighbors.add(new Neighbor(neighborTuple)); - break; - default: throw new IllegalStateException("unknown node kind"); } } - Verify.verify((nodeKind == NodeKind.DATA && neighborsWithVectors != null) || - (nodeKind == NodeKind.INTERMEDIATE && neighbors != null)); + Verify.verify((nodeKind == NodeKind.DATA && neighbors != null) || + (nodeKind == NodeKind.INTERMEDIATE && neighborsWithVectors != null)); return nodeKind == NodeKind.DATA - ? new DataNode(nodeId, itemSlots) - : new IntermediateNode(nodeId, childSlots); + ? new DataNode(primaryKey, vector, neighbors) + : new IntermediateNode(primaryKey, vector, neighborsWithVectors); } @Nonnull diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java index d9bded06d4..34fe48d579 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java @@ -49,6 +49,12 @@ public IntermediateNode asIntermediateNode() { throw new IllegalStateException("this is not a data node"); } + @Nonnull + @Override + public NodeWithLayer withLayer(final int layer) { + return new NodeWithLayer<>(layer, this); + } + @Nonnull @Override public NodeKind getKind() { diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index eb7c8aa7ef..9a2bcf249b 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -526,70 +526,77 @@ private CompletableFuture fetchLeftmostPathToLeaf(@Nonnull final @Nullable final Tuple lastKey, @Nonnull final Predicate mbrPredicate, @Nonnull final BiPredicate suffixPredicate) { - final AtomicReference currentId = new AtomicReference<>(nodeId); - final List> toBeProcessed = Lists.newArrayList(); - final AtomicReference leafNode = new AtomicReference<>(null); - return AsyncUtil.whileTrue(() -> onReadListener.onAsyncRead(storageAdapter.fetchNode(readTransaction, currentId.get())) - .thenApply(node -> { - if (node == null) { - if (Arrays.equals(currentId.get(), rootId)) { - Verify.verify(leafNode.get() == null); - return false; - } - throw new IllegalStateException("unable to fetch node for scan"); - } - if (node.getKind() == NodeKind.INTERMEDIATE) { - final Iterable childSlots = ((IntermediateNode)node).getSlots(); - Deque toBeProcessedThisLevel = new ArrayDeque<>(); - for (final Iterator iterator = childSlots.iterator(); iterator.hasNext(); ) { - final ChildSlot childSlot = iterator.next(); - if (lastHilbertValue != null && - lastKey != null) { - final int hilbertValueAndKeyCompare = - childSlot.compareLargestHilbertValueAndKey(lastHilbertValue, lastKey); - if (hilbertValueAndKeyCompare < 0) { - // - // The (lastHilbertValue, lastKey) pair is larger than the - // (largestHilbertValue, largestKey) pair of the current child. Advance to the next - // child. - // - continue; + final AtomicReference> currentNodeWithLayer = + new AtomicReference<>(); + + storageAdapter.fetchEntryNode(readTransaction) + .thenApply(nodeWithLayer -> { + currentNodeWithLayer.set(nodeWithLayer); + + final List> toBeProcessed = Lists.newArrayList(); + final AtomicReference leafNode = new AtomicReference<>(null); + return AsyncUtil.whileTrue(() -> onReadListener.onAsyncRead(storageAdapter.fetchNode(readTransaction, currentId.get())) + .thenApply(node -> { + if (node == null) { + if (Arrays.equals(currentId.get(), rootId)) { + Verify.verify(leafNode.get() == null); + return false; + } + throw new IllegalStateException("unable to fetch node for scan"); } - } + if (node.getKind() == NodeKind.INTERMEDIATE) { + final Iterable childSlots = ((IntermediateNode)node).getSlots(); + Deque toBeProcessedThisLevel = new ArrayDeque<>(); + for (final Iterator iterator = childSlots.iterator(); iterator.hasNext(); ) { + final ChildSlot childSlot = iterator.next(); + if (lastHilbertValue != null && + lastKey != null) { + final int hilbertValueAndKeyCompare = + childSlot.compareLargestHilbertValueAndKey(lastHilbertValue, lastKey); + if (hilbertValueAndKeyCompare < 0) { + // + // The (lastHilbertValue, lastKey) pair is larger than the + // (largestHilbertValue, largestKey) pair of the current child. Advance to the next + // child. + // + continue; + } + } + + if (!mbrPredicate.test(childSlot.getMbr())) { + onReadListener.onChildNodeDiscard(childSlot); + continue; + } + + if (childSlot.suffixPredicateCanBeApplied()) { + if (!suffixPredicate.test(childSlot.getSmallestKeySuffix(), + childSlot.getLargestKeySuffix())) { + onReadListener.onChildNodeDiscard(childSlot); + continue; + } + } + + toBeProcessedThisLevel.addLast(childSlot); + iterator.forEachRemaining(toBeProcessedThisLevel::addLast); + } + toBeProcessed.add(toBeProcessedThisLevel); - if (!mbrPredicate.test(childSlot.getMbr())) { - onReadListener.onChildNodeDiscard(childSlot); - continue; - } + final ChildSlot nextChildSlot = resolveNextIdForFetch(toBeProcessed, mbrPredicate, + suffixPredicate, onReadListener); + if (nextChildSlot == null) { + return false; + } - if (childSlot.suffixPredicateCanBeApplied()) { - if (!suffixPredicate.test(childSlot.getSmallestKeySuffix(), - childSlot.getLargestKeySuffix())) { - onReadListener.onChildNodeDiscard(childSlot); - continue; + currentId.set(Objects.requireNonNull(nextChildSlot.getChildId())); + return true; + } else { + leafNode.set((DataNode)node); + return false; } - } - - toBeProcessedThisLevel.addLast(childSlot); - iterator.forEachRemaining(toBeProcessedThisLevel::addLast); - } - toBeProcessed.add(toBeProcessedThisLevel); - - final ChildSlot nextChildSlot = resolveNextIdForFetch(toBeProcessed, mbrPredicate, - suffixPredicate, onReadListener); - if (nextChildSlot == null) { - return false; - } - - currentId.set(Objects.requireNonNull(nextChildSlot.getChildId())); - return true; - } else { - leafNode.set((DataNode)node); - return false; - } - }), executor).thenApply(vignore -> leafNode.get() == null - ? TraversalState.end() - : TraversalState.of(toBeProcessed, leafNode.get())); + }), executor).thenApply(vignore -> leafNode.get() == null + ? TraversalState.end() + : TraversalState.of(toBeProcessed, leafNode.get())); + }); } /** diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java index 040653e4e0..ff1c494420 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java @@ -49,6 +49,12 @@ public IntermediateNode asIntermediateNode() { return this; } + @Nonnull + @Override + public NodeWithLayer withLayer(final int layer) { + return new NodeWithLayer<>(layer, this); + } + @Nonnull @Override public NodeKind getKind() { diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java index cdaa5ad0d1..dab2bac5c1 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java @@ -67,4 +67,7 @@ public interface Node { @Nonnull IntermediateNode asIntermediateNode(); + + @Nonnull + NodeWithLayer withLayer(int layer); } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeWithLayer.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeWithLayer.java new file mode 100644 index 0000000000..f3a3dd49a8 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeWithLayer.java @@ -0,0 +1,43 @@ +/* + * NodeWithLayer.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import javax.annotation.Nonnull; + +class NodeWithLayer { + private final int layer; + @Nonnull + private final Node node; + + public NodeWithLayer(final int layer, @Nonnull final Node node) { + this.layer = layer; + this.node = node; + } + + public int getLayer() { + return layer; + } + + @Nonnull + public Node getNode() { + return node; + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java index 6e6ecde9d7..e6e8e1cf8f 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java @@ -82,7 +82,7 @@ interface StorageAdapter { @Nonnull OnReadListener getOnReadListener(); - CompletableFuture> fetchEntryNode(@Nonnull Transaction transaction); + CompletableFuture> fetchEntryNode(@Nonnull ReadTransaction readTransaction); /** * Insert a new entry into the node index if configuration indicates we should maintain such an index. From 8bdf410950527c7da6ec0afbf6e59b8a202a5f52 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Sat, 26 Jul 2025 12:34:20 +0200 Subject: [PATCH 03/34] save point -- in the middle of just mess --- .../async/hnsw/AbstractStorageAdapter.java | 17 +- .../async/hnsw/ByNodeStorageAdapter.java | 40 ++- .../foundationdb/async/hnsw/DataNode.java | 22 +- .../apple/foundationdb/async/hnsw/HNSW.java | 249 ++++++++---------- .../async/hnsw/IntermediateNode.java | 26 +- .../apple/foundationdb/async/hnsw/Metric.java | 33 +++ .../apple/foundationdb/async/hnsw/Node.java | 11 +- .../async/hnsw/NodeKeyWithLayer.java | 45 ++++ .../hnsw/NodeKeyWithLayerAndDistance.java | 38 +++ .../async/hnsw/OnReadListener.java | 5 +- .../async/hnsw/StorageAdapter.java | 24 +- 11 files changed, 317 insertions(+), 193 deletions(-) create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayer.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayerAndDistance.java diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java index 57c8f3d730..d1eb8c6d37 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java @@ -173,12 +173,16 @@ public void deleteFromNodeIndexIfNecessary(@Nonnull final Transaction transactio @Nonnull @Override - public CompletableFuture fetchNode(@Nonnull final ReadTransaction transaction, @Nonnull final byte[] nodeId) { - return getOnWriteListener().onAsyncReadForWrite(fetchNodeInternal(transaction, nodeId).thenApply(this::checkNode)); + public CompletableFuture> fetchNode(@Nonnull final Node.NodeCreator creator, + @Nonnull final ReadTransaction readTransaction, + int layer, @Nonnull Tuple primaryKey) { + return fetchNodeInternal(creator, readTransaction, layer, primaryKey).thenApply(this::checkNode); } @Nonnull - protected abstract CompletableFuture fetchNodeInternal(@Nonnull ReadTransaction transaction, @Nonnull byte[] nodeId); + protected abstract CompletableFuture> fetchNodeInternal(@Nonnull Node.NodeCreator creator, + @Nonnull ReadTransaction readTransaction, + int layer, @Nonnull Tuple primaryKey); /** * Method to perform basic invariant check(s) on a newly-fetched node. @@ -190,12 +194,7 @@ public CompletableFuture fetchNode(@Nonnull final ReadTransaction transact * @return the node that was passed in */ @Nullable - private N checkNode(@Nullable final N node) { - if (node != null && (node.size() < getConfig().getMinM() || node.size() > getConfig().getMaxM())) { - if (!node.isRoot()) { - throw new IllegalStateException("packing of non-root is out of valid range"); - } - } + private NodeWithLayer checkNode(@Nullable final NodeWithLayer node) { return node; } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java index 63a4df1905..426465f4df 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java @@ -75,22 +75,44 @@ public ByNodeStorageAdapter(@Nonnull final HNSW.Config config, @Nonnull final Su } @Override - public CompletableFuture> fetchEntryNode(@Nonnull final ReadTransaction readTransaction) { + public CompletableFuture fetchEntryNodeKey(@Nonnull final ReadTransaction readTransaction) { final byte[] key = getEntryNodeSubspace().pack(); return readTransaction.get(key) .thenApply(valueBytes -> { if (valueBytes == null) { - throw new IllegalStateException("cannot fetch entry point"); + return null; // not a single node in the index } final Tuple entryTuple = Tuple.fromBytes(valueBytes); final int lMax = (int)entryTuple.getLong(0); - final Node node = nodeFromTuple(entryTuple.getNestedTuple(1)); + final Tuple primaryKey = entryTuple.getNestedTuple(1); + final OnReadListener onReadListener = getOnReadListener(); + onReadListener.onKeyValueRead(key, valueBytes); + return new NodeKeyWithLayer(lMax, primaryKey); + }); + } + + @Nonnull + @Override + protected CompletableFuture> fetchNodeInternal(@Nonnull final Node.NodeCreator creator, + @Nonnull final ReadTransaction readTransaction, + final int layer, + @Nonnull final Tuple primaryKey) { + final byte[] key = getDataSubspace().pack(Tuple.from(layer, primaryKey)); + + return readTransaction.get(key) + .thenApply(valueBytes -> { + if (valueBytes == null) { + throw new IllegalStateException("cannot fetch node"); + } + + final Tuple nodeTuple = Tuple.fromBytes(valueBytes); + final Node node = nodeFromTuple(creator, nodeTuple); final OnReadListener onReadListener = getOnReadListener(); onReadListener.onNodeRead(node); - onReadListener.onKeyValueRead(node, key, valueBytes); - return node.withLayer(lMax); + onReadListener.onKeyValueRead(key, valueBytes); + return node.withLayer(layer); }); } @@ -153,7 +175,8 @@ public CompletableFuture fetchNodeInternal(@Nonnull final ReadTransaction } @Nonnull - private Node nodeFromTuple(@Nonnull final Tuple tuple) { + private Node nodeFromTuple(@Nonnull final Node.NodeCreator creator, + @Nonnull final Tuple tuple) { final NodeKind nodeKind = NodeKind.fromSerializedNodeKind((byte)tuple.getLong(0)); final Tuple primaryKey = tuple.getNestedTuple(1); final Tuple vectorTuple = tuple.getNestedTuple(2); @@ -200,9 +223,8 @@ private Node nodeFromTuple(@Nonnull final Tuple tuple) { Verify.verify((nodeKind == NodeKind.DATA && neighbors != null) || (nodeKind == NodeKind.INTERMEDIATE && neighborsWithVectors != null)); - return nodeKind == NodeKind.DATA - ? new DataNode(primaryKey, vector, neighbors) - : new IntermediateNode(primaryKey, vector, neighborsWithVectors); + return creator.create(nodeKind, primaryKey, vector, + nodeKind == NodeKind.DATA ? neighbors : neighborsWithVectors); } @Nonnull diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java index 34fe48d579..011d7761a1 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java @@ -22,6 +22,7 @@ import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; +import com.google.common.base.Verify; import com.google.common.collect.Lists; import javax.annotation.Nonnull; @@ -37,6 +38,12 @@ public DataNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector vec super(primaryKey, vector, neighbors); } + @Nonnull + @Override + public NodeKind getKind() { + return NodeKind.DATA; + } + @Nonnull @Override public DataNode asDataNode() { @@ -55,9 +62,18 @@ public NodeWithLayer withLayer(final int layer) { return new NodeWithLayer<>(layer, this); } - @Nonnull @Override - public NodeKind getKind() { - return NodeKind.DATA; + public NodeCreator sameCreator() { + return DataNode::creator; + } + + @Nonnull + @SuppressWarnings("unchecked") + public static Node creator(@Nonnull final NodeKind nodeKind, + @Nonnull final Tuple primaryKey, + @Nonnull final Vector vector, + @Nonnull final List neighbors) { + Verify.verify(nodeKind == NodeKind.INTERMEDIATE); + return new DataNode(primaryKey, vector, (List)neighbors); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index 9a2bcf249b..a4fef2b289 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -31,6 +31,7 @@ import com.apple.foundationdb.subspace.Subspace; import com.apple.foundationdb.tuple.Tuple; import com.apple.foundationdb.tuple.TupleHelpers; +import com.christianheina.langx.half4j.Half; import com.google.common.base.Preconditions; import com.google.common.base.Verify; import com.google.common.collect.ImmutableList; @@ -42,6 +43,7 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; +import java.lang.reflect.Array; import java.math.BigInteger; import java.util.ArrayDeque; import java.util.Arrays; @@ -250,6 +252,10 @@ public boolean isStoreHilbertValues() { return storeHilbertValues; } + public Metric getMetric() { + return Metric.euclideanMetric(); + } + public ConfigBuilder toBuilder() { return new ConfigBuilder(useNodeSlotIndex, minM, maxM, splitS, storage, storeHilbertValues); } @@ -504,101 +510,88 @@ public AsyncIterator scan(@Nonnull final ReadTransaction readTransacti } /** - * Returns the left-most path from a given node id to a leaf node containing items as a {@link TraversalState}. - * The term left-most used here is defined by comparing {@code (largestHilbertValue, largestKey)} when - * comparing nodes (the left one being the smaller, the right one being the greater). - * @param readTransaction the transaction to use - * @param nodeId node id to start from. This may be the actual root of the tree or some other node within the tree. - * @param lastHilbertValue hilbert value serving as a watermark to return only items that are larger than the - * {@code (lastHilbertValue, lastKey)} pair - * @param lastKey key serving as a watermark to return only items that are larger than the - * {@code (lastHilbertValue, lastKey)} pair - * @param mbrPredicate a predicate on an mbr {@link Rectangle}. This predicate is evaluated on the way down to the - * leaf node. - * @param suffixPredicate predicate to be invoked on a range of suffixes - * @return a {@link TraversalState} of the left-most path from {@code nodeId} to a {@link DataNode} whose - * {@link Node}s all pass the mbr predicate test. + * TODO. */ @Nonnull - private CompletableFuture fetchLeftmostPathToLeaf(@Nonnull final ReadTransaction readTransaction, - @Nonnull final byte[] nodeId, - @Nullable final BigInteger lastHilbertValue, - @Nullable final Tuple lastKey, - @Nonnull final Predicate mbrPredicate, - @Nonnull final BiPredicate suffixPredicate) { - final AtomicReference> currentNodeWithLayer = - new AtomicReference<>(); - - storageAdapter.fetchEntryNode(readTransaction) - .thenApply(nodeWithLayer -> { - currentNodeWithLayer.set(nodeWithLayer); - - final List> toBeProcessed = Lists.newArrayList(); - final AtomicReference leafNode = new AtomicReference<>(null); - return AsyncUtil.whileTrue(() -> onReadListener.onAsyncRead(storageAdapter.fetchNode(readTransaction, currentId.get())) - .thenApply(node -> { - if (node == null) { - if (Arrays.equals(currentId.get(), rootId)) { - Verify.verify(leafNode.get() == null); - return false; - } - throw new IllegalStateException("unable to fetch node for scan"); - } - if (node.getKind() == NodeKind.INTERMEDIATE) { - final Iterable childSlots = ((IntermediateNode)node).getSlots(); - Deque toBeProcessedThisLevel = new ArrayDeque<>(); - for (final Iterator iterator = childSlots.iterator(); iterator.hasNext(); ) { - final ChildSlot childSlot = iterator.next(); - if (lastHilbertValue != null && - lastKey != null) { - final int hilbertValueAndKeyCompare = - childSlot.compareLargestHilbertValueAndKey(lastHilbertValue, lastKey); - if (hilbertValueAndKeyCompare < 0) { - // - // The (lastHilbertValue, lastKey) pair is larger than the - // (largestHilbertValue, largestKey) pair of the current child. Advance to the next - // child. - // - continue; - } - } - - if (!mbrPredicate.test(childSlot.getMbr())) { - onReadListener.onChildNodeDiscard(childSlot); - continue; - } - - if (childSlot.suffixPredicateCanBeApplied()) { - if (!suffixPredicate.test(childSlot.getSmallestKeySuffix(), - childSlot.getLargestKeySuffix())) { - onReadListener.onChildNodeDiscard(childSlot); - continue; - } - } - - toBeProcessedThisLevel.addLast(childSlot); - iterator.forEachRemaining(toBeProcessedThisLevel::addLast); - } - toBeProcessed.add(toBeProcessedThisLevel); + private CompletableFuture nearestDropNodeKeyOnLayer0(@Nonnull final ReadTransaction readTransaction, + @Nonnull final Vector queryVector) { + return storageAdapter.fetchEntryNodeKey(readTransaction) + .thenApply(entryNodeKeyWithLayer -> + entryNodeKeyWithLayer == null + ? null + : new NodeKeyWithLayerAndDistance(entryNodeKeyWithLayer.getLayer(), entryNodeKeyWithLayer.getPrimaryKey(), + Double.POSITIVE_INFINITY)) + .thenCompose(entryNodeKeyWithLayerAndDistance -> { + if (entryNodeKeyWithLayerAndDistance == null) { + return CompletableFuture.completedFuture(null); // not a single node in the index + } - final ChildSlot nextChildSlot = resolveNextIdForFetch(toBeProcessed, mbrPredicate, - suffixPredicate, onReadListener); - if (nextChildSlot == null) { - return false; - } + final AtomicReference currentNodeKeyWithLayerReference = + new AtomicReference<>(entryNodeKeyWithLayerAndDistance); - currentId.set(Objects.requireNonNull(nextChildSlot.getChildId())); - return true; - } else { - leafNode.set((DataNode)node); - return false; - } - }), executor).thenApply(vignore -> leafNode.get() == null - ? TraversalState.end() - : TraversalState.of(toBeProcessed, leafNode.get())); + return AsyncUtil.whileTrue(() -> + nearestDropNodeKey(readTransaction, entryNodeKeyWithLayerAndDistance, queryVector) + .thenApply(nodeKeyWithLayerAndDistance -> { + currentNodeKeyWithLayerReference.set(nodeKeyWithLayerAndDistance); + return nodeKeyWithLayerAndDistance.getLayer() > 0; + }), executor).thenApply(ignored -> currentNodeKeyWithLayerReference.get()); }); } + /** + * TODO. + */ + @Nonnull + private CompletableFuture nearestDropNodeKey(@Nonnull final ReadTransaction readTransaction, + @Nonnull final NodeKeyWithLayerAndDistance entryNodeKey, + @Nonnull final Vector queryVector) { + final var layer = entryNodeKey.getLayer(); + Verify.verify(layer > 0); + final Metric metric = getConfig().getMetric(); + final AtomicReference currentNodeKeyReference = + new AtomicReference<>(entryNodeKey); + + return AsyncUtil.whileTrue(() -> onReadListener.onAsyncRead( + storageAdapter.fetchNode(IntermediateNode::creator, readTransaction, + layer, currentNodeKeyReference.get().getPrimaryKey())) + .thenApply(nodeWithLayer -> { + if (nodeWithLayer == null) { + throw new IllegalStateException("unable to fetch node"); + } + final IntermediateNode node = nodeWithLayer.getNode().asIntermediateNode(); + final List neighbors = node.getNeighbors(); + + final NodeKeyWithLayerAndDistance currentNodeKey = currentNodeKeyReference.get(); + + double minDistance = + currentNodeKey.getDistance() == Double.POSITIVE_INFINITY + ? Vector.comparativeDistance(metric, node.getVector(), queryVector) + : currentNodeKey.getDistance(); + + NeighborWithVector nearestNeighbor = null; + for (final NeighborWithVector neighbor : neighbors) { + final double distance = + Vector.comparativeDistance(metric, neighbor.getVector(), queryVector); + if (distance < minDistance) { + minDistance = distance; + nearestNeighbor = neighbor; + } + } + + if (nearestNeighbor == null) { + currentNodeKeyReference.set( + new NodeKeyWithLayerAndDistance(layer - 1, currentNodeKey.getPrimaryKey(), + minDistance)); + return false; + } + + currentNodeKeyReference.set( + new NodeKeyWithLayerAndDistance(layer, nearestNeighbor.getPrimaryKey(), + minDistance)); + return true; + }), executor).thenApply(ignored -> currentNodeKeyReference.get()); + } + /** * Returns the next left-most path from a given {@link TraversalState} to a leaf node containing items as * a {@link TraversalState}. The term left-most used here is defined by comparing @@ -1548,51 +1541,23 @@ private CompletableFuture fetchUpdatePathToLeaf(@Nonnull final Transac } /** - * Method to fetch the siblings of a given node. The node passed in must not be the root node and must be linked up - * to its parent. The parent already has information obout the children ids. This method (through the slot - * information of the node passed in) can then determine adjacent nodes. - * @param transaction the transaction to use - * @param node the node to fetch siblings for - * @return a completable future containing a list of {@link Node}s that contain the {@link Config#getSplitS()} - * number of siblings (where the node passed in is counted as a sibling) if that many siblings exist. In - * the case (i.e. for a small root node) where there are not enough siblings we return the maximum possible - * number of siblings. The returned sibling nodes are returned in Hilbert value order and contain the node - * passed in at the correct position in the returned list. The siblings will also attempt to hug the nodes - * passed in as good as possible meaning that we attempt to return the node passed in as middle-most element - * of the returned list. + * TODO. */ @Nonnull - private CompletableFuture> fetchSiblings(@Nonnull final Transaction transaction, - @Nonnull final Node node) { + @SuppressWarnings("unchecked") + private CompletableFuture>> fetchNeighborNodes(@Nonnull final Transaction transaction, + @Nonnull final NodeWithLayer nodeWithLayer) { // this deque is only modified by once upon creation - final ArrayDeque toBeProcessed = new ArrayDeque<>(); + final ArrayDeque toBeProcessed = new ArrayDeque<>(); final List> working = Lists.newArrayList(); - final int numSiblings = config.getSplitS(); - final Node[] siblings = new Node[numSiblings]; + final Node node = nodeWithLayer.getNode(); + final List neighbors = node.getNeighbors(); + final AtomicInteger neighborIndex = new AtomicInteger(0); + final NodeWithLayer[] neighborNodeArray = + (NodeWithLayer[])Array.newInstance(NodeWithLayer.class, neighbors.size()); - // - // Do some acrobatics to find the best start/end positions for the siblings. Take into account how many - // are warranted, if the node that was passed occupies a slot in its parent node that is touching the end or the - // beginning of the parent's slots, and the total number of slots in the parent of the node that was - // passed in. - // - final IntermediateNode parentNode = Objects.requireNonNull(node.getParentNode()); - int slotIndexInParent = node.getSlotIndexInParent(); - int start = slotIndexInParent - numSiblings / 2; - int end = start + numSiblings; - if (start < 0) { - start = 0; - end = numSiblings; - } else if (end > parentNode.size()) { - end = parentNode.size(); - start = end - numSiblings; - } - - // because lambdas - final int minSibling = start; - - for (int i = start; i < end; i++) { - toBeProcessed.addLast(parentNode.getSlot(i).getChildId()); + for (final N neighbor : neighbors) { + toBeProcessed.addLast(neighbor.getPrimaryKey()); } // Fetch all sibling nodes (in parallel if possible). @@ -1600,31 +1565,25 @@ private CompletableFuture> fetchSiblings(@Nonnull final Transaction t working.removeIf(CompletableFuture::isDone); while (working.size() <= MAX_CONCURRENT_READS) { - final int index = numSiblings - toBeProcessed.size(); - final byte[] currentId = toBeProcessed.pollFirst(); - if (currentId == null) { + final Tuple currentNeighborKey = toBeProcessed.pollFirst(); + if (currentNeighborKey == null) { break; } - final int slotIndex = minSibling + index; - if (slotIndex != slotIndexInParent) { - working.add(storageAdapter.fetchNode(transaction, currentId) - .thenAccept(siblingNode -> { - Objects.requireNonNull(siblingNode); - siblingNode.linkToParent(parentNode, slotIndex); - siblings[index] = siblingNode; - })); - } else { - // put node in the list of siblings -- even though node is strictly speaking not a sibling of itself - siblings[index] = node; - } + final int index = neighborIndex.getAndIncrement(); + + working.add(storageAdapter.fetchNode(node.sameCreator(), transaction, nodeWithLayer.getLayer(), currentNeighborKey) + .thenAccept(resultNode -> { + Objects.requireNonNull(resultNode); + neighborNodeArray[index] = resultNode; + })); } if (working.isEmpty()) { return AsyncUtil.READY_FALSE; } - return AsyncUtil.whenAny(working).thenApply(v -> true); - }, executor).thenApply(vignore -> Lists.newArrayList(siblings)); + return AsyncUtil.whenAny(working).thenApply(ignored -> true); + }, executor).thenApply(ignored -> Lists.newArrayList(neighborNodeArray)); } /** diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java index ff1c494420..689148eb54 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java @@ -22,14 +22,13 @@ import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; +import com.google.common.base.Verify; import javax.annotation.Nonnull; import java.util.List; /** - * An intermediate node of the R-tree. An intermediate node holds the holds information about its children nodes that - * be intermediate nodes or leaf nodes. The secondary attributes such as {@code largestHilbertValue}, - * {@code largestKey} can be derived (and recomputed) if the children of this node are available to be introspected. + * TODO. */ class IntermediateNode extends AbstractNode { public IntermediateNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, @@ -37,6 +36,12 @@ public IntermediateNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector withLayer(final int layer) { return new NodeWithLayer<>(layer, this); } - @Nonnull @Override - public NodeKind getKind() { - return NodeKind.INTERMEDIATE; + public NodeCreator sameCreator() { + return IntermediateNode::creator; + } + + @Nonnull + @SuppressWarnings("unchecked") + public static Node creator(@Nonnull final NodeKind nodeKind, + @Nonnull final Tuple primaryKey, + @Nonnull final Vector vector, + @Nonnull final List neighbors) { + Verify.verify(nodeKind == NodeKind.INTERMEDIATE); + return new IntermediateNode(primaryKey, vector, (List)neighbors); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java index afa0e14fc5..f59f15bf22 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java @@ -20,7 +20,15 @@ package com.apple.foundationdb.async.hnsw; +import javax.annotation.Nonnull; + public interface Metric { + ManhattanMetric MANHATTAN_METRIC = new ManhattanMetric(); + EuclideanMetric EUCLIDEAN_METRIC = new EuclideanMetric(); + EuclideanSquareMetric EUCLIDEAN_SQUARE_METRIC = new EuclideanSquareMetric(); + CosineMetric COSINE_METRIC = new CosineMetric(); + DotProductMetric DOT_PRODUCT_METRIC = new DotProductMetric(); + double distance(Double[] vector1, Double[] vector2); default double comparativeDistance(Double[] vector1, Double[] vector2) { @@ -46,6 +54,31 @@ private static void validate(Double[] vector1, Double[] vector2) { } } + @Nonnull + static ManhattanMetric manhattanMetric() { + return Metric.MANHATTAN_METRIC; + } + + @Nonnull + static EuclideanMetric euclideanMetric() { + return Metric.EUCLIDEAN_METRIC; + } + + @Nonnull + static EuclideanSquareMetric euclideanSquareMetric() { + return Metric.EUCLIDEAN_SQUARE_METRIC; + } + + @Nonnull + static CosineMetric cosineMetric() { + return Metric.COSINE_METRIC; + } + + @Nonnull + static DotProductMetric dotProductMetric() { + return Metric.DOT_PRODUCT_METRIC; + } + class ManhattanMetric implements Metric { @Override public double distance(final Double[] vector1, final Double[] vector2) { diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java index dab2bac5c1..c5d0f71fb6 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java @@ -25,6 +25,7 @@ import com.google.errorprone.annotations.CanIgnoreReturnValue; import javax.annotation.Nonnull; +import java.util.List; /** * TODO. @@ -38,7 +39,7 @@ public interface Node { Vector getVector(); @Nonnull - Iterable getNeighbors(); + List getNeighbors(); @Nonnull N getNeighbor(int index); @@ -70,4 +71,12 @@ public interface Node { @Nonnull NodeWithLayer withLayer(int layer); + + NodeCreator sameCreator(); + + @FunctionalInterface + interface NodeCreator { + Node create(@Nonnull NodeKind nodeKind, @Nonnull Tuple primaryKey, @Nonnull Vector vector, + @Nonnull List neighbors); + } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayer.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayer.java new file mode 100644 index 0000000000..33f833dd20 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayer.java @@ -0,0 +1,45 @@ +/* + * NodeWithLayer.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.tuple.Tuple; + +import javax.annotation.Nonnull; + +class NodeKeyWithLayer { + private final int layer; + @Nonnull + private final Tuple primaryKey; + + public NodeKeyWithLayer(final int layer, @Nonnull final Tuple primaryKey) { + this.layer = layer; + this.primaryKey = primaryKey; + } + + public int getLayer() { + return layer; + } + + @Nonnull + public Tuple getPrimaryKey() { + return primaryKey; + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayerAndDistance.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayerAndDistance.java new file mode 100644 index 0000000000..a733e033f9 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayerAndDistance.java @@ -0,0 +1,38 @@ +/* + * NodeWithLayer.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.tuple.Tuple; + +import javax.annotation.Nonnull; + +class NodeKeyWithLayerAndDistance extends NodeKeyWithLayer { + private final double distance; + + public NodeKeyWithLayerAndDistance(final int layer, @Nonnull final Tuple primaryKey, final double distance) { + super(layer, primaryKey); + this.distance = distance; + } + + public double getDistance() { + return distance; + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java index 98e74bf8f6..da12f8199d 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java @@ -34,7 +34,7 @@ default void onSlotIndexEntryRead(@Nonnull final byte[] key) { // nothing } - default CompletableFuture onAsyncRead(@Nonnull CompletableFuture future) { + default CompletableFuture> onAsyncRead(@Nonnull CompletableFuture> future) { return future; } @@ -42,8 +42,7 @@ default void onNodeRead(@Nonnull Node node) { // nothing } - default void onKeyValueRead(@Nonnull Node node, - @Nonnull byte[] key, + default void onKeyValueRead(@Nonnull byte[] key, @Nonnull byte[] value) { // nothing } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java index e6e8e1cf8f..fc30302990 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java @@ -29,7 +29,6 @@ import javax.annotation.Nullable; import java.math.BigInteger; import java.util.List; -import java.util.UUID; import java.util.concurrent.CompletableFuture; /** @@ -82,7 +81,13 @@ interface StorageAdapter { @Nonnull OnReadListener getOnReadListener(); - CompletableFuture> fetchEntryNode(@Nonnull ReadTransaction readTransaction); + CompletableFuture fetchEntryNodeKey(@Nonnull ReadTransaction readTransaction); + + @Nonnull + CompletableFuture> fetchNode(@Nonnull Node.NodeCreator creator, + @Nonnull ReadTransaction readTransaction, + int layer, + @Nonnull Tuple primaryKey); /** * Insert a new entry into the node index if configuration indicates we should maintain such an index. @@ -147,19 +152,4 @@ interface StorageAdapter { CompletableFuture scanNodeIndexAndFetchNode(@Nonnull ReadTransaction transaction, int level, @Nonnull BigInteger hilbertValue, @Nonnull Tuple key, boolean isInsertUpdate); - - /** - * Method to fetch the data needed to construct a {@link Node}. Note that a node on disk is represented by its - * slots. Each slot is represented by a key/value pair in FDB. Each key (common for both leaf and intermediate - * nodes) starts with an 8-byte node id (which is usually a serialized {@link UUID}) followed by one byte which - * indicates the {@link NodeKind} of node the slot belongs to. - * - * @param transaction the transaction to use - * @param nodeId the node id we should use - * - * @return A completable future containing the {@link Node} that was fetched from the database once completed. - * The node may be an object of {@link DataNode} or of {@link IntermediateNode}. - */ - @Nonnull - CompletableFuture fetchNode(@Nonnull ReadTransaction transaction, @Nonnull byte[] nodeId); } From 5d07bce0416a43a9c74daec4d1ab7f753c837ce0 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Sat, 26 Jul 2025 19:20:40 +0200 Subject: [PATCH 04/34] save point -- in the middle of just mess --- .../foundationdb/async/hnsw/AbstractNode.java | 12 +- .../async/hnsw/ByNodeStorageAdapter.java | 117 +++++++------- .../foundationdb/async/hnsw/DataNode.java | 19 ++- ...WithLayerAndDistance.java => Element.java} | 21 ++- ...WithLayer.java => EntryPointAndLayer.java} | 13 +- .../foundationdb/async/hnsw/GreedyResult.java | 55 +++++++ .../apple/foundationdb/async/hnsw/HNSW.java | 146 +++++++++++++----- .../async/hnsw/IntermediateNode.java | 9 +- .../apple/foundationdb/async/hnsw/Node.java | 6 +- .../async/hnsw/StorageAdapter.java | 2 +- 10 files changed, 275 insertions(+), 125 deletions(-) rename fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/{NodeKeyWithLayerAndDistance.java => Element.java} (66%) rename fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/{NodeKeyWithLayer.java => EntryPointAndLayer.java} (76%) create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyResult.java diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java index 87e03b393b..70387939ad 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java @@ -21,7 +21,6 @@ package com.apple.foundationdb.async.hnsw; import com.apple.foundationdb.tuple.Tuple; -import com.christianheina.langx.half4j.Half; import com.google.common.collect.ImmutableList; import javax.annotation.Nonnull; @@ -35,16 +34,12 @@ abstract class AbstractNode implements Node { @Nonnull private final Tuple primaryKey; - @Nonnull - private final Vector vector; - @Nonnull private final List neighbors; - protected AbstractNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, + protected AbstractNode(@Nonnull final Tuple primaryKey, @Nonnull final List neighbors) { this.primaryKey = primaryKey; - this.vector = vector; this.neighbors = ImmutableList.copyOf(neighbors); } @@ -54,11 +49,6 @@ public Tuple getPrimaryKey() { return primaryKey; } - @Nonnull - public Vector getVector() { - return vector; - } - @Nonnull @Override public List getNeighbors() { diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java index 426465f4df..786d924e2b 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java @@ -75,7 +75,7 @@ public ByNodeStorageAdapter(@Nonnull final HNSW.Config config, @Nonnull final Su } @Override - public CompletableFuture fetchEntryNodeKey(@Nonnull final ReadTransaction readTransaction) { + public CompletableFuture fetchEntryNodeKey(@Nonnull final ReadTransaction readTransaction) { final byte[] key = getEntryNodeSubspace().pack(); return readTransaction.get(key) @@ -83,13 +83,14 @@ public CompletableFuture fetchEntryNodeKey(@Nonnull final Read if (valueBytes == null) { return null; // not a single node in the index } + final OnReadListener onReadListener = getOnReadListener(); + onReadListener.onKeyValueRead(key, valueBytes); final Tuple entryTuple = Tuple.fromBytes(valueBytes); final int lMax = (int)entryTuple.getLong(0); final Tuple primaryKey = entryTuple.getNestedTuple(1); - final OnReadListener onReadListener = getOnReadListener(); - onReadListener.onKeyValueRead(key, valueBytes); - return new NodeKeyWithLayer(lMax, primaryKey); + final Tuple vectorTuple = entryTuple.getNestedTuple(2); + return new EntryPointAndLayer(lMax, primaryKey, vectorFromTuple(vectorTuple)); }); } @@ -156,77 +157,79 @@ private Tuple toTuple(@Nonnull final Node node) { return Tuple.from(node.getKind().getSerialized(), slotTuples); } - @Nonnull - @Override - public CompletableFuture fetchNodeInternal(@Nonnull final ReadTransaction transaction, - @Nonnull final byte[] nodeId) { - final byte[] key = packWithSubspace(nodeId); - return transaction.get(key) - .thenApply(valueBytes -> { - if (valueBytes == null) { - return null; - } - final Node node = nodeFromTuple(nodeId, Tuple.fromBytes(valueBytes)); - final OnReadListener onReadListener = getOnReadListener(); - onReadListener.onNodeRead(node); - onReadListener.onKeyValueRead(node, key, valueBytes); - return node; - }); - } - @Nonnull private Node nodeFromTuple(@Nonnull final Node.NodeCreator creator, @Nonnull final Tuple tuple) { final NodeKind nodeKind = NodeKind.fromSerializedNodeKind((byte)tuple.getLong(0)); final Tuple primaryKey = tuple.getNestedTuple(1); - final Tuple vectorTuple = tuple.getNestedTuple(2); - final Tuple neighborsTuple = tuple.getNestedTuple(3); + final Tuple vectorTuple; + final Tuple neighborsTuple; + + switch (nodeKind) { + case DATA: + vectorTuple = tuple.getNestedTuple(2); + neighborsTuple = tuple.getNestedTuple(3); + return dataNodeFromTuples(creator, primaryKey, vectorTuple, neighborsTuple); + case INTERMEDIATE: + neighborsTuple = tuple.getNestedTuple(3); + return intermediateNodeFromTuples(creator, primaryKey, neighborsTuple); + default: + throw new IllegalStateException("unknown node kind"); + } + } - final Half[] vectorHalfs = new Half[vectorTuple.size()]; - for (int i = 0; i < vectorTuple.size(); i ++) { - vectorHalfs[i] = Half.shortBitsToHalf(shortFromBytes(vectorTuple.getBytes(i))); + @Nonnull + private Node dataNodeFromTuples(@Nonnull final Node.NodeCreator creator, + @Nonnull final Tuple primaryKey, + @Nonnull final Tuple vectorTuple, + @Nonnull final Tuple neighborsTuple) { + final Vector vector = vectorFromTuple(vectorTuple); + + List neighbors = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); + + for (final Object neighborObject : neighborsTuple) { + final Tuple neighborTuple = (Tuple)neighborObject; + neighbors.add(new Neighbor(neighborTuple)); } - final Vector.HalfVector vector = new Vector.HalfVector(vectorHalfs); - List neighborsWithVectors = null; + + return creator.create(NodeKind.DATA, primaryKey, vector, neighbors); + } + + @Nonnull + private Node intermediateNodeFromTuples(@Nonnull final Node.NodeCreator creator, + @Nonnull final Tuple primaryKey, + @Nonnull final Tuple neighborsTuple) { + List neighborsWithVectors = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); Half[] neighborVectorHalfs = null; - List neighbors = null; for (final Object neighborObject : neighborsTuple) { final Tuple neighborTuple = (Tuple)neighborObject; - switch (nodeKind) { - case DATA: - if (neighbors == null) { - neighbors = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); - } - neighbors.add(new Neighbor(neighborTuple)); - break; - - case INTERMEDIATE: - final Tuple neighborPrimaryKey = neighborTuple.getNestedTuple(0); - final Tuple neighborVectorTuple = neighborTuple.getNestedTuple(1); - if (neighborsWithVectors == null) { - neighborsWithVectors = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); - neighborVectorHalfs = new Half[neighborVectorTuple.size()]; - } - - for (int i = 0; i < neighborVectorTuple.size(); i ++) { - neighborVectorHalfs[i] = Half.shortBitsToHalf(shortFromBytes(neighborVectorTuple.getBytes(i))); - } - neighborsWithVectors.add(new NeighborWithVector(neighborPrimaryKey, new Vector.HalfVector(neighborVectorHalfs))); - break; + final Tuple neighborPrimaryKey = neighborTuple.getNestedTuple(0); + final Tuple neighborVectorTuple = neighborTuple.getNestedTuple(1); + if (neighborVectorHalfs == null) { + neighborVectorHalfs = new Half[neighborVectorTuple.size()]; + } - default: - throw new IllegalStateException("unknown node kind"); + for (int i = 0; i < neighborVectorTuple.size(); i ++) { + neighborVectorHalfs[i] = Half.shortBitsToHalf(shortFromBytes(neighborVectorTuple.getBytes(i))); } + neighborsWithVectors.add(new NeighborWithVector(neighborPrimaryKey, new Vector.HalfVector(neighborVectorHalfs))); } - Verify.verify((nodeKind == NodeKind.DATA && neighbors != null) || - (nodeKind == NodeKind.INTERMEDIATE && neighborsWithVectors != null)); + return creator.create(NodeKind.INTERMEDIATE, primaryKey, null, neighborsWithVectors); + } - return creator.create(nodeKind, primaryKey, vector, - nodeKind == NodeKind.DATA ? neighbors : neighborsWithVectors); + @Nonnull + private Vector vectorFromTuple(final Tuple vectorTuple) { + final Half[] vectorHalfs = new Half[vectorTuple.size()]; + for (int i = 0; i < vectorTuple.size(); i ++) { + vectorHalfs[i] = Half.shortBitsToHalf(shortFromBytes(vectorTuple.getBytes(i))); + } + return new Vector.HalfVector(vectorHalfs); } + + @Nonnull @Override public > AbstractChangeSet diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java index 011d7761a1..d99afcbdb3 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java @@ -23,19 +23,23 @@ import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; import com.google.common.base.Verify; -import com.google.common.collect.Lists; import javax.annotation.Nonnull; import javax.annotation.Nullable; import java.util.List; +import java.util.Objects; /** - * A leaf node of the R-tree. A leaf node holds the actual data in {@link ItemSlot}s. + * TODO. */ class DataNode extends AbstractNode { + @Nonnull + private final Vector vector; + public DataNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, @Nonnull final List neighbors) { - super(primaryKey, vector, neighbors); + super(primaryKey, neighbors); + this.vector = vector; } @Nonnull @@ -44,6 +48,11 @@ public NodeKind getKind() { return NodeKind.DATA; } + @Nonnull + public Vector getVector() { + return vector; + } + @Nonnull @Override public DataNode asDataNode() { @@ -71,9 +80,9 @@ public NodeCreator sameCreator() { @SuppressWarnings("unchecked") public static Node creator(@Nonnull final NodeKind nodeKind, @Nonnull final Tuple primaryKey, - @Nonnull final Vector vector, + @Nullable final Vector vector, @Nonnull final List neighbors) { Verify.verify(nodeKind == NodeKind.INTERMEDIATE); - return new DataNode(primaryKey, vector, (List)neighbors); + return new DataNode(primaryKey, Objects.requireNonNull(vector), (List)neighbors); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayerAndDistance.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Element.java similarity index 66% rename from fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayerAndDistance.java rename to fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Element.java index a733e033f9..53e1390687 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayerAndDistance.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Element.java @@ -23,16 +23,31 @@ import com.apple.foundationdb.tuple.Tuple; import javax.annotation.Nonnull; +import java.util.List; -class NodeKeyWithLayerAndDistance extends NodeKeyWithLayer { +class Element { + @Nonnull + private final Tuple primaryKey; private final double distance; - public NodeKeyWithLayerAndDistance(final int layer, @Nonnull final Tuple primaryKey, final double distance) { - super(layer, primaryKey); + public Element(@Nonnull final Tuple primaryKey, final double distance) { + this.primaryKey = primaryKey; this.distance = distance; } + @Nonnull + public Tuple getPrimaryKey() { + return primaryKey; + } + public double getDistance() { return distance; } + + @Nonnull + public static Iterable primaryKeys(@Nonnull List elements) { + return () -> elements.stream() + .map(Element::getPrimaryKey) + .iterator(); + } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayer.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/EntryPointAndLayer.java similarity index 76% rename from fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayer.java rename to fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/EntryPointAndLayer.java index 33f833dd20..f621916079 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayer.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/EntryPointAndLayer.java @@ -21,17 +21,21 @@ package com.apple.foundationdb.async.hnsw; import com.apple.foundationdb.tuple.Tuple; +import com.christianheina.langx.half4j.Half; import javax.annotation.Nonnull; -class NodeKeyWithLayer { +class EntryPointAndLayer { private final int layer; @Nonnull private final Tuple primaryKey; + @Nonnull + private final Vector vector; - public NodeKeyWithLayer(final int layer, @Nonnull final Tuple primaryKey) { + public EntryPointAndLayer(final int layer, @Nonnull final Tuple primaryKey, @Nonnull final Vector vector) { this.layer = layer; this.primaryKey = primaryKey; + this.vector = vector; } public int getLayer() { @@ -42,4 +46,9 @@ public int getLayer() { public Tuple getPrimaryKey() { return primaryKey; } + + @Nonnull + public Vector getVector() { + return vector; + } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyResult.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyResult.java new file mode 100644 index 0000000000..0a66c8bca7 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyResult.java @@ -0,0 +1,55 @@ +/* + * NodeWithLayer.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.tuple.Tuple; + +import javax.annotation.Nonnull; + +class GreedyResult { + private final int layer; + @Nonnull + private final Tuple primaryKey; + private final double distance; + + public GreedyResult(final int layer, @Nonnull final Tuple primaryKey, final double distance) { + this.layer = layer; + this.primaryKey = primaryKey; + this.distance = distance; + } + + public int getLayer() { + return layer; + } + + @Nonnull + public Tuple getPrimaryKey() { + return primaryKey; + } + + public double getDistance() { + return distance; + } + + public Element toElement() { + return new Element(getPrimaryKey(), getDistance()); + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index a4fef2b289..7f26e16740 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -37,6 +37,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterables; import com.google.common.collect.Lists; +import com.google.common.collect.Sets; import com.google.errorprone.annotations.CanIgnoreReturnValue; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -48,13 +49,16 @@ import java.util.ArrayDeque; import java.util.Arrays; import java.util.Collections; +import java.util.Comparator; import java.util.Deque; import java.util.Iterator; import java.util.List; import java.util.NoSuchElementException; import java.util.Objects; +import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; +import java.util.concurrent.PriorityBlockingQueue; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiPredicate; @@ -512,29 +516,37 @@ public AsyncIterator scan(@Nonnull final ReadTransaction readTransacti /** * TODO. */ + @SuppressWarnings("checkstyle:MethodName") // method name introduced by paper @Nonnull - private CompletableFuture nearestDropNodeKeyOnLayer0(@Nonnull final ReadTransaction readTransaction, - @Nonnull final Vector queryVector) { + private CompletableFuture kNearestNeighborsSearch(@Nonnull final ReadTransaction readTransaction, + @Nonnull final Vector queryVector) { return storageAdapter.fetchEntryNodeKey(readTransaction) - .thenApply(entryNodeKeyWithLayer -> - entryNodeKeyWithLayer == null - ? null - : new NodeKeyWithLayerAndDistance(entryNodeKeyWithLayer.getLayer(), entryNodeKeyWithLayer.getPrimaryKey(), - Double.POSITIVE_INFINITY)) - .thenCompose(entryNodeKeyWithLayerAndDistance -> { - if (entryNodeKeyWithLayerAndDistance == null) { + .thenCompose(entryPointAndLayer -> { + if (entryPointAndLayer == null) { return CompletableFuture.completedFuture(null); // not a single node in the index } - final AtomicReference currentNodeKeyWithLayerReference = - new AtomicReference<>(entryNodeKeyWithLayerAndDistance); + final Metric metric = getConfig().getMetric(); - return AsyncUtil.whileTrue(() -> - nearestDropNodeKey(readTransaction, entryNodeKeyWithLayerAndDistance, queryVector) - .thenApply(nodeKeyWithLayerAndDistance -> { - currentNodeKeyWithLayerReference.set(nodeKeyWithLayerAndDistance); - return nodeKeyWithLayerAndDistance.getLayer() > 0; - }), executor).thenApply(ignored -> currentNodeKeyWithLayerReference.get()); + final var entryState = new GreedyResult(entryPointAndLayer.getLayer(), + entryPointAndLayer.getPrimaryKey(), + Vector.comparativeDistance(metric, entryPointAndLayer.getVector(), queryVector)); + final AtomicReference greedyResultReference = + new AtomicReference<>(entryState); + + if (entryPointAndLayer.getLayer() == 0) { + // entry data points to a node in layer 0 directly + return CompletableFuture.completedFuture(entryState); + } + + return AsyncUtil.whileTrue(() -> { + final var greedyIn = greedyResultReference.get(); + return greedySearchLayer(readTransaction, greedyIn.toElement(), greedyIn.getLayer(), queryVector) + .thenApply(greedyResult -> { + greedyResultReference.set(greedyResult); + return greedyResult.getLayer() > 0; + }); + }, executor).thenApply(ignored -> greedyResultReference.get()); }); } @@ -542,18 +554,18 @@ private CompletableFuture nearestDropNodeKeyOnLayer * TODO. */ @Nonnull - private CompletableFuture nearestDropNodeKey(@Nonnull final ReadTransaction readTransaction, - @Nonnull final NodeKeyWithLayerAndDistance entryNodeKey, - @Nonnull final Vector queryVector) { - final var layer = entryNodeKey.getLayer(); + private CompletableFuture greedySearchLayer(@Nonnull final ReadTransaction readTransaction, + @Nonnull final Element entryElement, + final int layer, + @Nonnull final Vector queryVector) { Verify.verify(layer > 0); final Metric metric = getConfig().getMetric(); - final AtomicReference currentNodeKeyReference = - new AtomicReference<>(entryNodeKey); + final AtomicReference greedyStateReference = + new AtomicReference<>(new GreedyResult(layer, entryElement.getPrimaryKey(), entryElement.getDistance())); return AsyncUtil.whileTrue(() -> onReadListener.onAsyncRead( storageAdapter.fetchNode(IntermediateNode::creator, readTransaction, - layer, currentNodeKeyReference.get().getPrimaryKey())) + layer, greedyStateReference.get().getPrimaryKey())) .thenApply(nodeWithLayer -> { if (nodeWithLayer == null) { throw new IllegalStateException("unable to fetch node"); @@ -561,12 +573,70 @@ private CompletableFuture nearestDropNodeKey(@Nonnu final IntermediateNode node = nodeWithLayer.getNode().asIntermediateNode(); final List neighbors = node.getNeighbors(); - final NodeKeyWithLayerAndDistance currentNodeKey = currentNodeKeyReference.get(); + final GreedyResult currentNodeKey = greedyStateReference.get(); + double minDistance = currentNodeKey.getDistance(); + + NeighborWithVector nearestNeighbor = null; + for (final NeighborWithVector neighbor : neighbors) { + final double distance = + Vector.comparativeDistance(metric, neighbor.getVector(), queryVector); + if (distance < minDistance) { + minDistance = distance; + nearestNeighbor = neighbor; + } + } + + if (nearestNeighbor == null) { + greedyStateReference.set( + new GreedyResult(layer - 1, currentNodeKey.getPrimaryKey(), minDistance)); + return false; + } + + greedyStateReference.set( + new GreedyResult(layer, nearestNeighbor.getPrimaryKey(), + minDistance)); + return true; + }), executor).thenApply(ignored -> greedyStateReference.get()); + } + + /** + * TODO. + */ + @Nonnull + private CompletableFuture searchLayer(@Nonnull final ReadTransaction readTransaction, + @Nonnull final List entryElements, + final int layer, + @Nonnull final Vector queryVector) { + final Set visited = Sets.newHashSet(Element.primaryKeys(entryElements)); + final PriorityBlockingQueue candidates = + new PriorityBlockingQueue<>(entryElements.size(), + Comparator.comparing(Element::getDistance)); + candidates.addAll(entryElements); + final PriorityBlockingQueue nearestNeighbors = + new PriorityBlockingQueue<>(entryElements.size(), + Comparator.comparing(Element::getDistance).reversed()); + nearestNeighbors.addAll(entryElements); + + if (candidates.isEmpty()) { + return CompletableFuture.completedFuture(null); // TODO + } + + final Metric metric = getConfig().getMetric(); + final AtomicReference greedyStateReference = + new AtomicReference<>(entryStates); + + return AsyncUtil.whileTrue(() -> + fetchNeighborNodes(IntermediateNode::creator, readTransaction, + layer, greedyStateReference.get().getPrimaryKey()) + .thenApply(nodeWithLayer -> { + if (nodeWithLayer == null) { + throw new IllegalStateException("unable to fetch node"); + } + final IntermediateNode node = nodeWithLayer.getNode().asIntermediateNode(); + final List neighbors = node.getNeighbors(); - double minDistance = - currentNodeKey.getDistance() == Double.POSITIVE_INFINITY - ? Vector.comparativeDistance(metric, node.getVector(), queryVector) - : currentNodeKey.getDistance(); + final GreedyResult currentNodeKey = greedyStateReference.get(); + double minDistance = currentNodeKey.getDistance(); NeighborWithVector nearestNeighbor = null; for (final NeighborWithVector neighbor : neighbors) { @@ -579,17 +649,16 @@ private CompletableFuture nearestDropNodeKey(@Nonnu } if (nearestNeighbor == null) { - currentNodeKeyReference.set( - new NodeKeyWithLayerAndDistance(layer - 1, currentNodeKey.getPrimaryKey(), - minDistance)); + greedyStateReference.set( + new GreedyResult(layer - 1, currentNodeKey.getPrimaryKey(), minDistance)); return false; } - currentNodeKeyReference.set( - new NodeKeyWithLayerAndDistance(layer, nearestNeighbor.getPrimaryKey(), + greedyStateReference.set( + new GreedyResult(layer, nearestNeighbor.getPrimaryKey(), minDistance)); return true; - }), executor).thenApply(ignored -> currentNodeKeyReference.get()); + }), executor).thenApply(ignored -> greedyStateReference.get()); } /** @@ -1546,11 +1615,11 @@ private CompletableFuture fetchUpdatePathToLeaf(@Nonnull final Transac @Nonnull @SuppressWarnings("unchecked") private CompletableFuture>> fetchNeighborNodes(@Nonnull final Transaction transaction, - @Nonnull final NodeWithLayer nodeWithLayer) { + @Nonnull final Node node, + final int layer) { // this deque is only modified by once upon creation final ArrayDeque toBeProcessed = new ArrayDeque<>(); final List> working = Lists.newArrayList(); - final Node node = nodeWithLayer.getNode(); final List neighbors = node.getNeighbors(); final AtomicInteger neighborIndex = new AtomicInteger(0); final NodeWithLayer[] neighborNodeArray = @@ -1572,7 +1641,8 @@ private CompletableFuture>> fetchNeig final int index = neighborIndex.getAndIncrement(); - working.add(storageAdapter.fetchNode(node.sameCreator(), transaction, nodeWithLayer.getLayer(), currentNeighborKey) + working.add(onReadListener.onAsyncRead(storageAdapter.fetchNode(node.sameCreator(), transaction, layer, + currentNeighborKey)) .thenAccept(resultNode -> { Objects.requireNonNull(resultNode); neighborNodeArray[index] = resultNode; diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java index 689148eb54..e16bdde8d0 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java @@ -25,15 +25,16 @@ import com.google.common.base.Verify; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import java.util.List; /** * TODO. */ class IntermediateNode extends AbstractNode { - public IntermediateNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, + public IntermediateNode(@Nonnull final Tuple primaryKey, @Nonnull final List neighbors) { - super(primaryKey, vector, neighbors); + super(primaryKey, neighbors); } @Nonnull @@ -69,9 +70,9 @@ public NodeCreator sameCreator() { @SuppressWarnings("unchecked") public static Node creator(@Nonnull final NodeKind nodeKind, @Nonnull final Tuple primaryKey, - @Nonnull final Vector vector, + @Nullable final Vector vector, @Nonnull final List neighbors) { Verify.verify(nodeKind == NodeKind.INTERMEDIATE); - return new IntermediateNode(primaryKey, vector, (List)neighbors); + return new IntermediateNode(primaryKey, (List)neighbors); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java index c5d0f71fb6..c0d8462903 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java @@ -25,6 +25,7 @@ import com.google.errorprone.annotations.CanIgnoreReturnValue; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import java.util.List; /** @@ -35,9 +36,6 @@ public interface Node { @Nonnull Tuple getPrimaryKey(); - @Nonnull - Vector getVector(); - @Nonnull List getNeighbors(); @@ -76,7 +74,7 @@ public interface Node { @FunctionalInterface interface NodeCreator { - Node create(@Nonnull NodeKind nodeKind, @Nonnull Tuple primaryKey, @Nonnull Vector vector, + Node create(@Nonnull NodeKind nodeKind, @Nonnull Tuple primaryKey, @Nullable Vector vector, @Nonnull List neighbors); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java index fc30302990..d564b20095 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java @@ -81,7 +81,7 @@ interface StorageAdapter { @Nonnull OnReadListener getOnReadListener(); - CompletableFuture fetchEntryNodeKey(@Nonnull ReadTransaction readTransaction); + CompletableFuture fetchEntryNodeKey(@Nonnull ReadTransaction readTransaction); @Nonnull CompletableFuture> fetchNode(@Nonnull Node.NodeCreator creator, From d9078e208fe9fb4cd02c1e07ca13efb40dcc9870 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Sat, 26 Jul 2025 22:34:15 +0200 Subject: [PATCH 05/34] save point -- in the middle of just mess and mess --- .../foundationdb/async/hnsw/AbstractNode.java | 12 +- .../async/hnsw/AbstractStorageAdapter.java | 14 +- .../async/hnsw/ByNodeStorageAdapter.java | 40 +- .../foundationdb/async/hnsw/DataNode.java | 24 +- .../foundationdb/async/hnsw/GreedyResult.java | 4 +- .../apple/foundationdb/async/hnsw/HNSW.java | 1921 +---------------- .../async/hnsw/IntermediateNode.java | 22 +- .../apple/foundationdb/async/hnsw/Node.java | 23 +- .../hnsw/{Element.java => NodeReference.java} | 33 +- ...or.java => NodeReferenceWithDistance.java} | 35 +- ...ctor.java => NodeReferenceWithVector.java} | 6 +- .../async/hnsw/NodeWithLayer.java | 43 - .../async/hnsw/OnReadListener.java | 2 +- .../async/hnsw/StorageAdapter.java | 8 +- .../apple/foundationdb/async/hnsw/Vector.java | 16 + 15 files changed, 236 insertions(+), 1967 deletions(-) rename fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/{Element.java => NodeReference.java} (61%) rename fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/{Neighbor.java => NodeReferenceWithDistance.java} (50%) rename fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/{NeighborWithVector.java => NodeReferenceWithVector.java} (85%) delete mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeWithLayer.java diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java index 70387939ad..9d5c1ff95c 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java @@ -28,17 +28,17 @@ /** * TODO. - * @param node type class. + * @param node type class. */ -abstract class AbstractNode implements Node { +abstract class AbstractNode implements Node { @Nonnull private final Tuple primaryKey; @Nonnull - private final List neighbors; + private final List neighbors; protected AbstractNode(@Nonnull final Tuple primaryKey, - @Nonnull final List neighbors) { + @Nonnull final List neighbors) { this.primaryKey = primaryKey; this.neighbors = ImmutableList.copyOf(neighbors); } @@ -51,13 +51,13 @@ public Tuple getPrimaryKey() { @Nonnull @Override - public List getNeighbors() { + public List getNeighbors() { return neighbors; } @Nonnull @Override - public N getNeighbor(final int index) { + public R getNeighbor(final int index) { return neighbors.get(index); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java index d1eb8c6d37..98ba205405 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java @@ -173,16 +173,16 @@ public void deleteFromNodeIndexIfNecessary(@Nonnull final Transaction transactio @Nonnull @Override - public CompletableFuture> fetchNode(@Nonnull final Node.NodeCreator creator, - @Nonnull final ReadTransaction readTransaction, - int layer, @Nonnull Tuple primaryKey) { + public CompletableFuture> fetchNode(@Nonnull final Node.NodeCreator creator, + @Nonnull final ReadTransaction readTransaction, + int layer, @Nonnull Tuple primaryKey) { return fetchNodeInternal(creator, readTransaction, layer, primaryKey).thenApply(this::checkNode); } @Nonnull - protected abstract CompletableFuture> fetchNodeInternal(@Nonnull Node.NodeCreator creator, - @Nonnull ReadTransaction readTransaction, - int layer, @Nonnull Tuple primaryKey); + protected abstract CompletableFuture> fetchNodeInternal(@Nonnull Node.NodeCreator creator, + @Nonnull ReadTransaction readTransaction, + int layer, @Nonnull Tuple primaryKey); /** * Method to perform basic invariant check(s) on a newly-fetched node. @@ -194,7 +194,7 @@ protected abstract CompletableFuture> fetc * @return the node that was passed in */ @Nullable - private NodeWithLayer checkNode(@Nullable final NodeWithLayer node) { + private Node checkNode(@Nullable final Node node) { return node; } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java index 786d924e2b..1ec1a12fa4 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java @@ -96,10 +96,10 @@ public CompletableFuture fetchEntryNodeKey(@Nonnull final Re @Nonnull @Override - protected CompletableFuture> fetchNodeInternal(@Nonnull final Node.NodeCreator creator, - @Nonnull final ReadTransaction readTransaction, - final int layer, - @Nonnull final Tuple primaryKey) { + protected CompletableFuture> fetchNodeInternal(@Nonnull final Node.NodeCreator creator, + @Nonnull final ReadTransaction readTransaction, + final int layer, + @Nonnull final Tuple primaryKey) { final byte[] key = getDataSubspace().pack(Tuple.from(layer, primaryKey)); return readTransaction.get(key) @@ -109,11 +109,11 @@ protected CompletableFuture> fetchNodeInte } final Tuple nodeTuple = Tuple.fromBytes(valueBytes); - final Node node = nodeFromTuple(creator, nodeTuple); + final Node node = nodeFromTuple(creator, nodeTuple); final OnReadListener onReadListener = getOnReadListener(); onReadListener.onNodeRead(node); onReadListener.onKeyValueRead(key, valueBytes); - return node.withLayer(layer); + return node; }); } @@ -158,8 +158,8 @@ private Tuple toTuple(@Nonnull final Node node) { } @Nonnull - private Node nodeFromTuple(@Nonnull final Node.NodeCreator creator, - @Nonnull final Tuple tuple) { + private Node nodeFromTuple(@Nonnull final Node.NodeCreator creator, + @Nonnull final Tuple tuple) { final NodeKind nodeKind = NodeKind.fromSerializedNodeKind((byte)tuple.getLong(0)); final Tuple primaryKey = tuple.getNestedTuple(1); final Tuple vectorTuple; @@ -179,27 +179,27 @@ private Node nodeFromTuple(@Nonnull final Node.NodeCreat } @Nonnull - private Node dataNodeFromTuples(@Nonnull final Node.NodeCreator creator, - @Nonnull final Tuple primaryKey, - @Nonnull final Tuple vectorTuple, - @Nonnull final Tuple neighborsTuple) { + private Node dataNodeFromTuples(@Nonnull final Node.NodeCreator creator, + @Nonnull final Tuple primaryKey, + @Nonnull final Tuple vectorTuple, + @Nonnull final Tuple neighborsTuple) { final Vector vector = vectorFromTuple(vectorTuple); - List neighbors = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); + List nodeReferences = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); for (final Object neighborObject : neighborsTuple) { final Tuple neighborTuple = (Tuple)neighborObject; - neighbors.add(new Neighbor(neighborTuple)); + nodeReferences.add(new NodeReference(neighborTuple)); } - return creator.create(NodeKind.DATA, primaryKey, vector, neighbors); + return creator.create(NodeKind.DATA, primaryKey, vector, nodeReferences); } @Nonnull - private Node intermediateNodeFromTuples(@Nonnull final Node.NodeCreator creator, - @Nonnull final Tuple primaryKey, - @Nonnull final Tuple neighborsTuple) { - List neighborsWithVectors = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); + private Node intermediateNodeFromTuples(@Nonnull final Node.NodeCreator creator, + @Nonnull final Tuple primaryKey, + @Nonnull final Tuple neighborsTuple) { + List neighborsWithVectors = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); Half[] neighborVectorHalfs = null; for (final Object neighborObject : neighborsTuple) { @@ -213,7 +213,7 @@ private Node intermediateNodeFromTuples(@Nonnull final N for (int i = 0; i < neighborVectorTuple.size(); i ++) { neighborVectorHalfs[i] = Half.shortBitsToHalf(shortFromBytes(neighborVectorTuple.getBytes(i))); } - neighborsWithVectors.add(new NeighborWithVector(neighborPrimaryKey, new Vector.HalfVector(neighborVectorHalfs))); + neighborsWithVectors.add(new NodeReferenceWithVector(neighborPrimaryKey, new Vector.HalfVector(neighborVectorHalfs))); } return creator.create(NodeKind.INTERMEDIATE, primaryKey, null, neighborsWithVectors); diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java index d99afcbdb3..feb2f79cb0 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java @@ -32,13 +32,13 @@ /** * TODO. */ -class DataNode extends AbstractNode { +class DataNode extends AbstractNode { @Nonnull private final Vector vector; public DataNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, - @Nonnull final List neighbors) { - super(primaryKey, neighbors); + @Nonnull final List nodeReferences) { + super(primaryKey, nodeReferences); this.vector = vector; } @@ -65,24 +65,18 @@ public IntermediateNode asIntermediateNode() { throw new IllegalStateException("this is not a data node"); } - @Nonnull - @Override - public NodeWithLayer withLayer(final int layer) { - return new NodeWithLayer<>(layer, this); - } - @Override - public NodeCreator sameCreator() { + public NodeCreator sameCreator() { return DataNode::creator; } @Nonnull @SuppressWarnings("unchecked") - public static Node creator(@Nonnull final NodeKind nodeKind, - @Nonnull final Tuple primaryKey, - @Nullable final Vector vector, - @Nonnull final List neighbors) { + public static Node creator(@Nonnull final NodeKind nodeKind, + @Nonnull final Tuple primaryKey, + @Nullable final Vector vector, + @Nonnull final List neighbors) { Verify.verify(nodeKind == NodeKind.INTERMEDIATE); - return new DataNode(primaryKey, Objects.requireNonNull(vector), (List)neighbors); + return new DataNode(primaryKey, Objects.requireNonNull(vector), (List)neighbors); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyResult.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyResult.java index 0a66c8bca7..b104848603 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyResult.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyResult.java @@ -49,7 +49,7 @@ public double getDistance() { return distance; } - public Element toElement() { - return new Element(getPrimaryKey(), getDistance()); + public NodeReferenceWithDistance toElement() { + return new NodeReferenceWithDistance(getPrimaryKey(), getDistance()); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index 7f26e16740..5a970fd2ab 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -20,23 +20,17 @@ package com.apple.foundationdb.async.hnsw; -import com.apple.foundationdb.Database; import com.apple.foundationdb.ReadTransaction; -import com.apple.foundationdb.Transaction; -import com.apple.foundationdb.TransactionContext; import com.apple.foundationdb.annotation.API; -import com.apple.foundationdb.annotation.SpotBugsSuppressWarnings; import com.apple.foundationdb.async.AsyncIterator; import com.apple.foundationdb.async.AsyncUtil; import com.apple.foundationdb.subspace.Subspace; import com.apple.foundationdb.tuple.Tuple; -import com.apple.foundationdb.tuple.TupleHelpers; import com.christianheina.langx.half4j.Half; import com.google.common.base.Preconditions; import com.google.common.base.Verify; -import com.google.common.collect.ImmutableList; -import com.google.common.collect.Iterables; import com.google.common.collect.Lists; +import com.google.common.collect.Maps; import com.google.common.collect.Sets; import com.google.errorprone.annotations.CanIgnoreReturnValue; import org.slf4j.Logger; @@ -47,13 +41,9 @@ import java.lang.reflect.Array; import java.math.BigInteger; import java.util.ArrayDeque; -import java.util.Arrays; -import java.util.Collections; import java.util.Comparator; -import java.util.Deque; -import java.util.Iterator; import java.util.List; -import java.util.NoSuchElementException; +import java.util.Map; import java.util.Objects; import java.util.Set; import java.util.concurrent.CompletableFuture; @@ -65,7 +55,6 @@ import java.util.function.Function; import java.util.function.Predicate; import java.util.function.Supplier; -import java.util.stream.Collectors; /** * TODO. @@ -150,8 +139,6 @@ public class HNSW { @Nonnull private final Config config; @Nonnull - private final Function hilbertValueFunction; - @Nonnull private final Supplier nodeIdSupplier; @Nonnull private final OnWriteListener onWriteListener; @@ -181,7 +168,7 @@ public enum Storage { @Nonnull private StorageAdapter newStorageAdapter(@Nonnull final Config config, @Nonnull final Subspace subspace, @Nonnull final Subspace nodeSlotIndexSubspace, - @Nonnull final Function hilbertValueFunction, + @Nonnull final Function hilbertValueFunction, @Nonnull final OnWriteListener onWriteListener, @Nonnull final OnReadListener onReadListener) { return storageAdapterCreator.create(config, subspace, nodeSlotIndexSubspace, @@ -194,7 +181,7 @@ private StorageAdapter newStorageAdapter(@Nonnull final Config config, @Nonnull */ private interface StorageAdapterCreator { StorageAdapter create(@Nonnull Config config, @Nonnull Subspace subspace, @Nonnull Subspace nodeSlotIndexSubspace, - @Nonnull Function hilbertValueFunction, + @Nonnull Function hilbertValueFunction, @Nonnull OnWriteListener onWriteListener, @Nonnull OnReadListener onReadListener); } @@ -374,10 +361,10 @@ public static ConfigBuilder newConfigBuilder() { * @param subspace the subspace where the r-tree is stored * @param secondarySubspace the subspace where the node index (if used is stored) * @param executor an executor to use when running asynchronous tasks - * @param hilbertValueFunction function to compute the Hilbert value from a {@link Point} + * @param hilbertValueFunction function to compute the Hilbert value from a {@link NodeReferenceWithDistance} */ public HNSW(@Nonnull final Subspace subspace, @Nonnull final Subspace secondarySubspace, - @Nonnull final Executor executor, @Nonnull final Function hilbertValueFunction) { + @Nonnull final Executor executor, @Nonnull final Function hilbertValueFunction) { this(subspace, secondarySubspace, executor, DEFAULT_CONFIG, hilbertValueFunction, NodeHelpers::newRandomNodeId, OnWriteListener.NOOP, OnReadListener.NOOP); } @@ -388,14 +375,14 @@ public HNSW(@Nonnull final Subspace subspace, @Nonnull final Subspace secondaryS * @param nodeSlotIndexSubspace the subspace where the node index (if used is stored) * @param executor an executor to use when running asynchronous tasks * @param config configuration to use - * @param hilbertValueFunction function to compute the Hilbert value for a {@link Point} + * @param hilbertValueFunction function to compute the Hilbert value for a {@link NodeReferenceWithDistance} * @param nodeIdSupplier supplier to be invoked when new nodes are created * @param onWriteListener an on-write listener to be called after writes take place * @param onReadListener an on-read listener to be called after reads take place */ public HNSW(@Nonnull final Subspace subspace, @Nonnull final Subspace nodeSlotIndexSubspace, @Nonnull final Executor executor, @Nonnull final Config config, - @Nonnull final Function hilbertValueFunction, + @Nonnull final Function hilbertValueFunction, @Nonnull final Supplier nodeIdSupplier, @Nonnull final OnWriteListener onWriteListener, @Nonnull final OnReadListener onReadListener) { @@ -555,29 +542,29 @@ private CompletableFuture kNearestNeighborsSearch(@Nonnull final R */ @Nonnull private CompletableFuture greedySearchLayer(@Nonnull final ReadTransaction readTransaction, - @Nonnull final Element entryElement, + @Nonnull final NodeReferenceWithDistance entryNeighbor, final int layer, @Nonnull final Vector queryVector) { Verify.verify(layer > 0); final Metric metric = getConfig().getMetric(); final AtomicReference greedyStateReference = - new AtomicReference<>(new GreedyResult(layer, entryElement.getPrimaryKey(), entryElement.getDistance())); + new AtomicReference<>(new GreedyResult(layer, entryNeighbor.getPrimaryKey(), entryNeighbor.getDistance())); return AsyncUtil.whileTrue(() -> onReadListener.onAsyncRead( storageAdapter.fetchNode(IntermediateNode::creator, readTransaction, layer, greedyStateReference.get().getPrimaryKey())) - .thenApply(nodeWithLayer -> { - if (nodeWithLayer == null) { + .thenApply(node -> { + if (node == null) { throw new IllegalStateException("unable to fetch node"); } - final IntermediateNode node = nodeWithLayer.getNode().asIntermediateNode(); - final List neighbors = node.getNeighbors(); + final IntermediateNode intermediateNode = node.asIntermediateNode(); + final List neighbors = intermediateNode.getNeighbors(); final GreedyResult currentNodeKey = greedyStateReference.get(); double minDistance = currentNodeKey.getDistance(); - NeighborWithVector nearestNeighbor = null; - for (final NeighborWithVector neighbor : neighbors) { + NodeReferenceWithVector nearestNeighbor = null; + for (final NodeReferenceWithVector neighbor : neighbors) { final double distance = Vector.comparativeDistance(metric, neighbor.getVector(), queryVector); if (distance < minDistance) { @@ -603,19 +590,75 @@ private CompletableFuture greedySearchLayer(@Nonnull final ReadTra * TODO. */ @Nonnull - private CompletableFuture searchLayer(@Nonnull final ReadTransaction readTransaction, - @Nonnull final List entryElements, - final int layer, - @Nonnull final Vector queryVector) { - final Set visited = Sets.newHashSet(Element.primaryKeys(entryElements)); - final PriorityBlockingQueue candidates = - new PriorityBlockingQueue<>(entryElements.size(), - Comparator.comparing(Element::getDistance)); - candidates.addAll(entryElements); - final PriorityBlockingQueue nearestNeighbors = - new PriorityBlockingQueue<>(entryElements.size(), - Comparator.comparing(Element::getDistance).reversed()); - nearestNeighbors.addAll(entryElements); + private CompletableFuture searchLayer(@Nonnull Node.NodeCreator creator, + @Nonnull final ReadTransaction readTransaction, + @Nonnull final List entryNeighbors, + final int layer, + final int efSearch, + @Nonnull final Vector queryVector) { + final Set visited = Sets.newConcurrentHashSet(NodeReference.primaryKeys(entryNeighbors)); + final PriorityBlockingQueue candidates = + new PriorityBlockingQueue<>(entryNeighbors.size(), + Comparator.comparing(NodeReferenceWithDistance::getDistance)); + candidates.addAll(entryNeighbors); + final PriorityBlockingQueue nearestNeighbors = + new PriorityBlockingQueue<>(entryNeighbors.size(), + Comparator.comparing(NodeReferenceWithDistance::getDistance).reversed()); + nearestNeighbors.addAll(entryNeighbors); + final Map> nodeCache = Maps.newConcurrentMap(); + + final Metric metric = getConfig().getMetric(); + + return AsyncUtil.whileTrue(() -> { + if (candidates.isEmpty()) { + return false; + } + + final NodeReferenceWithDistance candidate = candidates.poll(); + final NodeReferenceWithDistance furthestNeighbor = Objects.requireNonNull(nearestNeighbors.peek()); + + if (candidate.getDistance() > furthestNeighbor.getDistance()) { + return false; + } + + final CompletableFuture> candidateNodeFuture; + final Node cachedCandidateNode = nodeCache.get(candidate.getPrimaryKey()); + if (cachedCandidateNode != null) { + candidateNodeFuture = CompletableFuture.completedFuture(cachedCandidateNode); + } else { + candidateNodeFuture = + storageAdapter.fetchNode(creator, readTransaction, layer, candidate.getPrimaryKey()) + .thenApply(candidateNode -> nodeCache.put(candidate.getPrimaryKey(), candidateNode)); + } + + candidateNodeFuture + .thenCompose(candidateNode -> fetchNodes(creator, readTransaction, + NodeReference.primaryKeys(candidateNode.getNeighbors()), layer) + .thenApply(neighbors -> { + + for (final Node neighbor : neighbors) { + + } + })) + .thenApply(neighbors -> { + for (final Node neighbor : neighbors) { + if (visited.contains(neighbor.getPrimaryKey())) { + continue; + } + visited.add(neighbor.getPrimaryKey()); + final NodeReferenceWithDistance furthestNeighbor1 = + Objects.requireNonNull(nearestNeighbors.peek()); + + final + Vector.comparativeDistance(metric, neighbor.) + + } + }) + + return AsyncUtil.whileTrue(() -> { + fetchNodes(creator, readTransaction, candidateReference.get().) + }); + }).thenApply(ignored -> null); // TODO if (candidates.isEmpty()) { return CompletableFuture.completedFuture(null); // TODO @@ -625,21 +668,23 @@ private CompletableFuture searchLayer(@Nonnull final ReadTransacti final AtomicReference greedyStateReference = new AtomicReference<>(entryStates); + fetchNodes(creator, readTransaction, ) + return AsyncUtil.whileTrue(() -> - fetchNeighborNodes(IntermediateNode::creator, readTransaction, - layer, greedyStateReference.get().getPrimaryKey()) - .thenApply(nodeWithLayer -> { + fetchNodes(IntermediateNode::creator, readTransaction, + layer, greedyStateReference.get().getPrimaryKey()) + .thenApply(nodeWithLayer -> { if (nodeWithLayer == null) { throw new IllegalStateException("unable to fetch node"); } final IntermediateNode node = nodeWithLayer.getNode().asIntermediateNode(); - final List neighbors = node.getNeighbors(); + final List neighbors = node.getNeighbors(); final GreedyResult currentNodeKey = greedyStateReference.get(); double minDistance = currentNodeKey.getDistance(); - NeighborWithVector nearestNeighbor = null; - for (final NeighborWithVector neighbor : neighbors) { + NodeReferenceWithVector nearestNeighbor = null; + for (final NodeReferenceWithVector neighbor : neighbors) { final double distance = Vector.comparativeDistance(metric, neighbor.getVector(), queryVector); if (distance < minDistance) { @@ -661,973 +706,24 @@ private CompletableFuture searchLayer(@Nonnull final ReadTransacti }), executor).thenApply(ignored -> greedyStateReference.get()); } - /** - * Returns the next left-most path from a given {@link TraversalState} to a leaf node containing items as - * a {@link TraversalState}. The term left-most used here is defined by comparing - * {@code (largestHilbertValue, largestKey)} when comparing nodes (the left one being the smaller, the right one - * being the greater). - * @param readTransaction the transaction to use - * @param traversalState traversal state to start from. The initial traversal state is always obtained by initially - * calling {@link #fetchLeftmostPathToLeaf(ReadTransaction, byte[], BigInteger, Tuple, Predicate, BiPredicate)}. - * @param mbrPredicate a predicate on an mbr {@link Rectangle}. This predicate is evaluated for each node that - * is processed. - * @return a {@link TraversalState} of the left-most path from {@code nodeId} to a {@link DataNode} whose - * {@link Node}s all pass the mbr predicate test. - */ - @Nonnull - private CompletableFuture fetchNextPathToLeaf(@Nonnull final ReadTransaction readTransaction, - @Nonnull final TraversalState traversalState, - @Nullable final BigInteger lastHilbertValue, - @Nullable final Tuple lastKey, - @Nonnull final Predicate mbrPredicate, - @Nonnull final BiPredicate suffixPredicate) { - - final List> toBeProcessed = traversalState.getToBeProcessed(); - final AtomicReference leafNode = new AtomicReference<>(null); - - return AsyncUtil.whileTrue(() -> { - final ChildSlot nextChildSlot = resolveNextIdForFetch(toBeProcessed, mbrPredicate, suffixPredicate, - onReadListener); - if (nextChildSlot == null) { - return AsyncUtil.READY_FALSE; - } - - // fetch the left-most path rooted at the current child to its left-most leaf and concatenate the paths - return fetchLeftmostPathToLeaf(readTransaction, nextChildSlot.getChildId(), lastHilbertValue, - lastKey, mbrPredicate, suffixPredicate) - .thenApply(nestedTraversalState -> { - if (nestedTraversalState.isEnd()) { - // no more data in this subtree - return true; - } - // combine the traversal states - leafNode.set(nestedTraversalState.getCurrentLeafNode()); - toBeProcessed.addAll(nestedTraversalState.getToBeProcessed()); - return false; - }); - }, executor).thenApply(v -> leafNode.get() == null - ? TraversalState.end() - : TraversalState.of(toBeProcessed, leafNode.get())); - } - - /** - * Return the next {@link ChildSlot} that needs to be processed given a list of deques that need to be processed - * as part of the current scan. - * @param toBeProcessed list of deques - * @param mbrPredicate a predicate on an mbr {@link Rectangle} - * @param suffixPredicate a predicate that is tested if applicable on the key suffix - * @return The next child slot that needs to be processed or {@code null} if there is no next child slot. - * As a side effect of calling this method the child slot is removed from {@code toBeProcessed}. - */ - @Nullable - @SuppressWarnings("PMD.AvoidBranchingStatementAsLastInLoop") - private static ChildSlot resolveNextIdForFetch(@Nonnull final List> toBeProcessed, - @Nonnull final Predicate mbrPredicate, - @Nonnull final BiPredicate suffixPredicate, - @Nonnull final OnReadListener onReadListener) { - for (int level = toBeProcessed.size() - 1; level >= 0; level--) { - final Deque toBeProcessedThisLevel = toBeProcessed.get(level); - - while (!toBeProcessedThisLevel.isEmpty()) { - final ChildSlot childSlot = toBeProcessedThisLevel.pollFirst(); - if (!mbrPredicate.test(childSlot.getMbr())) { - onReadListener.onChildNodeDiscard(childSlot); - continue; - } - if (childSlot.suffixPredicateCanBeApplied()) { - if (!suffixPredicate.test(childSlot.getSmallestKeySuffix(), - childSlot.getLargestKeySuffix())) { - onReadListener.onChildNodeDiscard(childSlot); - continue; - } - } - toBeProcessed.subList(level + 1, toBeProcessed.size()).clear(); - return childSlot; - } - } - return null; - } - - // - // Insert/Update path - // - - - /** - * Method to insert an object/item into the R-tree. The item is treated unique per its point in space as well as its - * additional key that is also passed in. The Hilbert value of the point is passed in as to allow the caller to - * compute Hilbert values themselves. Note that there is a bijective mapping between point and Hilbert - * value which allows us to recompute point from Hilbert value as well as Hilbert value from point. We currently - * treat point and Hilbert value independent, however, they are redundant and not independent at all. The implication - * is that we do not have to store both point and Hilbert value (but we currently do). - * @param tc transaction context - * @param point the point to be used in space - * @param keySuffix the additional key to be stored with the item - * @param value the additional value to be stored with the item - * @return a completable future that completes when the insert is completed - */ - @Nonnull - public CompletableFuture insertOrUpdate(@Nonnull final TransactionContext tc, - @Nonnull final Point point, - @Nonnull final Tuple keySuffix, - @Nonnull final Tuple value) { - final BigInteger hilbertValue = hilbertValueFunction.apply(point); - final Tuple itemKey = Tuple.from(point.getCoordinates(), keySuffix); - - // - // Get to the leaf node we need to start the insert from and then call the appropriate method to perform - // the actual insert/update. - // - return tc.runAsync(transaction -> fetchPathForModification(transaction, hilbertValue, itemKey, true) - .thenCompose(leafNode -> { - if (leafNode == null) { - leafNode = new DataNode(rootId, Lists.newArrayList()); - } - return insertOrUpdateSlot(transaction, leafNode, point, hilbertValue, itemKey, value); - })); - } - - /** - * Inserts a new slot into the {@link DataNode} passed in or updates an existing slot of the {@link DataNode} passed - * in. - * @param transaction transaction - * @param targetNode leaf node that is the target of this insert or update - * @param point the point to be used in space - * @param hilbertValue the hilbert value of the point - * @param key the additional key to be stored with the item - * @param value the additional value to be stored with the item - * @return a completable future that completes when the insert/update is completed - */ - @Nonnull - private CompletableFuture insertOrUpdateSlot(@Nonnull final Transaction transaction, - @Nonnull final DataNode targetNode, - @Nonnull final Point point, - @Nonnull final BigInteger hilbertValue, - @Nonnull final Tuple key, - @Nonnull final Tuple value) { - Verify.verify(targetNode.size() <= config.getMaxM()); - - final AtomicInteger level = new AtomicInteger(0); - final ItemSlot newSlot = new ItemSlot(hilbertValue, point, key, value); - final AtomicInteger insertSlotIndex = new AtomicInteger(findInsertUpdateItemSlotIndex(targetNode, hilbertValue, key)); - if (insertSlotIndex.get() < 0) { - // just update the slot with the potentially new value - storageAdapter.writeLeafNodeSlot(transaction, targetNode, newSlot); - return AsyncUtil.DONE; - } - - // - // This is an insert. - // - - final AtomicReference currentNode = new AtomicReference<>(targetNode); - final AtomicReference parentSlot = new AtomicReference<>(newSlot); - - // - // Inch our way upwards in the tree to perform the necessary adjustments. What needs to be done next - // is informed by the result of the current operation: - // 1. A split happened; we need to insert a new slot into the parent node -- prime current node and - // current slot and continue. - // 2. The slot was inserted but mbrs, largest Hilbert Values and largest Keys need to be adjusted upwards. - // 3. We are done as no further adjustments are necessary. - // - return AsyncUtil.whileTrue(() -> { - final NodeSlot currentNewSlot = parentSlot.get(); - - if (currentNewSlot != null) { - return insertSlotIntoTargetNode(transaction, level.get(), hilbertValue, key, currentNode.get(), currentNewSlot, insertSlotIndex.get()) - .thenApply(nodeOrAdjust -> { - if (currentNode.get().isRoot()) { - return false; - } - currentNode.set(currentNode.get().getParentNode()); - parentSlot.set(nodeOrAdjust.getSlotInParent()); - insertSlotIndex.set(nodeOrAdjust.getSplitNode() == null ? -1 : nodeOrAdjust.getSplitNode().getSlotIndexInParent()); - level.incrementAndGet(); - return nodeOrAdjust.getSplitNode() != null || nodeOrAdjust.parentNeedsAdjustment(); - }); - } else { - // adjustment only - return updateSlotsAndAdjustNode(transaction, level.get(), hilbertValue, key, currentNode.get(), true) - .thenApply(nodeOrAdjust -> { - Verify.verify(nodeOrAdjust.getSlotInParent() == null); - if (currentNode.get().isRoot()) { - return false; - } - currentNode.set(currentNode.get().getParentNode()); - level.incrementAndGet(); - return nodeOrAdjust.parentNeedsAdjustment(); - }); - } - }, executor); - } - - /** - * Insert a new slot into the target node passed in. - * @param transaction transaction - * @param level the current level of target node, {@code 0} indicating the leaf level - * @param hilbertValue the Hilbert Value of the record that is being inserted - * @param key the key of the record that is being inserted - * @param targetNode target node - * @param newSlot new slot - * @param slotIndexInTargetNode The index of the new slot that we should use when inserting the new slot. While - * this information can be computed from the other arguments passed in, the caller already knows this - * information; we can avoid searching for the proper spot on our own. - * @return a completable future that when completed indicates what needs to be done next (see {@link NodeOrAdjust}). - */ - @Nonnull - private CompletableFuture insertSlotIntoTargetNode(@Nonnull final Transaction transaction, - final int level, - @Nonnull final BigInteger hilbertValue, - @Nonnull final Tuple key, - @Nonnull final Node targetNode, - @Nonnull final NodeSlot newSlot, - final int slotIndexInTargetNode) { - if (targetNode.size() < config.getMaxM()) { - // enough space left in target - - if (logger.isTraceEnabled()) { - logger.trace("regular insert without splitting; node={}; size={}", targetNode, targetNode.size()); - } - targetNode.insertSlot(storageAdapter, level - 1, slotIndexInTargetNode, newSlot); - - if (targetNode.getKind() == NodeKind.INTERMEDIATE) { - // - // If this is an insert for an intermediate node, the child node referred to by newSlot - // is a split node from a lower level meaning a split has happened on a lower level and the - // participating siblings of that split have potentially changed. - // - storageAdapter.writeNodes(transaction, Collections.singletonList(targetNode)); - } else { - // if this is an insert for a leaf node we can just write the slot - Verify.verify(targetNode.getKind() == NodeKind.LEAF); - storageAdapter.writeLeafNodeSlot(transaction, (DataNode)targetNode, (ItemSlot)newSlot); - } - - // node has left some space -- indicate that we are done splitting at the current node - if (!targetNode.isRoot()) { - return fetchParentNodeIfNecessary(transaction, targetNode, level, hilbertValue, key, true) - .thenApply(ignored -> adjustSlotInParent(targetNode, level) - ? NodeOrAdjust.ADJUST - : NodeOrAdjust.NONE); - } - - // no split and no adjustment - return CompletableFuture.completedFuture(NodeOrAdjust.NONE); - } else { - // - // If this is the root we need to grow the tree taller by splitting the root to get a new root - // with two children each containing half of the slots previously contained by the old root node. - // - if (targetNode.isRoot()) { - if (logger.isTraceEnabled()) { - logger.trace("splitting root node; size={}", targetNode.size()); - } - // temporarily overfill the old root node - targetNode.insertSlot(storageAdapter, level - 1, slotIndexInTargetNode, newSlot); - - splitRootNode(transaction, level, targetNode); - return CompletableFuture.completedFuture(NodeOrAdjust.NONE); - } - - // - // Node is full -- borrow some space from the siblings if possible. The paper does overflow handling and - // node splitting separately -- we do it in one path. - // - final CompletableFuture> siblings = - fetchParentNodeIfNecessary(transaction, targetNode, level, hilbertValue, key, true) - .thenCompose(ignored -> - fetchSiblings(transaction, targetNode)); - - return siblings.thenApply(siblingNodes -> { - int numSlots = - Math.toIntExact(siblingNodes - .stream() - .mapToLong(Node::size) - .sum()); - - // First determine if we actually need to split; create the split node if we do; for the remainder of - // this method splitNode != null <=> we are splitting; otherwise we handle overflow. - final Node splitNode; - final List newSiblingNodes; - if (numSlots == siblingNodes.size() * config.getMaxM()) { - if (logger.isTraceEnabled()) { - logger.trace("splitting node; node={}, siblings={}", - targetNode, - siblingNodes.stream().map(Node::toString) - .collect(Collectors.joining(","))); - } - splitNode = targetNode.newOfSameKind(nodeIdSupplier.get()); - // link this split node to become the last node of the siblings - splitNode.linkToParent(Objects.requireNonNull(targetNode.getParentNode()), - siblingNodes.get(siblingNodes.size() - 1).getSlotIndexInParent() + 1); - newSiblingNodes = Lists.newArrayList(siblingNodes); - newSiblingNodes.add(splitNode); - } else { - if (logger.isTraceEnabled()) { - logger.trace("handling overflow; node={}, numSlots={}, siblings={}", - targetNode, - numSlots, - siblingNodes.stream().map(Node::toString) - .collect(Collectors.joining(","))); - } - splitNode = null; - newSiblingNodes = siblingNodes; - } - - // temporarily overfill targetNode - numSlots++; - targetNode.insertSlot(storageAdapter, level - 1, slotIndexInTargetNode, newSlot); - - // sibling nodes are in hilbert value order - final Iterator slotIterator = - siblingNodes - .stream() - .flatMap(Node::slotsStream) - .iterator(); - - // - // Distribute all slots (including the new one which is now at its correct position among its brethren) - // across all siblings (which includes the targetNode and (if we are splitting) the splitNode). - // At the end of this modification all siblings have and (almost) equal count of slots that is - // guaranteed to be between minM and maxM. - // - - final int base = numSlots / newSiblingNodes.size(); - int rest = numSlots % newSiblingNodes.size(); - - List> newNodeSlotLists = Lists.newArrayList(); - List currentNodeSlots = Lists.newArrayList(); - while (slotIterator.hasNext()) { - final NodeSlot slot = slotIterator.next(); - currentNodeSlots.add(slot); - if (currentNodeSlots.size() == base + (rest > 0 ? 1 : 0)) { - if (rest > 0) { - // one fewer to distribute - rest--; - } - - newNodeSlotLists.add(currentNodeSlots); - currentNodeSlots = Lists.newArrayList(); - } - } - - Verify.verify(newSiblingNodes.size() == newNodeSlotLists.size()); - - final Iterator newSiblingNodesIterator = newSiblingNodes.iterator(); - final Iterator> newNodeSlotsIterator = newNodeSlotLists.iterator(); - - // assign slots to nodes - while (newSiblingNodesIterator.hasNext()) { - final Node newSiblingNode = newSiblingNodesIterator.next(); - Verify.verify(newNodeSlotsIterator.hasNext()); - final List newNodeSlots = newNodeSlotsIterator.next(); - newSiblingNode.moveOutAllSlots(storageAdapter); - newSiblingNode.moveInSlots(storageAdapter, newNodeSlots); - } - - // update nodes - storageAdapter.writeNodes(transaction, newSiblingNodes); - - // - // Adjust the parent's slot information in memory only; we'll write it in the next iteration when - // we go one level up. - // - for (final Node siblingNode : siblingNodes) { - adjustSlotInParent(siblingNode, level); - } - - if (splitNode == null) { - // didn't split -- just continue adjusting - return NodeOrAdjust.ADJUST; - } - - // - // Manufacture a new slot for the splitNode; the caller will then use that slot to insert it into the - // parent. - // - final NodeSlot firstSlotOfSplitNode = splitNode.getSlot(0); - final NodeSlot lastSlotOfSplitNode = splitNode.getSlot(splitNode.size() - 1); - return new NodeOrAdjust( - new ChildSlot(firstSlotOfSplitNode.getSmallestHilbertValue(), firstSlotOfSplitNode.getSmallestKey(), - lastSlotOfSplitNode.getLargestHilbertValue(), lastSlotOfSplitNode.getLargestKey(), - splitNode.getId(), NodeHelpers.computeMbr(splitNode.getSlots())), - splitNode, true); - }); - } - } - - /** - * Split the root node. This method first creates two nodes {@code left} and {@code right}. The root node, - * whose ID is always a string of {@code 0x00}, contains some number {@code n} of slots. {@code n / 2} slots of those - * {@code n} slots are moved to {@code left}, the rest to {@code right}. The root node is then updated to have two - * children: {@code left} and {@code right}. All three nodes are then updated in the database. - * @param transaction transaction to use - * @param level the level counting starting at {@code 0} indicating the leaf level increasing upwards - * @param oldRootNode the old root node - */ - private void splitRootNode(@Nonnull final Transaction transaction, - final int level, - @Nonnull final Node oldRootNode) { - final Node leftNode = oldRootNode.newOfSameKind(nodeIdSupplier.get()); - final Node rightNode = oldRootNode.newOfSameKind(nodeIdSupplier.get()); - final int leftSize = oldRootNode.size() / 2; - final List leftSlots = ImmutableList.copyOf(oldRootNode.getSlots(0, leftSize)); - leftNode.moveInSlots(storageAdapter, leftSlots); - final int rightSize = oldRootNode.size() - leftSize; - final List rightSlots = ImmutableList.copyOf(oldRootNode.getSlots(leftSize, leftSize + rightSize)); - rightNode.moveInSlots(storageAdapter, rightSlots); - - final NodeSlot firstSlotOfLeftNode = leftSlots.get(0); - final NodeSlot lastSlotOfLeftNode = leftSlots.get(leftSlots.size() - 1); - final NodeSlot firstSlotOfRightNode = rightSlots.get(0); - final NodeSlot lastSlotOfRightNode = rightSlots.get(rightSlots.size() - 1); - - final ChildSlot leftChildSlot = new ChildSlot(firstSlotOfLeftNode.getSmallestHilbertValue(), firstSlotOfLeftNode.getSmallestKey(), - lastSlotOfLeftNode.getLargestHilbertValue(), lastSlotOfLeftNode.getLargestKey(), - leftNode.getId(), NodeHelpers.computeMbr(leftNode.getSlots())); - final ChildSlot rightChildSlot = new ChildSlot(firstSlotOfRightNode.getSmallestHilbertValue(), firstSlotOfRightNode.getSmallestKey(), - lastSlotOfRightNode.getLargestHilbertValue(), lastSlotOfRightNode.getLargestKey(), - rightNode.getId(), NodeHelpers.computeMbr(rightNode.getSlots())); - - oldRootNode.moveOutAllSlots(storageAdapter); - final IntermediateNode newRootNode = new IntermediateNode(rootId) - .insertSlot(storageAdapter, level, 0, leftChildSlot) - .insertSlot(storageAdapter, level, 1, rightChildSlot); - - storageAdapter.writeNodes(transaction, Lists.newArrayList(oldRootNode, newRootNode, leftNode, rightNode)); - } - - // Delete Path - - /** - * Method to delete from the R-tree. The item is treated unique per its point in space as well as its - * additional key that is passed in. - * @param tc transaction context - * @param point the point - * @param keySuffix the additional key to be stored with the item - * @return a completable future that completes when the delete operation is completed - */ - @Nonnull - public CompletableFuture delete(@Nonnull final TransactionContext tc, - @Nonnull final Point point, - @Nonnull final Tuple keySuffix) { - final BigInteger hilbertValue = hilbertValueFunction.apply(point); - final Tuple itemKey = Tuple.from(point.getCoordinates(), keySuffix); - - // - // Get to the leaf node we need to start the delete operation from and then call the appropriate method to - // perform the actual delete. - // - return tc.runAsync(transaction -> fetchPathForModification(transaction, hilbertValue, itemKey, false) - .thenCompose(leafNode -> { - if (leafNode == null) { - return AsyncUtil.DONE; - } - return deleteSlotIfExists(transaction, leafNode, hilbertValue, itemKey); - })); - } - - /** - * Deletes a slot from the {@link DataNode} passed or exits if the slot could not be found in the target node. - * in. - * @param transaction transaction - * @param targetNode leaf node that is the target of this delete operation - * @param hilbertValue the hilbert value of the point - * @param key the additional key to be stored with the item - * @return a completable future that completes when the delete is completed - */ - @Nonnull - private CompletableFuture deleteSlotIfExists(@Nonnull final Transaction transaction, - @Nonnull final DataNode targetNode, - @Nonnull final BigInteger hilbertValue, - @Nonnull final Tuple key) { - Verify.verify(targetNode.size() <= config.getMaxM()); - - final AtomicInteger level = new AtomicInteger(0); - final AtomicInteger deleteSlotIndex = new AtomicInteger(findDeleteItemSlotIndex(targetNode, hilbertValue, key)); - if (deleteSlotIndex.get() < 0) { - // - // The slot was not found meaning that the item was not found and that means we don't have to do anything - // here. - // - return AsyncUtil.DONE; - } - - // - // We found the slot and therefore the item. - // - - final NodeSlot deleteSlot = targetNode.getSlot(deleteSlotIndex.get()); - final AtomicReference currentNode = new AtomicReference<>(targetNode); - final AtomicReference parentSlot = new AtomicReference<>(deleteSlot); - - // - // Inch our way upwards in the tree to perform the necessary adjustments. What needs to be done next - // is informed by the result of the current operation: - // 1. A fuse happened; we need to delete an existing slot from the parent node -- prime current node and - // current slot and continue. - // 2. The slot was deleted but mbrs, largest Hilbert Values and largest Keys need to be adjusted upwards. - // 3. We are done as no further adjustments are necessary. - // - return AsyncUtil.whileTrue(() -> { - final NodeSlot currentDeleteSlot = parentSlot.get(); - - if (currentDeleteSlot != null) { - return deleteSlotFromTargetNode(transaction, level.get(), hilbertValue, key, currentNode.get(), currentDeleteSlot, deleteSlotIndex.get()) - .thenApply(nodeOrAdjust -> { - if (currentNode.get().isRoot()) { - return false; - } - currentNode.set(currentNode.get().getParentNode()); - parentSlot.set(nodeOrAdjust.getSlotInParent()); - deleteSlotIndex.set(nodeOrAdjust.getTombstoneNode() == null ? -1 : nodeOrAdjust.getTombstoneNode().getSlotIndexInParent()); - level.incrementAndGet(); - return nodeOrAdjust.getTombstoneNode() != null || nodeOrAdjust.parentNeedsAdjustment(); - }); - } else { - // adjustment only - return updateSlotsAndAdjustNode(transaction, level.get(), hilbertValue, key, currentNode.get(), false) - .thenApply(nodeOrAdjust -> { - Verify.verify(nodeOrAdjust.getSlotInParent() == null); - if (currentNode.get().isRoot()) { - return false; - } - currentNode.set(currentNode.get().getParentNode()); - level.incrementAndGet(); - return nodeOrAdjust.parentNeedsAdjustment(); - }); - } - }, executor); - } - - /** - * Delete and existing slot from the target node passed in. - * @param transaction transaction - * @param level the current level of target node, {@code 0} indicating the leaf level - * @param hilbertValue the Hilbert Value of the record that is being deleted - * @param key the key of the record that is being deleted - * @param targetNode target node - * @param deleteSlot existing slot that is to be deleted - * @param slotIndexInTargetNode The index of the new slot that we should use when inserting the new slot. While - * this information can be computed from the other arguments passed in, the caller already knows this - * information; we can avoid searching for the proper spot on our own. - * @return a completable future that when completed indicates what needs to be done next (see {@link NodeOrAdjust}). - */ - @Nonnull - private CompletableFuture deleteSlotFromTargetNode(@Nonnull final Transaction transaction, - final int level, - final BigInteger hilbertValue, - final Tuple key, - @Nonnull final Node targetNode, - @Nonnull final NodeSlot deleteSlot, - final int slotIndexInTargetNode) { - // - // We need to keep the number of slots per node between minM <= size() <= maxM unless this is the root node. - // - if (targetNode.isRoot() || targetNode.size() > config.getMinM()) { - if (logger.isTraceEnabled()) { - logger.trace("regular delete; node={}; size={}", targetNode, targetNode.size()); - } - targetNode.deleteSlot(storageAdapter, level - 1, slotIndexInTargetNode); - - if (targetNode.getKind() == NodeKind.INTERMEDIATE) { - // - // If this node is the root and the root node is an intermediate node, then it should at least have two - // children. - // - Verify.verify(!targetNode.isRoot() || targetNode.size() >= 2); - // - // If this is a delete operation within an intermediate node, the slot being deleted results from a - // fuse operation meaning a fuse has occurred on a lower level and the participating siblings of that split have - // potentially changed. - // - storageAdapter.writeNodes(transaction, Collections.singletonList(targetNode)); - } else { - Verify.verify(targetNode.getKind() == NodeKind.LEAF); - storageAdapter.clearLeafNodeSlot(transaction, (DataNode)targetNode, (ItemSlot)deleteSlot); - } - - // node is not under-flowing -- indicate that we are done fusing at the current node - if (!targetNode.isRoot()) { - return fetchParentNodeIfNecessary(transaction, targetNode, level, hilbertValue, key, false) - .thenApply(ignored -> adjustSlotInParent(targetNode, level) - ? NodeOrAdjust.ADJUST - : NodeOrAdjust.NONE); - } - - // no fuse and no adjustment - return CompletableFuture.completedFuture(NodeOrAdjust.NONE); // no fuse and no adjustment - } else { - // - // Node is under min-capacity -- borrow some children/items from the siblings if possible. - // - final CompletableFuture> siblings = - fetchParentNodeIfNecessary(transaction, targetNode, level, hilbertValue, key, false) - .thenCompose(ignored -> fetchSiblings(transaction, targetNode)); - - return siblings.thenApply(siblingNodes -> { - int numSlots = - Math.toIntExact(siblingNodes - .stream() - .mapToLong(Node::size) - .sum()); - - final Node tombstoneNode; - final List newSiblingNodes; - if (numSlots == siblingNodes.size() * config.getMinM()) { - if (logger.isTraceEnabled()) { - logger.trace("fusing nodes; node={}, siblings={}", - targetNode, - siblingNodes.stream().map(Node::toString).collect(Collectors.joining(","))); - } - tombstoneNode = siblingNodes.get(siblingNodes.size() - 1); - newSiblingNodes = siblingNodes.subList(0, siblingNodes.size() - 1); - } else { - if (logger.isTraceEnabled()) { - logger.trace("handling underflow; node={}, numSlots={}, siblings={}", - targetNode, - numSlots, - siblingNodes.stream().map(Node::toString).collect(Collectors.joining(","))); - } - tombstoneNode = null; - newSiblingNodes = siblingNodes; - } - - // temporarily underfill targetNode - numSlots--; - targetNode.deleteSlot(storageAdapter, level - 1, slotIndexInTargetNode); - - // sibling nodes are in hilbert value order - final Iterator slotIterator = - siblingNodes - .stream() - .flatMap(Node::slotsStream) - .iterator(); - - // - // Distribute all slots (excluding the one we want to delete) across all siblings (which also excludes - // the targetNode and (if we are fusing) the tombstoneNode). - // At the end of this modification all siblings have and (almost) equal count of slots that is - // guaranteed to be between minM and maxM. - // - - final int base = numSlots / newSiblingNodes.size(); - int rest = numSlots % newSiblingNodes.size(); - - List> newNodeSlotLists = Lists.newArrayList(); - List currentNodeSlots = Lists.newArrayList(); - while (slotIterator.hasNext()) { - final NodeSlot slot = slotIterator.next(); - currentNodeSlots.add(slot); - if (currentNodeSlots.size() == base + (rest > 0 ? 1 : 0)) { - if (rest > 0) { - // one fewer to distribute - rest--; - } - - newNodeSlotLists.add(currentNodeSlots); - currentNodeSlots = Lists.newArrayList(); - } - } - - Verify.verify(newSiblingNodes.size() == newNodeSlotLists.size()); - - if (tombstoneNode != null) { - // remove the slots for the tombstone node and update - tombstoneNode.moveOutAllSlots(storageAdapter); - storageAdapter.writeNodes(transaction, Collections.singletonList(tombstoneNode)); - } - - final Iterator newSiblingNodesIterator = newSiblingNodes.iterator(); - final Iterator> newNodeSlotsIterator = newNodeSlotLists.iterator(); - - // assign the slots to the appropriate nodes - while (newSiblingNodesIterator.hasNext()) { - final Node newSiblingNode = newSiblingNodesIterator.next(); - Verify.verify(newNodeSlotsIterator.hasNext()); - final List newNodeSlots = newNodeSlotsIterator.next(); - newSiblingNode.moveOutAllSlots(storageAdapter); - newSiblingNode.moveInSlots(storageAdapter, newNodeSlots); - } - - final IntermediateNode parentNode = Objects.requireNonNull(targetNode.getParentNode()); - if (parentNode.isRoot() && parentNode.size() == 2 && tombstoneNode != null) { - // - // The parent node (root) would only have one child after this delete. - // We shrink the tree by removing the root and making the last remaining sibling the root. - // - final Node toBePromotedNode = Iterables.getOnlyElement(newSiblingNodes); - promoteNodeToRoot(transaction, level, parentNode, toBePromotedNode); - return NodeOrAdjust.NONE; - } - - storageAdapter.writeNodes(transaction, newSiblingNodes); - - for (final Node newSiblingNode : newSiblingNodes) { - adjustSlotInParent(newSiblingNode, level); - } - - if (tombstoneNode == null) { - // - // We only handled underfill (and didn't need to fuse) but still need to continue adjusting - // mbrs, largest Hilbert values, and largest keys upward the tree. - // - return NodeOrAdjust.ADJUST; - } - - // - // We need to signal that the current operation ended in a fuse, and we need to delete the slot for - // the tombstoneNode one level higher. - // - return new NodeOrAdjust(parentNode.getSlot(tombstoneNode.getSlotIndexInParent()), - tombstoneNode, true); - }); - } - } - - /** - * Promote the given node to become the new root node. The node that is passed only changes its node id but retains - * all of it slots. This operation is the opposite of {@link #splitRootNode(Transaction, int, Node)} which can be - * invoked by the insert code path. - * @param transaction transaction - * @param level the level counting starting at {@code 0} indicating the leaf level increasing upwards - * @param oldRootNode the old root node - * @param toBePromotedNode node to be promoted. - */ - private void promoteNodeToRoot(final @Nonnull Transaction transaction, final int level, final IntermediateNode oldRootNode, - final Node toBePromotedNode) { - oldRootNode.deleteAllSlots(storageAdapter, level); - - // hold on to the slots of the to-be-promoted node -- copy them as moveOutAllSlots() will mutate the slot list - final List newRootSlots = ImmutableList.copyOf(toBePromotedNode.getSlots()); - toBePromotedNode.moveOutAllSlots(storageAdapter); - final Node newRootNode = toBePromotedNode.newOfSameKind(rootId).moveInSlots(storageAdapter, newRootSlots); - - // We need to update the node and the new root node in order to clear out the existing slots of the pre-promoted - // node. - storageAdapter.writeNodes(transaction, ImmutableList.of(oldRootNode, newRootNode, toBePromotedNode)); - } - - // - // Helper methods that may be called from more than one code path. - // - - /** - * Updates (persists) the slots for a target node and then computes the necessary adjustments in its parent - * node (without persisting those). - * @param transaction the transaction to use - * @param level the current level of target node, {@code 0} indicating the leaf level - * @param targetNode the target node - * @return A future containing either {@link NodeOrAdjust#NONE} if no further adjustments need to be persisted or - * {@link NodeOrAdjust#ADJUST} if the slots of the parent node of the target node need to be adjusted as - * well. - */ - @Nonnull - private CompletableFuture updateSlotsAndAdjustNode(@Nonnull final Transaction transaction, - final int level, - @Nonnull final BigInteger hilbertValue, - @Nonnull final Tuple key, - @Nonnull final Node targetNode, - final boolean isInsertUpdate) { - storageAdapter.writeNodes(transaction, Collections.singletonList(targetNode)); - - if (targetNode.isRoot()) { - return CompletableFuture.completedFuture(NodeOrAdjust.NONE); - } - - return fetchParentNodeIfNecessary(transaction, targetNode, level, hilbertValue, key, isInsertUpdate) - .thenApply(ignored -> adjustSlotInParent(targetNode, level) - ? NodeOrAdjust.ADJUST - : NodeOrAdjust.NONE); - } - - /** - * Updates the target node's mbr, largest Hilbert value as well its largest key in the target node's parent slot. - * @param targetNode target node - * @return {@code true} if any attributes of the target slot were modified, {@code false} otherwise. This will - * inform the caller if modifications need to be persisted and/or if the parent node itseld=f needs to be - * adjusted as well. - */ - private boolean adjustSlotInParent(@Nonnull final Node targetNode, final int level) { - Preconditions.checkArgument(!targetNode.isRoot()); - boolean slotHasChanged; - final IntermediateNode parentNode = Objects.requireNonNull(targetNode.getParentNode()); - final int slotIndexInParent = targetNode.getSlotIndexInParent(); - final ChildSlot childSlot = parentNode.getSlot(slotIndexInParent); - final Rectangle newMbr = NodeHelpers.computeMbr(targetNode.getSlots()); - slotHasChanged = !childSlot.getMbr().equals(newMbr); - final NodeSlot firstSlotOfTargetNode = targetNode.getSlot(0); - slotHasChanged |= !childSlot.getSmallestHilbertValue().equals(firstSlotOfTargetNode.getSmallestHilbertValue()); - slotHasChanged |= !childSlot.getSmallestKey().equals(firstSlotOfTargetNode.getSmallestKey()); - final NodeSlot lastSlotOfTargetNode = targetNode.getSlot(targetNode.size() - 1); - slotHasChanged |= !childSlot.getLargestHilbertValue().equals(lastSlotOfTargetNode.getLargestHilbertValue()); - slotHasChanged |= !childSlot.getLargestKey().equals(lastSlotOfTargetNode.getLargestKey()); - - if (slotHasChanged) { - parentNode.updateSlot(storageAdapter, level, slotIndexInParent, - new ChildSlot(firstSlotOfTargetNode.getSmallestHilbertValue(), firstSlotOfTargetNode.getSmallestKey(), - lastSlotOfTargetNode.getLargestHilbertValue(), lastSlotOfTargetNode.getLargestKey(), childSlot.getChildId(), - newMbr)); - } - return slotHasChanged; - } - - @Nonnull - private CompletableFuture fetchPathForModification(@Nonnull final Transaction transaction, - @Nonnull final BigInteger hilbertValue, - @Nonnull final Tuple key, - final boolean isInsertUpdate) { - if (config.isUseNodeSlotIndex()) { - return scanIndexAndFetchLeafNode(transaction, hilbertValue, key, isInsertUpdate); - } else { - return fetchUpdatePathToLeaf(transaction, hilbertValue, key, isInsertUpdate); - } - } - - @Nonnull - private CompletableFuture scanIndexAndFetchLeafNode(@Nonnull final ReadTransaction transaction, - @Nonnull final BigInteger hilbertValue, - @Nonnull final Tuple key, - final boolean isInsertUpdate) { - return storageAdapter.scanNodeIndexAndFetchNode(transaction, 0, hilbertValue, key, isInsertUpdate) - .thenApply(node -> { - Verify.verify(node == null || - (node.getKind() == NodeKind.LEAF && node instanceof DataNode)); - return (DataNode)node; - }); - } - - @Nonnull - private CompletableFuture scanIndexAndFetchIntermediateNode(@Nonnull final ReadTransaction transaction, - final int level, - @Nonnull final BigInteger hilbertValue, - @Nonnull final Tuple key, - final boolean isInsertUpdate) { - Verify.verify(level > 0); - return storageAdapter.scanNodeIndexAndFetchNode(transaction, level, hilbertValue, key, isInsertUpdate) - .thenApply(node -> { - // - // Note that there is no non-error scenario where node can be null here; either the node is - // not in the node slot index but is the root node which has already been resolved and fetched OR - // this node is a legitimate parent node of a node we know must exist as level > 0. If node were - // null here, it would mean that there is a node that is not the root but its parent is not in - // the R-tree. - // - Verify.verify(node.getKind() == NodeKind.INTERMEDIATE && node instanceof IntermediateNode); - return (IntermediateNode)node; - }); - } - - @Nonnull - private CompletableFuture fetchParentNodeIfNecessary(@Nonnull final ReadTransaction transaction, - @Nonnull final Node node, - final int level, - @Nonnull final BigInteger hilbertValue, - @Nonnull final Tuple key, - final boolean isInsertUpdate) { - Verify.verify(!node.isRoot()); - final IntermediateNode linkedParentNode = node.getParentNode(); - if (linkedParentNode != null) { - return CompletableFuture.completedFuture(linkedParentNode); - } - - Verify.verify(getConfig().isUseNodeSlotIndex()); - return scanIndexAndFetchIntermediateNode(transaction, level + 1, hilbertValue, key, isInsertUpdate) - .thenApply(parentNode -> { - final int slotIndexInParent = findChildSlotIndex(parentNode, node.getId()); - Verify.verify(slotIndexInParent >= 0); - node.linkToParent(parentNode, slotIndexInParent); - return parentNode; - }); - } - - /** - * Method to fetch the update path of a given {@code (hilbertValue, key)} pair. The update path is a {@link DataNode} - * and all its parent nodes to the root node. The caller can invoke {@link Node#getParentNode()} to navigate to - * all nodes in the update path starting from the {@link DataNode} that is returned. The {@link DataNode} that is - * returned may or may not already contain a slot for the {@code (hilbertValue, key)} pair passed in. This logic is - * invoked for insert, updates, as well as delete operations. If it is used for insert and the item is not yet - * part of the leaf node, the leaf node that is returned can be understood as the correct place to insert the item - * in question. - * @param transaction the transaction to use - * @param hilbertValue the Hilbert value to look for - * @param key the key to look for - * @param isInsertUpdate is this call part of and index/update operation or a delete operation - * @return A completable future containing a {@link DataNode} and by extension (through {@link Node#getParentNode()}) - * all intermediate nodes up to the root node that may get affected by an insert, update, or delete - * of the specified item. - */ - @Nonnull - private CompletableFuture fetchUpdatePathToLeaf(@Nonnull final Transaction transaction, - @Nonnull final BigInteger hilbertValue, - @Nonnull final Tuple key, - final boolean isInsertUpdate) { - final AtomicReference parentNode = new AtomicReference<>(null); - final AtomicInteger slotInParent = new AtomicInteger(-1); - final AtomicReference currentId = new AtomicReference<>(rootId); - final AtomicReference leafNode = new AtomicReference<>(null); - return AsyncUtil.whileTrue(() -> storageAdapter.fetchNode(transaction, currentId.get()) - .thenApply(node -> { - if (node == null) { - if (Arrays.equals(currentId.get(), rootId)) { - Verify.verify(leafNode.get() == null); - return false; - } - throw new IllegalStateException("unable to fetch node for insert or update"); - } - if (parentNode.get() != null) { - node.linkToParent(parentNode.get(), slotInParent.get()); - } - if (node.getKind() == NodeKind.INTERMEDIATE) { - final IntermediateNode intermediateNode = (IntermediateNode)node; - final int slotIndex = findChildSlotIndex(intermediateNode, hilbertValue, key, isInsertUpdate); - if (slotIndex < 0) { - Verify.verify(!isInsertUpdate); - // - // This is for a delete operation and we were unable to find a child that covers - // the Hilbert Value/key to be deleted - return false; - } - - parentNode.set(intermediateNode); - slotInParent.set(slotIndex); - final ChildSlot childSlot = intermediateNode.getSlot(slotIndex); - currentId.set(childSlot.getChildId()); - return true; - } else { - leafNode.set((DataNode)node); - return false; - } - }), executor) - .thenApply(ignored -> { - final DataNode node = leafNode.get(); - if (logger.isTraceEnabled()) { - logger.trace("update path; path={}", NodeHelpers.nodeIdPath(node)); - } - return node; - }); - } - /** * TODO. */ @Nonnull @SuppressWarnings("unchecked") - private CompletableFuture>> fetchNeighborNodes(@Nonnull final Transaction transaction, - @Nonnull final Node node, - final int layer) { + private CompletableFuture>> fetchNodes(@Nonnull final Node.NodeCreator creator, + @Nonnull final ReadTransaction readTransaction, + @Nonnull final Iterable primaryKeys, + final int layer) { // this deque is only modified by once upon creation final ArrayDeque toBeProcessed = new ArrayDeque<>(); + for (final var primaryKey : primaryKeys) { + toBeProcessed.addLast(primaryKey); + } final List> working = Lists.newArrayList(); - final List neighbors = node.getNeighbors(); final AtomicInteger neighborIndex = new AtomicInteger(0); - final NodeWithLayer[] neighborNodeArray = - (NodeWithLayer[])Array.newInstance(NodeWithLayer.class, neighbors.size()); - - for (final N neighbor : neighbors) { - toBeProcessed.addLast(neighbor.getPrimaryKey()); - } + final Node[] neighborNodeArray = + (Node[])Array.newInstance(Node.class, toBeProcessed.size()); // Fetch all sibling nodes (in parallel if possible). return AsyncUtil.whileTrue(() -> { @@ -1641,7 +737,7 @@ private CompletableFuture>> fetchNeig final int index = neighborIndex.getAndIncrement(); - working.add(onReadListener.onAsyncRead(storageAdapter.fetchNode(node.sameCreator(), transaction, layer, + working.add(onReadListener.onAsyncRead(storageAdapter.fetchNode(creator, readTransaction, layer, currentNeighborKey)) .thenAccept(resultNode -> { Objects.requireNonNull(resultNode); @@ -1656,816 +752,5 @@ private CompletableFuture>> fetchNeig }, executor).thenApply(ignored -> Lists.newArrayList(neighborNodeArray)); } - /** - * Method to compute the depth of this R-tree. - * @param transactionContext transaction context to be used - * @return the depth of the R-tree - */ - public int depth(@Nonnull final TransactionContext transactionContext) { - // - // find the number of levels in this tree - // - Node node = - transactionContext.run(tr -> fetchUpdatePathToLeaf(tr, BigInteger.ONE, new Tuple(), true).join()); - if (node == null) { - logger.trace("R-tree is empty."); - return 0; - } - - int numLevels = 1; - while (node.getParentNode() != null) { - numLevels ++; - node = node.getParentNode(); - } - Verify.verify(node.isRoot(), "end of update path should be the root"); - logger.trace("numLevels = {}", numLevels); - return numLevels; - } - - /** - * Method to validate the Hilbert R-tree. - * @param db the database to use - */ - public void validate(@Nonnull final Database db) { - validate(db, Integer.MAX_VALUE); - } - - /** - * Method to validate the Hilbert R-tree. - * @param db the database to use - * @param maxNumNodesToBeValidated a maximum number of nodes this call should attempt to validate - */ - public void validate(@Nonnull final Database db, - final int maxNumNodesToBeValidated) { - - ArrayDeque toBeProcessed = new ArrayDeque<>(); - toBeProcessed.addLast(new ValidationTraversalState(depth(db) - 1, null, rootId)); - - while (!toBeProcessed.isEmpty()) { - db.run(tr -> validate(tr, maxNumNodesToBeValidated, toBeProcessed).join()); - } - } - - /** - * Method to validate the Hilbert R-tree. - * @param transaction the transaction to use - * @param maxNumNodesToBeValidated a maximum number of nodes this call should attempt to validate - * @param toBeProcessed a deque with node information that still needs to be processed - * @return a completable future that completes successfully with the current deque of to-be-processed nodes if the - * portion of the tree that was validated is in fact valid, completes with failure otherwise - */ - @Nonnull - private CompletableFuture> validate(@Nonnull final Transaction transaction, - final int maxNumNodesToBeValidated, - @Nonnull final ArrayDeque toBeProcessed) { - final AtomicInteger numNodesEnqueued = new AtomicInteger(0); - final List>> working = Lists.newArrayList(); - - // Fetch the entire tree. - return AsyncUtil.whileTrue(() -> { - final Iterator>> workingIterator = working.iterator(); - while (workingIterator.hasNext()) { - final CompletableFuture> nextFuture = workingIterator.next(); - if (nextFuture.isDone()) { - toBeProcessed.addAll(nextFuture.join()); - workingIterator.remove(); - } - } - - while (working.size() <= MAX_CONCURRENT_READS && numNodesEnqueued.get() < maxNumNodesToBeValidated) { - final ValidationTraversalState currentValidationTraversalState = toBeProcessed.pollFirst(); - if (currentValidationTraversalState == null) { - break; - } - - final IntermediateNode parentNode = currentValidationTraversalState.getParentNode(); - final int level = currentValidationTraversalState.getLevel(); - final ChildSlot childSlotInParentNode; - final int slotIndexInParent; - if (parentNode != null) { - int slotIndex; - ChildSlot childSlot = null; - for (slotIndex = 0; slotIndex < parentNode.size(); slotIndex++) { - childSlot = parentNode.getSlot(slotIndex); - if (Arrays.equals(childSlot.getChildId(), currentValidationTraversalState.getChildId())) { - break; - } - } - - if (slotIndex == parentNode.size()) { - throw new IllegalStateException("child slot not found in parent for child node"); - } else { - childSlotInParentNode = childSlot; - slotIndexInParent = slotIndex; - } - } else { - childSlotInParentNode = null; - slotIndexInParent = -1; - } - - final CompletableFuture fetchedNodeFuture = - onReadListener.onAsyncRead(storageAdapter.fetchNode(transaction, currentValidationTraversalState.getChildId()) - .thenApply(node -> { - if (parentNode != null) { - Objects.requireNonNull(node); - node.linkToParent(parentNode, slotIndexInParent); - } - return node; - }) - .thenCompose(childNode -> { - if (parentNode != null && getConfig().isUseNodeSlotIndex()) { - final var childSlot = parentNode.getSlot(slotIndexInParent); - return storageAdapter.scanNodeIndexAndFetchNode(transaction, level, - childSlot.getLargestHilbertValue(), childSlot.getLargestKey(), false) - .thenApply(nodeFromIndex -> { - Objects.requireNonNull(nodeFromIndex); - if (!Arrays.equals(nodeFromIndex.getId(), childNode.getId())) { - logger.warn("corrupt node slot index at level {}, parentNode = {}", level, parentNode); - throw new IllegalStateException("corrupt node index"); - } - return childNode; - }); - } - return CompletableFuture.completedFuture(childNode); - })); - working.add(fetchedNodeFuture.thenApply(childNode -> { - if (childNode == null) { - // Starting at root node but root node was not fetched since the R-tree has no entries. - return ImmutableList.of(); - } - childNode.validate(); - childNode.validateParentNode(parentNode, childSlotInParentNode); - - // add all children to the to be processed queue - if (childNode.getKind() == NodeKind.INTERMEDIATE) { - return ((IntermediateNode)childNode).getSlots() - .stream() - .map(childSlot -> new ValidationTraversalState(level - 1, - (IntermediateNode)childNode, childSlot.getChildId())) - .collect(ImmutableList.toImmutableList()); - } else { - return ImmutableList.of(); - } - })); - numNodesEnqueued.addAndGet(1); - } - - if (working.isEmpty()) { - return AsyncUtil.READY_FALSE; - } - return AsyncUtil.whenAny(working).thenApply(v -> true); - }, executor).thenApply(vignore -> toBeProcessed); - } - - /** - * Method to find the appropriate child slot index for a given Hilbert value and key. This method is used - * to find the proper slot indexes for the insert/update path and for the delete path. Note that if - * {@code (largestHilbertValue, largestKey)} of the last child is less than {@code (hilbertValue, key)}, we insert - * through the last child as we treat the (non-existing) next item as {@code (infinity, infinity)}. - * @param intermediateNode the intermediate node to search - * @param hilbertValue hilbert value - * @param key key - * @param isInsertUpdate indicator if the caller - * @return the 0-based slot index that corresponds to the given {@code (hilbertValue, key)} pair {@code p} if a slot - * covers that pair. If such a slot cannot be found while a new record is inserted, slot {@code 0} is - * returned if that slot is compared larger than {@code p}, the last slot ({@code size - 1}) if that slot is - * compared smaller than {@code p}. If, on the contrary, a record is deleted and a slot covering {@code p} - * cannot be found, this method returns {@code -1}. - */ - private static int findChildSlotIndex(@Nonnull final IntermediateNode intermediateNode, - @Nonnull final BigInteger hilbertValue, - @Nonnull final Tuple key, - final boolean isInsertUpdate) { - Verify.verify(!intermediateNode.isEmpty()); - - if (!isInsertUpdate) { - // make sure that the node covers the Hilbert Value/key we would like to delete - final ChildSlot firstChildSlot = intermediateNode.getSlot(0); - - final int compare = NodeSlot.compareHilbertValueKeyPair(firstChildSlot.getSmallestHilbertValue(), firstChildSlot.getSmallestKey(), - hilbertValue, key); - if (compare > 0) { - // child smallest HV/key > target HV/key - return -1; - } - } - - for (int slotIndex = 0; slotIndex < intermediateNode.size(); slotIndex++) { - final ChildSlot childSlot = intermediateNode.getSlot(slotIndex); - - // - // Choose subtree with the minimum Hilbert value that is greater than the target - // Hilbert value. If there is no such subtree, i.e. the target Hilbert value is the - // largest Hilbert value, we choose the largest one in the current node. - // - final int compare = NodeSlot.compareHilbertValueKeyPair(childSlot.getLargestHilbertValue(), childSlot.getLargestKey(), hilbertValue, key); - if (compare >= 0) { - // child largest HV/key > target HV/key - return slotIndex; - } - } - - // - // This is an intermediate node; we insert through the last child, but return -1 if this is for a delete - // operation. - return isInsertUpdate ? intermediateNode.size() - 1 : - 1; - } - - /** - * Method to find the appropriate child slot index for a given child it. - * @param parentNode the intermediate node to search - * @param childId the child id to search for - * @return if found the 0-based slot index that corresponds to slot using holding the given {@code childId}; - * {@code -1} otherwise - */ - private static int findChildSlotIndex(@Nonnull final IntermediateNode parentNode, @Nonnull final byte[] childId) { - for (int slotIndex = 0; slotIndex < parentNode.size(); slotIndex++) { - final ChildSlot childSlot = parentNode.getSlot(slotIndex); - - if (Arrays.equals(childSlot.getChildId(), childId)) { - return slotIndex; - } - } - return -1; - } - - /** - * Method to find the appropriate item slot index for a given Hilbert value and key. This method is used - * to find the proper item slot index for the insert/update path. - * @param leafNode the leaf node to search - * @param hilbertValue hilbert value - * @param key key - * @return {@code -1} if the item specified by {@code (hilbertValue, key)} already exists in {@code leafNode}; - * the 0-based slot index that represents the insertion point index of the given {@code (hilbertValue, key)} - * pair, otherwise - */ - private static int findInsertUpdateItemSlotIndex(@Nonnull final DataNode leafNode, - @Nonnull final BigInteger hilbertValue, - @Nonnull final Tuple key) { - for (int slotIndex = 0; slotIndex < leafNode.size(); slotIndex++) { - final ItemSlot slot = leafNode.getSlot(slotIndex); - - final int compare = NodeSlot.compareHilbertValueKeyPair(slot.getHilbertValue(), slot.getKey(), hilbertValue, key); - if (compare == 0) { - return -1; - } - - if (compare > 0) { - return slotIndex; - } - } - - return leafNode.size(); - } - - /** - * Method to find the appropriate item slot index for a given Hilbert value and key. This method is used - * to find the proper item slot index for the delete path. - * @param leafNode the leaf node to search - * @param hilbertValue hilbert value - * @param key key - * @return {@code -1} if the item specified by {@code (hilbertValue, key)} does not exist in {@code leafNode}; - * the 0-based slot index that corresponds to the slot for the given {@code (hilbertValue, key)} - * pair, otherwise - */ - private static int findDeleteItemSlotIndex(@Nonnull final DataNode leafNode, - @Nonnull final BigInteger hilbertValue, - @Nonnull final Tuple key) { - for (int slotIndex = 0; slotIndex < leafNode.size(); slotIndex++) { - final ItemSlot slot = leafNode.getSlot(slotIndex); - - final int compare = NodeSlot.compareHilbertValueKeyPair(slot.getHilbertValue(), slot.getKey(), hilbertValue, key); - if (compare == 0) { - return slotIndex; - } - - if (compare > 0) { - return -1; - } - } - - return -1; - } - - /** - * Traversal state of a scan over the tree. A scan consists of an initial walk to the left-most applicable leaf node - * potentially containing items relevant to the scan. The caller then consumes that leaf node and advances to the - * next leaf node that is relevant to the scan. The notion of next emerges using the order defined by the - * composite {@code (hilbertValue, key)} for items in leaf nodes and {@code (largestHilbertValue, largestKey)} in - * intermediate nodes. The traversal state captures the node ids that still have to be processed on each discovered - * level in order to fulfill the requirements of the scan operation. - */ - private static class TraversalState { - @Nullable - private final List> toBeProcessed; - - @Nullable - private final DataNode currentLeafNode; - - private TraversalState(@Nullable final List> toBeProcessed, @Nullable final DataNode currentLeafNode) { - this.toBeProcessed = toBeProcessed; - this.currentLeafNode = currentLeafNode; - } - - @Nonnull - public List> getToBeProcessed() { - return Objects.requireNonNull(toBeProcessed); - } - - @Nonnull - public DataNode getCurrentLeafNode() { - return Objects.requireNonNull(currentLeafNode); - } - - public boolean isEnd() { - return currentLeafNode == null; - } - - public static TraversalState of(@Nonnull final List> toBeProcessed, @Nonnull final DataNode currentLeafNode) { - return new TraversalState(toBeProcessed, currentLeafNode); - } - - public static TraversalState end() { - return new TraversalState(null, null); - } - } - - /** - * An {@link AsyncIterator} over the leaf nodes that represent the result of a scan over the tree. This iterator - * interfaces with the scan logic - * (see {@link #fetchLeftmostPathToLeaf(ReadTransaction, byte[], BigInteger, Tuple, Predicate, BiPredicate)} and - * {@link #fetchNextPathToLeaf(ReadTransaction, TraversalState, BigInteger, Tuple, Predicate, BiPredicate)}) and wraps - * intermediate {@link TraversalState}s created by these methods. - */ - private class LeafIterator implements AsyncIterator { - @Nonnull - private final ReadTransaction readTransaction; - @Nonnull - private final byte[] rootId; - @Nullable - private final BigInteger lastHilbertValue; - @Nullable - private final Tuple lastKey; - @Nonnull - private final Predicate mbrPredicate; - @Nonnull - private final BiPredicate suffixKeyPredicate; - - @Nullable - private TraversalState currentState; - @Nullable - private CompletableFuture nextStateFuture; - - @SpotBugsSuppressWarnings("EI_EXPOSE_REP2") - public LeafIterator(@Nonnull final ReadTransaction readTransaction, @Nonnull final byte[] rootId, - @Nullable final BigInteger lastHilbertValue, @Nullable final Tuple lastKey, - @Nonnull final Predicate mbrPredicate, @Nonnull final BiPredicate suffixKeyPredicate) { - Preconditions.checkArgument((lastHilbertValue == null && lastKey == null) || - (lastHilbertValue != null && lastKey != null)); - this.readTransaction = readTransaction; - this.rootId = rootId; - this.lastHilbertValue = lastHilbertValue; - this.lastKey = lastKey; - this.mbrPredicate = mbrPredicate; - this.suffixKeyPredicate = suffixKeyPredicate; - this.currentState = null; - this.nextStateFuture = null; - } - - @Override - public CompletableFuture onHasNext() { - if (nextStateFuture == null) { - if (currentState == null) { - nextStateFuture = fetchLeftmostPathToLeaf(readTransaction, rootId, lastHilbertValue, lastKey, - mbrPredicate, suffixKeyPredicate); - } else { - nextStateFuture = fetchNextPathToLeaf(readTransaction, currentState, lastHilbertValue, lastKey, - mbrPredicate, suffixKeyPredicate); - } - } - return nextStateFuture.thenApply(traversalState -> !traversalState.isEnd()); - } - - @Override - public boolean hasNext() { - return onHasNext().join(); - } - - @Override - public DataNode next() { - if (hasNext()) { - // underlying has already completed - currentState = Objects.requireNonNull(nextStateFuture).join(); - nextStateFuture = null; - return currentState.getCurrentLeafNode(); - } - throw new NoSuchElementException("called next() on exhausted iterator"); - } - - @Override - public void cancel() { - if (nextStateFuture != null) { - nextStateFuture.cancel(false); - } - } - } - - /** - * Iterator for iterating the items contained in the leaf nodes produced by an underlying {@link LeafIterator}. - * This iterator is the async equivalent of - * {@code Streams.stream(leafIterator).flatMap(leafNode -> leafNode.getItems().stream()).toIterator()}. - */ - public static class ItemSlotIterator implements AsyncIterator { - @Nonnull - private final AsyncIterator leafIterator; - @Nullable - private DataNode currentLeafNode; - @Nullable - private Iterator currenLeafItemsIterator; - - private ItemSlotIterator(@Nonnull final AsyncIterator leafIterator) { - this.leafIterator = leafIterator; - this.currentLeafNode = null; - this.currenLeafItemsIterator = null; - } - - @Override - public CompletableFuture onHasNext() { - if (currenLeafItemsIterator != null && currenLeafItemsIterator.hasNext()) { - return CompletableFuture.completedFuture(true); - } - // we know that each leaf has items (or if it doesn't it is the root; we are done if there are no items - return leafIterator.onHasNext() - .thenApply(hasNext -> { - if (hasNext) { - this.currentLeafNode = leafIterator.next(); - this.currenLeafItemsIterator = currentLeafNode.getSlots().iterator(); - return currenLeafItemsIterator.hasNext(); - } - return false; - }); - } - - @Override - public boolean hasNext() { - return onHasNext().join(); - } - - @Override - public ItemSlot next() { - if (hasNext()) { - return Objects.requireNonNull(currenLeafItemsIterator).next(); - } - throw new NoSuchElementException("called next() on exhausted iterator"); - } - - @Override - public void cancel() { - leafIterator.cancel(); - } - } - - /** - * Class to signal the caller of insert/update/delete code paths what the next action in that path should be. - * The indicated action is either another insert/delete on a higher level in the tree, further adjustments of - * secondary attributes on a higher level in the tree, or an indication that the insert/update/delete path is done - * with all necessary modifications. - */ - private static class NodeOrAdjust { - public static final NodeOrAdjust NONE = new NodeOrAdjust(null, null, false); - public static final NodeOrAdjust ADJUST = new NodeOrAdjust(null, null, true); - - @Nullable - private final ChildSlot slotInParent; - @Nullable - private final Node node; - - private final boolean parentNeedsAdjustment; - - private NodeOrAdjust(@Nullable final ChildSlot slotInParent, @Nullable final Node node, final boolean parentNeedsAdjustment) { - Verify.verify((slotInParent == null && node == null) || - (slotInParent != null && node != null)); - this.slotInParent = slotInParent; - this.node = node; - this.parentNeedsAdjustment = parentNeedsAdjustment; - } - - @Nullable - public ChildSlot getSlotInParent() { - return slotInParent; - } - - @Nullable - public Node getSplitNode() { - return node; - } - - @Nullable - public Node getTombstoneNode() { - return node; - } - - public boolean parentNeedsAdjustment() { - return parentNeedsAdjustment; - } - } - - /** - * Helper class for the traversal of nodes during tree validation. - */ - private static class ValidationTraversalState { - final int level; - @Nullable - private final IntermediateNode parentNode; - @Nonnull - private final byte[] childId; - - public ValidationTraversalState(final int level, @Nullable final IntermediateNode parentNode, @Nonnull final byte[] childId) { - this.level = level; - this.parentNode = parentNode; - this.childId = childId; - } - - public int getLevel() { - return level; - } - - @Nullable - public IntermediateNode getParentNode() { - return parentNode; - } - - @Nonnull - public byte[] getChildId() { - return childId; - } - } - - /** - * Class to capture an N-dimensional point. It wraps a {@link Tuple} mostly due to proximity with its serialization - * format and provides helpers for Euclidean operations. Note that the coordinates used here do not need to be - * numbers. - */ - public static class Point { - @Nonnull - private final Tuple coordinates; - - public Point(@Nonnull final Tuple coordinates) { - Preconditions.checkArgument(!coordinates.isEmpty()); - this.coordinates = coordinates; - } - - @Nonnull - public Tuple getCoordinates() { - return coordinates; - } - - public int getNumDimensions() { - return coordinates.size(); - } - - @Nullable - public Object getCoordinate(final int dimension) { - return coordinates.get(dimension); - } - - @Nullable - public Number getCoordinateAsNumber(final int dimension) { - return (Number)getCoordinate(dimension); - } - - @Override - public boolean equals(final Object o) { - if (this == o) { - return true; - } - if (!(o instanceof Point)) { - return false; - } - final Point point = (Point)o; - return TupleHelpers.equals(coordinates, point.coordinates); - } - - @Override - public int hashCode() { - return coordinates.hashCode(); - } - - @Nonnull - @Override - public String toString() { - return coordinates.toString(); - } - } - - /** - * Class to capture an N-dimensional rectangle/cube/hypercube. It wraps a {@link Tuple} mostly due to proximity - * with its serialization format and provides helpers for Euclidean operations. Note that the coordinates used here - * do not need to be numbers. - */ - public static class Rectangle { - /** - * A tuple that holds the coordinates of this N-dimensional rectangle. The layout is defined as - * {@code (low1, low2, ..., lowN, high1, high2, ..., highN}. Note that we don't use nested {@link Tuple}s for - * space-saving reasons (when the tuple is serialized). - */ - @Nonnull - private final Tuple ranges; - - public Rectangle(final Tuple ranges) { - Preconditions.checkArgument(!ranges.isEmpty() && ranges.size() % 2 == 0); - this.ranges = ranges; - } - - public int getNumDimensions() { - return ranges.size() >> 1; - } - - @Nonnull - public Tuple getRanges() { - return ranges; - } - - @Nonnull - public Object getLow(final int dimension) { - return ranges.get(dimension); - } - - @Nonnull - public Object getHigh(final int dimension) { - return ranges.get((ranges.size() >> 1) + dimension); - } - - @Nonnull - public BigInteger area() { - BigInteger currentArea = BigInteger.ONE; - for (int d = 0; d < getNumDimensions(); d++) { - currentArea = currentArea.multiply(BigInteger.valueOf(((Number)getHigh(d)).longValue() - ((Number)getLow(d)).longValue())); - } - return currentArea; - } - - @Nonnull - public Rectangle unionWith(@Nonnull final Point point) { - Preconditions.checkArgument(getNumDimensions() == point.getNumDimensions()); - boolean isModified = false; - Object[] ranges = new Object[getNumDimensions() << 1]; - - for (int d = 0; d < getNumDimensions(); d++) { - final Object coordinate = point.getCoordinate(d); - final Tuple coordinateTuple = Tuple.from(coordinate); - final Object low = getLow(d); - final Tuple lowTuple = Tuple.from(low); - if (TupleHelpers.compare(coordinateTuple, lowTuple) < 0) { - ranges[d] = coordinate; - isModified = true; - } else { - ranges[d] = low; - } - - final Object high = getHigh(d); - final Tuple highTuple = Tuple.from(high); - if (TupleHelpers.compare(coordinateTuple, highTuple) > 0) { - ranges[getNumDimensions() + d] = coordinate; - isModified = true; - } else { - ranges[getNumDimensions() + d] = high; - } - } - - if (!isModified) { - return this; - } - - return new Rectangle(Tuple.from(ranges)); - } - - @Nonnull - public Rectangle unionWith(@Nonnull final Rectangle other) { - Preconditions.checkArgument(getNumDimensions() == other.getNumDimensions()); - boolean isModified = false; - Object[] ranges = new Object[getNumDimensions() << 1]; - - for (int d = 0; d < getNumDimensions(); d++) { - final Object otherLow = other.getLow(d); - final Tuple otherLowTuple = Tuple.from(otherLow); - final Object otherHigh = other.getHigh(d); - final Tuple otherHighTuple = Tuple.from(otherHigh); - - final Object low = getLow(d); - final Tuple lowTuple = Tuple.from(low); - if (TupleHelpers.compare(otherLowTuple, lowTuple) < 0) { - ranges[d] = otherLow; - isModified = true; - } else { - ranges[d] = low; - } - final Object high = getHigh(d); - final Tuple highTuple = Tuple.from(high); - if (TupleHelpers.compare(otherHighTuple, highTuple) > 0) { - ranges[getNumDimensions() + d] = otherHigh; - isModified = true; - } else { - ranges[getNumDimensions() + d] = high; - } - } - - if (!isModified) { - return this; - } - - return new Rectangle(Tuple.from(ranges)); - } - - public boolean isOverlapping(@Nonnull final Rectangle other) { - Preconditions.checkArgument(getNumDimensions() == other.getNumDimensions()); - - for (int d = 0; d < getNumDimensions(); d++) { - final Tuple otherLowTuple = Tuple.from(other.getLow(d)); - final Tuple otherHighTuple = Tuple.from(other.getHigh(d)); - - final Tuple lowTuple = Tuple.from(getLow(d)); - final Tuple highTuple = Tuple.from(getHigh(d)); - - if (TupleHelpers.compare(highTuple, otherLowTuple) < 0 || - TupleHelpers.compare(lowTuple, otherHighTuple) > 0) { - return false; - } - } - return true; - } - - public boolean contains(@Nonnull final Point point) { - Preconditions.checkArgument(getNumDimensions() == point.getNumDimensions()); - - for (int d = 0; d < getNumDimensions(); d++) { - final Tuple otherTuple = Tuple.from(point.getCoordinate(d)); - - final Tuple lowTuple = Tuple.from(getLow(d)); - final Tuple highTuple = Tuple.from(getHigh(d)); - - if (TupleHelpers.compare(highTuple, otherTuple) < 0 || - TupleHelpers.compare(lowTuple, otherTuple) > 0) { - return false; - } - } - return true; - } - - @Override - public boolean equals(final Object o) { - if (this == o) { - return true; - } - if (!(o instanceof Rectangle)) { - return false; - } - final Rectangle rectangle = (Rectangle)o; - return TupleHelpers.equals(ranges, rectangle.ranges); - } - - @Override - public int hashCode() { - return ranges.hashCode(); - } - - @Nonnull - public String toPlotString() { - final StringBuilder builder = new StringBuilder(); - for (int d = 0; d < getNumDimensions(); d++) { - builder.append(((Number)getLow(d)).longValue()); - if (d + 1 < getNumDimensions()) { - builder.append(","); - } - } - - builder.append(","); - - for (int d = 0; d < getNumDimensions(); d++) { - builder.append(((Number)getHigh(d)).longValue()); - if (d + 1 < getNumDimensions()) { - builder.append(","); - } - } - return builder.toString(); - } - - @Nonnull - @Override - public String toString() { - return ranges.toString(); - } - - @Nonnull - public static Rectangle fromPoint(@Nonnull final Point point) { - final Object[] mbrRanges = new Object[point.getNumDimensions() * 2]; - for (int d = 0; d < point.getNumDimensions(); d++) { - final Object coordinate = point.getCoordinate(d); - mbrRanges[d] = coordinate; - mbrRanges[point.getNumDimensions() + d] = coordinate; - } - return new Rectangle(Tuple.from(mbrRanges)); - } - } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java index e16bdde8d0..a1de1a5c2a 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java @@ -31,9 +31,9 @@ /** * TODO. */ -class IntermediateNode extends AbstractNode { +class IntermediateNode extends AbstractNode { public IntermediateNode(@Nonnull final Tuple primaryKey, - @Nonnull final List neighbors) { + @Nonnull final List neighbors) { super(primaryKey, neighbors); } @@ -55,24 +55,18 @@ public IntermediateNode asIntermediateNode() { return this; } - @Nonnull - @Override - public NodeWithLayer withLayer(final int layer) { - return new NodeWithLayer<>(layer, this); - } - @Override - public NodeCreator sameCreator() { + public NodeCreator sameCreator() { return IntermediateNode::creator; } @Nonnull @SuppressWarnings("unchecked") - public static Node creator(@Nonnull final NodeKind nodeKind, - @Nonnull final Tuple primaryKey, - @Nullable final Vector vector, - @Nonnull final List neighbors) { + public static Node creator(@Nonnull final NodeKind nodeKind, + @Nonnull final Tuple primaryKey, + @Nullable final Vector vector, + @Nonnull final List neighbors) { Verify.verify(nodeKind == NodeKind.INTERMEDIATE); - return new IntermediateNode(primaryKey, (List)neighbors); + return new IntermediateNode(primaryKey, (List)neighbors); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java index c0d8462903..d3f24e66b6 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java @@ -30,29 +30,29 @@ /** * TODO. - * @param neighbor type + * @param neighbor type */ -public interface Node { +public interface Node { @Nonnull Tuple getPrimaryKey(); @Nonnull - List getNeighbors(); + List getNeighbors(); @Nonnull - N getNeighbor(int index); + R getNeighbor(int index); @CanIgnoreReturnValue @Nonnull - Node insert(@Nonnull StorageAdapter storageAdapter, int level, int slotIndex, @Nonnull NodeSlot slot); + Node insert(@Nonnull StorageAdapter storageAdapter, int level, int slotIndex, @Nonnull NodeSlot slot); @CanIgnoreReturnValue @Nonnull - Node update(@Nonnull StorageAdapter storageAdapter, int level, int slotIndex, @Nonnull NodeSlot updatedSlot); + Node update(@Nonnull StorageAdapter storageAdapter, int level, int slotIndex, @Nonnull NodeSlot updatedSlot); @CanIgnoreReturnValue @Nonnull - Node delete(@Nonnull StorageAdapter storageAdapter, int level, int slotIndex); + Node delete(@Nonnull StorageAdapter storageAdapter, int level, int slotIndex); /** * Return the kind of the node, i.e. {@link NodeKind#DATA} or {@link NodeKind#INTERMEDIATE}. @@ -67,14 +67,11 @@ public interface Node { @Nonnull IntermediateNode asIntermediateNode(); - @Nonnull - NodeWithLayer withLayer(int layer); - - NodeCreator sameCreator(); + NodeCreator sameCreator(); @FunctionalInterface - interface NodeCreator { + interface NodeCreator { Node create(@Nonnull NodeKind nodeKind, @Nonnull Tuple primaryKey, @Nullable Vector vector, - @Nonnull List neighbors); + @Nonnull List neighbors); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Element.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReference.java similarity index 61% rename from fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Element.java rename to fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReference.java index 53e1390687..c395f4796c 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Element.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReference.java @@ -1,5 +1,5 @@ /* - * NodeWithLayer.java + * NodeReference.java * * This source file is part of the FoundationDB open source project * @@ -24,15 +24,14 @@ import javax.annotation.Nonnull; import java.util.List; +import java.util.Objects; -class Element { +public class NodeReference { @Nonnull private final Tuple primaryKey; - private final double distance; - public Element(@Nonnull final Tuple primaryKey, final double distance) { + public NodeReference(@Nonnull final Tuple primaryKey) { this.primaryKey = primaryKey; - this.distance = distance; } @Nonnull @@ -40,14 +39,24 @@ public Tuple getPrimaryKey() { return primaryKey; } - public double getDistance() { - return distance; - } - @Nonnull - public static Iterable primaryKeys(@Nonnull List elements) { - return () -> elements.stream() - .map(Element::getPrimaryKey) + public static Iterable primaryKeys(@Nonnull List neighbors) { + return () -> neighbors.stream() + .map(NodeReference::getPrimaryKey) .iterator(); } + + @Override + public boolean equals(final Object o) { + if (!(o instanceof NodeReference)) { + return false; + } + final NodeReference that = (NodeReference)o; + return Objects.equals(primaryKey, that.primaryKey); + } + + @Override + public int hashCode() { + return Objects.hashCode(primaryKey); + } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Neighbor.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithDistance.java similarity index 50% rename from fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Neighbor.java rename to fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithDistance.java index 9f3107ddcb..4408bf0f4e 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Neighbor.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithDistance.java @@ -1,5 +1,5 @@ /* - * Neighbor.java + * NodeReferenceWithDistance.java * * This source file is part of the FoundationDB open source project * @@ -23,17 +23,34 @@ import com.apple.foundationdb.tuple.Tuple; import javax.annotation.Nonnull; +import java.util.Objects; -public class Neighbor { - @Nonnull - private final Tuple primaryKey; +class NodeReferenceWithDistance extends NodeReference { + private final double distance; - public Neighbor(@Nonnull final Tuple primaryKey) { - this.primaryKey = primaryKey; + public NodeReferenceWithDistance(@Nonnull final Tuple primaryKey, final double distance) { + super(primaryKey); + this.distance = distance; } - @Nonnull - public Tuple getPrimaryKey() { - return primaryKey; + public double getDistance() { + return distance; + } + + @Override + public boolean equals(final Object o) { + if (!(o instanceof NodeReferenceWithDistance)) { + return false; + } + if (!super.equals(o)) { + return false; + } + final NodeReferenceWithDistance that = (NodeReferenceWithDistance)o; + return Double.compare(distance, that.distance) == 0; + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), distance); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NeighborWithVector.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithVector.java similarity index 85% rename from fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NeighborWithVector.java rename to fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithVector.java index a531cb3e7c..ead8135b6a 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NeighborWithVector.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithVector.java @@ -1,5 +1,5 @@ /* - * Neighbor.java + * NodeReferenceWithVector.java * * This source file is part of the FoundationDB open source project * @@ -25,11 +25,11 @@ import javax.annotation.Nonnull; -public class NeighborWithVector extends Neighbor { +public class NodeReferenceWithVector extends NodeReference { @Nonnull private final Vector vector; - public NeighborWithVector(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector) { + public NodeReferenceWithVector(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector) { super(primaryKey); this.vector = vector; } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeWithLayer.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeWithLayer.java deleted file mode 100644 index f3a3dd49a8..0000000000 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeWithLayer.java +++ /dev/null @@ -1,43 +0,0 @@ -/* - * NodeWithLayer.java - * - * This source file is part of the FoundationDB open source project - * - * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors - * - * Licensed 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 com.apple.foundationdb.async.hnsw; - -import javax.annotation.Nonnull; - -class NodeWithLayer { - private final int layer; - @Nonnull - private final Node node; - - public NodeWithLayer(final int layer, @Nonnull final Node node) { - this.layer = layer; - this.node = node; - } - - public int getLayer() { - return layer; - } - - @Nonnull - public Node getNode() { - return node; - } -} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java index da12f8199d..711dc578ad 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java @@ -34,7 +34,7 @@ default void onSlotIndexEntryRead(@Nonnull final byte[] key) { // nothing } - default CompletableFuture> onAsyncRead(@Nonnull CompletableFuture> future) { + default CompletableFuture> onAsyncRead(@Nonnull CompletableFuture> future) { return future; } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java index d564b20095..cea4a56f3c 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java @@ -84,10 +84,10 @@ interface StorageAdapter { CompletableFuture fetchEntryNodeKey(@Nonnull ReadTransaction readTransaction); @Nonnull - CompletableFuture> fetchNode(@Nonnull Node.NodeCreator creator, - @Nonnull ReadTransaction readTransaction, - int layer, - @Nonnull Tuple primaryKey); + CompletableFuture> fetchNode(@Nonnull Node.NodeCreator creator, + @Nonnull ReadTransaction readTransaction, + int layer, + @Nonnull Tuple primaryKey); /** * Insert a new entry into the node index if configuration indicates we should maintain such an index. diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java index d8482e2ac4..d88bcda5c1 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java @@ -24,6 +24,8 @@ import com.google.common.base.Suppliers; import javax.annotation.Nonnull; +import java.util.Arrays; +import java.util.Objects; import java.util.function.Supplier; /** @@ -53,6 +55,20 @@ public R[] getData() { @Nonnull public abstract DoubleVector toDoubleVector(); + @Override + public boolean equals(final Object o) { + if (!(o instanceof Vector)) { + return false; + } + final Vector vector = (Vector)o; + return Objects.deepEquals(data, vector.data); + } + + @Override + public int hashCode() { + return Arrays.hashCode(data); + } + public static class HalfVector extends Vector { @Nonnull private final Supplier toDoubleVectorSupplier; From 3e63761d165516aadb229f9589eaf1faf63e4ce7 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Sun, 27 Jul 2025 12:22:11 +0200 Subject: [PATCH 06/34] save point -- read path almost done --- .../{GreedyResult.java => GreedyState.java} | 6 +- .../apple/foundationdb/async/hnsw/HNSW.java | 270 ++++++++++-------- .../async/hnsw/NodeReference.java | 6 +- .../foundationdb/async/hnsw/SearchResult.java | 66 +++++ 4 files changed, 226 insertions(+), 122 deletions(-) rename fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/{GreedyResult.java => GreedyState.java} (87%) create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/SearchResult.java diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyResult.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyState.java similarity index 87% rename from fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyResult.java rename to fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyState.java index b104848603..97b7c30941 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyResult.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyState.java @@ -24,13 +24,13 @@ import javax.annotation.Nonnull; -class GreedyResult { +class GreedyState { private final int layer; @Nonnull private final Tuple primaryKey; private final double distance; - public GreedyResult(final int layer, @Nonnull final Tuple primaryKey, final double distance) { + public GreedyState(final int layer, @Nonnull final Tuple primaryKey, final double distance) { this.layer = layer; this.primaryKey = primaryKey; this.distance = distance; @@ -49,7 +49,7 @@ public double getDistance() { return distance; } - public NodeReferenceWithDistance toElement() { + public NodeReferenceWithDistance toNodeReferenceWithDistance() { return new NodeReferenceWithDistance(getPrimaryKey(), getDistance()); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index 5a970fd2ab..e6580b1786 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -29,6 +29,9 @@ import com.christianheina.langx.half4j.Half; import com.google.common.base.Preconditions; import com.google.common.base.Verify; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Iterables; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.google.common.collect.Sets; @@ -38,7 +41,6 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; -import java.lang.reflect.Array; import java.math.BigInteger; import java.util.ArrayDeque; import java.util.Comparator; @@ -51,6 +53,7 @@ import java.util.concurrent.PriorityBlockingQueue; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiFunction; import java.util.function.BiPredicate; import java.util.function.Function; import java.util.function.Predicate; @@ -505,8 +508,8 @@ public AsyncIterator scan(@Nonnull final ReadTransaction readTransacti */ @SuppressWarnings("checkstyle:MethodName") // method name introduced by paper @Nonnull - private CompletableFuture kNearestNeighborsSearch(@Nonnull final ReadTransaction readTransaction, - @Nonnull final Vector queryVector) { + private CompletableFuture kNearestNeighborsSearch(@Nonnull final ReadTransaction readTransaction, + @Nonnull final Vector queryVector) { return storageAdapter.fetchEntryNodeKey(readTransaction) .thenCompose(entryPointAndLayer -> { if (entryPointAndLayer == null) { @@ -515,10 +518,10 @@ private CompletableFuture kNearestNeighborsSearch(@Nonnull final R final Metric metric = getConfig().getMetric(); - final var entryState = new GreedyResult(entryPointAndLayer.getLayer(), + final var entryState = new GreedyState(entryPointAndLayer.getLayer(), entryPointAndLayer.getPrimaryKey(), Vector.comparativeDistance(metric, entryPointAndLayer.getVector(), queryVector)); - final AtomicReference greedyResultReference = + final AtomicReference greedyResultReference = new AtomicReference<>(entryState); if (entryPointAndLayer.getLayer() == 0) { @@ -528,10 +531,11 @@ private CompletableFuture kNearestNeighborsSearch(@Nonnull final R return AsyncUtil.whileTrue(() -> { final var greedyIn = greedyResultReference.get(); - return greedySearchLayer(readTransaction, greedyIn.toElement(), greedyIn.getLayer(), queryVector) - .thenApply(greedyResult -> { - greedyResultReference.set(greedyResult); - return greedyResult.getLayer() > 0; + return greedySearchLayer(readTransaction, greedyIn.toNodeReferenceWithDistance(), + greedyIn.getLayer(), queryVector) + .thenApply(greedyState -> { + greedyResultReference.set(greedyState); + return greedyState.getLayer() > 0; }); }, executor).thenApply(ignored -> greedyResultReference.get()); }); @@ -541,14 +545,14 @@ private CompletableFuture kNearestNeighborsSearch(@Nonnull final R * TODO. */ @Nonnull - private CompletableFuture greedySearchLayer(@Nonnull final ReadTransaction readTransaction, - @Nonnull final NodeReferenceWithDistance entryNeighbor, - final int layer, - @Nonnull final Vector queryVector) { + private CompletableFuture greedySearchLayer(@Nonnull final ReadTransaction readTransaction, + @Nonnull final NodeReferenceWithDistance entryNeighbor, + final int layer, + @Nonnull final Vector queryVector) { Verify.verify(layer > 0); final Metric metric = getConfig().getMetric(); - final AtomicReference greedyStateReference = - new AtomicReference<>(new GreedyResult(layer, entryNeighbor.getPrimaryKey(), entryNeighbor.getDistance())); + final AtomicReference greedyStateReference = + new AtomicReference<>(new GreedyState(layer, entryNeighbor.getPrimaryKey(), entryNeighbor.getDistance())); return AsyncUtil.whileTrue(() -> onReadListener.onAsyncRead( storageAdapter.fetchNode(IntermediateNode::creator, readTransaction, @@ -560,7 +564,7 @@ private CompletableFuture greedySearchLayer(@Nonnull final ReadTra final IntermediateNode intermediateNode = node.asIntermediateNode(); final List neighbors = intermediateNode.getNeighbors(); - final GreedyResult currentNodeKey = greedyStateReference.get(); + final GreedyState currentNodeKey = greedyStateReference.get(); double minDistance = currentNodeKey.getDistance(); NodeReferenceWithVector nearestNeighbor = null; @@ -575,12 +579,12 @@ private CompletableFuture greedySearchLayer(@Nonnull final ReadTra if (nearestNeighbor == null) { greedyStateReference.set( - new GreedyResult(layer - 1, currentNodeKey.getPrimaryKey(), minDistance)); + new GreedyState(layer - 1, currentNodeKey.getPrimaryKey(), minDistance)); return false; } greedyStateReference.set( - new GreedyResult(layer, nearestNeighbor.getPrimaryKey(), + new GreedyState(layer, nearestNeighbor.getPrimaryKey(), minDistance)); return true; }), executor).thenApply(ignored -> greedyStateReference.get()); @@ -590,12 +594,12 @@ private CompletableFuture greedySearchLayer(@Nonnull final ReadTra * TODO. */ @Nonnull - private CompletableFuture searchLayer(@Nonnull Node.NodeCreator creator, - @Nonnull final ReadTransaction readTransaction, - @Nonnull final List entryNeighbors, - final int layer, - final int efSearch, - @Nonnull final Vector queryVector) { + private CompletableFuture> searchLayer(@Nonnull Node.NodeCreator creator, + @Nonnull final ReadTransaction readTransaction, + @Nonnull final List entryNeighbors, + final int layer, + final int efSearch, + @Nonnull final Vector queryVector) { final Set visited = Sets.newConcurrentHashSet(NodeReference.primaryKeys(entryNeighbors)); final PriorityBlockingQueue candidates = new PriorityBlockingQueue<>(entryNeighbors.size(), @@ -605,105 +609,135 @@ private CompletableFuture searchLayer(@N new PriorityBlockingQueue<>(entryNeighbors.size(), Comparator.comparing(NodeReferenceWithDistance::getDistance).reversed()); nearestNeighbors.addAll(entryNeighbors); - final Map> nodeCache = Maps.newConcurrentMap(); - + final Map> nodeCache = Maps.newConcurrentMap(); final Metric metric = getConfig().getMetric(); return AsyncUtil.whileTrue(() -> { if (candidates.isEmpty()) { - return false; + return AsyncUtil.READY_FALSE; } final NodeReferenceWithDistance candidate = candidates.poll(); final NodeReferenceWithDistance furthestNeighbor = Objects.requireNonNull(nearestNeighbors.peek()); if (candidate.getDistance() > furthestNeighbor.getDistance()) { - return false; - } - - final CompletableFuture> candidateNodeFuture; - final Node cachedCandidateNode = nodeCache.get(candidate.getPrimaryKey()); - if (cachedCandidateNode != null) { - candidateNodeFuture = CompletableFuture.completedFuture(cachedCandidateNode); - } else { - candidateNodeFuture = - storageAdapter.fetchNode(creator, readTransaction, layer, candidate.getPrimaryKey()) - .thenApply(candidateNode -> nodeCache.put(candidate.getPrimaryKey(), candidateNode)); + return AsyncUtil.READY_FALSE; } - candidateNodeFuture - .thenCompose(candidateNode -> fetchNodes(creator, readTransaction, - NodeReference.primaryKeys(candidateNode.getNeighbors()), layer) - .thenApply(neighbors -> { - - for (final Node neighbor : neighbors) { - + return fetchNodeIfNotCached(creator, readTransaction, layer, candidate, nodeCache) + .thenApply(candidateNode -> + Iterables.filter(candidateNode.getNeighbors(), + neighbor -> !visited.contains(neighbor.getPrimaryKey()))) + .thenCompose(neighborReferences -> fetchSomeNeighbors(creator, readTransaction, + layer, neighborReferences, nodeCache)) + .thenApply(neighborReferences -> { + for (final NodeReferenceWithVector current : neighborReferences) { + visited.add(current.getPrimaryKey()); + final double furthestDistance = + Objects.requireNonNull(nearestNeighbors.peek()).getDistance(); + + final double currentDistance = + Vector.comparativeDistance(metric, current.getVector(), queryVector); + if (currentDistance < furthestDistance || nearestNeighbors.size() < efSearch) { + final NodeReferenceWithDistance currentWithDistance = + new NodeReferenceWithDistance(current.getPrimaryKey(), currentDistance); + candidates.add(currentWithDistance); + nearestNeighbors.add(currentWithDistance); + if (nearestNeighbors.size() > efSearch) { + nearestNeighbors.poll(); } - })) - .thenApply(neighbors -> { - for (final Node neighbor : neighbors) { - if (visited.contains(neighbor.getPrimaryKey())) { - continue; } - visited.add(neighbor.getPrimaryKey()); - final NodeReferenceWithDistance furthestNeighbor1 = - Objects.requireNonNull(nearestNeighbors.peek()); - - final - Vector.comparativeDistance(metric, neighbor.) - } - }) + return true; + }); + }).thenCompose(ignored -> fetchResultsIfNecessary(creator, readTransaction, layer, nearestNeighbors, + nodeCache)); + } - return AsyncUtil.whileTrue(() -> { - fetchNodes(creator, readTransaction, candidateReference.get().) - }); - }).thenApply(ignored -> null); // TODO + /** + * TODO. + */ + @Nonnull + private CompletableFuture> fetchNodeIfNotCached(@Nonnull final Node.NodeCreator creator, + @Nonnull final ReadTransaction readTransaction, + final int layer, + @Nonnull final NodeReference nodeReference, + @Nonnull final Map> nodeCache) { + return fetchNodeIfNecessaryAndApply(creator, readTransaction, layer, nodeReference, + nR -> nodeCache.get(nR.getPrimaryKey()), + (ignored, node) -> node); + } - if (candidates.isEmpty()) { - return CompletableFuture.completedFuture(null); // TODO + /** + * TODO. + */ + @Nonnull + private CompletableFuture fetchNodeIfNecessaryAndApply(@Nonnull final Node.NodeCreator creator, + @Nonnull final ReadTransaction readTransaction, + final int layer, + @Nonnull final R nodeReference, + @Nonnull final Function fetchBypassFunction, + @Nonnull final BiFunction, U> biMapFunction) { + final U bypass = fetchBypassFunction.apply(nodeReference); + if (bypass != null) { + return CompletableFuture.completedFuture(bypass); } - final Metric metric = getConfig().getMetric(); - final AtomicReference greedyStateReference = - new AtomicReference<>(entryStates); - - fetchNodes(creator, readTransaction, ) + return onReadListener.onAsyncRead( + storageAdapter.fetchNode(creator, readTransaction, layer, nodeReference.getPrimaryKey())) + .thenApply(node -> biMapFunction.apply(nodeReference, node)); + } - return AsyncUtil.whileTrue(() -> - fetchNodes(IntermediateNode::creator, readTransaction, - layer, greedyStateReference.get().getPrimaryKey()) - .thenApply(nodeWithLayer -> { - if (nodeWithLayer == null) { - throw new IllegalStateException("unable to fetch node"); + /** + * TODO. + */ + @Nonnull + private CompletableFuture> fetchSomeNeighbors(@Nonnull final Node.NodeCreator creator, + @Nonnull final ReadTransaction readTransaction, + final int layer, + @Nonnull final Iterable neighborReferences, + @Nonnull final Map> nodeCache) { + return fetchSomeNodesAndApply(creator, readTransaction, layer, neighborReferences, + neighborReference -> { + if (neighborReference instanceof NodeReferenceWithVector) { + return (NodeReferenceWithVector)neighborReference; } - final IntermediateNode node = nodeWithLayer.getNode().asIntermediateNode(); - final List neighbors = node.getNeighbors(); - - final GreedyResult currentNodeKey = greedyStateReference.get(); - double minDistance = currentNodeKey.getDistance(); - - NodeReferenceWithVector nearestNeighbor = null; - for (final NodeReferenceWithVector neighbor : neighbors) { - final double distance = - Vector.comparativeDistance(metric, neighbor.getVector(), queryVector); - if (distance < minDistance) { - minDistance = distance; - nearestNeighbor = neighbor; - } + final Node neighborNode = nodeCache.get(neighborReference.getPrimaryKey()); + if (neighborNode == null) { + return null; } + return new NodeReferenceWithVector(neighborReference.getPrimaryKey(), neighborNode.asDataNode().getVector()); + }, + (neighborReference, neighborNode) -> + new NodeReferenceWithVector(neighborReference.getPrimaryKey(), neighborNode.asDataNode().getVector())); + } - if (nearestNeighbor == null) { - greedyStateReference.set( - new GreedyResult(layer - 1, currentNodeKey.getPrimaryKey(), minDistance)); - return false; + /** + * TODO. + */ + @Nonnull + private CompletableFuture> fetchResultsIfNecessary(@Nonnull final Node.NodeCreator creator, + @Nonnull final ReadTransaction readTransaction, + final int layer, + @Nonnull final Iterable nodeReferences, + @Nonnull final Map> nodeCache) { + return fetchSomeNodesAndApply(creator, readTransaction, layer, nodeReferences, + nodeReference -> { + final Node node = nodeCache.get(nodeReference.getPrimaryKey()); + if (node == null) { + return null; } - - greedyStateReference.set( - new GreedyResult(layer, nearestNeighbor.getPrimaryKey(), - minDistance)); - return true; - }), executor).thenApply(ignored -> greedyStateReference.get()); + return new SearchResult.NodeReferenceWithNode<>(nodeReference, node); + }, + SearchResult.NodeReferenceWithNode::new) + .thenApply(nodeReferencesWithNodes -> { + final ImmutableMap.Builder> nodeMapBuilder = + ImmutableMap.builder(); + for (final SearchResult.NodeReferenceWithNode nodeReferenceWithNode : nodeReferencesWithNodes) { + nodeMapBuilder.put(nodeReferenceWithNode.getNodeReferenceWithDistance(), nodeReferenceWithNode.getNode()); + } + return new SearchResult<>(layer, nodeMapBuilder.build()); + }); } /** @@ -711,34 +745,34 @@ private CompletableFuture searchLayer(@N */ @Nonnull @SuppressWarnings("unchecked") - private CompletableFuture>> fetchNodes(@Nonnull final Node.NodeCreator creator, - @Nonnull final ReadTransaction readTransaction, - @Nonnull final Iterable primaryKeys, - final int layer) { + private CompletableFuture> fetchSomeNodesAndApply(@Nonnull final Node.NodeCreator creator, + @Nonnull final ReadTransaction readTransaction, + final int layer, + @Nonnull final Iterable nodeReferences, + @Nonnull final Function fetchBypassFunction, + @Nonnull final BiFunction, U> biMapFunction) { // this deque is only modified by once upon creation - final ArrayDeque toBeProcessed = new ArrayDeque<>(); - for (final var primaryKey : primaryKeys) { - toBeProcessed.addLast(primaryKey); + final ArrayDeque toBeProcessed = new ArrayDeque<>(); + for (final var nodeReference : nodeReferences) { + toBeProcessed.addLast(nodeReference); } final List> working = Lists.newArrayList(); final AtomicInteger neighborIndex = new AtomicInteger(0); - final Node[] neighborNodeArray = - (Node[])Array.newInstance(Node.class, toBeProcessed.size()); + final Object[] neighborNodeArray = new Object[toBeProcessed.size()]; - // Fetch all sibling nodes (in parallel if possible). return AsyncUtil.whileTrue(() -> { working.removeIf(CompletableFuture::isDone); while (working.size() <= MAX_CONCURRENT_READS) { - final Tuple currentNeighborKey = toBeProcessed.pollFirst(); - if (currentNeighborKey == null) { + final R currentNeighborReference = toBeProcessed.pollFirst(); + if (currentNeighborReference == null) { break; } final int index = neighborIndex.getAndIncrement(); - working.add(onReadListener.onAsyncRead(storageAdapter.fetchNode(creator, readTransaction, layer, - currentNeighborKey)) + working.add(fetchNodeIfNecessaryAndApply(creator, readTransaction, layer, + currentNeighborReference, fetchBypassFunction, biMapFunction) .thenAccept(resultNode -> { Objects.requireNonNull(resultNode); neighborNodeArray[index] = resultNode; @@ -749,8 +783,12 @@ private CompletableFuture>> fetchNodes(@N return AsyncUtil.READY_FALSE; } return AsyncUtil.whenAny(working).thenApply(ignored -> true); - }, executor).thenApply(ignored -> Lists.newArrayList(neighborNodeArray)); + }, executor).thenApply(ignored -> { + final ImmutableList.Builder resultBuilder = ImmutableList.builder(); + for (final Object o : neighborNodeArray) { + resultBuilder.add((U)o); + } + return resultBuilder.build(); + }); } - - } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReference.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReference.java index c395f4796c..dd7d1680d5 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReference.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReference.java @@ -21,9 +21,9 @@ package com.apple.foundationdb.async.hnsw; import com.apple.foundationdb.tuple.Tuple; +import com.google.common.collect.Streams; import javax.annotation.Nonnull; -import java.util.List; import java.util.Objects; public class NodeReference { @@ -40,8 +40,8 @@ public Tuple getPrimaryKey() { } @Nonnull - public static Iterable primaryKeys(@Nonnull List neighbors) { - return () -> neighbors.stream() + public static Iterable primaryKeys(@Nonnull Iterable neighbors) { + return () -> Streams.stream(neighbors) .map(NodeReference::getPrimaryKey) .iterator(); } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/SearchResult.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/SearchResult.java new file mode 100644 index 0000000000..af66ca7e4b --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/SearchResult.java @@ -0,0 +1,66 @@ +/* + * NodeWithLayer.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import javax.annotation.Nonnull; +import java.util.Map; + +class SearchResult { + private final int layer; + @Nonnull + private final Map> nodeMap; + + public SearchResult(final int layer, @Nonnull final Map> nodeMap) { + this.layer = layer; + this.nodeMap = nodeMap; + } + + public int getLayer() { + return layer; + } + + @Nonnull + public Map> getNodeMap() { + return nodeMap; + } + + public static class NodeReferenceWithNode { + @Nonnull + private final NodeReferenceWithDistance nodeReferenceWithDistance; + @Nonnull + private final Node node; + + public NodeReferenceWithNode(@Nonnull final NodeReferenceWithDistance nodeReferenceWithDistance, @Nonnull final Node node) { + this.nodeReferenceWithDistance = nodeReferenceWithDistance; + this.node = node; + } + + @Nonnull + public NodeReferenceWithDistance getNodeReferenceWithDistance() { + return nodeReferenceWithDistance; + } + + @Nonnull + public Node getNode() { + return node; + } + } +} From 102c4adff8694729d2abf66a502deb1d8910106b Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Sun, 27 Jul 2025 18:12:54 +0200 Subject: [PATCH 07/34] save point -- read path done --- .../apple/foundationdb/async/hnsw/HNSW.java | 21 +++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index e6580b1786..4ed82af6a7 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -508,8 +508,9 @@ public AsyncIterator scan(@Nonnull final ReadTransaction readTransacti */ @SuppressWarnings("checkstyle:MethodName") // method name introduced by paper @Nonnull - private CompletableFuture kNearestNeighborsSearch(@Nonnull final ReadTransaction readTransaction, - @Nonnull final Vector queryVector) { + private CompletableFuture> kNearestNeighborsSearch(@Nonnull final ReadTransaction readTransaction, + final int efSearch, + @Nonnull final Vector queryVector) { return storageAdapter.fetchEntryNodeKey(readTransaction) .thenCompose(entryPointAndLayer -> { if (entryPointAndLayer == null) { @@ -521,7 +522,7 @@ private CompletableFuture kNearestNeighborsSearch(@Nonnull final Re final var entryState = new GreedyState(entryPointAndLayer.getLayer(), entryPointAndLayer.getPrimaryKey(), Vector.comparativeDistance(metric, entryPointAndLayer.getVector(), queryVector)); - final AtomicReference greedyResultReference = + final AtomicReference greedyStateReference = new AtomicReference<>(entryState); if (entryPointAndLayer.getLayer() == 0) { @@ -530,14 +531,22 @@ private CompletableFuture kNearestNeighborsSearch(@Nonnull final Re } return AsyncUtil.whileTrue(() -> { - final var greedyIn = greedyResultReference.get(); + final var greedyIn = greedyStateReference.get(); return greedySearchLayer(readTransaction, greedyIn.toNodeReferenceWithDistance(), greedyIn.getLayer(), queryVector) .thenApply(greedyState -> { - greedyResultReference.set(greedyState); + greedyStateReference.set(greedyState); return greedyState.getLayer() > 0; }); - }, executor).thenApply(ignored -> greedyResultReference.get()); + }, executor).thenApply(ignored -> greedyStateReference.get()); + }).thenCompose(greedyState -> { + if (greedyState == null) { + return CompletableFuture.completedFuture(null); + } + + return searchLayer(DataNode::creator, readTransaction, + ImmutableList.of(greedyState.toNodeReferenceWithDistance()), 0, efSearch, + queryVector); }); } From eb7e54d4aef412352f34466f73e97bda937992d7 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Sun, 27 Jul 2025 22:51:04 +0200 Subject: [PATCH 08/34] some renamings --- .../async/hnsw/ByNodeStorageAdapter.java | 8 +++--- .../hnsw/{DataNode.java => CompactNode.java} | 20 +++++++------- .../apple/foundationdb/async/hnsw/HNSW.java | 26 +++++++++---------- ...ntermediateNode.java => InliningNode.java} | 22 ++++++++-------- .../apple/foundationdb/async/hnsw/Node.java | 6 ++--- .../foundationdb/async/hnsw/NodeKind.java | 4 +-- .../async/hnsw/StorageAdapter.java | 4 +-- 7 files changed, 45 insertions(+), 45 deletions(-) rename fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/{DataNode.java => CompactNode.java} (75%) rename fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/{IntermediateNode.java => InliningNode.java} (73%) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java index 1ec1a12fa4..54d8f4895b 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java @@ -118,13 +118,13 @@ protected CompletableFuture> fetchNodeInternal } @Override - public void writeLeafNodeSlot(@Nonnull final Transaction transaction, @Nonnull final DataNode node, + public void writeLeafNodeSlot(@Nonnull final Transaction transaction, @Nonnull final CompactNode node, @Nonnull final ItemSlot itemSlot) { persistNode(transaction, node); } @Override - public void clearLeafNodeSlot(@Nonnull final Transaction transaction, @Nonnull final DataNode node, + public void clearLeafNodeSlot(@Nonnull final Transaction transaction, @Nonnull final CompactNode node, @Nonnull final ItemSlot itemSlot) { persistNode(transaction, node); } @@ -170,7 +170,7 @@ private Node nodeFromTuple(@Nonnull final Node.Node vectorTuple = tuple.getNestedTuple(2); neighborsTuple = tuple.getNestedTuple(3); return dataNodeFromTuples(creator, primaryKey, vectorTuple, neighborsTuple); - case INTERMEDIATE: + case INLINING: neighborsTuple = tuple.getNestedTuple(3); return intermediateNodeFromTuples(creator, primaryKey, neighborsTuple); default: @@ -216,7 +216,7 @@ private Node intermediateNodeFromTuples(@Nonnull fi neighborsWithVectors.add(new NodeReferenceWithVector(neighborPrimaryKey, new Vector.HalfVector(neighborVectorHalfs))); } - return creator.create(NodeKind.INTERMEDIATE, primaryKey, null, neighborsWithVectors); + return creator.create(NodeKind.INLINING, primaryKey, null, neighborsWithVectors); } @Nonnull diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java similarity index 75% rename from fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java rename to fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java index feb2f79cb0..343b0e12f6 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DataNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java @@ -1,5 +1,5 @@ /* - * DataNode.java + * CompactNode.java * * This source file is part of the FoundationDB open source project * @@ -32,12 +32,12 @@ /** * TODO. */ -class DataNode extends AbstractNode { +class CompactNode extends AbstractNode { @Nonnull private final Vector vector; - public DataNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, - @Nonnull final List nodeReferences) { + public CompactNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, + @Nonnull final List nodeReferences) { super(primaryKey, nodeReferences); this.vector = vector; } @@ -55,19 +55,19 @@ public Vector getVector() { @Nonnull @Override - public DataNode asDataNode() { + public CompactNode asCompactNode() { return this; } @Nonnull @Override - public IntermediateNode asIntermediateNode() { - throw new IllegalStateException("this is not a data node"); + public InliningNode asInliningNode() { + throw new IllegalStateException("this is not an inlining node"); } @Override public NodeCreator sameCreator() { - return DataNode::creator; + return CompactNode::creator; } @Nonnull @@ -76,7 +76,7 @@ public static Node creator(@Nonnull final NodeKind nodeKind, @Nonnull final Tuple primaryKey, @Nullable final Vector vector, @Nonnull final List neighbors) { - Verify.verify(nodeKind == NodeKind.INTERMEDIATE); - return new DataNode(primaryKey, Objects.requireNonNull(vector), (List)neighbors); + Verify.verify(nodeKind == NodeKind.INLINING); + return new CompactNode(primaryKey, Objects.requireNonNull(vector), (List)neighbors); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index 4ed82af6a7..087844f1f2 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -498,7 +498,7 @@ public AsyncIterator scan(@Nonnull final ReadTransaction readTransacti @Nonnull final BiPredicate suffixKeyPredicate) { Preconditions.checkArgument((lastHilbertValue == null && lastKey == null) || (lastHilbertValue != null && lastKey != null)); - AsyncIterator leafIterator = + AsyncIterator leafIterator = new LeafIterator(readTransaction, rootId, lastHilbertValue, lastKey, mbrPredicate, suffixKeyPredicate); return new ItemSlotIterator(leafIterator); } @@ -544,7 +544,7 @@ private CompletableFuture> kNearestNeighborsSearch(@ return CompletableFuture.completedFuture(null); } - return searchLayer(DataNode::creator, readTransaction, + return searchLayer(CompactNode::creator, readTransaction, ImmutableList.of(greedyState.toNodeReferenceWithDistance()), 0, efSearch, queryVector); }); @@ -564,14 +564,14 @@ private CompletableFuture greedySearchLayer(@Nonnull final ReadTran new AtomicReference<>(new GreedyState(layer, entryNeighbor.getPrimaryKey(), entryNeighbor.getDistance())); return AsyncUtil.whileTrue(() -> onReadListener.onAsyncRead( - storageAdapter.fetchNode(IntermediateNode::creator, readTransaction, + storageAdapter.fetchNode(InliningNode::creator, readTransaction, layer, greedyStateReference.get().getPrimaryKey())) .thenApply(node -> { if (node == null) { throw new IllegalStateException("unable to fetch node"); } - final IntermediateNode intermediateNode = node.asIntermediateNode(); - final List neighbors = intermediateNode.getNeighbors(); + final InliningNode inliningNode = node.asInliningNode(); + final List neighbors = inliningNode.getNeighbors(); final GreedyState currentNodeKey = greedyStateReference.get(); double minDistance = currentNodeKey.getDistance(); @@ -637,7 +637,7 @@ private CompletableFuture> searchLayer .thenApply(candidateNode -> Iterables.filter(candidateNode.getNeighbors(), neighbor -> !visited.contains(neighbor.getPrimaryKey()))) - .thenCompose(neighborReferences -> fetchSomeNeighbors(creator, readTransaction, + .thenCompose(neighborReferences -> neighborhood(creator, readTransaction, layer, neighborReferences, nodeCache)) .thenApply(neighborReferences -> { for (final NodeReferenceWithVector current : neighborReferences) { @@ -701,11 +701,11 @@ private CompletableFuture< * TODO. */ @Nonnull - private CompletableFuture> fetchSomeNeighbors(@Nonnull final Node.NodeCreator creator, - @Nonnull final ReadTransaction readTransaction, - final int layer, - @Nonnull final Iterable neighborReferences, - @Nonnull final Map> nodeCache) { + private CompletableFuture> neighborhood(@Nonnull final Node.NodeCreator creator, + @Nonnull final ReadTransaction readTransaction, + final int layer, + @Nonnull final Iterable neighborReferences, + @Nonnull final Map> nodeCache) { return fetchSomeNodesAndApply(creator, readTransaction, layer, neighborReferences, neighborReference -> { if (neighborReference instanceof NodeReferenceWithVector) { @@ -715,10 +715,10 @@ private CompletableFuture - new NodeReferenceWithVector(neighborReference.getPrimaryKey(), neighborNode.asDataNode().getVector())); + new NodeReferenceWithVector(neighborReference.getPrimaryKey(), neighborNode.asCompactNode().getVector())); } /** diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java similarity index 73% rename from fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java rename to fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java index a1de1a5c2a..6619a2003e 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/IntermediateNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java @@ -1,5 +1,5 @@ /* - * IntermediateNode.java + * InliningNode.java * * This source file is part of the FoundationDB open source project * @@ -31,33 +31,33 @@ /** * TODO. */ -class IntermediateNode extends AbstractNode { - public IntermediateNode(@Nonnull final Tuple primaryKey, - @Nonnull final List neighbors) { +class InliningNode extends AbstractNode { + public InliningNode(@Nonnull final Tuple primaryKey, + @Nonnull final List neighbors) { super(primaryKey, neighbors); } @Nonnull @Override public NodeKind getKind() { - return NodeKind.INTERMEDIATE; + return NodeKind.INLINING; } @Nonnull @Override - public DataNode asDataNode() { - throw new IllegalStateException("this is not a data node"); + public CompactNode asCompactNode() { + throw new IllegalStateException("this is not a compact node"); } @Nonnull @Override - public IntermediateNode asIntermediateNode() { + public InliningNode asInliningNode() { return this; } @Override public NodeCreator sameCreator() { - return IntermediateNode::creator; + return InliningNode::creator; } @Nonnull @@ -66,7 +66,7 @@ public static Node creator(@Nonnull final NodeKind node @Nonnull final Tuple primaryKey, @Nullable final Vector vector, @Nonnull final List neighbors) { - Verify.verify(nodeKind == NodeKind.INTERMEDIATE); - return new IntermediateNode(primaryKey, (List)neighbors); + Verify.verify(nodeKind == NodeKind.INLINING); + return new InliningNode(primaryKey, (List)neighbors); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java index d3f24e66b6..8c10378d31 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java @@ -55,17 +55,17 @@ public interface Node { Node delete(@Nonnull StorageAdapter storageAdapter, int level, int slotIndex); /** - * Return the kind of the node, i.e. {@link NodeKind#DATA} or {@link NodeKind#INTERMEDIATE}. + * Return the kind of the node, i.e. {@link NodeKind#DATA} or {@link NodeKind#INLINING}. * @return the kind of this node as a {@link NodeKind} */ @Nonnull NodeKind getKind(); @Nonnull - DataNode asDataNode(); + CompactNode asCompactNode(); @Nonnull - IntermediateNode asIntermediateNode(); + InliningNode asInliningNode(); NodeCreator sameCreator(); diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKind.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKind.java index 98f0c1adfd..0a9f6031e6 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKind.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKind.java @@ -29,7 +29,7 @@ */ public enum NodeKind { DATA((byte)0x00), - INTERMEDIATE((byte)0x01); + INLINING((byte)0x01); private final byte serialized; @@ -49,7 +49,7 @@ static NodeKind fromSerializedNodeKind(byte serializedNodeKind) { nodeKind = NodeKind.DATA; break; case 0x01: - nodeKind = NodeKind.INTERMEDIATE; + nodeKind = NodeKind.INLINING; break; default: throw new IllegalArgumentException("unknown node kind"); diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java index cea4a56f3c..c2c34fa669 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java @@ -114,7 +114,7 @@ CompletableFuture> fetchNode(@Nonnull Node.Nod * @param node node whose slot to persist * @param itemSlot the node slot to persist */ - void writeLeafNodeSlot(@Nonnull Transaction transaction, @Nonnull DataNode node, @Nonnull ItemSlot itemSlot); + void writeLeafNodeSlot(@Nonnull Transaction transaction, @Nonnull CompactNode node, @Nonnull ItemSlot itemSlot); /** * Clear out a leaf node slot. @@ -123,7 +123,7 @@ CompletableFuture> fetchNode(@Nonnull Node.Nod * @param node node whose slot is cleared out * @param itemSlot the node slot to clear out */ - void clearLeafNodeSlot(@Nonnull Transaction transaction, @Nonnull DataNode node, @Nonnull ItemSlot itemSlot); + void clearLeafNodeSlot(@Nonnull Transaction transaction, @Nonnull CompactNode node, @Nonnull ItemSlot itemSlot); /** * Method to (re-)persist a list of nodes passed in. From 66547ca37c61577cedf3f4dab93ea25ffbcc3b7c Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Mon, 28 Jul 2025 18:21:57 +0200 Subject: [PATCH 09/34] started on the writing path --- .../foundationdb/async/MoreAsyncUtil.java | 22 + .../async/hnsw/AbstractStorageAdapter.java | 88 +--- .../async/hnsw/ByNodeStorageAdapter.java | 301 +---------- .../foundationdb/async/hnsw/CompactNode.java | 49 +- ...tAndLayer.java => EntryNodeReference.java} | 21 +- .../apple/foundationdb/async/hnsw/HNSW.java | 473 +++++++++--------- .../foundationdb/async/hnsw/InliningNode.java | 46 +- .../apple/foundationdb/async/hnsw/Metric.java | 30 ++ .../apple/foundationdb/async/hnsw/Node.java | 79 ++- .../{GreedyState.java => NodeFactory.java} | 34 +- .../foundationdb/async/hnsw/NodeKind.java | 4 +- .../async/hnsw/OnWriteListener.java | 2 +- .../foundationdb/async/hnsw/SearchResult.java | 8 +- .../async/hnsw/StorageAdapter.java | 104 ++-- .../apple/foundationdb/async/hnsw/Vector.java | 5 + 15 files changed, 511 insertions(+), 755 deletions(-) rename fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/{EntryPointAndLayer.java => EntryNodeReference.java} (68%) rename fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/{GreedyState.java => NodeFactory.java} (55%) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java index 563dec11a6..f6f0c999e5 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java @@ -42,9 +42,13 @@ import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.ThreadFactory; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiConsumer; import java.util.function.BiFunction; import java.util.function.Function; +import java.util.function.IntFunction; +import java.util.function.IntPredicate; +import java.util.function.IntUnaryOperator; import java.util.function.Predicate; import java.util.function.Supplier; @@ -1051,6 +1055,24 @@ public static CompletableFuture swallowException(@Nonnull CompletableFutur return result; } + public static CompletableFuture forLoop(int startI, @Nonnull final IntPredicate conditionPredicate, + @Nonnull final IntUnaryOperator stepFunction, + @Nonnull final IntFunction> body, + @Nonnull final Executor executor) { + final AtomicInteger loopVariableAtomic = new AtomicInteger(startI); + return AsyncUtil.whileTrue(() -> { + final int loopVariable = loopVariableAtomic.get(); + if (!conditionPredicate.test(loopVariable)) { + return AsyncUtil.READY_FALSE; + } + return body.apply(loopVariable) + .thenApply(ignored -> { + loopVariableAtomic.set(stepFunction.applyAsInt(loopVariable)); + return true; + }); + }, executor); + } + /** * A {@code Boolean} function that is always true. * @param the type of the (ignored) argument to the function diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java index 98ba205405..383c22f27c 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java @@ -21,17 +21,12 @@ package com.apple.foundationdb.async.hnsw; import com.apple.foundationdb.ReadTransaction; -import com.apple.foundationdb.Transaction; import com.apple.foundationdb.subspace.Subspace; import com.apple.foundationdb.tuple.Tuple; import javax.annotation.Nonnull; import javax.annotation.Nullable; -import java.math.BigInteger; -import java.util.List; -import java.util.Objects; import java.util.concurrent.CompletableFuture; -import java.util.function.Function; /** * Implementations and attributes common to all concrete implementations of {@link StorageAdapter}. @@ -53,8 +48,6 @@ abstract class AbstractStorageAdapter implements StorageAdapter { private final Subspace dataSubspace; protected AbstractStorageAdapter(@Nonnull final HNSW.Config config, @Nonnull final Subspace subspace, - @Nonnull final Subspace nodeSlotIndexSubspace, - @Nonnull final Function hilbertValueFunction, @Nonnull final OnWriteListener onWriteListener, @Nonnull final OnReadListener onReadListener) { this.config = config; @@ -108,79 +101,16 @@ public OnReadListener getOnReadListener() { return onReadListener; } - @Override - public void writeNodes(@Nonnull final Transaction transaction, @Nonnull final List nodes) { - for (final Node node : nodes) { - writeNode(transaction, node); - } - } - - protected void writeNode(@Nonnull final Transaction transaction, @Nonnull final Node node) { - final Node.ChangeSet changeSet = node.getChangeSet(); - if (changeSet == null) { - return; - } - - changeSet.apply(transaction); - getOnWriteListener().onNodeWritten(node); - } - - @Nonnull - public byte[] packWithSubspace(final byte[] key) { - return getSubspace().pack(key); - } - - @Nonnull - public byte[] packWithSubspace(final Tuple tuple) { - return getSubspace().pack(tuple); - } - - @Nonnull - @Override - public CompletableFuture scanNodeIndexAndFetchNode(@Nonnull final ReadTransaction transaction, - final int level, - @Nonnull final BigInteger hilbertValue, - @Nonnull final Tuple key, - final boolean isInsertUpdate) { - Objects.requireNonNull(nodeSlotIndexAdapter); - return nodeSlotIndexAdapter.scanIndexForNodeId(transaction, level, hilbertValue, key, isInsertUpdate) - .thenCompose(nodeId -> nodeId == null - ? CompletableFuture.completedFuture(null) - : fetchNode(transaction, nodeId)); - } - - @Override - public void insertIntoNodeIndexIfNecessary(@Nonnull final Transaction transaction, final int level, - @Nonnull final NodeSlot nodeSlot) { - if (!getConfig().isUseNodeSlotIndex() || !(nodeSlot instanceof ChildSlot)) { - return; - } - - Objects.requireNonNull(nodeSlotIndexAdapter); - nodeSlotIndexAdapter.writeChildSlot(transaction, level, (ChildSlot)nodeSlot); - } - - @Override - public void deleteFromNodeIndexIfNecessary(@Nonnull final Transaction transaction, final int level, - @Nonnull final NodeSlot nodeSlot) { - if (!getConfig().isUseNodeSlotIndex() || !(nodeSlot instanceof ChildSlot)) { - return; - } - - Objects.requireNonNull(nodeSlotIndexAdapter); - nodeSlotIndexAdapter.clearChildSlot(transaction, level, (ChildSlot)nodeSlot); - } - @Nonnull @Override - public CompletableFuture> fetchNode(@Nonnull final Node.NodeCreator creator, + public CompletableFuture> fetchNode(@Nonnull final NodeFactory nodeFactory, @Nonnull final ReadTransaction readTransaction, int layer, @Nonnull Tuple primaryKey) { - return fetchNodeInternal(creator, readTransaction, layer, primaryKey).thenApply(this::checkNode); + return fetchNodeInternal(nodeFactory, readTransaction, layer, primaryKey).thenApply(this::checkNode); } @Nonnull - protected abstract CompletableFuture> fetchNodeInternal(@Nonnull Node.NodeCreator creator, + protected abstract CompletableFuture> fetchNodeInternal(@Nonnull NodeFactory nodeFactory, @Nonnull ReadTransaction readTransaction, int layer, @Nonnull Tuple primaryKey); @@ -197,16 +127,4 @@ protected abstract CompletableFuture> fetchNod private Node checkNode(@Nullable final Node node) { return node; } - - @Nonnull - abstract > AbstractChangeSet - newInsertChangeSet(@Nonnull N node, int level, @Nonnull List insertedSlots); - - @Nonnull - abstract > AbstractChangeSet - newUpdateChangeSet(@Nonnull N node, int level, @Nonnull S originalSlot, @Nonnull S updatedSlot); - - @Nonnull - abstract > AbstractChangeSet - newDeleteChangeSet(@Nonnull N node, int level, @Nonnull List deletedSlots); } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java index 54d8f4895b..6b0d0ff898 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java @@ -24,58 +24,22 @@ import com.apple.foundationdb.Transaction; import com.apple.foundationdb.subspace.Subspace; import com.apple.foundationdb.tuple.Tuple; -import com.christianheina.langx.half4j.Half; -import com.google.common.base.Verify; -import com.google.common.collect.ImmutableList; -import com.google.common.collect.Lists; -import com.google.common.collect.Streams; import javax.annotation.Nonnull; -import java.math.BigInteger; -import java.util.List; import java.util.concurrent.CompletableFuture; -import java.util.function.Function; /** - * Storage adapter that represents each node as a single key/value pair in the database. That paradigm - * greatly simplifies how nodes, node slots, and operations on these data structures are managed. Also, a - * node that is serialized into a key/value pair using this storage adapter leaves a significantly smaller - * footprint when compared to the multitude of key/value pairs that are serialized/deserialized using - * {@link BySlotStorageAdapter}. - *
- * These advantages are offset by the key disadvantage of having to read/deserialize and the serialize/write - * the entire node to realize/persist a minute change to one of its slots. That, in turn, may cause a higher - * likelihood of conflicting with another transaction. - *
- * Each node is serialized as follows (each {@code (thing)} is a tuple containing {@code thing}): - * 
- * {@code
- *    key: nodeId: byte[16]
- *    value: Tuple(nodeKind: Long, slotList: (slot1, ..., slotn])
- *           slot:
- *               for leaf nodes:
- *                   (hilbertValue: BigInteger, itemKey: (point: Tuple(d1, ..., dk),
- *                    keySuffix: (...)),
- *                    value: (...))
- *               for intermediate nodes:
- *                   (smallestHV: BigInteger, smallestKey: (...),
- *                    largestHV: BigInteger, largestKey: (...),
- *                    childId: byte[16],
- *                    mbr: (minD1, minD2, ..., minDk, maxD1, maxD2, ..., maxDk)))
- * }
- *
+ * TODO. */ class ByNodeStorageAdapter extends AbstractStorageAdapter implements StorageAdapter { public ByNodeStorageAdapter(@Nonnull final HNSW.Config config, @Nonnull final Subspace subspace, - @Nonnull final Subspace nodeSlotIndexSubspace, - @Nonnull final Function hilbertValueFunction, @Nonnull final OnWriteListener onWriteListener, @Nonnull final OnReadListener onReadListener) { - super(config, subspace, nodeSlotIndexSubspace, hilbertValueFunction, onWriteListener, onReadListener); + super(config, subspace, onWriteListener, onReadListener); } @Override - public CompletableFuture fetchEntryNodeKey(@Nonnull final ReadTransaction readTransaction) { + public CompletableFuture fetchEntryNodeReference(@Nonnull final ReadTransaction readTransaction) { final byte[] key = getEntryNodeSubspace().pack(); return readTransaction.get(key) @@ -90,13 +54,24 @@ public CompletableFuture fetchEntryNodeKey(@Nonnull final Re final int lMax = (int)entryTuple.getLong(0); final Tuple primaryKey = entryTuple.getNestedTuple(1); final Tuple vectorTuple = entryTuple.getNestedTuple(2); - return new EntryPointAndLayer(lMax, primaryKey, vectorFromTuple(vectorTuple)); + return new EntryNodeReference(primaryKey, StorageAdapter.vectorFromTuple(vectorTuple), lMax); }); } + + @Override + public void writeEntryNodeReference(@Nonnull final Transaction transaction, + @Nonnull final EntryNodeReference entryNodeReference) { + transaction.set(getEntryNodeSubspace().pack(), + Tuple.from(entryNodeReference.getLayer(), + entryNodeReference.getPrimaryKey(), + StorageAdapter.tupleFromVector(entryNodeReference.getVector())).pack()); + } + + @Nonnull @Override - protected CompletableFuture> fetchNodeInternal(@Nonnull final Node.NodeCreator creator, + protected CompletableFuture> fetchNodeInternal(@Nonnull final NodeFactory nodeFactory, @Nonnull final ReadTransaction readTransaction, final int layer, @Nonnull final Tuple primaryKey) { @@ -109,7 +84,7 @@ protected CompletableFuture> fetchNodeInternal } final Tuple nodeTuple = Tuple.fromBytes(valueBytes); - final Node node = nodeFromTuple(creator, nodeTuple); + final Node node = Node.nodeFromTuples(nodeFactory, primaryKey, nodeTuple); final OnReadListener onReadListener = getOnReadListener(); onReadListener.onNodeRead(node); onReadListener.onKeyValueRead(key, valueBytes); @@ -118,242 +93,10 @@ protected CompletableFuture> fetchNodeInternal } @Override - public void writeLeafNodeSlot(@Nonnull final Transaction transaction, @Nonnull final CompactNode node, - @Nonnull final ItemSlot itemSlot) { - persistNode(transaction, node); - } - - @Override - public void clearLeafNodeSlot(@Nonnull final Transaction transaction, @Nonnull final CompactNode node, - @Nonnull final ItemSlot itemSlot) { - persistNode(transaction, node); - } - - private void persistNode(@Nonnull final Transaction transaction, @Nonnull final Node node) { - final byte[] packedKey = packWithSubspace(node.getId()); - - if (node.isEmpty()) { - // this can only happen when we just deleted the last slot; delete the entire node - transaction.clear(packedKey); - getOnWriteListener().onKeyCleared(node, packedKey); - } else { - // updateNodeIndexIfNecessary(transaction, level, node); - final byte[] packedValue = toTuple(node).pack(); - transaction.set(packedKey, packedValue); - getOnWriteListener().onKeyValueWritten(node, packedKey, packedValue); - } - } - - @Nonnull - private Tuple toTuple(@Nonnull final Node node) { - final HNSW.Config config = getConfig(); - final List slotTuples = Lists.newArrayListWithExpectedSize(node.size()); - for (final NodeSlot nodeSlot : node.getSlots()) { - final Tuple slotTuple = Tuple.fromStream( - Streams.concat(nodeSlot.getSlotKey(config.isStoreHilbertValues()).getItems().stream(), - nodeSlot.getSlotValue().getItems().stream())); - slotTuples.add(slotTuple); - } - return Tuple.from(node.getKind().getSerialized(), slotTuples); - } - - @Nonnull - private Node nodeFromTuple(@Nonnull final Node.NodeCreator creator, - @Nonnull final Tuple tuple) { - final NodeKind nodeKind = NodeKind.fromSerializedNodeKind((byte)tuple.getLong(0)); - final Tuple primaryKey = tuple.getNestedTuple(1); - final Tuple vectorTuple; - final Tuple neighborsTuple; - - switch (nodeKind) { - case DATA: - vectorTuple = tuple.getNestedTuple(2); - neighborsTuple = tuple.getNestedTuple(3); - return dataNodeFromTuples(creator, primaryKey, vectorTuple, neighborsTuple); - case INLINING: - neighborsTuple = tuple.getNestedTuple(3); - return intermediateNodeFromTuples(creator, primaryKey, neighborsTuple); - default: - throw new IllegalStateException("unknown node kind"); - } - } - - @Nonnull - private Node dataNodeFromTuples(@Nonnull final Node.NodeCreator creator, - @Nonnull final Tuple primaryKey, - @Nonnull final Tuple vectorTuple, - @Nonnull final Tuple neighborsTuple) { - final Vector vector = vectorFromTuple(vectorTuple); - - List nodeReferences = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); - - for (final Object neighborObject : neighborsTuple) { - final Tuple neighborTuple = (Tuple)neighborObject; - nodeReferences.add(new NodeReference(neighborTuple)); - } - - return creator.create(NodeKind.DATA, primaryKey, vector, nodeReferences); - } - - @Nonnull - private Node intermediateNodeFromTuples(@Nonnull final Node.NodeCreator creator, - @Nonnull final Tuple primaryKey, - @Nonnull final Tuple neighborsTuple) { - List neighborsWithVectors = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); - Half[] neighborVectorHalfs = null; - - for (final Object neighborObject : neighborsTuple) { - final Tuple neighborTuple = (Tuple)neighborObject; - final Tuple neighborPrimaryKey = neighborTuple.getNestedTuple(0); - final Tuple neighborVectorTuple = neighborTuple.getNestedTuple(1); - if (neighborVectorHalfs == null) { - neighborVectorHalfs = new Half[neighborVectorTuple.size()]; - } - - for (int i = 0; i < neighborVectorTuple.size(); i ++) { - neighborVectorHalfs[i] = Half.shortBitsToHalf(shortFromBytes(neighborVectorTuple.getBytes(i))); - } - neighborsWithVectors.add(new NodeReferenceWithVector(neighborPrimaryKey, new Vector.HalfVector(neighborVectorHalfs))); - } - - return creator.create(NodeKind.INLINING, primaryKey, null, neighborsWithVectors); - } - - @Nonnull - private Vector vectorFromTuple(final Tuple vectorTuple) { - final Half[] vectorHalfs = new Half[vectorTuple.size()]; - for (int i = 0; i < vectorTuple.size(); i ++) { - vectorHalfs[i] = Half.shortBitsToHalf(shortFromBytes(vectorTuple.getBytes(i))); - } - return new Vector.HalfVector(vectorHalfs); - } - - - - @Nonnull - @Override - public > AbstractChangeSet - newInsertChangeSet(@Nonnull final N node, final int level, @Nonnull final List insertedSlots) { - return new InsertChangeSet<>(node, level, insertedSlots); - } - - @Nonnull - @Override - public > AbstractChangeSet - newUpdateChangeSet(@Nonnull final N node, final int level, - @Nonnull final S originalSlot, @Nonnull final S updatedSlot) { - return new UpdateChangeSet<>(node, level, originalSlot, updatedSlot); - } - - @Nonnull - @Override - public > AbstractChangeSet - newDeleteChangeSet(@Nonnull final N node, final int level, @Nonnull final List deletedSlots) { - return new DeleteChangeSet<>(node, level, deletedSlots); - } - - private class InsertChangeSet> extends AbstractChangeSet { - @Nonnull - private final List insertedSlots; - - public InsertChangeSet(@Nonnull final N node, final int level, @Nonnull final List insertedSlots) { - super(node.getChangeSet(), node, level); - this.insertedSlots = ImmutableList.copyOf(insertedSlots); - } - - @Override - public void apply(@Nonnull final Transaction transaction) { - super.apply(transaction); - - // - // If this change set is the first, we persist the node, don't persist the node otherwise. This is a - // performance optimization to avoid writing and rewriting the node for each change set in the chain - // of change sets. - // - if (getPreviousChangeSet() == null) { - persistNode(transaction, getNode()); - } - if (isUpdateNodeSlotIndex()) { - for (final S insertedSlot : insertedSlots) { - insertIntoNodeIndexIfNecessary(transaction, getLevel(), insertedSlot); - } - } - } - } - - private class UpdateChangeSet> extends AbstractChangeSet { - @Nonnull - private final S originalSlot; - @Nonnull - private final S updatedSlot; - - public UpdateChangeSet(@Nonnull final N node, final int level, @Nonnull final S originalSlot, - @Nonnull final S updatedSlot) { - super(node.getChangeSet(), node, level); - this.originalSlot = originalSlot; - this.updatedSlot = updatedSlot; - } - - @Override - public void apply(@Nonnull final Transaction transaction) { - super.apply(transaction); - - // - // If this change set is the first, we persist the node, don't persist the node otherwise. This is a - // performance optimization to avoid writing and rewriting the node for each change set in the chain - // of change sets. - // - if (getPreviousChangeSet() == null) { - persistNode(transaction, getNode()); - } - if (isUpdateNodeSlotIndex()) { - deleteFromNodeIndexIfNecessary(transaction, getLevel(), originalSlot); - insertIntoNodeIndexIfNecessary(transaction, getLevel(), updatedSlot); - } - } - } - - private class DeleteChangeSet> extends AbstractChangeSet { - @Nonnull - private final List deletedSlots; - - public DeleteChangeSet(@Nonnull final N node, final int level, @Nonnull final List deletedSlots) { - super(node.getChangeSet(), node, level); - this.deletedSlots = ImmutableList.copyOf(deletedSlots); - } - - @Override - public void apply(@Nonnull final Transaction transaction) { - super.apply(transaction); - - // - // If this change set is the first, we persist the node, don't persist the node otherwise. This is a - // performance optimization to avoid writing and rewriting the node for each change set in the chain - // of change sets. - // - if (getPreviousChangeSet() == null) { - persistNode(transaction, getNode()); - } - if (isUpdateNodeSlotIndex()) { - for (final S deletedSlot : deletedSlots) { - deleteFromNodeIndexIfNecessary(transaction, getLevel(), deletedSlot); - } - } - } - } - - private short shortFromBytes(byte[] bytes) { - Verify.verify(bytes.length == 2); - int high = bytes[0] & 0xFF; // Convert to unsigned int - int low = bytes[1] & 0xFF; - - return (short) ((high << 8) | low); - } - - private byte[] bytesFromShort(short value) { - byte[] result = new byte[2]; - result[0] = (byte) ((value >> 8) & 0xFF); // high byte first - result[1] = (byte) (value & 0xFF); // low byte second - return result; + public void writeNode(@Nonnull Transaction transaction, @Nonnull final Node node, + final int layer) { + final byte[] key = getDataSubspace().pack(Tuple.from(layer, node.getPrimaryKey())); + transaction.set(key, node.toTuple().pack()); + getOnWriteListener().onNodeWritten(node); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java index 343b0e12f6..2bf332751a 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java @@ -23,6 +23,7 @@ import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; import com.google.common.base.Verify; +import com.google.common.collect.Lists; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -32,7 +33,24 @@ /** * TODO. */ -class CompactNode extends AbstractNode { +public class CompactNode extends AbstractNode { + @Nonnull + private static final NodeFactory FACTORY = new NodeFactory<>() { + @SuppressWarnings("unchecked") + @Nonnull + @Override + public Node create(@Nonnull final NodeKind nodeKind, @Nonnull final Tuple primaryKey, @Nullable final Vector vector, @Nonnull final List neighbors) { + Verify.verify(nodeKind == NodeKind.COMPACT); + return new CompactNode(primaryKey, Objects.requireNonNull(vector), (List)neighbors); + } + + @Nonnull + @Override + public NodeKind getNodeKind() { + return NodeKind.COMPACT; + } + }; + @Nonnull private final Vector vector; @@ -45,7 +63,7 @@ public CompactNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector @Nonnull @Override public NodeKind getKind() { - return NodeKind.DATA; + return NodeKind.COMPACT; } @Nonnull @@ -66,17 +84,26 @@ public InliningNode asInliningNode() { } @Override - public NodeCreator sameCreator() { - return CompactNode::creator; + public NodeFactory sameCreator() { + return CompactNode.factory(); + } + + @Nonnull + @Override + public Tuple toTuple() { + final List nodeItems = Lists.newArrayListWithExpectedSize(4); + nodeItems.add(NodeKind.COMPACT.getSerialized()); + nodeItems.add(StorageAdapter.tupleFromVector(getVector())); + final List neighborItems = Lists.newArrayListWithExpectedSize(getNeighbors().size()); + for (final NodeReference nodeReference : getNeighbors()) { + neighborItems.add(nodeReference.getPrimaryKey()); + } + nodeItems.add(Tuple.fromList(neighborItems)); + return Tuple.fromList(nodeItems); } @Nonnull - @SuppressWarnings("unchecked") - public static Node creator(@Nonnull final NodeKind nodeKind, - @Nonnull final Tuple primaryKey, - @Nullable final Vector vector, - @Nonnull final List neighbors) { - Verify.verify(nodeKind == NodeKind.INLINING); - return new CompactNode(primaryKey, Objects.requireNonNull(vector), (List)neighbors); + public static NodeFactory factory() { + return FACTORY; } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/EntryPointAndLayer.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/EntryNodeReference.java similarity index 68% rename from fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/EntryPointAndLayer.java rename to fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/EntryNodeReference.java index f621916079..3cbbb3c2dd 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/EntryPointAndLayer.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/EntryNodeReference.java @@ -25,30 +25,15 @@ import javax.annotation.Nonnull; -class EntryPointAndLayer { +class EntryNodeReference extends NodeReferenceWithVector { private final int layer; - @Nonnull - private final Tuple primaryKey; - @Nonnull - private final Vector vector; - public EntryPointAndLayer(final int layer, @Nonnull final Tuple primaryKey, @Nonnull final Vector vector) { + public EntryNodeReference(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, final int layer) { + super(primaryKey, vector); this.layer = layer; - this.primaryKey = primaryKey; - this.vector = vector; } public int getLayer() { return layer; } - - @Nonnull - public Tuple getPrimaryKey() { - return primaryKey; - } - - @Nonnull - public Vector getVector() { - return vector; - } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index 087844f1f2..5b9400cd3b 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -21,9 +21,11 @@ package com.apple.foundationdb.async.hnsw; import com.apple.foundationdb.ReadTransaction; +import com.apple.foundationdb.Transaction; import com.apple.foundationdb.annotation.API; import com.apple.foundationdb.async.AsyncIterator; import com.apple.foundationdb.async.AsyncUtil; +import com.apple.foundationdb.async.MoreAsyncUtil; import com.apple.foundationdb.subspace.Subspace; import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; @@ -47,6 +49,8 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Queue; +import java.util.Random; import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; @@ -66,72 +70,14 @@ public class HNSW { private static final Logger logger = LoggerFactory.getLogger(HNSW.class); - /** - * root id. The root id is always only zeros. - */ - static final byte[] rootId = new byte[16]; - public static final int MAX_CONCURRENT_READS = 16; - - /** - * Indicator if we should maintain a secondary node index consisting of hilbet value and key to speed up - * update/deletes. - */ - public static final boolean DEFAULT_USE_NODE_SLOT_INDEX = false; - - /** - * The minimum number of slots a node has (if not the root node). {@code M} should be chosen in a way that the - * minimum is half of the maximum. That in turn guarantees that overflow/underflow handling can be performed without - * causing further underflow/overflow. - */ - public static final int DEFAULT_MIN_M = 16; - /** - * The maximum number of slots a node has. This value is derived from {@link #DEFAULT_MIN_M}. - */ - public static final int DEFAULT_MAX_M = 2 * DEFAULT_MIN_M; - - /** - * The magic split number. We split {@code S} to {@code S + 1} nodes while inserting data and fuse - * {@code S + 1} to {@code S} nodes while deleting data. Academically, 2-to-3 splits and 3-to-2 fuses - * seem to yield the best results. Please be aware of the following constraints: - *
    - *
  1. When splitting {@code S} to {@code S + 1} nodes, we re-distribute the children of {@code S} nodes - * into {@code S + 1} nodes which may cause an underflow if {@code S} and {@code M} are not set carefully with - * respect to each other. Example: {@code MIN_M = 25}, {@code MAX_M = 32}, {@code S = 2}, two nodes at - * already at maximum capacity containing a combined total of 64 children when a new child is inserted. - * We split the two nodes into three as indicated by {@code S = 2}. We have 65 children but there is no way - * of distributing them among three nodes such that none of them underflows. This constraint can be - * formulated as {@code S * MAX_M / (S + 1) >= MIN_M}.
    2. - *
  2. When fusing {@code S + 1} to {@code S} nodes, we re-distribute the children of {@code S + 1} nodes - * into {@code S + 1} nodes which may cause an overflow if {@code S} and {@code M} are not set carefully with - * respect to each other. Example: {@code MIN_M = 25}, {@code MAX_M = 32}, {@code S = 2}, three nodes at - * already at minimum capacity containing a combined total of 75 children when a child is deleted. - * We fuse the three nodes into two as indicated by {@code S = 2}. We have 75 children but there is no way - * of distributing them among two nodes such that none of them overflows. This constraint can be formulated as - * {@code (S + 1) * MIN_M / S <= MAX_M}.
    3. - *
- * Both constraints are in fact the same constraint and can be written as {@code MAX_M / MIN_M >= (S + 1) / S}. - */ - public static final int DEFAULT_S = 2; - - /** - * Default storage layout. Can be either {@code BY_SLOT} or {@code BY_NODE}. {@code BY_SLOT} encodes all information - * pertaining to a {@link NodeSlot} as one key/value pair in the database; {@code BY_NODE} encodes all information - * pertaining to a {@link Node} as one key/value pair in the database. While {@code BY_SLOT} avoids conflicts as - * most inserts/updates only need to update one slot, it is by far less compact as some information is stored - * in a normalized fashion and therefore repeated multiple times (i.e. node identifiers, etc.). {@code BY_NODE} - * inlines slot information into the node leading to a more size-efficient layout of the data. That advantage is - * offset by a higher likelihood of conflicts. - */ @Nonnull - public static final Storage DEFAULT_STORAGE = Storage.BY_NODE; - - /** - * Indicator if Hilbert values should be stored or not with the data (in leaf nodes). A Hilbert value can always - * be recomputed from the point. - */ - public static final boolean DEFAULT_STORE_HILBERT_VALUES = true; - + public static final Random DEFAULT_RANDOM = new Random(0L); + @Nonnull + public static final Metric DEFAULT_METRIC = new Metric.EuclideanMetric(); + public static final int DEFAULT_M = 48; + public static final int DEFAULT_EF_SEARCH = 64; + public static final int DEFAULT_EF_CONSTRUCTION = 100; @Nonnull public static final Config DEFAULT_CONFIG = new Config(); @@ -142,8 +88,6 @@ public class HNSW { @Nonnull private final Config config; @Nonnull - private final Supplier nodeIdSupplier; - @Nonnull private final OnWriteListener onWriteListener; @Nonnull private final OnReadListener onReadListener; @@ -152,10 +96,6 @@ public class HNSW { * Different kinds of storage layouts. */ public enum Storage { - /** - * Every node slot is serialized as a key/value pair in FDB. - */ - BY_SLOT(BySlotStorageAdapter::new), /** * Every node with all its slots is serialized as one key/value pair. */ @@ -171,11 +111,9 @@ public enum Storage { @Nonnull private StorageAdapter newStorageAdapter(@Nonnull final Config config, @Nonnull final Subspace subspace, @Nonnull final Subspace nodeSlotIndexSubspace, - @Nonnull final Function hilbertValueFunction, @Nonnull final OnWriteListener onWriteListener, @Nonnull final OnReadListener onReadListener) { - return storageAdapterCreator.create(config, subspace, nodeSlotIndexSubspace, - hilbertValueFunction, onWriteListener, onReadListener); + return storageAdapterCreator.create(config, subspace, nodeSlotIndexSubspace, onWriteListener, onReadListener); } } @@ -184,7 +122,6 @@ private StorageAdapter newStorageAdapter(@Nonnull final Config config, @Nonnull */ private interface StorageAdapterCreator { StorageAdapter create(@Nonnull Config config, @Nonnull Subspace subspace, @Nonnull Subspace nodeSlotIndexSubspace, - @Nonnull Function hilbertValueFunction, @Nonnull OnWriteListener onWriteListener, @Nonnull OnReadListener onReadListener); } @@ -193,72 +130,62 @@ StorageAdapter create(@Nonnull Config config, @Nonnull Subspace subspace, @Nonnu * Configuration settings for a {@link HNSW}. */ public static class Config { - private final boolean useNodeSlotIndex; - private final int minM; - private final int maxM; - private final int splitS; @Nonnull - private final Storage storage; - - private final boolean storeHilbertValues; + private final Random random; + @Nonnull + private final Metric metric; + private final int m; + private final int efSearch; + private final int efConstruction; protected Config() { - this.useNodeSlotIndex = DEFAULT_USE_NODE_SLOT_INDEX; - this.minM = DEFAULT_MIN_M; - this.maxM = DEFAULT_MAX_M; - this.splitS = DEFAULT_S; - this.storage = DEFAULT_STORAGE; - this.storeHilbertValues = DEFAULT_STORE_HILBERT_VALUES; + this.random = DEFAULT_RANDOM; + this.metric = DEFAULT_METRIC; + this.m = DEFAULT_M; + this.efSearch = DEFAULT_EF_SEARCH; + this.efConstruction = DEFAULT_EF_CONSTRUCTION; } - protected Config(final boolean useNodeSlotIndex, final int minM, final int maxM, final int splitS, - @Nonnull final Storage storage, final boolean storeHilbertValues) { - this.useNodeSlotIndex = useNodeSlotIndex; - this.minM = minM; - this.maxM = maxM; - this.splitS = splitS; - this.storage = storage; - this.storeHilbertValues = storeHilbertValues; + protected Config(@Nonnull final Random random, @Nonnull final Metric metric, final int m, + final int efSearch, final int efConstruction) { + this.random = random; + this.metric = metric; + this.m = m; + this.efSearch = efSearch; + this.efConstruction = efConstruction; } - public boolean isUseNodeSlotIndex() { - return useNodeSlotIndex; - } - - public int getMinM() { - return minM; - } - - public int getMaxM() { - return maxM; + @Nonnull + public Random getRandom() { + return random; } - public int getSplitS() { - return splitS; + @Nonnull + public Metric getMetric() { + return metric; } - @Nonnull - public Storage getStorage() { - return storage; + public int getM() { + return m; } - public boolean isStoreHilbertValues() { - return storeHilbertValues; + public int getEfSearch() { + return efSearch; } - public Metric getMetric() { - return Metric.euclideanMetric(); + public int getEfConstruction() { + return efConstruction; } + @Nonnull public ConfigBuilder toBuilder() { - return new ConfigBuilder(useNodeSlotIndex, minM, maxM, splitS, storage, storeHilbertValues); + return new ConfigBuilder(getRandom(), getMetric(), getM(), getEfSearch(), getEfConstruction()); } @Override + @Nonnull public String toString() { - return storage + ", M=" + minM + "-" + maxM + ", S=" + splitS + - (useNodeSlotIndex ? ", slotIndex" : "") + - (storeHilbertValues ? ", storeHV" : ""); + return "Config[M=" + getM() + " , metric=" + getMetric() + "]"; } } @@ -269,84 +196,78 @@ public String toString() { */ @CanIgnoreReturnValue public static class ConfigBuilder { - private boolean useNodeSlotIndex = DEFAULT_USE_NODE_SLOT_INDEX; - private int minM = DEFAULT_MIN_M; - private int maxM = DEFAULT_MAX_M; - private int splitS = DEFAULT_S; @Nonnull - private Storage storage = DEFAULT_STORAGE; - private boolean storeHilbertValues = DEFAULT_STORE_HILBERT_VALUES; + private Random random = DEFAULT_RANDOM; + @Nonnull + private Metric metric = DEFAULT_METRIC; + private int m = DEFAULT_M; + private int efSearch = DEFAULT_EF_SEARCH; + private int efConstruction = DEFAULT_EF_CONSTRUCTION; public ConfigBuilder() { } - public ConfigBuilder(final boolean useNodeSlotIndex, final int minM, final int maxM, final int splitS, - @Nonnull final Storage storage, final boolean storeHilbertValues) { - this.useNodeSlotIndex = useNodeSlotIndex; - this.minM = minM; - this.maxM = maxM; - this.splitS = splitS; - this.storage = storage; - this.storeHilbertValues = storeHilbertValues; + public ConfigBuilder(@Nonnull Random random, @Nonnull final Metric metric, final int m, + final int efSearch, final int efConstruction) { + this.random = random; + this.metric = metric; + this.m = m; + this.efSearch = efSearch; + this.efConstruction = efConstruction; } - public int getMinM() { - return minM; + @Nonnull + public Random getRandom() { + return random; } - public ConfigBuilder setMinM(final int minM) { - this.minM = minM; + @Nonnull + public ConfigBuilder setRandom(@Nonnull final Random random) { + this.random = random; return this; } - public int getMaxM() { - return maxM; + @Nonnull + public Metric getMetric() { + return metric; } - public ConfigBuilder setMaxM(final int maxM) { - this.maxM = maxM; + @Nonnull + public ConfigBuilder setMetric(@Nonnull final Metric metric) { + this.metric = metric; return this; } - public int getSplitS() { - return splitS; - } - - public ConfigBuilder setSplitS(final int splitS) { - this.splitS = splitS; - return this; + public int getM() { + return m; } @Nonnull - public Storage getStorage() { - return storage; - } - - public ConfigBuilder setStorage(@Nonnull final Storage storage) { - this.storage = storage; + public ConfigBuilder setM(final int m) { + this.m = m; return this; } - public boolean isStoreHilbertValues() { - return storeHilbertValues; + public int getEfSearch() { + return efSearch; } - public ConfigBuilder setStoreHilbertValues(final boolean storeHilbertValues) { - this.storeHilbertValues = storeHilbertValues; + public ConfigBuilder setEfSearch(final int efSearch) { + this.efSearch = efSearch; return this; } - public boolean isUseNodeSlotIndex() { - return useNodeSlotIndex; + public int getEfConstruction() { + return efConstruction; } - public ConfigBuilder setUseNodeSlotIndex(final boolean useNodeSlotIndex) { - this.useNodeSlotIndex = useNodeSlotIndex; + public ConfigBuilder setEfConstruction(final int efConstruction) { + this.efConstruction = efConstruction; return this; } public Config build() { - return new Config(isUseNodeSlotIndex(), getMinM(), getMaxM(), getSplitS(), getStorage(), isStoreHilbertValues()); + return new Config(getRandom(), getMetric(), getM(), getEfSearch(), getEfConstruction()); } } @@ -364,11 +285,10 @@ public static ConfigBuilder newConfigBuilder() { * @param subspace the subspace where the r-tree is stored * @param secondarySubspace the subspace where the node index (if used is stored) * @param executor an executor to use when running asynchronous tasks - * @param hilbertValueFunction function to compute the Hilbert value from a {@link NodeReferenceWithDistance} */ public HNSW(@Nonnull final Subspace subspace, @Nonnull final Subspace secondarySubspace, - @Nonnull final Executor executor, @Nonnull final Function hilbertValueFunction) { - this(subspace, secondarySubspace, executor, DEFAULT_CONFIG, hilbertValueFunction, NodeHelpers::newRandomNodeId, + @Nonnull final Executor executor) { + this(subspace, secondarySubspace, executor, DEFAULT_CONFIG, OnWriteListener.NOOP, OnReadListener.NOOP); } @@ -378,15 +298,11 @@ public HNSW(@Nonnull final Subspace subspace, @Nonnull final Subspace secondaryS * @param nodeSlotIndexSubspace the subspace where the node index (if used is stored) * @param executor an executor to use when running asynchronous tasks * @param config configuration to use - * @param hilbertValueFunction function to compute the Hilbert value for a {@link NodeReferenceWithDistance} - * @param nodeIdSupplier supplier to be invoked when new nodes are created * @param onWriteListener an on-write listener to be called after writes take place * @param onReadListener an on-read listener to be called after reads take place */ public HNSW(@Nonnull final Subspace subspace, @Nonnull final Subspace nodeSlotIndexSubspace, @Nonnull final Executor executor, @Nonnull final Config config, - @Nonnull final Function hilbertValueFunction, - @Nonnull final Supplier nodeIdSupplier, @Nonnull final OnWriteListener onWriteListener, @Nonnull final OnReadListener onReadListener) { this.storageAdapter = config.getStorage() @@ -394,8 +310,6 @@ public HNSW(@Nonnull final Subspace subspace, @Nonnull final Subspace nodeSlotIn onReadListener); this.executor = executor; this.config = config; - this.hilbertValueFunction = hilbertValueFunction; - this.nodeIdSupplier = nodeIdSupplier; this.onWriteListener = onWriteListener; this.onReadListener = onReadListener; } @@ -509,9 +423,8 @@ public AsyncIterator scan(@Nonnull final ReadTransaction readTransacti @SuppressWarnings("checkstyle:MethodName") // method name introduced by paper @Nonnull private CompletableFuture> kNearestNeighborsSearch(@Nonnull final ReadTransaction readTransaction, - final int efSearch, @Nonnull final Vector queryVector) { - return storageAdapter.fetchEntryNodeKey(readTransaction) + return storageAdapter.fetchEntryNodeReference(readTransaction) .thenCompose(entryPointAndLayer -> { if (entryPointAndLayer == null) { return CompletableFuture.completedFuture(null); // not a single node in the index @@ -519,53 +432,73 @@ private CompletableFuture> kNearestNeighborsSearch(@ final Metric metric = getConfig().getMetric(); - final var entryState = new GreedyState(entryPointAndLayer.getLayer(), - entryPointAndLayer.getPrimaryKey(), - Vector.comparativeDistance(metric, entryPointAndLayer.getVector(), queryVector)); - final AtomicReference greedyStateReference = - new AtomicReference<>(entryState); + final NodeReferenceWithDistance entryState = + new NodeReferenceWithDistance(entryPointAndLayer.getPrimaryKey(), + Vector.comparativeDistance(metric, entryPointAndLayer.getVector(), queryVector)); if (entryPointAndLayer.getLayer() == 0) { // entry data points to a node in layer 0 directly return CompletableFuture.completedFuture(entryState); } - return AsyncUtil.whileTrue(() -> { - final var greedyIn = greedyStateReference.get(); - return greedySearchLayer(readTransaction, greedyIn.toNodeReferenceWithDistance(), - greedyIn.getLayer(), queryVector) + final AtomicInteger layerAtomic = new AtomicInteger(entryPointAndLayer.getLayer()); + final AtomicReference nodeReferenceAtomic = + new AtomicReference<>(entryState); + + return MoreAsyncUtil.forLoop(entryPointAndLayer.getLayer(), + layer -> layer > 0, + layer -> layer - 1, + layer -> { + final var greedyIn = nodeReferenceAtomic.get(); + return greedySearchLayer(InliningNode.factory(), readTransaction, greedyIn, layer, + queryVector) .thenApply(greedyState -> { - greedyStateReference.set(greedyState); - return greedyState.getLayer() > 0; + nodeReferenceAtomic.set(greedyState); + return null; }); - }, executor).thenApply(ignored -> greedyStateReference.get()); - }).thenCompose(greedyState -> { - if (greedyState == null) { + }, executor).thenApply(ignored -> nodeReferenceAtomic.get()); + }).thenCompose(nodeReference -> { + if (nodeReference == null) { return CompletableFuture.completedFuture(null); } - return searchLayer(CompactNode::creator, readTransaction, - ImmutableList.of(greedyState.toNodeReferenceWithDistance()), 0, efSearch, + return searchLayer(CompactNode.factory(), readTransaction, + ImmutableList.of(nodeReference), 0, config.getEfSearch(), queryVector); }); } + @Nonnull + private CompletableFuture greedySearchLayer(@Nonnull NodeFactory nodeFactory, + @Nonnull final ReadTransaction readTransaction, + @Nonnull final NodeReferenceWithDistance entryNeighbor, + final int layer, + @Nonnull final Vector queryVector) { + if (nodeFactory.getNodeKind() == NodeKind.INLINING) { + return greedySearchInliningLayer(readTransaction, entryNeighbor, layer, queryVector); + } else { + return searchLayer(nodeFactory, readTransaction, ImmutableList.of(entryNeighbor), layer, 1, queryVector) + .thenApply(searchResult -> Iterables.getOnlyElement(searchResult.getNodeMap().keySet())); + } + } + /** * TODO. */ @Nonnull - private CompletableFuture greedySearchLayer(@Nonnull final ReadTransaction readTransaction, - @Nonnull final NodeReferenceWithDistance entryNeighbor, - final int layer, - @Nonnull final Vector queryVector) { + private CompletableFuture greedySearchInliningLayer(@Nonnull final ReadTransaction readTransaction, + @Nonnull final NodeReferenceWithDistance entryNeighbor, + final int layer, + @Nonnull final Vector queryVector) { Verify.verify(layer > 0); final Metric metric = getConfig().getMetric(); - final AtomicReference greedyStateReference = - new AtomicReference<>(new GreedyState(layer, entryNeighbor.getPrimaryKey(), entryNeighbor.getDistance())); + final AtomicReference currentNodeReferenceAtomic = + new AtomicReference<>(new NodeReferenceWithDistance(entryNeighbor.getPrimaryKey(), + entryNeighbor.getDistance())); return AsyncUtil.whileTrue(() -> onReadListener.onAsyncRead( - storageAdapter.fetchNode(InliningNode::creator, readTransaction, - layer, greedyStateReference.get().getPrimaryKey())) + storageAdapter.fetchNode(InliningNode.factory(), readTransaction, + layer, currentNodeReferenceAtomic.get().getPrimaryKey())) .thenApply(node -> { if (node == null) { throw new IllegalStateException("unable to fetch node"); @@ -573,8 +506,8 @@ private CompletableFuture greedySearchLayer(@Nonnull final ReadTran final InliningNode inliningNode = node.asInliningNode(); final List neighbors = inliningNode.getNeighbors(); - final GreedyState currentNodeKey = greedyStateReference.get(); - double minDistance = currentNodeKey.getDistance(); + final NodeReferenceWithDistance currentNodeReference = currentNodeReferenceAtomic.get(); + double minDistance = currentNodeReference.getDistance(); NodeReferenceWithVector nearestNeighbor = null; for (final NodeReferenceWithVector neighbor : neighbors) { @@ -587,37 +520,36 @@ private CompletableFuture greedySearchLayer(@Nonnull final ReadTran } if (nearestNeighbor == null) { - greedyStateReference.set( - new GreedyState(layer - 1, currentNodeKey.getPrimaryKey(), minDistance)); + currentNodeReferenceAtomic.set( + new NodeReferenceWithDistance(currentNodeReference.getPrimaryKey(), minDistance)); return false; } - greedyStateReference.set( - new GreedyState(layer, nearestNeighbor.getPrimaryKey(), - minDistance)); + currentNodeReferenceAtomic.set( + new NodeReferenceWithDistance(nearestNeighbor.getPrimaryKey(), minDistance)); return true; - }), executor).thenApply(ignored -> greedyStateReference.get()); + }), executor).thenApply(ignored -> currentNodeReferenceAtomic.get()); } /** * TODO. */ @Nonnull - private CompletableFuture> searchLayer(@Nonnull Node.NodeCreator creator, + private CompletableFuture> searchLayer(@Nonnull NodeFactory nodeFactory, @Nonnull final ReadTransaction readTransaction, @Nonnull final List entryNeighbors, final int layer, final int efSearch, @Nonnull final Vector queryVector) { final Set visited = Sets.newConcurrentHashSet(NodeReference.primaryKeys(entryNeighbors)); - final PriorityBlockingQueue candidates = - new PriorityBlockingQueue<>(entryNeighbors.size(), + final Queue candidates = + new PriorityBlockingQueue<>(config.getM(), Comparator.comparing(NodeReferenceWithDistance::getDistance)); candidates.addAll(entryNeighbors); - final PriorityBlockingQueue nearestNeighbors = - new PriorityBlockingQueue<>(entryNeighbors.size(), + final Queue furthestNeighbors = + new PriorityBlockingQueue<>(config.getM(), Comparator.comparing(NodeReferenceWithDistance::getDistance).reversed()); - nearestNeighbors.addAll(entryNeighbors); + furthestNeighbors.addAll(entryNeighbors); final Map> nodeCache = Maps.newConcurrentMap(); final Metric metric = getConfig().getMetric(); @@ -627,39 +559,39 @@ private CompletableFuture> searchLayer } final NodeReferenceWithDistance candidate = candidates.poll(); - final NodeReferenceWithDistance furthestNeighbor = Objects.requireNonNull(nearestNeighbors.peek()); + final NodeReferenceWithDistance furthestNeighbor = Objects.requireNonNull(furthestNeighbors.peek()); if (candidate.getDistance() > furthestNeighbor.getDistance()) { return AsyncUtil.READY_FALSE; } - return fetchNodeIfNotCached(creator, readTransaction, layer, candidate, nodeCache) + return fetchNodeIfNotCached(nodeFactory, readTransaction, layer, candidate, nodeCache) .thenApply(candidateNode -> Iterables.filter(candidateNode.getNeighbors(), neighbor -> !visited.contains(neighbor.getPrimaryKey()))) - .thenCompose(neighborReferences -> neighborhood(creator, readTransaction, + .thenCompose(neighborReferences -> fetchNeighborhood(nodeFactory, readTransaction, layer, neighborReferences, nodeCache)) .thenApply(neighborReferences -> { for (final NodeReferenceWithVector current : neighborReferences) { visited.add(current.getPrimaryKey()); final double furthestDistance = - Objects.requireNonNull(nearestNeighbors.peek()).getDistance(); + Objects.requireNonNull(furthestNeighbors.peek()).getDistance(); final double currentDistance = Vector.comparativeDistance(metric, current.getVector(), queryVector); - if (currentDistance < furthestDistance || nearestNeighbors.size() < efSearch) { + if (currentDistance < furthestDistance || furthestNeighbors.size() < efSearch) { final NodeReferenceWithDistance currentWithDistance = new NodeReferenceWithDistance(current.getPrimaryKey(), currentDistance); candidates.add(currentWithDistance); - nearestNeighbors.add(currentWithDistance); - if (nearestNeighbors.size() > efSearch) { - nearestNeighbors.poll(); + furthestNeighbors.add(currentWithDistance); + if (furthestNeighbors.size() > efSearch) { + furthestNeighbors.poll(); } } } return true; }); - }).thenCompose(ignored -> fetchResultsIfNecessary(creator, readTransaction, layer, nearestNeighbors, + }).thenCompose(ignored -> fetchResultsIfNecessary(nodeFactory, readTransaction, layer, furthestNeighbors, nodeCache)); } @@ -667,12 +599,12 @@ private CompletableFuture> searchLayer * TODO. */ @Nonnull - private CompletableFuture> fetchNodeIfNotCached(@Nonnull final Node.NodeCreator creator, + private CompletableFuture> fetchNodeIfNotCached(@Nonnull final NodeFactory nodeFactory, @Nonnull final ReadTransaction readTransaction, final int layer, @Nonnull final NodeReference nodeReference, @Nonnull final Map> nodeCache) { - return fetchNodeIfNecessaryAndApply(creator, readTransaction, layer, nodeReference, + return fetchNodeIfNecessaryAndApply(nodeFactory, readTransaction, layer, nodeReference, nR -> nodeCache.get(nR.getPrimaryKey()), (ignored, node) -> node); } @@ -681,7 +613,7 @@ private CompletableFuture> fetchNodeIfNotCache * TODO. */ @Nonnull - private CompletableFuture fetchNodeIfNecessaryAndApply(@Nonnull final Node.NodeCreator creator, + private CompletableFuture fetchNodeIfNecessaryAndApply(@Nonnull final NodeFactory nodeFactory, @Nonnull final ReadTransaction readTransaction, final int layer, @Nonnull final R nodeReference, @@ -693,7 +625,7 @@ private CompletableFuture< } return onReadListener.onAsyncRead( - storageAdapter.fetchNode(creator, readTransaction, layer, nodeReference.getPrimaryKey())) + storageAdapter.fetchNode(nodeFactory, readTransaction, layer, nodeReference.getPrimaryKey())) .thenApply(node -> biMapFunction.apply(nodeReference, node)); } @@ -701,12 +633,12 @@ private CompletableFuture< * TODO. */ @Nonnull - private CompletableFuture> neighborhood(@Nonnull final Node.NodeCreator creator, - @Nonnull final ReadTransaction readTransaction, - final int layer, - @Nonnull final Iterable neighborReferences, - @Nonnull final Map> nodeCache) { - return fetchSomeNodesAndApply(creator, readTransaction, layer, neighborReferences, + private CompletableFuture> fetchNeighborhood(@Nonnull final NodeFactory nodeFactory, + @Nonnull final ReadTransaction readTransaction, + final int layer, + @Nonnull final Iterable neighborReferences, + @Nonnull final Map> nodeCache) { + return fetchSomeNodesAndApply(nodeFactory, readTransaction, layer, neighborReferences, neighborReference -> { if (neighborReference instanceof NodeReferenceWithVector) { return (NodeReferenceWithVector)neighborReference; @@ -725,7 +657,7 @@ private CompletableFuture CompletableFuture> fetchResultsIfNecessary(@Nonnull final Node.NodeCreator creator, + private CompletableFuture> fetchResultsIfNecessary(@Nonnull final NodeFactory creator, @Nonnull final ReadTransaction readTransaction, final int layer, @Nonnull final Iterable nodeReferences, @@ -745,7 +677,7 @@ private CompletableFuture> fetchResult for (final SearchResult.NodeReferenceWithNode nodeReferenceWithNode : nodeReferencesWithNodes) { nodeMapBuilder.put(nodeReferenceWithNode.getNodeReferenceWithDistance(), nodeReferenceWithNode.getNode()); } - return new SearchResult<>(layer, nodeMapBuilder.build()); + return new SearchResult<>(nodeMapBuilder.build()); }); } @@ -754,7 +686,7 @@ private CompletableFuture> fetchResult */ @Nonnull @SuppressWarnings("unchecked") - private CompletableFuture> fetchSomeNodesAndApply(@Nonnull final Node.NodeCreator creator, + private CompletableFuture> fetchSomeNodesAndApply(@Nonnull final NodeFactory creator, @Nonnull final ReadTransaction readTransaction, final int layer, @Nonnull final Iterable nodeReferences, @@ -800,4 +732,79 @@ private CompletableFuture< return resultBuilder.build(); }); } + + @Nonnull + public CompletableFuture insert(@Nonnull final Transaction transaction, @Nonnull final Tuple primaryKey, + @Nonnull final Vector vector) { + final Metric metric = getConfig().getMetric(); + final Queue furthestNeighbors = + new PriorityBlockingQueue<>(config.getM(), + Comparator.comparing(NodeReferenceWithDistance::getDistance).reversed()); + + final int l = insertionLayer(getConfig().getRandom()); + + storageAdapter.fetchEntryNodeReference(transaction) + .thenApply(entryNodeReference -> { + if (entryNodeReference == null) { + // this is the first node + writeLonelyNodes(InliningNode.factory(), transaction, primaryKey, vector, l, 0); + storageAdapter.writeNode(transaction, + CompactNode.factory() + .create(NodeKind.COMPACT, primaryKey, vector, ImmutableList.of()), + 0); + storageAdapter.writeEntryNodeReference(transaction, + new EntryNodeReference(primaryKey, vector, l)); + } else { + final int entryNodeLayer = entryNodeReference.getLayer(); + if (l > entryNodeLayer) { + writeLonelyNodes(InliningNode.factory(), transaction, primaryKey, vector, l, entryNodeLayer); + storageAdapter.writeEntryNodeReference(transaction, + new EntryNodeReference(primaryKey, vector, l)); + } + } + return entryNodeReference; + }).thenCompose(entryNodeReference -> { + if (entryNodeReference == null) { + return AsyncUtil.DONE; + } + + final int lMax = entryNodeReference.getLayer(); + + final AtomicReference nodeReferenceAtomic = + new AtomicReference<>(new NodeReferenceWithDistance(entryNodeReference.getPrimaryKey(), + Vector.comparativeDistance(metric, entryNodeReference.getVector(), vector))); + MoreAsyncUtil.forLoop(lMax, + layer -> layer > l, + layer -> layer - 1, + layer -> greedySearchLayer(InliningNode.factory(), transaction, + nodeReferenceAtomic.get(), layer, vector) + .thenApply(nodeReference -> { + nodeReferenceAtomic.set(nodeReference); + return null; + }), executor); + + final NodeReferenceWithDistance nodeReference = nodeReferenceAtomic.get(); + + + + }) + } + + public void writeLonelyNodes(@Nonnull final NodeFactory nodeFactory, + @Nonnull final Transaction transaction, + @Nonnull final Tuple primaryKey, + @Nonnull final Vector vector, + final int highestLayerInclusive, + final int lowestLayerExclusive) { + for (int layer = highestLayerInclusive; layer > lowestLayerExclusive; layer --) { + storageAdapter.writeNode(transaction, + nodeFactory.create(nodeFactory.getNodeKind(), primaryKey, vector, ImmutableList.of()), layer); + } + } + + private int insertionLayer(@Nonnull final Random random) { + double lambda = 1.0 / Math.log(getConfig().getM()); + double u = 1.0 - random.nextDouble(); // Avoid log(0) + return (int) Math.floor(-Math.log(u) * lambda); + } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java index 6619a2003e..bfc4de8153 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java @@ -23,6 +23,7 @@ import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; import com.google.common.base.Verify; +import com.google.common.collect.Lists; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -32,6 +33,24 @@ * TODO. */ class InliningNode extends AbstractNode { + @Nonnull + private static final NodeFactory FACTORY = new NodeFactory<>() { + @SuppressWarnings("unchecked") + @Nonnull + @Override + public Node create(@Nonnull final NodeKind nodeKind, @Nonnull final Tuple primaryKey, + @Nullable final Vector vector, @Nonnull final List neighbors) { + Verify.verify(nodeKind == NodeKind.INLINING); + return new InliningNode(primaryKey, (List)neighbors); + } + + @Nonnull + @Override + public NodeKind getNodeKind() { + return NodeKind.INLINING; + } + }; + public InliningNode(@Nonnull final Tuple primaryKey, @Nonnull final List neighbors) { super(primaryKey, neighbors); @@ -56,17 +75,26 @@ public InliningNode asInliningNode() { } @Override - public NodeCreator sameCreator() { - return InliningNode::creator; + public NodeFactory sameCreator() { + return InliningNode.factory(); + } + + @Nonnull + @Override + public Tuple toTuple() { + final List nodeItems = Lists.newArrayListWithExpectedSize(3); + nodeItems.add(NodeKind.INLINING.getSerialized()); + final List neighborItems = Lists.newArrayListWithExpectedSize(getNeighbors().size()); + for (final NodeReferenceWithVector nodeReference : getNeighbors()) { + neighborItems.add(Tuple.from(nodeReference.getPrimaryKey(), + StorageAdapter.tupleFromVector(nodeReference.getVector()))); + } + nodeItems.add(Tuple.fromList(neighborItems)); + return Tuple.fromList(nodeItems); } @Nonnull - @SuppressWarnings("unchecked") - public static Node creator(@Nonnull final NodeKind nodeKind, - @Nonnull final Tuple primaryKey, - @Nullable final Vector vector, - @Nonnull final List neighbors) { - Verify.verify(nodeKind == NodeKind.INLINING); - return new InliningNode(primaryKey, (List)neighbors); + public static NodeFactory factory() { + return FACTORY; } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java index f59f15bf22..d3fc11c082 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java @@ -90,6 +90,12 @@ public double distance(final Double[] vector1, final Double[] vector2) { } return sumOfAbsDiffs; } + + @Override + @Nonnull + public String toString() { + return this.getClass().getSimpleName(); + } } class EuclideanMetric implements Metric { @@ -99,6 +105,12 @@ public double distance(final Double[] vector1, final Double[] vector2) { return Math.sqrt(EuclideanSquareMetric.distanceInternal(vector1, vector2)); } + + @Override + @Nonnull + public String toString() { + return this.getClass().getSimpleName(); + } } class EuclideanSquareMetric implements Metric { @@ -116,6 +128,12 @@ private static double distanceInternal(final Double[] vector1, final Double[] ve } return sumOfSquares; } + + @Override + @Nonnull + public String toString() { + return this.getClass().getSimpleName(); + } } class CosineMetric implements Metric { @@ -140,6 +158,12 @@ public double distance(final Double[] vector1, final Double[] vector2) { return 1.0d - dotProduct / (Math.sqrt(normA) * Math.sqrt(normB)); } + + @Override + @Nonnull + public String toString() { + return this.getClass().getSimpleName(); + } } class DotProductMetric implements Metric { @@ -158,5 +182,11 @@ public double comparativeDistance(final Double[] vector1, final Double[] vector2 } return -product; } + + @Override + @Nonnull + public String toString() { + return this.getClass().getSimpleName(); + } } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java index 8c10378d31..49f80e130d 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java @@ -22,10 +22,9 @@ import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; -import com.google.errorprone.annotations.CanIgnoreReturnValue; +import com.google.common.collect.Lists; import javax.annotation.Nonnull; -import javax.annotation.Nullable; import java.util.List; /** @@ -42,20 +41,8 @@ public interface Node { @Nonnull R getNeighbor(int index); - @CanIgnoreReturnValue - @Nonnull - Node insert(@Nonnull StorageAdapter storageAdapter, int level, int slotIndex, @Nonnull NodeSlot slot); - - @CanIgnoreReturnValue - @Nonnull - Node update(@Nonnull StorageAdapter storageAdapter, int level, int slotIndex, @Nonnull NodeSlot updatedSlot); - - @CanIgnoreReturnValue - @Nonnull - Node delete(@Nonnull StorageAdapter storageAdapter, int level, int slotIndex); - /** - * Return the kind of the node, i.e. {@link NodeKind#DATA} or {@link NodeKind#INLINING}. + * Return the kind of the node, i.e. {@link NodeKind#COMPACT} or {@link NodeKind#INLINING}. * @return the kind of this node as a {@link NodeKind} */ @Nonnull @@ -67,11 +54,63 @@ public interface Node { @Nonnull InliningNode asInliningNode(); - NodeCreator sameCreator(); + NodeFactory sameCreator(); + + @Nonnull + Tuple toTuple(); + + @Nonnull + static Node nodeFromTuples(@Nonnull final NodeFactory creator, + @Nonnull final Tuple primaryKey, + @Nonnull final Tuple valueTuple) { + final NodeKind nodeKind = NodeKind.fromSerializedNodeKind((byte)valueTuple.getLong(0)); + final Tuple vectorTuple; + final Tuple neighborsTuple; + + switch (nodeKind) { + case COMPACT: + vectorTuple = valueTuple.getNestedTuple(1); + neighborsTuple = valueTuple.getNestedTuple(2); + return compactNodeFromTuples(creator, primaryKey, vectorTuple, neighborsTuple); + case INLINING: + neighborsTuple = valueTuple.getNestedTuple(1); + return inliningNodeFromTuples(creator, primaryKey, neighborsTuple); + default: + throw new IllegalStateException("unknown node kind"); + } + } + + @Nonnull + static Node compactNodeFromTuples(@Nonnull final NodeFactory creator, + @Nonnull final Tuple primaryKey, + @Nonnull final Tuple vectorTuple, + @Nonnull final Tuple neighborsTuple) { + final Vector vector = StorageAdapter.vectorFromTuple(vectorTuple); + + List nodeReferences = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); + + for (final Object neighborObject : neighborsTuple) { + final Tuple neighborTuple = (Tuple)neighborObject; + nodeReferences.add(new NodeReference(neighborTuple)); + } + + return creator.create(NodeKind.COMPACT, primaryKey, vector, nodeReferences); + } + + @Nonnull + static Node inliningNodeFromTuples(@Nonnull final NodeFactory creator, + @Nonnull final Tuple primaryKey, + @Nonnull final Tuple neighborsTuple) { + List neighborsWithVectors = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); + + for (final Object neighborObject : neighborsTuple) { + final Tuple neighborTuple = (Tuple)neighborObject; + final Tuple neighborPrimaryKey = neighborTuple.getNestedTuple(0); + final Tuple neighborVectorTuple = neighborTuple.getNestedTuple(1); + neighborsWithVectors.add(new NodeReferenceWithVector(neighborPrimaryKey, + StorageAdapter.vectorFromTuple(neighborVectorTuple))); + } - @FunctionalInterface - interface NodeCreator { - Node create(@Nonnull NodeKind nodeKind, @Nonnull Tuple primaryKey, @Nullable Vector vector, - @Nonnull List neighbors); + return creator.create(NodeKind.INLINING, primaryKey, null, neighborsWithVectors); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyState.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeFactory.java similarity index 55% rename from fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyState.java rename to fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeFactory.java index 97b7c30941..51e3fb59b1 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/GreedyState.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeFactory.java @@ -1,5 +1,5 @@ /* - * NodeWithLayer.java + * NodeFactory.java * * This source file is part of the FoundationDB open source project * @@ -21,35 +21,17 @@ package com.apple.foundationdb.async.hnsw; import com.apple.foundationdb.tuple.Tuple; +import com.christianheina.langx.half4j.Half; import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.util.List; -class GreedyState { - private final int layer; +public interface NodeFactory { @Nonnull - private final Tuple primaryKey; - private final double distance; - - public GreedyState(final int layer, @Nonnull final Tuple primaryKey, final double distance) { - this.layer = layer; - this.primaryKey = primaryKey; - this.distance = distance; - } - - public int getLayer() { - return layer; - } + Node create(@Nonnull NodeKind nodeKind, @Nonnull Tuple primaryKey, @Nullable Vector vector, + @Nonnull List neighbors); @Nonnull - public Tuple getPrimaryKey() { - return primaryKey; - } - - public double getDistance() { - return distance; - } - - public NodeReferenceWithDistance toNodeReferenceWithDistance() { - return new NodeReferenceWithDistance(getPrimaryKey(), getDistance()); - } + NodeKind getNodeKind(); } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKind.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKind.java index 0a9f6031e6..13d71a1b9b 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKind.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKind.java @@ -28,7 +28,7 @@ * Enum to capture the kind of node. */ public enum NodeKind { - DATA((byte)0x00), + COMPACT((byte)0x00), INLINING((byte)0x01); private final byte serialized; @@ -46,7 +46,7 @@ static NodeKind fromSerializedNodeKind(byte serializedNodeKind) { final NodeKind nodeKind; switch (serializedNodeKind) { case 0x00: - nodeKind = NodeKind.DATA; + nodeKind = NodeKind.COMPACT; break; case 0x01: nodeKind = NodeKind.INLINING; diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java index a2b7ad3697..7d39e0ef0d 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java @@ -42,7 +42,7 @@ default CompletableFuture onAsyncReadForWrite(@Nonnull Compl return future; } - default void onNodeWritten(@Nonnull Node node) { + default void onNodeWritten(@Nonnull Node node) { // nothing } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/SearchResult.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/SearchResult.java index af66ca7e4b..0685930957 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/SearchResult.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/SearchResult.java @@ -24,19 +24,13 @@ import java.util.Map; class SearchResult { - private final int layer; @Nonnull private final Map> nodeMap; - public SearchResult(final int layer, @Nonnull final Map> nodeMap) { - this.layer = layer; + public SearchResult(@Nonnull final Map> nodeMap) { this.nodeMap = nodeMap; } - public int getLayer() { - return layer; - } - @Nonnull public Map> getNodeMap() { return nodeMap; diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java index c2c34fa669..dfa7d75dde 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java @@ -24,10 +24,12 @@ import com.apple.foundationdb.Transaction; import com.apple.foundationdb.subspace.Subspace; import com.apple.foundationdb.tuple.Tuple; +import com.christianheina.langx.half4j.Half; +import com.google.common.base.Verify; +import com.google.common.collect.Lists; import javax.annotation.Nonnull; import javax.annotation.Nullable; -import java.math.BigInteger; import java.util.List; import java.util.concurrent.CompletableFuture; @@ -35,7 +37,6 @@ * Storage adapter used for serialization and deserialization of nodes. */ interface StorageAdapter { - /** * Get the {@link HNSW.Config} associated with this storage adapter. * @return the configuration used by this storage adapter @@ -81,75 +82,50 @@ interface StorageAdapter { @Nonnull OnReadListener getOnReadListener(); - CompletableFuture fetchEntryNodeKey(@Nonnull ReadTransaction readTransaction); + CompletableFuture fetchEntryNodeReference(@Nonnull ReadTransaction readTransaction); + + void writeEntryNodeReference(@Nonnull final Transaction transaction, + @Nonnull final EntryNodeReference entryNodeReference); @Nonnull - CompletableFuture> fetchNode(@Nonnull Node.NodeCreator creator, + CompletableFuture> fetchNode(@Nonnull NodeFactory nodeFactory, @Nonnull ReadTransaction readTransaction, int layer, @Nonnull Tuple primaryKey); - /** - * Insert a new entry into the node index if configuration indicates we should maintain such an index. - * - * @param transaction the transaction to use - * @param level the level counting starting at {@code 0} indicating the leaf level increasing upwards - * @param nodeSlot the {@link NodeSlot} to be inserted - */ - void insertIntoNodeIndexIfNecessary(@Nonnull Transaction transaction, int level, @Nonnull NodeSlot nodeSlot); - - /** - * Deletes an entry from the node index if configuration indicates we should maintain such an index. - * - * @param transaction the transaction to use - * @param level the level counting starting at {@code 0} indicating the leaf level increasing upwards - * @param nodeSlot the {@link NodeSlot} to be deleted - */ - void deleteFromNodeIndexIfNecessary(@Nonnull Transaction transaction, int level, @Nonnull NodeSlot nodeSlot); - - /** - * Persist a node slot. - * - * @param transaction the transaction to use - * @param node node whose slot to persist - * @param itemSlot the node slot to persist - */ - void writeLeafNodeSlot(@Nonnull Transaction transaction, @Nonnull CompactNode node, @Nonnull ItemSlot itemSlot); + void writeNode(@Nonnull final Transaction transaction, @Nonnull final Node node, + int layer); - /** - * Clear out a leaf node slot. - * - * @param transaction the transaction to use - * @param node node whose slot is cleared out - * @param itemSlot the node slot to clear out - */ - void clearLeafNodeSlot(@Nonnull Transaction transaction, @Nonnull CompactNode node, @Nonnull ItemSlot itemSlot); - - /** - * Method to (re-)persist a list of nodes passed in. - * - * @param transaction the transaction to use - * @param nodes a list of nodes to be (re-persisted) - */ - void writeNodes(@Nonnull Transaction transaction, @Nonnull List nodes); + @Nonnull + static Vector vectorFromTuple(final Tuple vectorTuple) { + final Half[] vectorHalfs = new Half[vectorTuple.size()]; + for (int i = 0; i < vectorTuple.size(); i ++) { + vectorHalfs[i] = Half.shortBitsToHalf(shortFromBytes(vectorTuple.getBytes(i))); + } + return new Vector.HalfVector(vectorHalfs); + } - /** - * Scan the node slot index for the given Hilbert Value/key pair and return the appropriate {@link Node}. - * Note that this method requires a node slot index to be maintained. - * - * @param transaction the transaction to use - * @param level the level we should search counting upwards starting from level {@code 0} for the leaf node - * level. - * @param hilbertValue the Hilbert Value of the {@code (Hilbert Value, key)} pair to search for - * @param key the key of the {@code (Hilbert Value, key)} pair to search for - * @param isInsertUpdate a use case indicator determining if this search is going to be used for an - * update operation or a delete operation - * - * @return a future that when completed holds the appropriate {@link Node} or {@code null} if such a - * {@link Node} could not be found. - */ @Nonnull - CompletableFuture scanNodeIndexAndFetchNode(@Nonnull ReadTransaction transaction, int level, - @Nonnull BigInteger hilbertValue, @Nonnull Tuple key, - boolean isInsertUpdate); + static Tuple tupleFromVector(final Vector vector) { + final List vectorBytes = Lists.newArrayListWithExpectedSize(vector.size()); + for (int i = 0; i < vector.size(); i ++) { + vectorBytes.add(bytesFromShort(Half.halfToShortBits(vector.getComponent(i)))); + } + return Tuple.fromList(vectorBytes); + } + + static short shortFromBytes(byte[] bytes) { + Verify.verify(bytes.length == 2); + int high = bytes[0] & 0xFF; // Convert to unsigned int + int low = bytes[1] & 0xFF; + + return (short) ((high << 8) | low); + } + + static byte[] bytesFromShort(short value) { + byte[] result = new byte[2]; + result[0] = (byte) ((value >> 8) & 0xFF); // high byte first + result[1] = (byte) (value & 0xFF); // low byte second + return result; + } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java index d88bcda5c1..985b783ed8 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java @@ -44,6 +44,11 @@ public int size() { return data.length; } + @Nonnull + R getComponent(int dimension) { + return data[dimension]; + } + @Nonnull public R[] getData() { return data; From d1388b8e3288f7e21176999aaa6c009d67d374e1 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Mon, 28 Jul 2025 22:10:22 +0200 Subject: [PATCH 10/34] save point --- .../apple/foundationdb/async/hnsw/HNSW.java | 35 +++++++++++-------- .../apple/foundationdb/async/hnsw/Vector.java | 7 ++++ 2 files changed, 28 insertions(+), 14 deletions(-) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index 5b9400cd3b..14631a00cd 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -61,7 +61,6 @@ import java.util.function.BiPredicate; import java.util.function.Function; import java.util.function.Predicate; -import java.util.function.Supplier; /** * TODO. @@ -546,10 +545,10 @@ private CompletableFuture> searchLayer new PriorityBlockingQueue<>(config.getM(), Comparator.comparing(NodeReferenceWithDistance::getDistance)); candidates.addAll(entryNeighbors); - final Queue furthestNeighbors = + final Queue nearestNeighbors = new PriorityBlockingQueue<>(config.getM(), Comparator.comparing(NodeReferenceWithDistance::getDistance).reversed()); - furthestNeighbors.addAll(entryNeighbors); + nearestNeighbors.addAll(entryNeighbors); final Map> nodeCache = Maps.newConcurrentMap(); final Metric metric = getConfig().getMetric(); @@ -559,7 +558,7 @@ private CompletableFuture> searchLayer } final NodeReferenceWithDistance candidate = candidates.poll(); - final NodeReferenceWithDistance furthestNeighbor = Objects.requireNonNull(furthestNeighbors.peek()); + final NodeReferenceWithDistance furthestNeighbor = Objects.requireNonNull(nearestNeighbors.peek()); if (candidate.getDistance() > furthestNeighbor.getDistance()) { return AsyncUtil.READY_FALSE; @@ -575,23 +574,23 @@ private CompletableFuture> searchLayer for (final NodeReferenceWithVector current : neighborReferences) { visited.add(current.getPrimaryKey()); final double furthestDistance = - Objects.requireNonNull(furthestNeighbors.peek()).getDistance(); + Objects.requireNonNull(nearestNeighbors.peek()).getDistance(); final double currentDistance = Vector.comparativeDistance(metric, current.getVector(), queryVector); - if (currentDistance < furthestDistance || furthestNeighbors.size() < efSearch) { + if (currentDistance < furthestDistance || nearestNeighbors.size() < efSearch) { final NodeReferenceWithDistance currentWithDistance = new NodeReferenceWithDistance(current.getPrimaryKey(), currentDistance); candidates.add(currentWithDistance); - furthestNeighbors.add(currentWithDistance); - if (furthestNeighbors.size() > efSearch) { - furthestNeighbors.poll(); + nearestNeighbors.add(currentWithDistance); + if (nearestNeighbors.size() > efSearch) { + nearestNeighbors.poll(); } } } return true; }); - }).thenCompose(ignored -> fetchResultsIfNecessary(nodeFactory, readTransaction, layer, furthestNeighbors, + }).thenCompose(ignored -> fetchResultsIfNecessary(nodeFactory, readTransaction, layer, nearestNeighbors, nodeCache)); } @@ -606,7 +605,10 @@ private CompletableFuture> fetchNodeIfNotCache @Nonnull final Map> nodeCache) { return fetchNodeIfNecessaryAndApply(nodeFactory, readTransaction, layer, nodeReference, nR -> nodeCache.get(nR.getPrimaryKey()), - (ignored, node) -> node); + (nR, node) -> { + nodeCache.put(nR.getPrimaryKey(), node); + return node; + }); } /** @@ -649,8 +651,10 @@ private CompletableFuture - new NodeReferenceWithVector(neighborReference.getPrimaryKey(), neighborNode.asCompactNode().getVector())); + (neighborReference, neighborNode) -> { + nodeCache.put(neighborReference.getPrimaryKey(), neighborNode); + return new NodeReferenceWithVector(neighborReference.getPrimaryKey(), neighborNode.asCompactNode().getVector()); + }); } /** @@ -670,7 +674,10 @@ private CompletableFuture> fetchResult } return new SearchResult.NodeReferenceWithNode<>(nodeReference, node); }, - SearchResult.NodeReferenceWithNode::new) + (nodeReferenceWithDistance, node) -> { + nodeCache.put(nodeReferenceWithDistance.getPrimaryKey(), node); + return new SearchResult.NodeReferenceWithNode(nodeReferenceWithDistance, node); + }) .thenApply(nodeReferencesWithNodes -> { final ImmutableMap.Builder> nodeMapBuilder = ImmutableMap.builder(); diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java index 985b783ed8..5308ea3271 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java @@ -35,9 +35,12 @@ public abstract class Vector { @Nonnull protected R[] data; + @Nonnull + protected Supplier hashCodeSupplier; public Vector(@Nonnull final R[] data) { this.data = data; + this.hashCodeSupplier = Suppliers.memoize(this::computeHashCode); } public int size() { @@ -71,6 +74,10 @@ public boolean equals(final Object o) { @Override public int hashCode() { + return hashCodeSupplier.get(); + } + + private int computeHashCode() { return Arrays.hashCode(data); } From 871d92c76a3627d1a686041b59bca732befdcde4 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Tue, 29 Jul 2025 17:44:37 +0200 Subject: [PATCH 11/34] insert almost works --- .../async/hnsw/ByNodeStorageAdapter.java | 6 + .../apple/foundationdb/async/hnsw/HNSW.java | 386 ++++++++++++------ .../async/hnsw/NodeReferenceAndNode.java | 57 +++ .../async/hnsw/NodeReferenceWithDistance.java | 8 +- .../foundationdb/async/hnsw/SearchResult.java | 60 --- .../async/hnsw/StorageAdapter.java | 3 + 6 files changed, 337 insertions(+), 183 deletions(-) create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceAndNode.java delete mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/SearchResult.java diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java index 6b0d0ff898..df1457458a 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java @@ -38,6 +38,12 @@ public ByNodeStorageAdapter(@Nonnull final HNSW.Config config, @Nonnull final Su super(config, subspace, onWriteListener, onReadListener); } + @Nonnull + @Override + public NodeFactory getNodeFactory(final int layer) { + return layer > 0 ? InliningNode.factory() : CompactNode.factory(); + } + @Override public CompletableFuture fetchEntryNodeReference(@Nonnull final ReadTransaction readTransaction) { final byte[] key = getEntryNodeSubspace().pack(); diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index 14631a00cd..9b185465ad 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -23,16 +23,13 @@ import com.apple.foundationdb.ReadTransaction; import com.apple.foundationdb.Transaction; import com.apple.foundationdb.annotation.API; -import com.apple.foundationdb.async.AsyncIterator; import com.apple.foundationdb.async.AsyncUtil; import com.apple.foundationdb.async.MoreAsyncUtil; import com.apple.foundationdb.subspace.Subspace; import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; -import com.google.common.base.Preconditions; import com.google.common.base.Verify; import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; import com.google.common.collect.Iterables; import com.google.common.collect.Lists; import com.google.common.collect.Maps; @@ -42,9 +39,8 @@ import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; -import javax.annotation.Nullable; -import java.math.BigInteger; import java.util.ArrayDeque; +import java.util.Collection; import java.util.Comparator; import java.util.List; import java.util.Map; @@ -58,25 +54,27 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; -import java.util.function.BiPredicate; import java.util.function.Function; -import java.util.function.Predicate; /** * TODO. */ @API(API.Status.EXPERIMENTAL) +@SuppressWarnings("checkstyle:AbbreviationAsWordInName") public class HNSW { private static final Logger logger = LoggerFactory.getLogger(HNSW.class); public static final int MAX_CONCURRENT_READS = 16; - @Nonnull - public static final Random DEFAULT_RANDOM = new Random(0L); - @Nonnull - public static final Metric DEFAULT_METRIC = new Metric.EuclideanMetric(); - public static final int DEFAULT_M = 48; + @Nonnull public static final Random DEFAULT_RANDOM = new Random(0L); + @Nonnull public static final Metric DEFAULT_METRIC = new Metric.EuclideanMetric(); + public static final int DEFAULT_M = 16; + public static final int DEFAULT_M_MAX = DEFAULT_M; + public static final int DEFAULT_M_MAX_0 = 2 * DEFAULT_M; public static final int DEFAULT_EF_SEARCH = 64; - public static final int DEFAULT_EF_CONSTRUCTION = 100; + public static final int DEFAULT_EF_CONSTRUCTION = 200; + public static final boolean DEFAULT_EXTEND_CANDIDATES = false; + public static final boolean DEFAULT_KEEP_PRUNED_CONNECTIONS = false; + @Nonnull public static final Config DEFAULT_CONFIG = new Config(); @@ -128,30 +126,44 @@ StorageAdapter create(@Nonnull Config config, @Nonnull Subspace subspace, @Nonnu /** * Configuration settings for a {@link HNSW}. */ + @SuppressWarnings("checkstyle:MemberName") public static class Config { @Nonnull private final Random random; @Nonnull private final Metric metric; private final int m; + private final int mMax; + private final int mMax0; private final int efSearch; private final int efConstruction; + private final boolean extendCandidates; + private final boolean keepPrunedConnections; protected Config() { this.random = DEFAULT_RANDOM; this.metric = DEFAULT_METRIC; this.m = DEFAULT_M; + this.mMax = DEFAULT_M_MAX; + this.mMax0 = DEFAULT_M_MAX_0; this.efSearch = DEFAULT_EF_SEARCH; this.efConstruction = DEFAULT_EF_CONSTRUCTION; + this.extendCandidates = DEFAULT_EXTEND_CANDIDATES; + this.keepPrunedConnections = DEFAULT_KEEP_PRUNED_CONNECTIONS; } - protected Config(@Nonnull final Random random, @Nonnull final Metric metric, final int m, - final int efSearch, final int efConstruction) { + protected Config(@Nonnull final Random random, @Nonnull final Metric metric, final int m, final int mMax, + final int mMax0, final int efSearch, final int efConstruction, final boolean extendCandidates, + final boolean keepPrunedConnections) { this.random = random; this.metric = metric; this.m = m; + this.mMax = mMax; + this.mMax0 = mMax0; this.efSearch = efSearch; this.efConstruction = efConstruction; + this.extendCandidates = extendCandidates; + this.keepPrunedConnections = keepPrunedConnections; } @Nonnull @@ -168,6 +180,14 @@ public int getM() { return m; } + public int getMMax() { + return mMax; + } + + public int getMMax0() { + return mMax0; + } + public int getEfSearch() { return efSearch; } @@ -176,15 +196,27 @@ public int getEfConstruction() { return efConstruction; } + public boolean isExtendCandidates() { + return extendCandidates; + } + + public boolean isKeepPrunedConnections() { + return keepPrunedConnections; + } + @Nonnull public ConfigBuilder toBuilder() { - return new ConfigBuilder(getRandom(), getMetric(), getM(), getEfSearch(), getEfConstruction()); + return new ConfigBuilder(getRandom(), getMetric(), getM(), getMMax(), getMMax0(), getEfSearch(), + getEfConstruction(), isExtendCandidates(), isKeepPrunedConnections()); } @Override @Nonnull public String toString() { - return "Config[M=" + getM() + " , metric=" + getMetric() + "]"; + return "Config[metric=" + getMetric() + "M=" + getM() + " , MMax=" + getMMax() + " , MMax0=" + getMMax0() + + ", efSearch=" + getEfSearch() + ", efConstruction=" + getEfConstruction() + + ", isExtendCandidates=" + isExtendCandidates() + + ", isKeepPrunedConnections=" + isKeepPrunedConnections() + "]"; } } @@ -194,25 +226,35 @@ public String toString() { * @see #newConfigBuilder */ @CanIgnoreReturnValue + @SuppressWarnings("checkstyle:MemberName") public static class ConfigBuilder { @Nonnull private Random random = DEFAULT_RANDOM; @Nonnull private Metric metric = DEFAULT_METRIC; private int m = DEFAULT_M; + private int mMax; + private int mMax0; private int efSearch = DEFAULT_EF_SEARCH; private int efConstruction = DEFAULT_EF_CONSTRUCTION; + private boolean extendCandidates = DEFAULT_EXTEND_CANDIDATES; + private boolean keepPrunedConnections = DEFAULT_KEEP_PRUNED_CONNECTIONS; public ConfigBuilder() { } - public ConfigBuilder(@Nonnull Random random, @Nonnull final Metric metric, final int m, - final int efSearch, final int efConstruction) { + public ConfigBuilder(@Nonnull Random random, @Nonnull final Metric metric, final int m, final int mMax, + final int mMax0, final int efSearch, final int efConstruction, + final boolean extendCandidates, final boolean keepPrunedConnections) { this.random = random; this.metric = metric; this.m = m; + this.mMax = mMax; + this.mMax0 = mMax0; this.efSearch = efSearch; this.efConstruction = efConstruction; + this.extendCandidates = extendCandidates; + this.keepPrunedConnections = keepPrunedConnections; } @Nonnull @@ -247,6 +289,26 @@ public ConfigBuilder setM(final int m) { return this; } + public int getMMax() { + return mMax; + } + + @Nonnull + public ConfigBuilder setMMax(final int mMax) { + this.mMax = mMax; + return this; + } + + public int getMMax0() { + return mMax0; + } + + @Nonnull + public ConfigBuilder setMMax0(final int mMax0) { + this.mMax0 = mMax0; + return this; + } + public int getEfSearch() { return efSearch; } @@ -265,8 +327,27 @@ public ConfigBuilder setEfConstruction(final int efConstruction) { return this; } + public boolean isExtendCandidates() { + return extendCandidates; + } + + public ConfigBuilder setExtendCandidates(final boolean extendCandidates) { + this.extendCandidates = extendCandidates; + return this; + } + + public boolean isKeepPrunedConnections() { + return keepPrunedConnections; + } + + public ConfigBuilder setKeepPrunedConnections(final boolean keepPrunedConnections) { + this.keepPrunedConnections = keepPrunedConnections; + return this; + } + public Config build() { - return new Config(getRandom(), getMetric(), getM(), getEfSearch(), getEfConstruction()); + return new Config(getRandom(), getMetric(), getM(), getMMax(), getMMax0(), getEfSearch(), + getEfConstruction(), isExtendCandidates(), isKeepPrunedConnections()); } } @@ -362,67 +443,13 @@ public OnReadListener getOnReadListener() { // Read Path // - /** - * Perform a scan over the tree within the transaction passed in using a predicate that is also passed in to - * eliminate subtrees from the scan. This predicate may be stateful which allows for dynamic adjustments of the - * queried area while the scan is active. - *
- * A scan of the tree offers all items that pass the {@code mbrPredicate} test in Hilbert Value order using an - * {@link AsyncIterator}. The predicate that is passed in is applied to intermediate nodes as well as leaf nodes, - * but not to elements contained by a leaf node. The caller should filter out items in a downstream operation. - * A scan of the tree will not prefetch the next node before the items of the current node have been consumed. This - * guarantees that the semantics of the mbr predicate can be adapted in response to the items being consumed. - * (this allows for efficient scans for {@code ORDER BY x, y LIMIT n} queries). - * @param readTransaction the transaction to use - * @param mbrPredicate a predicate on an mbr {@link Rectangle} - * @param suffixKeyPredicate a predicate on the suffix key - * @return an {@link AsyncIterator} of {@link ItemSlot}s. - */ - @Nonnull - public AsyncIterator scan(@Nonnull final ReadTransaction readTransaction, - @Nonnull final Predicate mbrPredicate, - @Nonnull final BiPredicate suffixKeyPredicate) { - return scan(readTransaction, null, null, mbrPredicate, suffixKeyPredicate); - } - - /** - * Perform a scan over the tree within the transaction passed in using a predicate that is also passed in to - * eliminate subtrees from the scan. This predicate may be stateful which allows for dynamic adjustments of the - * queried area while the scan is active. - *
- * A scan of the tree offers all items that pass the {@code mbrPredicate} test in Hilbert Value order using an - * {@link AsyncIterator}. The predicate that is passed in is applied to intermediate nodes as well as leaf nodes, - * but not to elements contained in a leaf node. The caller should filter out items in a downstream operation. - * A scan of the tree will not prefetch the next node before the items of the current node have been consumed. This - * guarantees that the semantics of the mbr predicate can be adapted in response to the items being consumed. - * (this allows for efficient scans for {@code ORDER BY x, y LIMIT n} queries). - * @param readTransaction the transaction to use - * @param lastHilbertValue the last Hilbert value that was returned by a previous call to this method - * @param lastKey the last key that was returned by a previous call to this method - * @param mbrPredicate a predicate on an mbr {@link Rectangle} - * @param suffixKeyPredicate a predicate on the suffix key - * @return an {@link AsyncIterator} of {@link ItemSlot}s. - */ - @Nonnull - public AsyncIterator scan(@Nonnull final ReadTransaction readTransaction, - @Nullable final BigInteger lastHilbertValue, - @Nullable final Tuple lastKey, - @Nonnull final Predicate mbrPredicate, - @Nonnull final BiPredicate suffixKeyPredicate) { - Preconditions.checkArgument((lastHilbertValue == null && lastKey == null) || - (lastHilbertValue != null && lastKey != null)); - AsyncIterator leafIterator = - new LeafIterator(readTransaction, rootId, lastHilbertValue, lastKey, mbrPredicate, suffixKeyPredicate); - return new ItemSlotIterator(leafIterator); - } - /** * TODO. */ @SuppressWarnings("checkstyle:MethodName") // method name introduced by paper @Nonnull - private CompletableFuture> kNearestNeighborsSearch(@Nonnull final ReadTransaction readTransaction, - @Nonnull final Vector queryVector) { + private CompletableFuture>> kNearestNeighborsSearch(@Nonnull final ReadTransaction readTransaction, + @Nonnull final Vector queryVector) { return storageAdapter.fetchEntryNodeReference(readTransaction) .thenCompose(entryPointAndLayer -> { if (entryPointAndLayer == null) { @@ -433,6 +460,7 @@ private CompletableFuture> kNearestNeighborsSearch(@ final NodeReferenceWithDistance entryState = new NodeReferenceWithDistance(entryPointAndLayer.getPrimaryKey(), + entryPointAndLayer.getVector(), Vector.comparativeDistance(metric, entryPointAndLayer.getVector(), queryVector)); if (entryPointAndLayer.getLayer() == 0) { @@ -440,7 +468,6 @@ private CompletableFuture> kNearestNeighborsSearch(@ return CompletableFuture.completedFuture(entryState); } - final AtomicInteger layerAtomic = new AtomicInteger(entryPointAndLayer.getLayer()); final AtomicReference nodeReferenceAtomic = new AtomicReference<>(entryState); @@ -449,7 +476,7 @@ private CompletableFuture> kNearestNeighborsSearch(@ layer -> layer - 1, layer -> { final var greedyIn = nodeReferenceAtomic.get(); - return greedySearchLayer(InliningNode.factory(), readTransaction, greedyIn, layer, + return greedySearchLayer(storageAdapter.getNodeFactory(layer), readTransaction, greedyIn, layer, queryVector) .thenApply(greedyState -> { nodeReferenceAtomic.set(greedyState); @@ -461,9 +488,9 @@ private CompletableFuture> kNearestNeighborsSearch(@ return CompletableFuture.completedFuture(null); } - return searchLayer(CompactNode.factory(), readTransaction, - ImmutableList.of(nodeReference), 0, config.getEfSearch(), - queryVector); + return searchLayer(storageAdapter.getNodeFactory(0), readTransaction, + ImmutableList.of(nodeReference), 0, config.getEfSearch(), + Maps.newConcurrentMap(), queryVector); }); } @@ -476,8 +503,8 @@ private CompletableFuture g if (nodeFactory.getNodeKind() == NodeKind.INLINING) { return greedySearchInliningLayer(readTransaction, entryNeighbor, layer, queryVector); } else { - return searchLayer(nodeFactory, readTransaction, ImmutableList.of(entryNeighbor), layer, 1, queryVector) - .thenApply(searchResult -> Iterables.getOnlyElement(searchResult.getNodeMap().keySet())); + return searchLayer(nodeFactory, readTransaction, ImmutableList.of(entryNeighbor), layer, 1, Maps.newConcurrentMap(), queryVector) + .thenApply(searchResult -> Iterables.getOnlyElement(searchResult).getNodeReferenceWithDistance()); } } @@ -492,8 +519,7 @@ private CompletableFuture greedySearchInliningLayer(@ Verify.verify(layer > 0); final Metric metric = getConfig().getMetric(); final AtomicReference currentNodeReferenceAtomic = - new AtomicReference<>(new NodeReferenceWithDistance(entryNeighbor.getPrimaryKey(), - entryNeighbor.getDistance())); + new AtomicReference<>(entryNeighbor); return AsyncUtil.whileTrue(() -> onReadListener.onAsyncRead( storageAdapter.fetchNode(InliningNode.factory(), readTransaction, @@ -519,13 +545,12 @@ private CompletableFuture greedySearchInliningLayer(@ } if (nearestNeighbor == null) { - currentNodeReferenceAtomic.set( - new NodeReferenceWithDistance(currentNodeReference.getPrimaryKey(), minDistance)); return false; } currentNodeReferenceAtomic.set( - new NodeReferenceWithDistance(nearestNeighbor.getPrimaryKey(), minDistance)); + new NodeReferenceWithDistance(nearestNeighbor.getPrimaryKey(), nearestNeighbor.getVector(), + minDistance)); return true; }), executor).thenApply(ignored -> currentNodeReferenceAtomic.get()); } @@ -534,12 +559,13 @@ private CompletableFuture greedySearchInliningLayer(@ * TODO. */ @Nonnull - private CompletableFuture> searchLayer(@Nonnull NodeFactory nodeFactory, - @Nonnull final ReadTransaction readTransaction, - @Nonnull final List entryNeighbors, - final int layer, - final int efSearch, - @Nonnull final Vector queryVector) { + private CompletableFuture>> searchLayer(@Nonnull NodeFactory nodeFactory, + @Nonnull final ReadTransaction readTransaction, + @Nonnull final Collection entryNeighbors, + final int layer, + final int efSearch, + @Nonnull final Map> nodeCache, + @Nonnull final Vector queryVector) { final Set visited = Sets.newConcurrentHashSet(NodeReference.primaryKeys(entryNeighbors)); final Queue candidates = new PriorityBlockingQueue<>(config.getM(), @@ -549,7 +575,6 @@ private CompletableFuture> searchLayer new PriorityBlockingQueue<>(config.getM(), Comparator.comparing(NodeReferenceWithDistance::getDistance).reversed()); nearestNeighbors.addAll(entryNeighbors); - final Map> nodeCache = Maps.newConcurrentMap(); final Metric metric = getConfig().getMetric(); return AsyncUtil.whileTrue(() -> { @@ -580,7 +605,8 @@ private CompletableFuture> searchLayer Vector.comparativeDistance(metric, current.getVector(), queryVector); if (currentDistance < furthestDistance || nearestNeighbors.size() < efSearch) { final NodeReferenceWithDistance currentWithDistance = - new NodeReferenceWithDistance(current.getPrimaryKey(), currentDistance); + new NodeReferenceWithDistance(current.getPrimaryKey(), current.getVector(), + currentDistance); candidates.add(currentWithDistance); nearestNeighbors.add(currentWithDistance); if (nearestNeighbors.size() > efSearch) { @@ -590,8 +616,8 @@ private CompletableFuture> searchLayer } return true; }); - }).thenCompose(ignored -> fetchResultsIfNecessary(nodeFactory, readTransaction, layer, nearestNeighbors, - nodeCache)); + }).thenCompose(ignored -> + fetchSomeNodesIfNotCached(nodeFactory, readTransaction, layer, nearestNeighbors, nodeCache)); } /** @@ -661,30 +687,22 @@ private CompletableFuture CompletableFuture> fetchResultsIfNecessary(@Nonnull final NodeFactory creator, - @Nonnull final ReadTransaction readTransaction, - final int layer, - @Nonnull final Iterable nodeReferences, - @Nonnull final Map> nodeCache) { + private CompletableFuture>> fetchSomeNodesIfNotCached(@Nonnull final NodeFactory creator, + @Nonnull final ReadTransaction readTransaction, + final int layer, + @Nonnull final Iterable nodeReferences, + @Nonnull final Map> nodeCache) { return fetchSomeNodesAndApply(creator, readTransaction, layer, nodeReferences, nodeReference -> { final Node node = nodeCache.get(nodeReference.getPrimaryKey()); if (node == null) { return null; } - return new SearchResult.NodeReferenceWithNode<>(nodeReference, node); + return new NodeReferenceAndNode<>(nodeReference, node); }, (nodeReferenceWithDistance, node) -> { nodeCache.put(nodeReferenceWithDistance.getPrimaryKey(), node); - return new SearchResult.NodeReferenceWithNode(nodeReferenceWithDistance, node); - }) - .thenApply(nodeReferencesWithNodes -> { - final ImmutableMap.Builder> nodeMapBuilder = - ImmutableMap.builder(); - for (final SearchResult.NodeReferenceWithNode nodeReferenceWithNode : nodeReferencesWithNodes) { - nodeMapBuilder.put(nodeReferenceWithNode.getNodeReferenceWithDistance(), nodeReferenceWithNode.getNode()); - } - return new SearchResult<>(nodeMapBuilder.build()); + return new NodeReferenceAndNode<>(nodeReferenceWithDistance, node); }); } @@ -718,7 +736,6 @@ private CompletableFuture< } final int index = neighborIndex.getAndIncrement(); - working.add(fetchNodeIfNecessaryAndApply(creator, readTransaction, layer, currentNeighborReference, fetchBypassFunction, biMapFunction) .thenAccept(resultNode -> { @@ -744,13 +761,10 @@ private CompletableFuture< public CompletableFuture insert(@Nonnull final Transaction transaction, @Nonnull final Tuple primaryKey, @Nonnull final Vector vector) { final Metric metric = getConfig().getMetric(); - final Queue furthestNeighbors = - new PriorityBlockingQueue<>(config.getM(), - Comparator.comparing(NodeReferenceWithDistance::getDistance).reversed()); final int l = insertionLayer(getConfig().getRandom()); - storageAdapter.fetchEntryNodeReference(transaction) + return storageAdapter.fetchEntryNodeReference(transaction) .thenApply(entryNodeReference -> { if (entryNodeReference == null) { // this is the first node @@ -779,6 +793,7 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N final AtomicReference nodeReferenceAtomic = new AtomicReference<>(new NodeReferenceWithDistance(entryNodeReference.getPrimaryKey(), + entryNodeReference.getVector(), Vector.comparativeDistance(metric, entryNodeReference.getVector(), vector))); MoreAsyncUtil.forLoop(lMax, layer -> layer > l, @@ -790,11 +805,142 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N return null; }), executor); - final NodeReferenceWithDistance nodeReference = nodeReferenceAtomic.get(); + final AtomicReference> nearestNeighborsAtomic = + new AtomicReference<>(ImmutableList.of(nodeReferenceAtomic.get())); + + MoreAsyncUtil.forLoop(Math.min(lMax, l), + layer -> layer >= 0, + layer -> layer - 1, + layer -> insertIntoLayer(storageAdapter.getNodeFactory(layer), transaction, + nearestNeighborsAtomic.get(), layer, primaryKey, vector) + .thenCompose(nearestNeighbors -> { + nearestNeighborsAtomic.set(nearestNeighbors); + return AsyncUtil.DONE; + }), executor); + }).thenCompose(ignored -> AsyncUtil.DONE); + } + + private CompletableFuture> insertIntoLayer(@Nonnull final NodeFactory nodeFactory, + @Nonnull final Transaction transaction, + @Nonnull final List nearestNeighbors, + int layer, + @Nonnull final Tuple primaryKey, + @Nonnull final Vector vector) { + final Map> nodeCache = Maps.newConcurrentMap(); + final Metric metric = getConfig().getMetric(); + + return searchLayer(nodeFactory, transaction, + nearestNeighbors, layer, config.getEfConstruction(), nodeCache, vector) + .thenApply(searchResult -> { + final List references = NodeReferenceAndNode.getReferences(searchResult); + selectNeighbors(nodeFactory, transaction, searchResult, + layer, getConfig().getM(), getConfig().isExtendCandidates(), nodeCache, vector); + return ImmutableList.copyOf(references); + }); + } + private CompletableFuture>> selectNeighbors(@Nonnull final NodeFactory nodeFactory, + @Nonnull final ReadTransaction readTransaction, + @Nonnull final List> nearestNeighbors, + final int layer, + final int m, + final boolean isExtendCandidates, + @Nonnull final Map> nodeCache, + @Nonnull final Vector vector) { + return extendCandidatesIfNecessary(nodeFactory, readTransaction, nearestNeighbors, layer, isExtendCandidates, nodeCache, vector) + .thenApply(extendedCandidates -> { + final List selected = Lists.newArrayListWithExpectedSize(m); + final Queue candidates = + new PriorityBlockingQueue<>(config.getM(), + Comparator.comparing(NodeReferenceWithDistance::getDistance)); + candidates.addAll(extendedCandidates); + final Queue discardedCandidates = + getConfig().isKeepPrunedConnections() + ? new PriorityBlockingQueue<>(config.getM(), + Comparator.comparing(NodeReferenceWithDistance::getDistance)) + : null; - }) + final Metric metric = getConfig().getMetric(); + + while (!candidates.isEmpty() && selected.size() < m) { + final NodeReferenceWithDistance nearestCandidate = candidates.poll(); + boolean shouldSelect = true; + for (final NodeReferenceWithDistance alreadySelected : selected) { + if (Vector.comparativeDistance(metric, nearestCandidate.getVector(), + alreadySelected.getVector()) < nearestCandidate.getDistance()) { + shouldSelect = false; + break; + } + } + if (shouldSelect) { + selected.add(nearestCandidate); + } else if (discardedCandidates != null) { + discardedCandidates.add(nearestCandidate); + } + } + + if (discardedCandidates != null) { // isKeepPrunedConnections is set to true + while (!discardedCandidates.isEmpty() && selected.size() < m) { + selected.add(discardedCandidates.poll()); + } + } + + return ImmutableList.copyOf(selected); + }).thenCompose(selectedNeighbors -> + fetchSomeNodesIfNotCached(nodeFactory, readTransaction, layer, selectedNeighbors, nodeCache)) + } + + private CompletableFuture> extendCandidatesIfNecessary(@Nonnull final NodeFactory nodeFactory, + @Nonnull final ReadTransaction readTransaction, + @Nonnull final List> candidates, + int layer, + boolean isExtendCandidates, + @Nonnull final Map> nodeCache, + @Nonnull final Vector vector) { + if (isExtendCandidates) { + final Metric metric = getConfig().getMetric(); + + final Set candidatesSeen = Sets.newConcurrentHashSet(); + for (final NodeReferenceAndNode candidate : candidates) { + candidatesSeen.add(candidate.getNode().getPrimaryKey()); + } + + final ImmutableList.Builder neighborsOfCandidatesBuilder = ImmutableList.builder(); + for (final NodeReferenceAndNode candidate : candidates) { + for (final N neighbor : candidate.getNode().getNeighbors()) { + final Tuple neighborPrimaryKey = neighbor.getPrimaryKey(); + if (!candidatesSeen.contains(neighborPrimaryKey)) { + candidatesSeen.add(neighborPrimaryKey); + neighborsOfCandidatesBuilder.add(neighbor); + } + } + } + + final Iterable neighborsOfCandidates = neighborsOfCandidatesBuilder.build(); + + return fetchNeighborhood(nodeFactory, readTransaction, layer, neighborsOfCandidates, nodeCache) + .thenApply(withVectors -> { + final ImmutableList.Builder extendedCandidatesBuilder = ImmutableList.builder(); + for (final NodeReferenceAndNode candidate : candidates) { + extendedCandidatesBuilder.add(candidate.getNodeReferenceWithDistance()); + } + + for (final NodeReferenceWithVector withVector : withVectors) { + final double distance = Vector.comparativeDistance(metric, withVector.getVector(), vector); + extendedCandidatesBuilder.add(new NodeReferenceWithDistance(withVector.getPrimaryKey(), + withVector.getVector(), distance)); + } + return extendedCandidatesBuilder.build(); + }); + } else { + final ImmutableList.Builder resultBuilder = ImmutableList.builder(); + for (final NodeReferenceAndNode candidate : candidates) { + resultBuilder.add(candidate.getNodeReferenceWithDistance()); + } + + return CompletableFuture.completedFuture(resultBuilder.build()); + } } public void writeLonelyNodes(@Nonnull final NodeFactory nodeFactory, diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceAndNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceAndNode.java new file mode 100644 index 0000000000..0b9ad99faa --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceAndNode.java @@ -0,0 +1,57 @@ +/* + * NodeReferenceAndNode.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.google.common.collect.ImmutableList; + +import javax.annotation.Nonnull; +import java.util.List; + +public class NodeReferenceAndNode { + @Nonnull + private final NodeReferenceWithDistance nodeReferenceWithDistance; + @Nonnull + private final Node node; + + public NodeReferenceAndNode(@Nonnull final NodeReferenceWithDistance nodeReferenceWithDistance, @Nonnull final Node node) { + this.nodeReferenceWithDistance = nodeReferenceWithDistance; + this.node = node; + } + + @Nonnull + public NodeReferenceWithDistance getNodeReferenceWithDistance() { + return nodeReferenceWithDistance; + } + + @Nonnull + public Node getNode() { + return node; + } + + @Nonnull + public static List getReferences(@Nonnull List> referencesAndNodes) { + final ImmutableList.Builder referencesBuilder = ImmutableList.builder(); + for (final NodeReferenceAndNode referenceWithNode : referencesAndNodes) { + referencesBuilder.add(referenceWithNode.getNodeReferenceWithDistance()); + } + return referencesBuilder.build(); + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithDistance.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithDistance.java index 4408bf0f4e..96a3a23720 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithDistance.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithDistance.java @@ -21,15 +21,17 @@ package com.apple.foundationdb.async.hnsw; import com.apple.foundationdb.tuple.Tuple; +import com.christianheina.langx.half4j.Half; import javax.annotation.Nonnull; import java.util.Objects; -class NodeReferenceWithDistance extends NodeReference { +class NodeReferenceWithDistance extends NodeReferenceWithVector { private final double distance; - public NodeReferenceWithDistance(@Nonnull final Tuple primaryKey, final double distance) { - super(primaryKey); + public NodeReferenceWithDistance(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, + final double distance) { + super(primaryKey, vector); this.distance = distance; } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/SearchResult.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/SearchResult.java deleted file mode 100644 index 0685930957..0000000000 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/SearchResult.java +++ /dev/null @@ -1,60 +0,0 @@ -/* - * NodeWithLayer.java - * - * This source file is part of the FoundationDB open source project - * - * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors - * - * Licensed 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 com.apple.foundationdb.async.hnsw; - -import javax.annotation.Nonnull; -import java.util.Map; - -class SearchResult { - @Nonnull - private final Map> nodeMap; - - public SearchResult(@Nonnull final Map> nodeMap) { - this.nodeMap = nodeMap; - } - - @Nonnull - public Map> getNodeMap() { - return nodeMap; - } - - public static class NodeReferenceWithNode { - @Nonnull - private final NodeReferenceWithDistance nodeReferenceWithDistance; - @Nonnull - private final Node node; - - public NodeReferenceWithNode(@Nonnull final NodeReferenceWithDistance nodeReferenceWithDistance, @Nonnull final Node node) { - this.nodeReferenceWithDistance = nodeReferenceWithDistance; - this.node = node; - } - - @Nonnull - public NodeReferenceWithDistance getNodeReferenceWithDistance() { - return nodeReferenceWithDistance; - } - - @Nonnull - public Node getNode() { - return node; - } - } -} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java index dfa7d75dde..9dcadcaa87 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java @@ -82,6 +82,9 @@ interface StorageAdapter { @Nonnull OnReadListener getOnReadListener(); + @Nonnull + NodeFactory getNodeFactory(final int layer); + CompletableFuture fetchEntryNodeReference(@Nonnull ReadTransaction readTransaction); void writeEntryNodeReference(@Nonnull final Transaction transaction, From dae05d552075db77c8a878ed3d5cb688a02b354a Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Wed, 30 Jul 2025 12:52:48 +0200 Subject: [PATCH 12/34] insert almost works --- .../foundationdb/async/MoreAsyncUtil.java | 50 +++- .../foundationdb/async/hnsw/AbstractNode.java | 12 +- .../async/hnsw/BaseNeighborsChangeSet.java | 53 ++++ .../foundationdb/async/hnsw/CompactNode.java | 11 +- .../async/hnsw/DeleteNeighborsChangeSet.java | 63 +++++ .../apple/foundationdb/async/hnsw/HNSW.java | 232 ++++++++++++------ .../foundationdb/async/hnsw/InliningNode.java | 11 +- .../async/hnsw/InsertNeighborsChangeSet.java | 60 +++++ .../async/hnsw/NeighborsChangeSet.java | 39 +++ .../apple/foundationdb/async/hnsw/Node.java | 14 +- .../foundationdb/async/hnsw/NodeFactory.java | 2 +- .../async/hnsw/NodeReferenceAndNode.java | 12 +- 12 files changed, 462 insertions(+), 97 deletions(-) create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/BaseNeighborsChangeSet.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NeighborsChangeSet.java diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java index f6f0c999e5..85ea6ad045 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java @@ -23,6 +23,8 @@ import com.apple.foundationdb.annotation.API; import com.apple.foundationdb.util.LoggableException; import com.google.common.base.Suppliers; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Lists; import com.google.common.util.concurrent.ThreadFactoryBuilder; import javax.annotation.Nonnull; @@ -33,6 +35,7 @@ import java.util.Iterator; import java.util.List; import java.util.NoSuchElementException; +import java.util.Objects; import java.util.Queue; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; @@ -1055,7 +1058,7 @@ public static CompletableFuture swallowException(@Nonnull CompletableFutur return result; } - public static CompletableFuture forLoop(int startI, @Nonnull final IntPredicate conditionPredicate, + public static CompletableFuture forLoop(final int startI, @Nonnull final IntPredicate conditionPredicate, @Nonnull final IntUnaryOperator stepFunction, @Nonnull final IntFunction> body, @Nonnull final Executor executor) { @@ -1073,6 +1076,51 @@ public static CompletableFuture forLoop(int startI, @Nonnull final IntPred }, executor); } + @SuppressWarnings("unchecked") + public static CompletableFuture> forEach(@Nonnull final Iterable items, + @Nonnull final Function> body, + final int parallelism, + @Nonnull final Executor executor) { + // this deque is only modified by once upon creation + final ArrayDeque toBeProcessed = new ArrayDeque<>(); + for (final T item : items) { + toBeProcessed.addLast(item); + } + + final List> working = Lists.newArrayList(); + final AtomicInteger indexAtomic = new AtomicInteger(0); + final Object[] resultArray = new Object[toBeProcessed.size()]; + + return AsyncUtil.whileTrue(() -> { + working.removeIf(CompletableFuture::isDone); + + while (working.size() <= parallelism) { + final T currentItem = toBeProcessed.pollFirst(); + if (currentItem == null) { + break; + } + + final int index = indexAtomic.getAndIncrement(); + working.add(body.apply(currentItem) + .thenAccept(resultNode -> { + Objects.requireNonNull(resultNode); + resultArray[index] = resultNode; + })); + } + + if (working.isEmpty()) { + return AsyncUtil.READY_FALSE; + } + return AsyncUtil.whenAny(working).thenApply(ignored -> true); + }, executor).thenApply(ignored -> { + final ImmutableList.Builder resultBuilder = ImmutableList.builder(); + for (final Object o : resultArray) { + resultBuilder.add((U)o); + } + return resultBuilder.build(); + }); + } + /** * A {@code Boolean} function that is always true. * @param the type of the (ignored) argument to the function diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java index 9d5c1ff95c..aa062e8700 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractNode.java @@ -28,17 +28,17 @@ /** * TODO. - * @param node type class. + * @param node type class. */ -abstract class AbstractNode implements Node { +abstract class AbstractNode implements Node { @Nonnull private final Tuple primaryKey; @Nonnull - private final List neighbors; + private final List neighbors; protected AbstractNode(@Nonnull final Tuple primaryKey, - @Nonnull final List neighbors) { + @Nonnull final List neighbors) { this.primaryKey = primaryKey; this.neighbors = ImmutableList.copyOf(neighbors); } @@ -51,13 +51,13 @@ public Tuple getPrimaryKey() { @Nonnull @Override - public List getNeighbors() { + public List getNeighbors() { return neighbors; } @Nonnull @Override - public R getNeighbor(final int index) { + public N getNeighbor(final int index) { return neighbors.get(index); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/BaseNeighborsChangeSet.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/BaseNeighborsChangeSet.java new file mode 100644 index 0000000000..6b5ce44b84 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/BaseNeighborsChangeSet.java @@ -0,0 +1,53 @@ +/* + * InliningNode.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.Transaction; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.util.List; + +/** + * TODO. + */ +class BaseNeighborsChangeSet implements NeighborsChangeSet { + @Nonnull + private final NodeReferenceAndNode baseNode; + + public BaseNeighborsChangeSet(@Nonnull final NodeReferenceAndNode baseNode) { + this.baseNode = baseNode; + } + + @Nullable + public BaseNeighborsChangeSet getParent() { + return null; + } + + @Nonnull + public List merge() { + return baseNode.getNode().getNeighbors(); + } + + @Override + public void writeDelta(@Nonnull final Transaction transaction) { + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java index 2bf332751a..a213d70411 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java @@ -22,7 +22,6 @@ import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; -import com.google.common.base.Verify; import com.google.common.collect.Lists; import javax.annotation.Nonnull; @@ -39,8 +38,8 @@ public class CompactNode extends AbstractNode { @SuppressWarnings("unchecked") @Nonnull @Override - public Node create(@Nonnull final NodeKind nodeKind, @Nonnull final Tuple primaryKey, @Nullable final Vector vector, @Nonnull final List neighbors) { - Verify.verify(nodeKind == NodeKind.COMPACT); + public Node create(@Nonnull final Tuple primaryKey, @Nullable final Vector vector, + @Nonnull final List neighbors) { return new CompactNode(primaryKey, Objects.requireNonNull(vector), (List)neighbors); } @@ -60,6 +59,12 @@ public CompactNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector this.vector = vector; } + @Nonnull + @Override + public NodeReference getSelfReference(@Nullable final Vector vector) { + return new NodeReference(getPrimaryKey()); + } + @Nonnull @Override public NodeKind getKind() { diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java new file mode 100644 index 0000000000..c1ff4e73c3 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java @@ -0,0 +1,63 @@ +/* + * InliningNode.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.Transaction; +import com.apple.foundationdb.tuple.Tuple; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Iterables; + +import javax.annotation.Nonnull; +import java.util.Collection; +import java.util.Set; + +/** + * TODO. + */ +class DeleteNeighborsChangeSet implements NeighborsChangeSet { + @Nonnull + private final NeighborsChangeSet parent; + + @Nonnull + private final Set deletedNeighborsTuples; + + public DeleteNeighborsChangeSet(@Nonnull final NeighborsChangeSet parent, + @Nonnull final Collection deletedNeighborsTuples) { + this.parent = parent; + this.deletedNeighborsTuples = ImmutableSet.copyOf(deletedNeighborsTuples); + } + + @Nonnull + public NeighborsChangeSet getParent() { + return parent; + } + + @Nonnull + public Iterable merge() { + return Iterables.filter(getParent().merge(), + current -> !deletedNeighborsTuples.contains(current.getPrimaryKey())); + } + + @Override + public void writeDelta(@Nonnull final Transaction transaction) { + throw new UnsupportedOperationException("not implemented yet"); + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index 9b185465ad..f4e007a638 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -39,7 +39,6 @@ import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; -import java.util.ArrayDeque; import java.util.Collection; import java.util.Comparator; import java.util.List; @@ -51,7 +50,6 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; import java.util.concurrent.PriorityBlockingQueue; -import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import java.util.function.Function; @@ -64,7 +62,8 @@ public class HNSW { private static final Logger logger = LoggerFactory.getLogger(HNSW.class); - public static final int MAX_CONCURRENT_READS = 16; + public static final int MAX_CONCURRENT_NODE_READS = 16; + public static final int MAX_CONCURRENT_NEIGHBOR_FETCHES = 3; @Nonnull public static final Random DEFAULT_RANDOM = new Random(0L); @Nonnull public static final Metric DEFAULT_METRIC = new Metric.EuclideanMetric(); public static final int DEFAULT_M = 16; @@ -475,14 +474,15 @@ private CompletableFuture layer > 0, layer -> layer - 1, layer -> { - final var greedyIn = nodeReferenceAtomic.get(); - return greedySearchLayer(storageAdapter.getNodeFactory(layer), readTransaction, greedyIn, layer, - queryVector) - .thenApply(greedyState -> { - nodeReferenceAtomic.set(greedyState); - return null; - }); - }, executor).thenApply(ignored -> nodeReferenceAtomic.get()); + final var greedyIn = nodeReferenceAtomic.get(); + return greedySearchLayer(storageAdapter.getNodeFactory(layer), readTransaction, greedyIn, layer, + queryVector) + .thenApply(greedyState -> { + nodeReferenceAtomic.set(greedyState); + return null; + }); + }, executor) + .thenApply(ignored -> nodeReferenceAtomic.get()); }).thenCompose(nodeReference -> { if (nodeReference == null) { return CompletableFuture.completedFuture(null); @@ -717,49 +717,15 @@ private CompletableFuture< @Nonnull final Iterable nodeReferences, @Nonnull final Function fetchBypassFunction, @Nonnull final BiFunction, U> biMapFunction) { - // this deque is only modified by once upon creation - final ArrayDeque toBeProcessed = new ArrayDeque<>(); - for (final var nodeReference : nodeReferences) { - toBeProcessed.addLast(nodeReference); - } - final List> working = Lists.newArrayList(); - final AtomicInteger neighborIndex = new AtomicInteger(0); - final Object[] neighborNodeArray = new Object[toBeProcessed.size()]; - - return AsyncUtil.whileTrue(() -> { - working.removeIf(CompletableFuture::isDone); - - while (working.size() <= MAX_CONCURRENT_READS) { - final R currentNeighborReference = toBeProcessed.pollFirst(); - if (currentNeighborReference == null) { - break; - } - - final int index = neighborIndex.getAndIncrement(); - working.add(fetchNodeIfNecessaryAndApply(creator, readTransaction, layer, - currentNeighborReference, fetchBypassFunction, biMapFunction) - .thenAccept(resultNode -> { - Objects.requireNonNull(resultNode); - neighborNodeArray[index] = resultNode; - })); - } - - if (working.isEmpty()) { - return AsyncUtil.READY_FALSE; - } - return AsyncUtil.whenAny(working).thenApply(ignored -> true); - }, executor).thenApply(ignored -> { - final ImmutableList.Builder resultBuilder = ImmutableList.builder(); - for (final Object o : neighborNodeArray) { - resultBuilder.add((U)o); - } - return resultBuilder.build(); - }); + return MoreAsyncUtil.forEach(nodeReferences, + currentNeighborReference -> fetchNodeIfNecessaryAndApply(creator, readTransaction, layer, + currentNeighborReference, fetchBypassFunction, biMapFunction), MAX_CONCURRENT_NODE_READS, + getExecutor()); } @Nonnull - public CompletableFuture insert(@Nonnull final Transaction transaction, @Nonnull final Tuple primaryKey, - @Nonnull final Vector vector) { + public CompletableFuture insert(@Nonnull final Transaction transaction, @Nonnull final Tuple newPrimaryKey, + @Nonnull final Vector newVector) { final Metric metric = getConfig().getMetric(); final int l = insertionLayer(getConfig().getRandom()); @@ -768,19 +734,19 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N .thenApply(entryNodeReference -> { if (entryNodeReference == null) { // this is the first node - writeLonelyNodes(InliningNode.factory(), transaction, primaryKey, vector, l, 0); + writeLonelyNodes(InliningNode.factory(), transaction, newPrimaryKey, newVector, l, 0); storageAdapter.writeNode(transaction, CompactNode.factory() - .create(NodeKind.COMPACT, primaryKey, vector, ImmutableList.of()), + .create(newPrimaryKey, newVector, ImmutableList.of()), 0); storageAdapter.writeEntryNodeReference(transaction, - new EntryNodeReference(primaryKey, vector, l)); + new EntryNodeReference(newPrimaryKey, newVector, l)); } else { final int entryNodeLayer = entryNodeReference.getLayer(); if (l > entryNodeLayer) { - writeLonelyNodes(InliningNode.factory(), transaction, primaryKey, vector, l, entryNodeLayer); + writeLonelyNodes(InliningNode.factory(), transaction, newPrimaryKey, newVector, l, entryNodeLayer); storageAdapter.writeEntryNodeReference(transaction, - new EntryNodeReference(primaryKey, vector, l)); + new EntryNodeReference(newPrimaryKey, newVector, l)); } } return entryNodeReference; @@ -794,12 +760,12 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N final AtomicReference nodeReferenceAtomic = new AtomicReference<>(new NodeReferenceWithDistance(entryNodeReference.getPrimaryKey(), entryNodeReference.getVector(), - Vector.comparativeDistance(metric, entryNodeReference.getVector(), vector))); + Vector.comparativeDistance(metric, entryNodeReference.getVector(), newVector))); MoreAsyncUtil.forLoop(lMax, layer -> layer > l, layer -> layer - 1, layer -> greedySearchLayer(InliningNode.factory(), transaction, - nodeReferenceAtomic.get(), layer, vector) + nodeReferenceAtomic.get(), layer, newVector) .thenApply(nodeReference -> { nodeReferenceAtomic.set(nodeReference); return null; @@ -808,11 +774,11 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N final AtomicReference> nearestNeighborsAtomic = new AtomicReference<>(ImmutableList.of(nodeReferenceAtomic.get())); - MoreAsyncUtil.forLoop(Math.min(lMax, l), + return MoreAsyncUtil.forLoop(Math.min(lMax, l), layer -> layer >= 0, layer -> layer - 1, layer -> insertIntoLayer(storageAdapter.getNodeFactory(layer), transaction, - nearestNeighborsAtomic.get(), layer, primaryKey, vector) + nearestNeighborsAtomic.get(), layer, newPrimaryKey, newVector) .thenCompose(nearestNeighbors -> { nearestNeighborsAtomic.set(nearestNeighbors); return AsyncUtil.DONE; @@ -820,29 +786,151 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N }).thenCompose(ignored -> AsyncUtil.DONE); } + @Nonnull private CompletableFuture> insertIntoLayer(@Nonnull final NodeFactory nodeFactory, @Nonnull final Transaction transaction, @Nonnull final List nearestNeighbors, int layer, - @Nonnull final Tuple primaryKey, - @Nonnull final Vector vector) { + @Nonnull final Tuple newPrimaryKey, + @Nonnull final Vector newVector) { final Map> nodeCache = Maps.newConcurrentMap(); - final Metric metric = getConfig().getMetric(); return searchLayer(nodeFactory, transaction, - nearestNeighbors, layer, config.getEfConstruction(), nodeCache, vector) - .thenApply(searchResult -> { + nearestNeighbors, layer, config.getEfConstruction(), nodeCache, newVector) + .thenCompose(searchResult -> { final List references = NodeReferenceAndNode.getReferences(searchResult); - selectNeighbors(nodeFactory, transaction, searchResult, - layer, getConfig().getM(), getConfig().isExtendCandidates(), nodeCache, vector); - return ImmutableList.copyOf(references); + return selectNeighbors(nodeFactory, transaction, searchResult, + layer, getConfig().getM(), getConfig().isExtendCandidates(), nodeCache, newVector) + .thenCompose(selectedNeighbors -> { + final Node newNode = + nodeFactory.create(newPrimaryKey, newVector, NodeReferenceAndNode.getReferences(selectedNeighbors)); + + // create change sets for each selected neighbor and insert new node into them + final Map> neighborChangeSetMap = + Maps.newLinkedHashMap(); + for (final NodeReferenceAndNode selectedNeighbor : selectedNeighbors) { + final NeighborsChangeSet baseSet = + new BaseNeighborsChangeSet<>(selectedNeighbor); + final NeighborsChangeSet insertSet = + new InsertNeighborsChangeSet<>(baseSet, ImmutableList.of(newNode.getSelfReference(newVector))); + neighborChangeSetMap.put(selectedNeighbor.getNode().getPrimaryKey(), + insertSet); + } + + final int currentMMax = layer == 0 ? getConfig().getMMax0() : getConfig().getMMax(); + return MoreAsyncUtil.forEach(selectedNeighbors, + selectedNeighbor -> { + final Node selectedNeighborNode = selectedNeighbor.getNode(); + final NeighborsChangeSet changeSet = + Objects.requireNonNull(neighborChangeSetMap.get(selectedNeighborNode.getPrimaryKey())); + return pruneNeighborsIfNecessary(nodeFactory, transaction, selectedNeighbor, layer, + currentMMax, changeSet, nodeCache) + .thenApply(nodeReferencesAndNodes -> { + if (nodeReferencesAndNodes == null) { + return changeSet; + } + return resolveChangeSetFromNewNeighbors(changeSet, nodeReferencesAndNodes); + }); + }, MAX_CONCURRENT_NEIGHBOR_FETCHES, getExecutor()) + .thenApply(changeSets -> { + for (int i = 0; i < selectedNeighbors.size(); i++) { + final NodeReferenceAndNode selectedNeighbor = selectedNeighbors.get(i); + final NeighborsChangeSet changeSet = changeSets.get(i); + neighborChangeSetMap.put( + selectedNeighbor.getNodeReferenceWithDistance().getPrimaryKey(), + changeSet); + } + + // TODO write newVector and all change sets pertaining to neighbors + return ImmutableList.copyOf(references); + }); + }); }); } + private NeighborsChangeSet resolveChangeSetFromNewNeighbors(@Nonnull final NeighborsChangeSet beforeChangeSet, + @Nonnull final Iterable> afterNeighbors) { + final Map beforeNeighborsMap = Maps.newLinkedHashMap(); + for (final N n : beforeChangeSet.merge()) { + beforeNeighborsMap.put(n.getPrimaryKey(), n); + } + + final Map afterNeighborsMap = Maps.newLinkedHashMap(); + for (final NodeReferenceAndNode nodeReferenceAndNode : afterNeighbors) { + final NodeReferenceWithDistance nodeReferenceWithDistance = nodeReferenceAndNode.getNodeReferenceWithDistance(); + + afterNeighborsMap.put(nodeReferenceWithDistance.getPrimaryKey(), + nodeReferenceAndNode.getNode().getSelfReference(nodeReferenceWithDistance.getVector())); + } + + final ImmutableList.Builder toBeDeletedBuilder = ImmutableList.builder(); + for (final Map.Entry beforeNeighborEntry : beforeNeighborsMap.entrySet()) { + if (!afterNeighborsMap.containsKey(beforeNeighborEntry.getKey())) { + toBeDeletedBuilder.add(beforeNeighborEntry.getValue().getPrimaryKey()); + } + } + final List toBeDeleted = toBeDeletedBuilder.build(); + + final ImmutableList.Builder toBeInsertedBuilder = ImmutableList.builder(); + for (final Map.Entry afterNeighborEntry : afterNeighborsMap.entrySet()) { + if (!beforeNeighborsMap.containsKey(afterNeighborEntry.getKey())) { + toBeInsertedBuilder.add(afterNeighborEntry.getValue()); + } + } + final List toBeInserted = toBeInsertedBuilder.build(); + + NeighborsChangeSet changeSet = beforeChangeSet; + + if (!toBeDeleted.isEmpty()) { + changeSet = new DeleteNeighborsChangeSet<>(changeSet, toBeDeleted); + } + if (!toBeInserted.isEmpty()) { + changeSet = new InsertNeighborsChangeSet<>(changeSet, toBeInserted); + } + return changeSet; + } + + @Nonnull + private CompletableFuture>> pruneNeighborsIfNecessary(@Nonnull final NodeFactory nodeFactory, + @Nonnull final Transaction transaction, + @Nonnull final NodeReferenceAndNode selectedNeighbor, + int layer, + int mMax, + @Nonnull final NeighborsChangeSet neighborChangeSet, + @Nonnull final Map> nodeCache) { + final Metric metric = getConfig().getMetric(); + final Node selectedNeighborNode = selectedNeighbor.getNode(); + if (selectedNeighborNode.getNeighbors().size() < mMax) { + return CompletableFuture.completedFuture(null); + } else { + return fetchNeighborhood(nodeFactory, transaction, layer, neighborChangeSet.merge(), nodeCache) + .thenCompose(nodeReferenceWithVectors -> { + final ImmutableList.Builder nodeReferencesWithDistancesBuilder = + ImmutableList.builder(); + for (final NodeReferenceWithVector nodeReferenceWithVector : nodeReferenceWithVectors) { + final var vector = nodeReferenceWithVector.getVector(); + final double distance = + Vector.comparativeDistance(metric, vector, + selectedNeighbor.getNodeReferenceWithDistance().getVector()); + nodeReferencesWithDistancesBuilder.add( + new NodeReferenceWithDistance(nodeReferenceWithVector.getPrimaryKey(), + vector, distance)); + } + return fetchSomeNodesIfNotCached(nodeFactory, transaction, layer, + nodeReferencesWithDistancesBuilder.build(), nodeCache); + }) + .thenCompose(nodeReferencesAndNodes -> + selectNeighbors(nodeFactory, transaction, + nodeReferencesAndNodes, layer, + mMax, false, nodeCache, + selectedNeighbor.getNodeReferenceWithDistance().getVector())); + } + } + private CompletableFuture>> selectNeighbors(@Nonnull final NodeFactory nodeFactory, @Nonnull final ReadTransaction readTransaction, - @Nonnull final List> nearestNeighbors, + @Nonnull final Iterable> nearestNeighbors, final int layer, final int m, final boolean isExtendCandidates, @@ -888,12 +976,12 @@ private CompletableFuture return ImmutableList.copyOf(selected); }).thenCompose(selectedNeighbors -> - fetchSomeNodesIfNotCached(nodeFactory, readTransaction, layer, selectedNeighbors, nodeCache)) + fetchSomeNodesIfNotCached(nodeFactory, readTransaction, layer, selectedNeighbors, nodeCache)); } private CompletableFuture> extendCandidatesIfNecessary(@Nonnull final NodeFactory nodeFactory, @Nonnull final ReadTransaction readTransaction, - @Nonnull final List> candidates, + @Nonnull final Iterable> candidates, int layer, boolean isExtendCandidates, @Nonnull final Map> nodeCache, @@ -951,7 +1039,7 @@ public void writeLonelyNodes(@Nonnull final NodeFactor final int lowestLayerExclusive) { for (int layer = highestLayerInclusive; layer > lowestLayerExclusive; layer --) { storageAdapter.writeNode(transaction, - nodeFactory.create(nodeFactory.getNodeKind(), primaryKey, vector, ImmutableList.of()), layer); + nodeFactory.create(primaryKey, vector, ImmutableList.of()), layer); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java index bfc4de8153..f3d08a25c1 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java @@ -22,12 +22,12 @@ import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; -import com.google.common.base.Verify; import com.google.common.collect.Lists; import javax.annotation.Nonnull; import javax.annotation.Nullable; import java.util.List; +import java.util.Objects; /** * TODO. @@ -38,9 +38,8 @@ class InliningNode extends AbstractNode { @SuppressWarnings("unchecked") @Nonnull @Override - public Node create(@Nonnull final NodeKind nodeKind, @Nonnull final Tuple primaryKey, + public Node create(@Nonnull final Tuple primaryKey, @Nullable final Vector vector, @Nonnull final List neighbors) { - Verify.verify(nodeKind == NodeKind.INLINING); return new InliningNode(primaryKey, (List)neighbors); } @@ -56,6 +55,12 @@ public InliningNode(@Nonnull final Tuple primaryKey, super(primaryKey, neighbors); } + @Nonnull + @Override + public NodeReferenceWithVector getSelfReference(@Nullable final Vector vector) { + return new NodeReferenceWithVector(getPrimaryKey(), Objects.requireNonNull(vector)); + } + @Nonnull @Override public NodeKind getKind() { diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java new file mode 100644 index 0000000000..e91a319c9b --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java @@ -0,0 +1,60 @@ +/* + * InliningNode.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.Transaction; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; + +import javax.annotation.Nonnull; +import java.util.List; + +/** + * TODO. + */ +class InsertNeighborsChangeSet implements NeighborsChangeSet { + @Nonnull + private final NeighborsChangeSet parent; + + @Nonnull + private final List insertedNeighbors; + + public InsertNeighborsChangeSet(@Nonnull final NeighborsChangeSet parent, + @Nonnull final List insertedNeighbors) { + this.parent = parent; + this.insertedNeighbors = ImmutableList.copyOf(insertedNeighbors); + } + + @Nonnull + public NeighborsChangeSet getParent() { + return parent; + } + + @Nonnull + public Iterable merge() { + return Iterables.concat(getParent().merge(), insertedNeighbors); + } + + @Override + public void writeDelta(@Nonnull final Transaction transaction) { + throw new UnsupportedOperationException("not implemented yet"); + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NeighborsChangeSet.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NeighborsChangeSet.java new file mode 100644 index 0000000000..21f621670a --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NeighborsChangeSet.java @@ -0,0 +1,39 @@ +/* + * InliningNode.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.Transaction; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +/** + * TODO. + */ +interface NeighborsChangeSet { + @Nullable + NeighborsChangeSet getParent(); + + @Nonnull + Iterable merge(); + + void writeDelta(@Nonnull final Transaction transaction); +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java index 49f80e130d..aaa65dd21a 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java @@ -25,21 +25,25 @@ import com.google.common.collect.Lists; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import java.util.List; /** * TODO. - * @param neighbor type + * @param neighbor type */ -public interface Node { +public interface Node { @Nonnull Tuple getPrimaryKey(); @Nonnull - List getNeighbors(); + N getSelfReference(@Nullable final Vector vector); @Nonnull - R getNeighbor(int index); + List getNeighbors(); + + @Nonnull + N getNeighbor(int index); /** * Return the kind of the node, i.e. {@link NodeKind#COMPACT} or {@link NodeKind#INLINING}. @@ -54,7 +58,7 @@ public interface Node { @Nonnull InliningNode asInliningNode(); - NodeFactory sameCreator(); + NodeFactory sameCreator(); @Nonnull Tuple toTuple(); diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeFactory.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeFactory.java index 51e3fb59b1..321e3f53d8 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeFactory.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeFactory.java @@ -29,7 +29,7 @@ public interface NodeFactory { @Nonnull - Node create(@Nonnull NodeKind nodeKind, @Nonnull Tuple primaryKey, @Nullable Vector vector, + Node create(@Nonnull Tuple primaryKey, @Nullable Vector vector, @Nonnull List neighbors); @Nonnull diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceAndNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceAndNode.java index 0b9ad99faa..bbf74e864a 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceAndNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceAndNode.java @@ -25,13 +25,13 @@ import javax.annotation.Nonnull; import java.util.List; -public class NodeReferenceAndNode { +public class NodeReferenceAndNode { @Nonnull private final NodeReferenceWithDistance nodeReferenceWithDistance; @Nonnull - private final Node node; + private final Node node; - public NodeReferenceAndNode(@Nonnull final NodeReferenceWithDistance nodeReferenceWithDistance, @Nonnull final Node node) { + public NodeReferenceAndNode(@Nonnull final NodeReferenceWithDistance nodeReferenceWithDistance, @Nonnull final Node node) { this.nodeReferenceWithDistance = nodeReferenceWithDistance; this.node = node; } @@ -42,14 +42,14 @@ public NodeReferenceWithDistance getNodeReferenceWithDistance() { } @Nonnull - public Node getNode() { + public Node getNode() { return node; } @Nonnull - public static List getReferences(@Nonnull List> referencesAndNodes) { + public static List getReferences(@Nonnull List> referencesAndNodes) { final ImmutableList.Builder referencesBuilder = ImmutableList.builder(); - for (final NodeReferenceAndNode referenceWithNode : referencesAndNodes) { + for (final NodeReferenceAndNode referenceWithNode : referencesAndNodes) { referencesBuilder.add(referenceWithNode.getNodeReferenceWithDistance()); } return referencesBuilder.build(); From 3445c3d3a6a97433c6a3c046dc23e47e1c6dbb27 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Wed, 30 Jul 2025 20:29:53 +0200 Subject: [PATCH 13/34] read and insert path both code complete --- .../async/hnsw/AbstractChangeSet.java | 94 ------- .../async/hnsw/AbstractStorageAdapter.java | 45 ++-- .../async/hnsw/BaseNeighborsChangeSet.java | 16 +- .../async/hnsw/ByNodeStorageAdapter.java | 108 -------- .../foundationdb/async/hnsw/CompactNode.java | 20 -- .../async/hnsw/CompactStorageAdapter.java | 132 +++++++++ .../async/hnsw/DeleteNeighborsChangeSet.java | 21 +- .../apple/foundationdb/async/hnsw/HNSW.java | 255 ++++++++---------- .../foundationdb/async/hnsw/InliningNode.java | 23 +- .../async/hnsw/InliningStorageAdapter.java | 125 +++++++++ .../async/hnsw/InsertNeighborsChangeSet.java | 30 ++- .../async/hnsw/NeighborsChangeSet.java | 5 +- .../apple/foundationdb/async/hnsw/Node.java | 61 ----- .../foundationdb/async/hnsw/NodeHelpers.java | 30 --- .../async/hnsw/NodeReference.java | 14 +- .../async/hnsw/NodeReferenceWithVector.java | 6 + .../async/hnsw/OnReadListener.java | 6 +- .../async/hnsw/OnWriteListener.java | 25 +- .../async/hnsw/StorageAdapter.java | 110 +++++--- .../async/rtree/StorageAdapter.java | 1 - 20 files changed, 547 insertions(+), 580 deletions(-) delete mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractChangeSet.java delete mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractChangeSet.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractChangeSet.java deleted file mode 100644 index 1b11d58856..0000000000 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractChangeSet.java +++ /dev/null @@ -1,94 +0,0 @@ -/* - * AbstractChangeSet.java - * - * This source file is part of the FoundationDB open source project - * - * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors - * - * Licensed 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 com.apple.foundationdb.async.hnsw; - -import com.apple.foundationdb.Transaction; -import com.apple.foundationdb.async.hnsw.Node.ChangeSet; - -import javax.annotation.Nonnull; -import javax.annotation.Nullable; - -/** - * Abstract base implementations for all {@link ChangeSet}s. - * @param slot type class - * @param node type class (self type) - */ -public abstract class AbstractChangeSet> implements ChangeSet { - @Nullable - private final ChangeSet previousChangeSet; - - @Nonnull - private final N node; - - private final int level; - - AbstractChangeSet(@Nullable final ChangeSet previousChangeSet, @Nonnull final N node, final int level) { - this.previousChangeSet = previousChangeSet; - this.node = node; - this.level = level; - } - - @Override - public void apply(@Nonnull final Transaction transaction) { - if (previousChangeSet != null) { - previousChangeSet.apply(transaction); - } - } - - /** - * Previous change set in the chain of change sets. Can be {@code null} if there is no previous change set. - * @return the previous change set in the chain of change sets - */ - @Nullable - public ChangeSet getPreviousChangeSet() { - return previousChangeSet; - } - - /** - * The node this change set applies to. - * @return the node this change set applies to - */ - @Nonnull - public N getNode() { - return node; - } - - /** - * The level we should use when maintaining the node slot index. If {@code level < 0}, do not maintain the node slot - * index. - * @return the level used when maintaing the node slot index - */ - public int getLevel() { - return level; - } - - /** - * Returns whether this change set needs to also update the node slot index. There are scenarios where we - * do not need to update such an index in general. For instance, the user may not want to use such an index. - * In addition to that, there are change set implementations that should not update the index even if such and index - * is maintained in general. For instance, the moved-in slots were already persisted in the database before the - * move-in operation. We should not update the node slot index in such a case. - * @return {@code true} if we need to update the node slot index, {@code false} otherwise - */ - public boolean isUpdateNodeSlotIndex() { - return level >= 0; - } -} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java index 383c22f27c..a306b3c600 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java @@ -31,31 +31,29 @@ /** * Implementations and attributes common to all concrete implementations of {@link StorageAdapter}. */ -abstract class AbstractStorageAdapter implements StorageAdapter { - public static final byte SUBSPACE_PREFIX_ENTRY_NODE = 0x01; - public static final byte SUBSPACE_PREFIX_DATA = 0x02; - +abstract class AbstractStorageAdapter implements StorageAdapter { @Nonnull private final HNSW.Config config; @Nonnull + private final NodeFactory nodeFactory; + @Nonnull private final Subspace subspace; @Nonnull private final OnWriteListener onWriteListener; @Nonnull private final OnReadListener onReadListener; - private final Subspace entryNodeSubspace; private final Subspace dataSubspace; - protected AbstractStorageAdapter(@Nonnull final HNSW.Config config, @Nonnull final Subspace subspace, + protected AbstractStorageAdapter(@Nonnull final HNSW.Config config, @Nonnull final NodeFactory nodeFactory, + @Nonnull final Subspace subspace, @Nonnull final OnWriteListener onWriteListener, @Nonnull final OnReadListener onReadListener) { this.config = config; + this.nodeFactory = nodeFactory; this.subspace = subspace; this.onWriteListener = onWriteListener; this.onReadListener = onReadListener; - - this.entryNodeSubspace = subspace.subspace(Tuple.from(SUBSPACE_PREFIX_ENTRY_NODE)); this.dataSubspace = subspace.subspace(Tuple.from(SUBSPACE_PREFIX_DATA)); } @@ -65,22 +63,22 @@ public HNSW.Config getConfig() { return config; } - @Override @Nonnull - public Subspace getSubspace() { - return subspace; + @Override + public NodeFactory getNodeFactory() { + return nodeFactory; } - @Nullable + @Nonnull @Override - public Subspace getSecondarySubspace() { - return null; + public NodeKind getNodeKind() { + return getNodeFactory().getNodeKind(); } @Override @Nonnull - public Subspace getEntryNodeSubspace() { - return entryNodeSubspace; + public Subspace getSubspace() { + return subspace; } @Override @@ -103,28 +101,25 @@ public OnReadListener getOnReadListener() { @Nonnull @Override - public CompletableFuture> fetchNode(@Nonnull final NodeFactory nodeFactory, - @Nonnull final ReadTransaction readTransaction, - int layer, @Nonnull Tuple primaryKey) { - return fetchNodeInternal(nodeFactory, readTransaction, layer, primaryKey).thenApply(this::checkNode); + public CompletableFuture> fetchNode(@Nonnull final ReadTransaction readTransaction, + int layer, @Nonnull Tuple primaryKey) { + return fetchNodeInternal(readTransaction, layer, primaryKey).thenApply(this::checkNode); } @Nonnull - protected abstract CompletableFuture> fetchNodeInternal(@Nonnull NodeFactory nodeFactory, - @Nonnull ReadTransaction readTransaction, - int layer, @Nonnull Tuple primaryKey); + protected abstract CompletableFuture> fetchNodeInternal(@Nonnull ReadTransaction readTransaction, + int layer, @Nonnull Tuple primaryKey); /** * Method to perform basic invariant check(s) on a newly-fetched node. * * @param node the node to check - * @param the type param for the node in order for this method to not be lossy on the type of the node that * was passed in * * @return the node that was passed in */ @Nullable - private Node checkNode(@Nullable final Node node) { + private Node checkNode(@Nullable final Node node) { return node; } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/BaseNeighborsChangeSet.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/BaseNeighborsChangeSet.java index 6b5ce44b84..397e818ac1 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/BaseNeighborsChangeSet.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/BaseNeighborsChangeSet.java @@ -21,20 +21,23 @@ package com.apple.foundationdb.async.hnsw; import com.apple.foundationdb.Transaction; +import com.apple.foundationdb.tuple.Tuple; +import com.google.common.collect.ImmutableList; import javax.annotation.Nonnull; import javax.annotation.Nullable; import java.util.List; +import java.util.function.Predicate; /** * TODO. */ class BaseNeighborsChangeSet implements NeighborsChangeSet { @Nonnull - private final NodeReferenceAndNode baseNode; + private final List neighbors; - public BaseNeighborsChangeSet(@Nonnull final NodeReferenceAndNode baseNode) { - this.baseNode = baseNode; + public BaseNeighborsChangeSet(@Nonnull final List neighbors) { + this.neighbors = ImmutableList.copyOf(neighbors); } @Nullable @@ -44,10 +47,13 @@ public BaseNeighborsChangeSet getParent() { @Nonnull public List merge() { - return baseNode.getNode().getNeighbors(); + return neighbors; } @Override - public void writeDelta(@Nonnull final Transaction transaction) { + public void writeDelta(@Nonnull final InliningStorageAdapter storageAdapter, @Nonnull final Transaction transaction, + final int layer, @Nonnull final Node node, + @Nonnull final Predicate primaryKeyPredicate) { + // nothing to be written } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java deleted file mode 100644 index df1457458a..0000000000 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java +++ /dev/null @@ -1,108 +0,0 @@ -/* - * ByNodeStorageAdapter.java - * - * This source file is part of the FoundationDB open source project - * - * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors - * - * Licensed 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 com.apple.foundationdb.async.hnsw; - -import com.apple.foundationdb.ReadTransaction; -import com.apple.foundationdb.Transaction; -import com.apple.foundationdb.subspace.Subspace; -import com.apple.foundationdb.tuple.Tuple; - -import javax.annotation.Nonnull; -import java.util.concurrent.CompletableFuture; - -/** - * TODO. - */ -class ByNodeStorageAdapter extends AbstractStorageAdapter implements StorageAdapter { - public ByNodeStorageAdapter(@Nonnull final HNSW.Config config, @Nonnull final Subspace subspace, - @Nonnull final OnWriteListener onWriteListener, - @Nonnull final OnReadListener onReadListener) { - super(config, subspace, onWriteListener, onReadListener); - } - - @Nonnull - @Override - public NodeFactory getNodeFactory(final int layer) { - return layer > 0 ? InliningNode.factory() : CompactNode.factory(); - } - - @Override - public CompletableFuture fetchEntryNodeReference(@Nonnull final ReadTransaction readTransaction) { - final byte[] key = getEntryNodeSubspace().pack(); - - return readTransaction.get(key) - .thenApply(valueBytes -> { - if (valueBytes == null) { - return null; // not a single node in the index - } - final OnReadListener onReadListener = getOnReadListener(); - onReadListener.onKeyValueRead(key, valueBytes); - - final Tuple entryTuple = Tuple.fromBytes(valueBytes); - final int lMax = (int)entryTuple.getLong(0); - final Tuple primaryKey = entryTuple.getNestedTuple(1); - final Tuple vectorTuple = entryTuple.getNestedTuple(2); - return new EntryNodeReference(primaryKey, StorageAdapter.vectorFromTuple(vectorTuple), lMax); - }); - } - - - @Override - public void writeEntryNodeReference(@Nonnull final Transaction transaction, - @Nonnull final EntryNodeReference entryNodeReference) { - transaction.set(getEntryNodeSubspace().pack(), - Tuple.from(entryNodeReference.getLayer(), - entryNodeReference.getPrimaryKey(), - StorageAdapter.tupleFromVector(entryNodeReference.getVector())).pack()); - } - - - @Nonnull - @Override - protected CompletableFuture> fetchNodeInternal(@Nonnull final NodeFactory nodeFactory, - @Nonnull final ReadTransaction readTransaction, - final int layer, - @Nonnull final Tuple primaryKey) { - final byte[] key = getDataSubspace().pack(Tuple.from(layer, primaryKey)); - - return readTransaction.get(key) - .thenApply(valueBytes -> { - if (valueBytes == null) { - throw new IllegalStateException("cannot fetch node"); - } - - final Tuple nodeTuple = Tuple.fromBytes(valueBytes); - final Node node = Node.nodeFromTuples(nodeFactory, primaryKey, nodeTuple); - final OnReadListener onReadListener = getOnReadListener(); - onReadListener.onNodeRead(node); - onReadListener.onKeyValueRead(key, valueBytes); - return node; - }); - } - - @Override - public void writeNode(@Nonnull Transaction transaction, @Nonnull final Node node, - final int layer) { - final byte[] key = getDataSubspace().pack(Tuple.from(layer, node.getPrimaryKey())); - transaction.set(key, node.toTuple().pack()); - getOnWriteListener().onNodeWritten(node); - } -} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java index a213d70411..8dd86f622b 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java @@ -22,7 +22,6 @@ import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; -import com.google.common.collect.Lists; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -88,25 +87,6 @@ public InliningNode asInliningNode() { throw new IllegalStateException("this is not an inlining node"); } - @Override - public NodeFactory sameCreator() { - return CompactNode.factory(); - } - - @Nonnull - @Override - public Tuple toTuple() { - final List nodeItems = Lists.newArrayListWithExpectedSize(4); - nodeItems.add(NodeKind.COMPACT.getSerialized()); - nodeItems.add(StorageAdapter.tupleFromVector(getVector())); - final List neighborItems = Lists.newArrayListWithExpectedSize(getNeighbors().size()); - for (final NodeReference nodeReference : getNeighbors()) { - neighborItems.add(nodeReference.getPrimaryKey()); - } - nodeItems.add(Tuple.fromList(neighborItems)); - return Tuple.fromList(nodeItems); - } - @Nonnull public static NodeFactory factory() { return FACTORY; diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java new file mode 100644 index 0000000000..9da72759f4 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java @@ -0,0 +1,132 @@ +/* + * CompactStorageAdapter.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.ReadTransaction; +import com.apple.foundationdb.Transaction; +import com.apple.foundationdb.subspace.Subspace; +import com.apple.foundationdb.tuple.Tuple; +import com.christianheina.langx.half4j.Half; +import com.google.common.base.Verify; +import com.google.common.collect.Lists; + +import javax.annotation.Nonnull; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +/** + * TODO. + */ +class CompactStorageAdapter extends AbstractStorageAdapter implements StorageAdapter { + public CompactStorageAdapter(@Nonnull final HNSW.Config config, @Nonnull final NodeFactory nodeFactory, + @Nonnull final Subspace subspace, + @Nonnull final OnWriteListener onWriteListener, + @Nonnull final OnReadListener onReadListener) { + super(config, nodeFactory, subspace, onWriteListener, onReadListener); + } + + @Nonnull + @Override + public StorageAdapter asCompactStorageAdapter() { + return this; + } + + @Nonnull + @Override + public StorageAdapter asInliningStorageAdapter() { + throw new IllegalStateException("cannot call this method on a compact storage adapter"); + } + + @Nonnull + @Override + protected CompletableFuture> fetchNodeInternal(@Nonnull final ReadTransaction readTransaction, + final int layer, + @Nonnull final Tuple primaryKey) { + final byte[] key = getDataSubspace().pack(Tuple.from(layer, primaryKey)); + + return readTransaction.get(key) + .thenApply(valueBytes -> { + if (valueBytes == null) { + throw new IllegalStateException("cannot fetch node"); + } + + final Tuple nodeTuple = Tuple.fromBytes(valueBytes); + final Node node = nodeFromTuples(primaryKey, nodeTuple); + final OnReadListener onReadListener = getOnReadListener(); + onReadListener.onNodeRead(node); + onReadListener.onKeyValueRead(key, valueBytes); + return node; + }); + } + + @Nonnull + private Node nodeFromTuples(@Nonnull final Tuple primaryKey, + @Nonnull final Tuple valueTuple) { + final NodeKind nodeKind = NodeKind.fromSerializedNodeKind((byte)valueTuple.getLong(0)); + Verify.verify(nodeKind == NodeKind.COMPACT); + + final Tuple vectorTuple; + final Tuple neighborsTuple; + + vectorTuple = valueTuple.getNestedTuple(1); + neighborsTuple = valueTuple.getNestedTuple(2); + return compactNodeFromTuples(primaryKey, vectorTuple, neighborsTuple); + } + + @Nonnull + private Node compactNodeFromTuples(@Nonnull final Tuple primaryKey, + @Nonnull final Tuple vectorTuple, + @Nonnull final Tuple neighborsTuple) { + final Vector vector = StorageAdapter.vectorFromTuple(vectorTuple); + final List nodeReferences = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); + + for (final Object neighborObject : neighborsTuple) { + final Tuple neighborTuple = (Tuple)neighborObject; + nodeReferences.add(new NodeReference(neighborTuple)); + } + + return getNodeFactory().create(primaryKey, vector, nodeReferences); + } + + + @Override + public void writeNode(@Nonnull final Transaction transaction, @Nonnull final Node node, + final int layer, @Nonnull final NeighborsChangeSet neighborsChangeSet) { + final byte[] key = getDataSubspace().pack(Tuple.from(layer, node.getPrimaryKey())); + + final List nodeItems = Lists.newArrayListWithExpectedSize(4); + nodeItems.add(NodeKind.COMPACT.getSerialized()); + final CompactNode compactNode = node.asCompactNode(); + nodeItems.add(StorageAdapter.tupleFromVector(compactNode.getVector())); + + final Iterable neighbors = neighborsChangeSet.merge(); + + final List neighborItems = Lists.newArrayList(); + for (final NodeReference neighborReference : neighbors) { + neighborItems.add(neighborReference.getPrimaryKey()); + } + nodeItems.add(Tuple.fromList(neighborItems)); + final Tuple nodeTuple = Tuple.fromList(nodeItems); + + transaction.set(key, nodeTuple.pack()); + getOnWriteListener().onNodeWritten(layer, node); + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java index c1ff4e73c3..feee9fa979 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java @@ -28,6 +28,7 @@ import javax.annotation.Nonnull; import java.util.Collection; import java.util.Set; +import java.util.function.Predicate; /** * TODO. @@ -37,12 +38,12 @@ class DeleteNeighborsChangeSet implements NeighborsChan private final NeighborsChangeSet parent; @Nonnull - private final Set deletedNeighborsTuples; + private final Set deletedNeighborsPrimaryKeys; public DeleteNeighborsChangeSet(@Nonnull final NeighborsChangeSet parent, - @Nonnull final Collection deletedNeighborsTuples) { + @Nonnull final Collection deletedNeighborsPrimaryKeys) { this.parent = parent; - this.deletedNeighborsTuples = ImmutableSet.copyOf(deletedNeighborsTuples); + this.deletedNeighborsPrimaryKeys = ImmutableSet.copyOf(deletedNeighborsPrimaryKeys); } @Nonnull @@ -53,11 +54,19 @@ public NeighborsChangeSet getParent() { @Nonnull public Iterable merge() { return Iterables.filter(getParent().merge(), - current -> !deletedNeighborsTuples.contains(current.getPrimaryKey())); + current -> !deletedNeighborsPrimaryKeys.contains(current.getPrimaryKey())); } @Override - public void writeDelta(@Nonnull final Transaction transaction) { - throw new UnsupportedOperationException("not implemented yet"); + public void writeDelta(@Nonnull final InliningStorageAdapter storageAdapter, @Nonnull final Transaction transaction, + final int layer, @Nonnull final Node node, @Nonnull final Predicate tuplePredicate) { + getParent().writeDelta(storageAdapter, transaction, layer, node, + tuplePredicate.and(tuple -> !deletedNeighborsPrimaryKeys.contains(tuple))); + + for (final Tuple deletedNeighborPrimaryKey : deletedNeighborsPrimaryKeys) { + if (tuplePredicate.test(deletedNeighborPrimaryKey)) { + storageAdapter.deleteNeighbor(transaction, layer, node.asInliningNode(), deletedNeighborPrimaryKey); + } + } } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index f4e007a638..8b948aa459 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -78,7 +78,7 @@ public class HNSW { public static final Config DEFAULT_CONFIG = new Config(); @Nonnull - private final StorageAdapter storageAdapter; + private final Subspace subspace; @Nonnull private final Executor executor; @Nonnull @@ -88,40 +88,6 @@ public class HNSW { @Nonnull private final OnReadListener onReadListener; - /** - * Different kinds of storage layouts. - */ - public enum Storage { - /** - * Every node with all its slots is serialized as one key/value pair. - */ - BY_NODE(ByNodeStorageAdapter::new); - - @Nonnull - private final StorageAdapterCreator storageAdapterCreator; - - Storage(@Nonnull final StorageAdapterCreator storageAdapterCreator) { - this.storageAdapterCreator = storageAdapterCreator; - } - - @Nonnull - private StorageAdapter newStorageAdapter(@Nonnull final Config config, @Nonnull final Subspace subspace, - @Nonnull final Subspace nodeSlotIndexSubspace, - @Nonnull final OnWriteListener onWriteListener, - @Nonnull final OnReadListener onReadListener) { - return storageAdapterCreator.create(config, subspace, nodeSlotIndexSubspace, onWriteListener, onReadListener); - } - } - - /** - * Functional interface to create a {@link StorageAdapter}. - */ - private interface StorageAdapterCreator { - StorageAdapter create(@Nonnull Config config, @Nonnull Subspace subspace, @Nonnull Subspace nodeSlotIndexSubspace, - @Nonnull OnWriteListener onWriteListener, - @Nonnull OnReadListener onReadListener); - } - /** * Configuration settings for a {@link HNSW}. */ @@ -360,46 +326,30 @@ public static ConfigBuilder newConfigBuilder() { } /** - * Initialize a new R-tree with the default configuration. - * @param subspace the subspace where the r-tree is stored - * @param secondarySubspace the subspace where the node index (if used is stored) - * @param executor an executor to use when running asynchronous tasks + * TODO. */ - public HNSW(@Nonnull final Subspace subspace, @Nonnull final Subspace secondarySubspace, - @Nonnull final Executor executor) { - this(subspace, secondarySubspace, executor, DEFAULT_CONFIG, - OnWriteListener.NOOP, OnReadListener.NOOP); + public HNSW(@Nonnull final Subspace subspace, @Nonnull final Executor executor) { + this(subspace, executor, DEFAULT_CONFIG, OnWriteListener.NOOP, OnReadListener.NOOP); } /** - * Initialize a new R-tree. - * @param subspace the subspace where the r-tree is stored - * @param nodeSlotIndexSubspace the subspace where the node index (if used is stored) - * @param executor an executor to use when running asynchronous tasks - * @param config configuration to use - * @param onWriteListener an on-write listener to be called after writes take place - * @param onReadListener an on-read listener to be called after reads take place + * TODO. */ - public HNSW(@Nonnull final Subspace subspace, @Nonnull final Subspace nodeSlotIndexSubspace, + public HNSW(@Nonnull final Subspace subspace, @Nonnull final Executor executor, @Nonnull final Config config, @Nonnull final OnWriteListener onWriteListener, @Nonnull final OnReadListener onReadListener) { - this.storageAdapter = config.getStorage() - .newStorageAdapter(config, subspace, nodeSlotIndexSubspace, hilbertValueFunction, onWriteListener, - onReadListener); + this.subspace = subspace; this.executor = executor; this.config = config; this.onWriteListener = onWriteListener; this.onReadListener = onReadListener; } - /** - * Get the {@link StorageAdapter} used to manage this r-tree. - * @return r-tree subspace - */ + @Nonnull - StorageAdapter getStorageAdapter() { - return storageAdapter; + public Subspace getSubspace() { + return subspace; } /** @@ -447,9 +397,9 @@ public OnReadListener getOnReadListener() { */ @SuppressWarnings("checkstyle:MethodName") // method name introduced by paper @Nonnull - private CompletableFuture>> kNearestNeighborsSearch(@Nonnull final ReadTransaction readTransaction, - @Nonnull final Vector queryVector) { - return storageAdapter.fetchEntryNodeReference(readTransaction) + public CompletableFuture>> kNearestNeighborsSearch(@Nonnull final ReadTransaction readTransaction, + @Nonnull final Vector queryVector) { + return StorageAdapter.fetchEntryNodeReference(readTransaction, getSubspace(), getOnReadListener()) .thenCompose(entryPointAndLayer -> { if (entryPointAndLayer == null) { return CompletableFuture.completedFuture(null); // not a single node in the index @@ -462,20 +412,23 @@ private CompletableFuture nodeReferenceAtomic = new AtomicReference<>(entryState); - return MoreAsyncUtil.forLoop(entryPointAndLayer.getLayer(), + return MoreAsyncUtil.forLoop(entryLayer, layer -> layer > 0, layer -> layer - 1, layer -> { + final var storageAdapter = getStorageAdapterForLayer(layer); final var greedyIn = nodeReferenceAtomic.get(); - return greedySearchLayer(storageAdapter.getNodeFactory(layer), readTransaction, greedyIn, layer, + return greedySearchLayer(storageAdapter, readTransaction, greedyIn, layer, queryVector) .thenApply(greedyState -> { nodeReferenceAtomic.set(greedyState); @@ -488,22 +441,24 @@ private CompletableFuture CompletableFuture greedySearchLayer(@Nonnull NodeFactory nodeFactory, + private CompletableFuture greedySearchLayer(@Nonnull StorageAdapter storageAdapter, @Nonnull final ReadTransaction readTransaction, @Nonnull final NodeReferenceWithDistance entryNeighbor, final int layer, @Nonnull final Vector queryVector) { - if (nodeFactory.getNodeKind() == NodeKind.INLINING) { - return greedySearchInliningLayer(readTransaction, entryNeighbor, layer, queryVector); + if (storageAdapter.getNodeKind() == NodeKind.INLINING) { + return greedySearchInliningLayer(storageAdapter.asInliningStorageAdapter(), readTransaction, entryNeighbor, layer, queryVector); } else { - return searchLayer(nodeFactory, readTransaction, ImmutableList.of(entryNeighbor), layer, 1, Maps.newConcurrentMap(), queryVector) + return searchLayer(storageAdapter, readTransaction, ImmutableList.of(entryNeighbor), layer, 1, Maps.newConcurrentMap(), queryVector) .thenApply(searchResult -> Iterables.getOnlyElement(searchResult).getNodeReferenceWithDistance()); } } @@ -512,7 +467,8 @@ private CompletableFuture g * TODO. */ @Nonnull - private CompletableFuture greedySearchInliningLayer(@Nonnull final ReadTransaction readTransaction, + private CompletableFuture greedySearchInliningLayer(@Nonnull final StorageAdapter storageAdapter, + @Nonnull final ReadTransaction readTransaction, @Nonnull final NodeReferenceWithDistance entryNeighbor, final int layer, @Nonnull final Vector queryVector) { @@ -522,8 +478,7 @@ private CompletableFuture greedySearchInliningLayer(@ new AtomicReference<>(entryNeighbor); return AsyncUtil.whileTrue(() -> onReadListener.onAsyncRead( - storageAdapter.fetchNode(InliningNode.factory(), readTransaction, - layer, currentNodeReferenceAtomic.get().getPrimaryKey())) + storageAdapter.fetchNode(readTransaction, layer, currentNodeReferenceAtomic.get().getPrimaryKey())) .thenApply(node -> { if (node == null) { throw new IllegalStateException("unable to fetch node"); @@ -559,7 +514,7 @@ private CompletableFuture greedySearchInliningLayer(@ * TODO. */ @Nonnull - private CompletableFuture>> searchLayer(@Nonnull NodeFactory nodeFactory, + private CompletableFuture>> searchLayer(@Nonnull StorageAdapter storageAdapter, @Nonnull final ReadTransaction readTransaction, @Nonnull final Collection entryNeighbors, final int layer, @@ -589,11 +544,11 @@ private CompletableFuture return AsyncUtil.READY_FALSE; } - return fetchNodeIfNotCached(nodeFactory, readTransaction, layer, candidate, nodeCache) + return fetchNodeIfNotCached(storageAdapter, readTransaction, layer, candidate, nodeCache) .thenApply(candidateNode -> Iterables.filter(candidateNode.getNeighbors(), neighbor -> !visited.contains(neighbor.getPrimaryKey()))) - .thenCompose(neighborReferences -> fetchNeighborhood(nodeFactory, readTransaction, + .thenCompose(neighborReferences -> fetchNeighborhood(storageAdapter, readTransaction, layer, neighborReferences, nodeCache)) .thenApply(neighborReferences -> { for (final NodeReferenceWithVector current : neighborReferences) { @@ -617,19 +572,19 @@ private CompletableFuture return true; }); }).thenCompose(ignored -> - fetchSomeNodesIfNotCached(nodeFactory, readTransaction, layer, nearestNeighbors, nodeCache)); + fetchSomeNodesIfNotCached(storageAdapter, readTransaction, layer, nearestNeighbors, nodeCache)); } /** * TODO. */ @Nonnull - private CompletableFuture> fetchNodeIfNotCached(@Nonnull final NodeFactory nodeFactory, + private CompletableFuture> fetchNodeIfNotCached(@Nonnull final StorageAdapter storageAdapter, @Nonnull final ReadTransaction readTransaction, final int layer, @Nonnull final NodeReference nodeReference, @Nonnull final Map> nodeCache) { - return fetchNodeIfNecessaryAndApply(nodeFactory, readTransaction, layer, nodeReference, + return fetchNodeIfNecessaryAndApply(storageAdapter, readTransaction, layer, nodeReference, nR -> nodeCache.get(nR.getPrimaryKey()), (nR, node) -> { nodeCache.put(nR.getPrimaryKey(), node); @@ -641,7 +596,7 @@ private CompletableFuture> fetchNodeIfNotCache * TODO. */ @Nonnull - private CompletableFuture fetchNodeIfNecessaryAndApply(@Nonnull final NodeFactory nodeFactory, + private CompletableFuture fetchNodeIfNecessaryAndApply(@Nonnull final StorageAdapter storageAdapter, @Nonnull final ReadTransaction readTransaction, final int layer, @Nonnull final R nodeReference, @@ -653,7 +608,7 @@ private CompletableFuture< } return onReadListener.onAsyncRead( - storageAdapter.fetchNode(nodeFactory, readTransaction, layer, nodeReference.getPrimaryKey())) + storageAdapter.fetchNode(readTransaction, layer, nodeReference.getPrimaryKey())) .thenApply(node -> biMapFunction.apply(nodeReference, node)); } @@ -661,12 +616,12 @@ private CompletableFuture< * TODO. */ @Nonnull - private CompletableFuture> fetchNeighborhood(@Nonnull final NodeFactory nodeFactory, + private CompletableFuture> fetchNeighborhood(@Nonnull final StorageAdapter storageAdapter, @Nonnull final ReadTransaction readTransaction, final int layer, @Nonnull final Iterable neighborReferences, @Nonnull final Map> nodeCache) { - return fetchSomeNodesAndApply(nodeFactory, readTransaction, layer, neighborReferences, + return fetchSomeNodesAndApply(storageAdapter, readTransaction, layer, neighborReferences, neighborReference -> { if (neighborReference instanceof NodeReferenceWithVector) { return (NodeReferenceWithVector)neighborReference; @@ -687,12 +642,12 @@ private CompletableFuture CompletableFuture>> fetchSomeNodesIfNotCached(@Nonnull final NodeFactory creator, + private CompletableFuture>> fetchSomeNodesIfNotCached(@Nonnull final StorageAdapter storageAdapter, @Nonnull final ReadTransaction readTransaction, final int layer, @Nonnull final Iterable nodeReferences, @Nonnull final Map> nodeCache) { - return fetchSomeNodesAndApply(creator, readTransaction, layer, nodeReferences, + return fetchSomeNodesAndApply(storageAdapter, readTransaction, layer, nodeReferences, nodeReference -> { final Node node = nodeCache.get(nodeReference.getPrimaryKey()); if (node == null) { @@ -710,15 +665,14 @@ private CompletableFuture * TODO. */ @Nonnull - @SuppressWarnings("unchecked") - private CompletableFuture> fetchSomeNodesAndApply(@Nonnull final NodeFactory creator, + private CompletableFuture> fetchSomeNodesAndApply(@Nonnull final StorageAdapter storageAdapter, @Nonnull final ReadTransaction readTransaction, final int layer, @Nonnull final Iterable nodeReferences, @Nonnull final Function fetchBypassFunction, @Nonnull final BiFunction, U> biMapFunction) { return MoreAsyncUtil.forEach(nodeReferences, - currentNeighborReference -> fetchNodeIfNecessaryAndApply(creator, readTransaction, layer, + currentNeighborReference -> fetchNodeIfNecessaryAndApply(storageAdapter, readTransaction, layer, currentNeighborReference, fetchBypassFunction, biMapFunction), MAX_CONCURRENT_NODE_READS, getExecutor()); } @@ -730,23 +684,19 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N final int l = insertionLayer(getConfig().getRandom()); - return storageAdapter.fetchEntryNodeReference(transaction) + return StorageAdapter.fetchEntryNodeReference(transaction, getSubspace(), getOnReadListener()) .thenApply(entryNodeReference -> { if (entryNodeReference == null) { // this is the first node - writeLonelyNodes(InliningNode.factory(), transaction, newPrimaryKey, newVector, l, 0); - storageAdapter.writeNode(transaction, - CompactNode.factory() - .create(newPrimaryKey, newVector, ImmutableList.of()), - 0); - storageAdapter.writeEntryNodeReference(transaction, - new EntryNodeReference(newPrimaryKey, newVector, l)); + writeLonelyNodes(transaction, newPrimaryKey, newVector, l, -1); + StorageAdapter.writeEntryNodeReference(transaction, getSubspace(), + new EntryNodeReference(newPrimaryKey, newVector, l), getOnWriteListener()); } else { final int entryNodeLayer = entryNodeReference.getLayer(); if (l > entryNodeLayer) { - writeLonelyNodes(InliningNode.factory(), transaction, newPrimaryKey, newVector, l, entryNodeLayer); - storageAdapter.writeEntryNodeReference(transaction, - new EntryNodeReference(newPrimaryKey, newVector, l)); + writeLonelyNodes(transaction, newPrimaryKey, newVector, l, entryNodeLayer); + StorageAdapter.writeEntryNodeReference(transaction, getSubspace(), + new EntryNodeReference(newPrimaryKey, newVector, l), getOnWriteListener()); } } return entryNodeReference; @@ -764,12 +714,15 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N MoreAsyncUtil.forLoop(lMax, layer -> layer > l, layer -> layer - 1, - layer -> greedySearchLayer(InliningNode.factory(), transaction, - nodeReferenceAtomic.get(), layer, newVector) - .thenApply(nodeReference -> { - nodeReferenceAtomic.set(nodeReference); - return null; - }), executor); + layer -> { + final StorageAdapter storageAdapter = getStorageAdapterForLayer(layer); + return greedySearchLayer(storageAdapter, transaction, + nodeReferenceAtomic.get(), layer, newVector) + .thenApply(nodeReference -> { + nodeReferenceAtomic.set(nodeReference); + return null; + }); + }, executor); final AtomicReference> nearestNeighborsAtomic = new AtomicReference<>(ImmutableList.of(nodeReferenceAtomic.get())); @@ -777,17 +730,20 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N return MoreAsyncUtil.forLoop(Math.min(lMax, l), layer -> layer >= 0, layer -> layer - 1, - layer -> insertIntoLayer(storageAdapter.getNodeFactory(layer), transaction, - nearestNeighborsAtomic.get(), layer, newPrimaryKey, newVector) - .thenCompose(nearestNeighbors -> { - nearestNeighborsAtomic.set(nearestNeighbors); - return AsyncUtil.DONE; - }), executor); + layer -> { + final StorageAdapter storageAdapter = getStorageAdapterForLayer(layer); + return insertIntoLayer(storageAdapter, transaction, + nearestNeighborsAtomic.get(), layer, newPrimaryKey, newVector) + .thenCompose(nearestNeighbors -> { + nearestNeighborsAtomic.set(nearestNeighbors); + return AsyncUtil.DONE; + }); + }, executor); }).thenCompose(ignored -> AsyncUtil.DONE); } @Nonnull - private CompletableFuture> insertIntoLayer(@Nonnull final NodeFactory nodeFactory, + private CompletableFuture> insertIntoLayer(@Nonnull final StorageAdapter storageAdapter, @Nonnull final Transaction transaction, @Nonnull final List nearestNeighbors, int layer, @@ -795,23 +751,30 @@ private CompletableFuture newVector) { final Map> nodeCache = Maps.newConcurrentMap(); - return searchLayer(nodeFactory, transaction, + return searchLayer(storageAdapter, transaction, nearestNeighbors, layer, config.getEfConstruction(), nodeCache, newVector) .thenCompose(searchResult -> { final List references = NodeReferenceAndNode.getReferences(searchResult); - return selectNeighbors(nodeFactory, transaction, searchResult, + return selectNeighbors(storageAdapter, transaction, searchResult, layer, getConfig().getM(), getConfig().isExtendCandidates(), nodeCache, newVector) .thenCompose(selectedNeighbors -> { + final NodeFactory nodeFactory = storageAdapter.getNodeFactory(); + final Node newNode = - nodeFactory.create(newPrimaryKey, newVector, NodeReferenceAndNode.getReferences(selectedNeighbors)); + nodeFactory.create(newPrimaryKey, newVector, + NodeReferenceAndNode.getReferences(selectedNeighbors)); + + final NeighborsChangeSet newNodeChangeSet = + new InsertNeighborsChangeSet<>(new BaseNeighborsChangeSet<>(ImmutableList.of()), + newNode.getNeighbors()); // create change sets for each selected neighbor and insert new node into them final Map> neighborChangeSetMap = Maps.newLinkedHashMap(); for (final NodeReferenceAndNode selectedNeighbor : selectedNeighbors) { final NeighborsChangeSet baseSet = - new BaseNeighborsChangeSet<>(selectedNeighbor); + new BaseNeighborsChangeSet<>(selectedNeighbor.getNode().getNeighbors()); final NeighborsChangeSet insertSet = new InsertNeighborsChangeSet<>(baseSet, ImmutableList.of(newNode.getSelfReference(newVector))); neighborChangeSetMap.put(selectedNeighbor.getNode().getPrimaryKey(), @@ -824,8 +787,8 @@ layer, getConfig().getM(), getConfig().isExtendCandidates(), nodeCache, newVecto final Node selectedNeighborNode = selectedNeighbor.getNode(); final NeighborsChangeSet changeSet = Objects.requireNonNull(neighborChangeSetMap.get(selectedNeighborNode.getPrimaryKey())); - return pruneNeighborsIfNecessary(nodeFactory, transaction, selectedNeighbor, layer, - currentMMax, changeSet, nodeCache) + return pruneNeighborsIfNecessary(storageAdapter, transaction, + selectedNeighbor, layer, currentMMax, changeSet, nodeCache) .thenApply(nodeReferencesAndNodes -> { if (nodeReferencesAndNodes == null) { return changeSet; @@ -842,7 +805,8 @@ layer, getConfig().getM(), getConfig().isExtendCandidates(), nodeCache, newVecto changeSet); } - // TODO write newVector and all change sets pertaining to neighbors + storageAdapter.writeNode(transaction, newNode, layer, newNodeChangeSet); + return ImmutableList.copyOf(references); }); }); @@ -892,7 +856,7 @@ private NeighborsChangeSet resolveChangeSetFromNewN } @Nonnull - private CompletableFuture>> pruneNeighborsIfNecessary(@Nonnull final NodeFactory nodeFactory, + private CompletableFuture>> pruneNeighborsIfNecessary(@Nonnull final StorageAdapter storageAdapter, @Nonnull final Transaction transaction, @Nonnull final NodeReferenceAndNode selectedNeighbor, int layer, @@ -904,7 +868,7 @@ private CompletableFuture if (selectedNeighborNode.getNeighbors().size() < mMax) { return CompletableFuture.completedFuture(null); } else { - return fetchNeighborhood(nodeFactory, transaction, layer, neighborChangeSet.merge(), nodeCache) + return fetchNeighborhood(storageAdapter, transaction, layer, neighborChangeSet.merge(), nodeCache) .thenCompose(nodeReferenceWithVectors -> { final ImmutableList.Builder nodeReferencesWithDistancesBuilder = ImmutableList.builder(); @@ -917,18 +881,18 @@ private CompletableFuture new NodeReferenceWithDistance(nodeReferenceWithVector.getPrimaryKey(), vector, distance)); } - return fetchSomeNodesIfNotCached(nodeFactory, transaction, layer, + return fetchSomeNodesIfNotCached(storageAdapter, transaction, layer, nodeReferencesWithDistancesBuilder.build(), nodeCache); }) .thenCompose(nodeReferencesAndNodes -> - selectNeighbors(nodeFactory, transaction, + selectNeighbors(storageAdapter, transaction, nodeReferencesAndNodes, layer, mMax, false, nodeCache, selectedNeighbor.getNodeReferenceWithDistance().getVector())); } } - private CompletableFuture>> selectNeighbors(@Nonnull final NodeFactory nodeFactory, + private CompletableFuture>> selectNeighbors(@Nonnull final StorageAdapter storageAdapter, @Nonnull final ReadTransaction readTransaction, @Nonnull final Iterable> nearestNeighbors, final int layer, @@ -936,7 +900,7 @@ private CompletableFuture final boolean isExtendCandidates, @Nonnull final Map> nodeCache, @Nonnull final Vector vector) { - return extendCandidatesIfNecessary(nodeFactory, readTransaction, nearestNeighbors, layer, isExtendCandidates, nodeCache, vector) + return extendCandidatesIfNecessary(storageAdapter, readTransaction, nearestNeighbors, layer, isExtendCandidates, nodeCache, vector) .thenApply(extendedCandidates -> { final List selected = Lists.newArrayListWithExpectedSize(m); final Queue candidates = @@ -976,10 +940,10 @@ private CompletableFuture return ImmutableList.copyOf(selected); }).thenCompose(selectedNeighbors -> - fetchSomeNodesIfNotCached(nodeFactory, readTransaction, layer, selectedNeighbors, nodeCache)); + fetchSomeNodesIfNotCached(storageAdapter, readTransaction, layer, selectedNeighbors, nodeCache)); } - private CompletableFuture> extendCandidatesIfNecessary(@Nonnull final NodeFactory nodeFactory, + private CompletableFuture> extendCandidatesIfNecessary(@Nonnull final StorageAdapter storageAdapter, @Nonnull final ReadTransaction readTransaction, @Nonnull final Iterable> candidates, int layer, @@ -1007,7 +971,7 @@ private CompletableFuture neighborsOfCandidates = neighborsOfCandidatesBuilder.build(); - return fetchNeighborhood(nodeFactory, readTransaction, layer, neighborsOfCandidates, nodeCache) + return fetchNeighborhood(storageAdapter, readTransaction, layer, neighborsOfCandidates, nodeCache) .thenApply(withVectors -> { final ImmutableList.Builder extendedCandidatesBuilder = ImmutableList.builder(); for (final NodeReferenceAndNode candidate : candidates) { @@ -1031,18 +995,35 @@ private CompletableFuture void writeLonelyNodes(@Nonnull final NodeFactory nodeFactory, - @Nonnull final Transaction transaction, - @Nonnull final Tuple primaryKey, - @Nonnull final Vector vector, - final int highestLayerInclusive, - final int lowestLayerExclusive) { + private void writeLonelyNodes(@Nonnull final Transaction transaction, + @Nonnull final Tuple primaryKey, + @Nonnull final Vector vector, + final int highestLayerInclusive, + final int lowestLayerExclusive) { for (int layer = highestLayerInclusive; layer > lowestLayerExclusive; layer --) { - storageAdapter.writeNode(transaction, - nodeFactory.create(primaryKey, vector, ImmutableList.of()), layer); + final StorageAdapter storageAdapter = getStorageAdapterForLayer(layer); + writeLonelyNodeOnLayer(storageAdapter, transaction, layer, primaryKey, vector); } } + private void writeLonelyNodeOnLayer(@Nonnull final StorageAdapter storageAdapter, + @Nonnull final Transaction transaction, + final int layer, + @Nonnull final Tuple primaryKey, + @Nonnull final Vector vector) { + storageAdapter.writeNode(transaction, + storageAdapter.getNodeFactory() + .create(primaryKey, vector, ImmutableList.of()), layer, + new BaseNeighborsChangeSet<>(ImmutableList.of())); + } + + @Nonnull + private StorageAdapter getStorageAdapterForLayer(final int layer) { + return layer > 0 + ? new InliningStorageAdapter(getConfig(), InliningNode.factory(), getSubspace(), getOnWriteListener(), getOnReadListener()) + : new CompactStorageAdapter(getConfig(), CompactNode.factory(), getSubspace(), getOnWriteListener(), getOnReadListener()); + } + private int insertionLayer(@Nonnull final Random random) { double lambda = 1.0 / Math.log(getConfig().getM()); double u = 1.0 - random.nextDouble(); // Avoid log(0) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java index f3d08a25c1..9967db60a8 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java @@ -22,7 +22,6 @@ import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; -import com.google.common.collect.Lists; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -39,7 +38,8 @@ class InliningNode extends AbstractNode { @Nonnull @Override public Node create(@Nonnull final Tuple primaryKey, - @Nullable final Vector vector, @Nonnull final List neighbors) { + @Nullable final Vector vector, + @Nonnull final List neighbors) { return new InliningNode(primaryKey, (List)neighbors); } @@ -79,25 +79,6 @@ public InliningNode asInliningNode() { return this; } - @Override - public NodeFactory sameCreator() { - return InliningNode.factory(); - } - - @Nonnull - @Override - public Tuple toTuple() { - final List nodeItems = Lists.newArrayListWithExpectedSize(3); - nodeItems.add(NodeKind.INLINING.getSerialized()); - final List neighborItems = Lists.newArrayListWithExpectedSize(getNeighbors().size()); - for (final NodeReferenceWithVector nodeReference : getNeighbors()) { - neighborItems.add(Tuple.from(nodeReference.getPrimaryKey(), - StorageAdapter.tupleFromVector(nodeReference.getVector()))); - } - nodeItems.add(Tuple.fromList(neighborItems)); - return Tuple.fromList(nodeItems); - } - @Nonnull public static NodeFactory factory() { return FACTORY; diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java new file mode 100644 index 0000000000..380eb48ef9 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java @@ -0,0 +1,125 @@ +/* + * CompactStorageAdapter.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.KeyValue; +import com.apple.foundationdb.Range; +import com.apple.foundationdb.ReadTransaction; +import com.apple.foundationdb.StreamingMode; +import com.apple.foundationdb.Transaction; +import com.apple.foundationdb.async.AsyncUtil; +import com.apple.foundationdb.subspace.Subspace; +import com.apple.foundationdb.tuple.Tuple; +import com.christianheina.langx.half4j.Half; +import com.google.common.collect.ImmutableList; + +import javax.annotation.Nonnull; +import java.util.concurrent.CompletableFuture; + +/** + * TODO. + */ +class InliningStorageAdapter extends AbstractStorageAdapter implements StorageAdapter { + public InliningStorageAdapter(@Nonnull final HNSW.Config config, + @Nonnull final NodeFactory nodeFactory, + @Nonnull final Subspace subspace, + @Nonnull final OnWriteListener onWriteListener, + @Nonnull final OnReadListener onReadListener) { + super(config, nodeFactory, subspace, onWriteListener, onReadListener); + } + + @Nonnull + @Override + public StorageAdapter asCompactStorageAdapter() { + throw new IllegalStateException("cannot call this method on an inlining storage adapter"); + } + + @Nonnull + @Override + public StorageAdapter asInliningStorageAdapter() { + return this; + } + + @Nonnull + protected CompletableFuture> fetchNodeInternal(@Nonnull final ReadTransaction readTransaction, + final int layer, + @Nonnull final Tuple primaryKey) { + final byte[] rangeKey = getNodeKey(layer, primaryKey); + + return AsyncUtil.collect(readTransaction.getRange(Range.startsWith(rangeKey), + ReadTransaction.ROW_LIMIT_UNLIMITED, false, StreamingMode.WANT_ALL), readTransaction.getExecutor()) + .thenApply(keyValues -> { + final OnReadListener onReadListener = getOnReadListener(); + + final ImmutableList.Builder nodeReferencesWithVectorBuilder = ImmutableList.builder(); + for (final KeyValue keyValue : keyValues) { + final byte[] key = keyValue.getKey(); + final byte[] value = keyValue.getValue(); + onReadListener.onKeyValueRead(key, value); + final Tuple neighborKeyTuple = Tuple.fromBytes(key); + final Tuple neighborValueTuple = Tuple.fromBytes(value); + + final Tuple neighborPrimaryKey = neighborKeyTuple.getNestedTuple(2); // neighbor primary key + final Vector neighborVector = StorageAdapter.vectorFromTuple(neighborValueTuple); // the entire value is the vector + nodeReferencesWithVectorBuilder.add(new NodeReferenceWithVector(neighborPrimaryKey, neighborVector)); + } + + final Node node = + getNodeFactory().create(primaryKey, null, nodeReferencesWithVectorBuilder.build()); + onReadListener.onNodeRead(node); + return node; + }); + } + + @Override + public void writeNode(@Nonnull final Transaction transaction, @Nonnull final Node node, + final int layer, @Nonnull final NeighborsChangeSet neighborsChangeSet) { + final InliningNode inliningNode = node.asInliningNode(); + + neighborsChangeSet.writeDelta(this, transaction, layer, inliningNode, t -> true); + getOnWriteListener().onNodeWritten(layer, node); + } + + @Nonnull + private byte[] getNodeKey(final int layer, @Nonnull final Tuple primaryKey) { + return getDataSubspace().pack(Tuple.from(layer, primaryKey)); + } + + public void writeNeighbor(@Nonnull final Transaction transaction, final int layer, + @Nonnull final Node node, @Nonnull final NodeReferenceWithVector neighbor) { + transaction.set(getNeighborKey(layer, node, neighbor.getPrimaryKey()), + StorageAdapter.tupleFromVector(neighbor.getVector()).pack()); + getOnWriteListener().onNeighborWritten(layer, node, neighbor); + } + + public void deleteNeighbor(@Nonnull final Transaction transaction, final int layer, + @Nonnull final Node node, @Nonnull final Tuple neighborPrimaryKey) { + transaction.clear(getNeighborKey(layer, node, neighborPrimaryKey)); + getOnWriteListener().onNeighborDeleted(layer, node, neighborPrimaryKey); + } + + @Nonnull + private byte[] getNeighborKey(final int layer, + @Nonnull final Node node, + @Nonnull final Tuple neighborPrimaryKey) { + return getDataSubspace().pack(Tuple.from(layer, node.getPrimaryKey(), neighborPrimaryKey)); + } +} diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java index e91a319c9b..f5473fd29e 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java @@ -21,11 +21,14 @@ package com.apple.foundationdb.async.hnsw; import com.apple.foundationdb.Transaction; -import com.google.common.collect.ImmutableList; +import com.apple.foundationdb.tuple.Tuple; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.Iterables; import javax.annotation.Nonnull; import java.util.List; +import java.util.Map; +import java.util.function.Predicate; /** * TODO. @@ -35,12 +38,17 @@ class InsertNeighborsChangeSet implements NeighborsChan private final NeighborsChangeSet parent; @Nonnull - private final List insertedNeighbors; + private final Map insertedNeighborsMap; public InsertNeighborsChangeSet(@Nonnull final NeighborsChangeSet parent, @Nonnull final List insertedNeighbors) { this.parent = parent; - this.insertedNeighbors = ImmutableList.copyOf(insertedNeighbors); + final ImmutableMap.Builder insertedNeighborsMapBuilder = ImmutableMap.builder(); + for (final N insertedNeighbor : insertedNeighbors) { + insertedNeighborsMapBuilder.put(insertedNeighbor.getPrimaryKey(), insertedNeighbor); + } + + this.insertedNeighborsMap = insertedNeighborsMapBuilder.build(); } @Nonnull @@ -50,11 +58,21 @@ public NeighborsChangeSet getParent() { @Nonnull public Iterable merge() { - return Iterables.concat(getParent().merge(), insertedNeighbors); + return Iterables.concat(getParent().merge(), insertedNeighborsMap.values()); } @Override - public void writeDelta(@Nonnull final Transaction transaction) { - throw new UnsupportedOperationException("not implemented yet"); + public void writeDelta(@Nonnull final InliningStorageAdapter storageAdapter, @Nonnull final Transaction transaction, + final int layer, @Nonnull final Node node, @Nonnull final Predicate tuplePredicate) { + getParent().writeDelta(storageAdapter, transaction, layer, node, + tuplePredicate.and(tuple -> !insertedNeighborsMap.containsKey(tuple))); + + for (final Map.Entry entry : insertedNeighborsMap.entrySet()) { + final Tuple primaryKey = entry.getKey(); + if (tuplePredicate.test(primaryKey)) { + storageAdapter.writeNeighbor(transaction, layer, node.asInliningNode(), + entry.getValue().asNodeReferenceWithVector()); + } + } } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NeighborsChangeSet.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NeighborsChangeSet.java index 21f621670a..b7f38ef1a7 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NeighborsChangeSet.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NeighborsChangeSet.java @@ -21,9 +21,11 @@ package com.apple.foundationdb.async.hnsw; import com.apple.foundationdb.Transaction; +import com.apple.foundationdb.tuple.Tuple; import javax.annotation.Nonnull; import javax.annotation.Nullable; +import java.util.function.Predicate; /** * TODO. @@ -35,5 +37,6 @@ interface NeighborsChangeSet { @Nonnull Iterable merge(); - void writeDelta(@Nonnull final Transaction transaction); + void writeDelta(@Nonnull InliningStorageAdapter storageAdapter, @Nonnull Transaction transaction, int layer, + @Nonnull Node node, @Nonnull Predicate primaryKeyPredicate); } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java index aaa65dd21a..4e7b08c301 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java @@ -22,7 +22,6 @@ import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; -import com.google.common.collect.Lists; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -57,64 +56,4 @@ public interface Node { @Nonnull InliningNode asInliningNode(); - - NodeFactory sameCreator(); - - @Nonnull - Tuple toTuple(); - - @Nonnull - static Node nodeFromTuples(@Nonnull final NodeFactory creator, - @Nonnull final Tuple primaryKey, - @Nonnull final Tuple valueTuple) { - final NodeKind nodeKind = NodeKind.fromSerializedNodeKind((byte)valueTuple.getLong(0)); - final Tuple vectorTuple; - final Tuple neighborsTuple; - - switch (nodeKind) { - case COMPACT: - vectorTuple = valueTuple.getNestedTuple(1); - neighborsTuple = valueTuple.getNestedTuple(2); - return compactNodeFromTuples(creator, primaryKey, vectorTuple, neighborsTuple); - case INLINING: - neighborsTuple = valueTuple.getNestedTuple(1); - return inliningNodeFromTuples(creator, primaryKey, neighborsTuple); - default: - throw new IllegalStateException("unknown node kind"); - } - } - - @Nonnull - static Node compactNodeFromTuples(@Nonnull final NodeFactory creator, - @Nonnull final Tuple primaryKey, - @Nonnull final Tuple vectorTuple, - @Nonnull final Tuple neighborsTuple) { - final Vector vector = StorageAdapter.vectorFromTuple(vectorTuple); - - List nodeReferences = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); - - for (final Object neighborObject : neighborsTuple) { - final Tuple neighborTuple = (Tuple)neighborObject; - nodeReferences.add(new NodeReference(neighborTuple)); - } - - return creator.create(NodeKind.COMPACT, primaryKey, vector, nodeReferences); - } - - @Nonnull - static Node inliningNodeFromTuples(@Nonnull final NodeFactory creator, - @Nonnull final Tuple primaryKey, - @Nonnull final Tuple neighborsTuple) { - List neighborsWithVectors = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); - - for (final Object neighborObject : neighborsTuple) { - final Tuple neighborTuple = (Tuple)neighborObject; - final Tuple neighborPrimaryKey = neighborTuple.getNestedTuple(0); - final Tuple neighborVectorTuple = neighborTuple.getNestedTuple(1); - neighborsWithVectors.add(new NodeReferenceWithVector(neighborPrimaryKey, - StorageAdapter.vectorFromTuple(neighborVectorTuple))); - } - - return creator.create(NodeKind.INLINING, primaryKey, null, neighborsWithVectors); - } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeHelpers.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeHelpers.java index 965f5742cc..1b81c9f847 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeHelpers.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeHelpers.java @@ -20,17 +20,7 @@ package com.apple.foundationdb.async.hnsw; -import com.google.common.collect.Lists; - import javax.annotation.Nonnull; -import javax.annotation.Nullable; -import java.nio.ByteBuffer; -import java.nio.ByteOrder; -import java.util.Collections; -import java.util.List; -import java.util.Objects; -import java.util.UUID; -import java.util.concurrent.atomic.AtomicLong; /** * Some helper methods for {@link Node}s. @@ -57,24 +47,4 @@ static String bytesToHex(byte[] bytes) { } return "0x" + new String(hexChars).replaceFirst("^0+(?!$)", ""); } - - /** - * Helper method to format the node ids of an insert/update path as a string. - * @param node a node that is usually linked up to its parents to form an insert/update path - * @return a {@link String} containing the string presentation of the insert/update path starting at {@code node} - */ - @Nonnull - static String nodeIdPath(@Nullable Node node) { - final List nodeIds = Lists.newArrayList(); - do { - if (node != null) { - nodeIds.add(bytesToHex(node.getId())); - node = node.getParentNode(); - } else { - nodeIds.add(""); - } - } while (node != null); - Collections.reverse(nodeIds); - return String.join(", ", nodeIds); - } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReference.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReference.java index dd7d1680d5..f057b199c5 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReference.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReference.java @@ -40,10 +40,8 @@ public Tuple getPrimaryKey() { } @Nonnull - public static Iterable primaryKeys(@Nonnull Iterable neighbors) { - return () -> Streams.stream(neighbors) - .map(NodeReference::getPrimaryKey) - .iterator(); + public NodeReferenceWithVector asNodeReferenceWithVector() { + throw new IllegalStateException("method should not be called"); } @Override @@ -59,4 +57,12 @@ public boolean equals(final Object o) { public int hashCode() { return Objects.hashCode(primaryKey); } + + @Nonnull + public static Iterable primaryKeys(@Nonnull Iterable neighbors) { + return () -> Streams.stream(neighbors) + .map(NodeReference::getPrimaryKey) + .iterator(); + } + } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithVector.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithVector.java index ead8135b6a..c3223aa200 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithVector.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithVector.java @@ -43,4 +43,10 @@ public Vector getVector() { public Vector getDoubleVector() { return vector.toDoubleVector(); } + + @Nonnull + @Override + public NodeReferenceWithVector asNodeReferenceWithVector() { + return this; + } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java index 711dc578ad..beb8530cd6 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java @@ -38,7 +38,7 @@ default CompletableFuture> onAsyncRead(@Nonnul return future; } - default void onNodeRead(@Nonnull Node node) { + default void onNodeRead(@Nonnull Node node) { // nothing } @@ -46,8 +46,4 @@ default void onKeyValueRead(@Nonnull byte[] key, @Nonnull byte[] value) { // nothing } - - default void onChildNodeDiscard(@Nonnull final ChildSlot childSlot) { - // nothing - } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java index 7d39e0ef0d..d57450d434 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java @@ -20,8 +20,9 @@ package com.apple.foundationdb.async.hnsw; +import com.apple.foundationdb.tuple.Tuple; + import javax.annotation.Nonnull; -import java.util.concurrent.CompletableFuture; /** * Function interface for a call back whenever we read the slots for a node. @@ -30,34 +31,20 @@ public interface OnWriteListener { OnWriteListener NOOP = new OnWriteListener() { }; - default void onSlotIndexEntryWritten(@Nonnull final byte[] key) { + default void onNodeWritten(final int layer, @Nonnull final Node node) { // nothing } - default void onSlotIndexEntryCleared(@Nonnull final byte[] key) { + default void onNeighborWritten(final int layer, @Nonnull final Node node, final NodeReference neighbor) { // nothing } - default CompletableFuture onAsyncReadForWrite(@Nonnull CompletableFuture future) { - return future; - } - - default void onNodeWritten(@Nonnull Node node) { + default void onNeighborDeleted(final int layer, @Nonnull final Node node, @Nonnull Tuple neighborPrimaryKey) { // nothing } - default void onKeyValueWritten(@Nonnull Node node, - @Nonnull byte[] key, + default void onKeyValueWritten(@Nonnull byte[] key, @Nonnull byte[] value) { // nothing } - - default void onNodeCleared(@Nonnull Node node) { - // nothing - } - - default void onKeyCleared(@Nonnull Node node, - @Nonnull byte[] key) { - // nothing - } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java index 9dcadcaa87..53e1a5750a 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java @@ -26,17 +26,17 @@ import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; import com.google.common.base.Verify; -import com.google.common.collect.Lists; import javax.annotation.Nonnull; -import javax.annotation.Nullable; -import java.util.List; import java.util.concurrent.CompletableFuture; /** * Storage adapter used for serialization and deserialization of nodes. */ -interface StorageAdapter { +interface StorageAdapter { + byte SUBSPACE_PREFIX_ENTRY_NODE = 0x01; + byte SUBSPACE_PREFIX_DATA = 0x02; + /** * Get the {@link HNSW.Config} associated with this storage adapter. * @return the configuration used by this storage adapter @@ -44,6 +44,18 @@ interface StorageAdapter { @Nonnull HNSW.Config getConfig(); + @Nonnull + NodeFactory getNodeFactory(); + + @Nonnull + NodeKind getNodeKind(); + + @Nonnull + StorageAdapter asCompactStorageAdapter(); + + @Nonnull + StorageAdapter asInliningStorageAdapter(); + /** * Get the subspace used to store this r-tree. * @@ -52,17 +64,6 @@ interface StorageAdapter { @Nonnull Subspace getSubspace(); - /** - * Get the subspace used to store a node slot index if in warranted by the {@link HNSW.Config}. - * - * @return secondary subspace or {@code null} if we do not maintain a node slot index - */ - @Nullable - Subspace getSecondarySubspace(); - - @Nonnull - Subspace getEntryNodeSubspace(); - @Nonnull Subspace getDataSubspace(); @@ -83,49 +84,84 @@ interface StorageAdapter { OnReadListener getOnReadListener(); @Nonnull - NodeFactory getNodeFactory(final int layer); - - CompletableFuture fetchEntryNodeReference(@Nonnull ReadTransaction readTransaction); + CompletableFuture> fetchNode(@Nonnull ReadTransaction readTransaction, + int layer, + @Nonnull Tuple primaryKey); - void writeEntryNodeReference(@Nonnull final Transaction transaction, - @Nonnull final EntryNodeReference entryNodeReference); + void writeNode(@Nonnull Transaction transaction, @Nonnull Node node, int layer, + @Nonnull NeighborsChangeSet changeSet); @Nonnull - CompletableFuture> fetchNode(@Nonnull NodeFactory nodeFactory, - @Nonnull ReadTransaction readTransaction, - int layer, - @Nonnull Tuple primaryKey); + static CompletableFuture fetchEntryNodeReference(@Nonnull final ReadTransaction readTransaction, + @Nonnull final Subspace subspace, + @Nonnull final OnReadListener onReadListener) { + final Subspace entryNodeSubspace = subspace.subspace(Tuple.from(SUBSPACE_PREFIX_ENTRY_NODE)); + final byte[] key = entryNodeSubspace.pack(); + + return readTransaction.get(key) + .thenApply(valueBytes -> { + if (valueBytes == null) { + return null; // not a single node in the index + } + onReadListener.onKeyValueRead(key, valueBytes); + + final Tuple entryTuple = Tuple.fromBytes(valueBytes); + final int lMax = (int)entryTuple.getLong(0); + final Tuple primaryKey = entryTuple.getNestedTuple(1); + final Tuple vectorTuple = entryTuple.getNestedTuple(2); + return new EntryNodeReference(primaryKey, StorageAdapter.vectorFromTuple(vectorTuple), lMax); + }); + } - void writeNode(@Nonnull final Transaction transaction, @Nonnull final Node node, - int layer); + static void writeEntryNodeReference(@Nonnull final Transaction transaction, + @Nonnull final Subspace subspace, + @Nonnull final EntryNodeReference entryNodeReference, + @Nonnull final OnWriteListener onWriteListener) { + final Subspace entryNodeSubspace = subspace.subspace(Tuple.from(SUBSPACE_PREFIX_ENTRY_NODE)); + final byte[] key = entryNodeSubspace.pack(); + final byte[] value = Tuple.from(entryNodeReference.getLayer(), + entryNodeReference.getPrimaryKey(), + StorageAdapter.tupleFromVector(entryNodeReference.getVector())).pack(); + transaction.set(key, + value); + onWriteListener.onKeyValueWritten(key, value); + } @Nonnull static Vector vectorFromTuple(final Tuple vectorTuple) { - final Half[] vectorHalfs = new Half[vectorTuple.size()]; - for (int i = 0; i < vectorTuple.size(); i ++) { - vectorHalfs[i] = Half.shortBitsToHalf(shortFromBytes(vectorTuple.getBytes(i))); + final byte[] vectorAsBytes = vectorTuple.getBytes(0); + final int bytesLength = vectorAsBytes.length; + Verify.verify(bytesLength % 2 == 0); + final int componentSize = bytesLength >>> 1; + final Half[] vectorHalfs = new Half[componentSize]; + for (int i = 0; i < componentSize; i ++) { + vectorHalfs[i] = Half.shortBitsToHalf(shortFromBytes(vectorAsBytes, i << 1)); } return new Vector.HalfVector(vectorHalfs); } @Nonnull + @SuppressWarnings("PrimitiveArrayArgumentToVarargsMethod") static Tuple tupleFromVector(final Vector vector) { - final List vectorBytes = Lists.newArrayListWithExpectedSize(vector.size()); + final byte[] vectorBytes = new byte[2 * vector.size()]; for (int i = 0; i < vector.size(); i ++) { - vectorBytes.add(bytesFromShort(Half.halfToShortBits(vector.getComponent(i)))); + final byte[] componentBytes = bytesFromShort(Half.halfToShortBits(vector.getComponent(i))); + final int indexTimesTwo = i << 1; + vectorBytes[indexTimesTwo] = componentBytes[0]; + vectorBytes[indexTimesTwo + 1] = componentBytes[1]; } - return Tuple.fromList(vectorBytes); + return Tuple.from(vectorBytes); } - static short shortFromBytes(byte[] bytes) { - Verify.verify(bytes.length == 2); - int high = bytes[0] & 0xFF; // Convert to unsigned int - int low = bytes[1] & 0xFF; + static short shortFromBytes(final byte[] bytes, final int offset) { + Verify.verify(offset % 2 == 0); + int high = bytes[offset] & 0xFF; // Convert to unsigned int + int low = bytes[offset + 1] & 0xFF; return (short) ((high << 8) | low); } - static byte[] bytesFromShort(short value) { + static byte[] bytesFromShort(final short value) { byte[] result = new byte[2]; result[0] = (byte) ((value >> 8) & 0xFF); // high byte first result[1] = (byte) (value & 0xFF); // low byte second diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/rtree/StorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/rtree/StorageAdapter.java index f60c17da63..2623cff1dc 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/rtree/StorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/rtree/StorageAdapter.java @@ -36,7 +36,6 @@ * Storage adapter used for serialization and deserialization of nodes. */ interface StorageAdapter { - /** * Get the {@link RTree.Config} associated with this storage adapter. * @return the configuration used by this storage adapter From 0a4b28293f24ac4ce1dbdf034536ac87a88702ad Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Thu, 31 Jul 2025 17:52:57 +0200 Subject: [PATCH 14/34] added logging --- .../async/hnsw/AbstractStorageAdapter.java | 18 ++ .../foundationdb/async/hnsw/CompactNode.java | 11 +- .../async/hnsw/CompactStorageAdapter.java | 20 +- .../async/hnsw/DeleteNeighborsChangeSet.java | 9 + .../apple/foundationdb/async/hnsw/HNSW.java | 89 +++++-- .../{NodeHelpers.java => HNSWHelpers.java} | 19 +- .../foundationdb/async/hnsw/InliningNode.java | 6 + .../async/hnsw/InliningStorageAdapter.java | 6 +- .../async/hnsw/InsertNeighborsChangeSet.java | 9 + .../async/hnsw/NodeReference.java | 6 +- .../async/hnsw/NodeReferenceWithVector.java | 7 + .../apple/foundationdb/async/hnsw/Vector.java | 20 +- .../foundationdb/async/rtree/NodeHelpers.java | 2 +- .../async/hnsw/HNSWModificationTest.java | 222 ++++++++++++++++++ .../indexes/VectorIndexTestBase.java | 8 +- .../src/test/proto/test_records_vector.proto | 2 +- 16 files changed, 411 insertions(+), 43 deletions(-) rename fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/{NodeHelpers.java => HNSWHelpers.java} (76%) create mode 100644 fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java index a306b3c600..5b074279eb 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java @@ -21,8 +21,11 @@ package com.apple.foundationdb.async.hnsw; import com.apple.foundationdb.ReadTransaction; +import com.apple.foundationdb.Transaction; import com.apple.foundationdb.subspace.Subspace; import com.apple.foundationdb.tuple.Tuple; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -32,6 +35,9 @@ * Implementations and attributes common to all concrete implementations of {@link StorageAdapter}. */ abstract class AbstractStorageAdapter implements StorageAdapter { + @Nonnull + private static final Logger logger = LoggerFactory.getLogger(AbstractStorageAdapter.class); + @Nonnull private final HNSW.Config config; @Nonnull @@ -122,4 +128,16 @@ protected abstract CompletableFuture> fetchNodeInternal(@Nonnull ReadTra private Node checkNode(@Nullable final Node node) { return node; } + + public void writeNode(@Nonnull Transaction transaction, @Nonnull Node node, int layer, + @Nonnull NeighborsChangeSet changeSet) { + writeNodeInternal(transaction, node, layer, changeSet); + if (logger.isDebugEnabled()) { + logger.debug("written node with key={} at layer={}", node.getPrimaryKey(), layer); + } + } + + protected abstract void writeNodeInternal(@Nonnull Transaction transaction, @Nonnull Node node, int layer, + @Nonnull NeighborsChangeSet changeSet); + } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java index 8dd86f622b..23f2667b63 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java @@ -53,8 +53,8 @@ public NodeKind getNodeKind() { private final Vector vector; public CompactNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, - @Nonnull final List nodeReferences) { - super(primaryKey, nodeReferences); + @Nonnull final List neighbors) { + super(primaryKey, neighbors); this.vector = vector; } @@ -91,4 +91,11 @@ public InliningNode asInliningNode() { public static NodeFactory factory() { return FACTORY; } + + @Override + public String toString() { + return "C[primaryKey=" + getPrimaryKey() + + ";vector=" + vector + + ";neighbors=" + getNeighbors() + "]"; + } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java index 9da72759f4..2ce21f7091 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java @@ -27,6 +27,8 @@ import com.christianheina.langx.half4j.Half; import com.google.common.base.Verify; import com.google.common.collect.Lists; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; import java.util.List; @@ -36,6 +38,9 @@ * TODO. */ class CompactStorageAdapter extends AbstractStorageAdapter implements StorageAdapter { + @Nonnull + private static final Logger logger = LoggerFactory.getLogger(CompactStorageAdapter.class); + public CompactStorageAdapter(@Nonnull final HNSW.Config config, @Nonnull final NodeFactory nodeFactory, @Nonnull final Subspace subspace, @Nonnull final OnWriteListener onWriteListener, @@ -98,8 +103,8 @@ private Node compactNodeFromTuples(@Nonnull final Tuple primaryKe final Vector vector = StorageAdapter.vectorFromTuple(vectorTuple); final List nodeReferences = Lists.newArrayListWithExpectedSize(neighborsTuple.size()); - for (final Object neighborObject : neighborsTuple) { - final Tuple neighborTuple = (Tuple)neighborObject; + for (int i = 0; i < neighborsTuple.size(); i ++) { + final Tuple neighborTuple = neighborsTuple.getNestedTuple(i); nodeReferences.add(new NodeReference(neighborTuple)); } @@ -108,11 +113,11 @@ private Node compactNodeFromTuples(@Nonnull final Tuple primaryKe @Override - public void writeNode(@Nonnull final Transaction transaction, @Nonnull final Node node, - final int layer, @Nonnull final NeighborsChangeSet neighborsChangeSet) { + public void writeNodeInternal(@Nonnull final Transaction transaction, @Nonnull final Node node, + final int layer, @Nonnull final NeighborsChangeSet neighborsChangeSet) { final byte[] key = getDataSubspace().pack(Tuple.from(layer, node.getPrimaryKey())); - final List nodeItems = Lists.newArrayListWithExpectedSize(4); + final List nodeItems = Lists.newArrayListWithExpectedSize(3); nodeItems.add(NodeKind.COMPACT.getSerialized()); final CompactNode compactNode = node.asCompactNode(); nodeItems.add(StorageAdapter.tupleFromVector(compactNode.getVector())); @@ -123,6 +128,11 @@ public void writeNode(@Nonnull final Transaction transaction, @Nonnull final Nod for (final NodeReference neighborReference : neighbors) { neighborItems.add(neighborReference.getPrimaryKey()); } + if (logger.isDebugEnabled()) { + logger.debug("written neighbors of primaryKey={}, oldSize={}, newSize={}", node.getPrimaryKey(), + node.getNeighbors().size(), neighborItems.size()); + } + nodeItems.add(Tuple.fromList(neighborItems)); final Tuple nodeTuple = Tuple.fromList(nodeItems); diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java index feee9fa979..0d4694b2fe 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java @@ -24,6 +24,8 @@ import com.apple.foundationdb.tuple.Tuple; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; import java.util.Collection; @@ -34,6 +36,9 @@ * TODO. */ class DeleteNeighborsChangeSet implements NeighborsChangeSet { + @Nonnull + private static final Logger logger = LoggerFactory.getLogger(DeleteNeighborsChangeSet.class); + @Nonnull private final NeighborsChangeSet parent; @@ -66,6 +71,10 @@ public void writeDelta(@Nonnull final InliningStorageAdapter storageAdapter, @No for (final Tuple deletedNeighborPrimaryKey : deletedNeighborsPrimaryKeys) { if (tuplePredicate.test(deletedNeighborPrimaryKey)) { storageAdapter.deleteNeighbor(transaction, layer, node.asInliningNode(), deletedNeighborPrimaryKey); + if (logger.isDebugEnabled()) { + logger.debug("deleted neighbor of primaryKey={} targeting primaryKey={}", node.getPrimaryKey(), + deletedNeighborPrimaryKey); + } } } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index 8b948aa459..b7cb4c3fd0 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -52,7 +52,9 @@ import java.util.concurrent.PriorityBlockingQueue; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; +import java.util.function.Consumer; import java.util.function.Function; +import java.util.stream.Collectors; /** * TODO. @@ -60,6 +62,7 @@ @API(API.Status.EXPERIMENTAL) @SuppressWarnings("checkstyle:AbbreviationAsWordInName") public class HNSW { + @Nonnull private static final Logger logger = LoggerFactory.getLogger(HNSW.class); public static final int MAX_CONCURRENT_NODE_READS = 16; @@ -198,8 +201,8 @@ public static class ConfigBuilder { @Nonnull private Metric metric = DEFAULT_METRIC; private int m = DEFAULT_M; - private int mMax; - private int mMax0; + private int mMax = DEFAULT_M_MAX; + private int mMax0 = DEFAULT_M_MAX_0; private int efSearch = DEFAULT_EF_SEARCH; private int efConstruction = DEFAULT_EF_CONSTRUCTION; private boolean extendCandidates = DEFAULT_EXTEND_CANDIDATES; @@ -398,6 +401,7 @@ public OnReadListener getOnReadListener() { @SuppressWarnings("checkstyle:MethodName") // method name introduced by paper @Nonnull public CompletableFuture>> kNearestNeighborsSearch(@Nonnull final ReadTransaction readTransaction, + final int efSearch, @Nonnull final Vector queryVector) { return StorageAdapter.fetchEntryNodeReference(readTransaction, getSubspace(), getOnReadListener()) .thenCompose(entryPointAndLayer -> { @@ -418,7 +422,6 @@ public CompletableFuture nodeReferenceAtomic = new AtomicReference<>(entryState); @@ -444,8 +447,8 @@ public CompletableFuture CompletableFuture return true; }); }).thenCompose(ignored -> - fetchSomeNodesIfNotCached(storageAdapter, readTransaction, layer, nearestNeighbors, nodeCache)); + fetchSomeNodesIfNotCached(storageAdapter, readTransaction, layer, nearestNeighbors, nodeCache)) + .thenApply(searchResult -> { + debug(l -> { + l.debug("searched layer={} for efSearch={} with result=={}", layer, efSearch, + searchResult.stream() + .map(nodeReferenceAndNode -> + "(primaryKey=" + nodeReferenceAndNode.getNodeReferenceWithDistance().getPrimaryKey() + + ",distance=" + nodeReferenceAndNode.getNodeReferenceWithDistance().getDistance() + ")") + .collect(Collectors.joining(","))); + }); + return searchResult; + }); } /** @@ -682,21 +696,24 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N @Nonnull final Vector newVector) { final Metric metric = getConfig().getMetric(); - final int l = insertionLayer(getConfig().getRandom()); + final int insertionLayer = insertionLayer(getConfig().getRandom()); + debug(l -> l.debug("new node with key={} selected to be inserted into layer={}", newPrimaryKey, insertionLayer)); return StorageAdapter.fetchEntryNodeReference(transaction, getSubspace(), getOnReadListener()) .thenApply(entryNodeReference -> { if (entryNodeReference == null) { // this is the first node - writeLonelyNodes(transaction, newPrimaryKey, newVector, l, -1); + writeLonelyNodes(transaction, newPrimaryKey, newVector, insertionLayer, -1); StorageAdapter.writeEntryNodeReference(transaction, getSubspace(), - new EntryNodeReference(newPrimaryKey, newVector, l), getOnWriteListener()); + new EntryNodeReference(newPrimaryKey, newVector, insertionLayer), getOnWriteListener()); + debug(l -> l.debug("written entry node reference with key={} on layer={}", newPrimaryKey, insertionLayer)); } else { final int entryNodeLayer = entryNodeReference.getLayer(); - if (l > entryNodeLayer) { - writeLonelyNodes(transaction, newPrimaryKey, newVector, l, entryNodeLayer); + if (insertionLayer > entryNodeLayer) { + writeLonelyNodes(transaction, newPrimaryKey, newVector, insertionLayer, entryNodeLayer); StorageAdapter.writeEntryNodeReference(transaction, getSubspace(), - new EntryNodeReference(newPrimaryKey, newVector, l), getOnWriteListener()); + new EntryNodeReference(newPrimaryKey, newVector, insertionLayer), getOnWriteListener()); + debug(l -> l.debug("written entry node reference with key={} on layer={}", newPrimaryKey, insertionLayer)); } } return entryNodeReference; @@ -706,13 +723,15 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N } final int lMax = entryNodeReference.getLayer(); + debug(l -> l.debug("entry node with key {} at layer {}", entryNodeReference.getPrimaryKey(), + lMax)); final AtomicReference nodeReferenceAtomic = new AtomicReference<>(new NodeReferenceWithDistance(entryNodeReference.getPrimaryKey(), entryNodeReference.getVector(), Vector.comparativeDistance(metric, entryNodeReference.getVector(), newVector))); MoreAsyncUtil.forLoop(lMax, - layer -> layer > l, + layer -> layer > insertionLayer, layer -> layer - 1, layer -> { final StorageAdapter storageAdapter = getStorageAdapterForLayer(layer); @@ -724,10 +743,15 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N }); }, executor); + debug(l -> { + final NodeReference nodeReference = nodeReferenceAtomic.get(); + l.debug("nearest entry point at lMax={} is at key={}", lMax, nodeReference.getPrimaryKey()); + }); + final AtomicReference> nearestNeighborsAtomic = new AtomicReference<>(ImmutableList.of(nodeReferenceAtomic.get())); - return MoreAsyncUtil.forLoop(Math.min(lMax, l), + return MoreAsyncUtil.forLoop(Math.min(lMax, insertionLayer), layer -> layer >= 0, layer -> layer - 1, layer -> { @@ -749,6 +773,7 @@ private CompletableFuture newVector) { + debug(l -> l.debug("begin insert key={} at layer={}", newPrimaryKey, layer)); final Map> nodeCache = Maps.newConcurrentMap(); return searchLayer(storageAdapter, transaction, @@ -756,8 +781,8 @@ private CompletableFuture { final List references = NodeReferenceAndNode.getReferences(searchResult); - return selectNeighbors(storageAdapter, transaction, searchResult, - layer, getConfig().getM(), getConfig().isExtendCandidates(), nodeCache, newVector) + return selectNeighbors(storageAdapter, transaction, searchResult, layer, getConfig().getM(), + getConfig().isExtendCandidates(), nodeCache, newVector) .thenCompose(selectedNeighbors -> { final NodeFactory nodeFactory = storageAdapter.getNodeFactory(); @@ -797,19 +822,19 @@ layer, getConfig().getM(), getConfig().isExtendCandidates(), nodeCache, newVecto }); }, MAX_CONCURRENT_NEIGHBOR_FETCHES, getExecutor()) .thenApply(changeSets -> { + storageAdapter.writeNode(transaction, newNode, layer, newNodeChangeSet); for (int i = 0; i < selectedNeighbors.size(); i++) { final NodeReferenceAndNode selectedNeighbor = selectedNeighbors.get(i); final NeighborsChangeSet changeSet = changeSets.get(i); - neighborChangeSetMap.put( - selectedNeighbor.getNodeReferenceWithDistance().getPrimaryKey(), - changeSet); + storageAdapter.writeNode(transaction, selectedNeighbor.getNode(), + layer, changeSet); } - - storageAdapter.writeNode(transaction, newNode, layer, newNodeChangeSet); - return ImmutableList.copyOf(references); }); }); + }).thenApply(nodeReferencesWithDistances -> { + debug(l -> l.debug("end insert key={} at layer={}", newPrimaryKey, layer)); + return nodeReferencesWithDistances; }); } @@ -940,7 +965,18 @@ private CompletableFuture return ImmutableList.copyOf(selected); }).thenCompose(selectedNeighbors -> - fetchSomeNodesIfNotCached(storageAdapter, readTransaction, layer, selectedNeighbors, nodeCache)); + fetchSomeNodesIfNotCached(storageAdapter, readTransaction, layer, selectedNeighbors, nodeCache)) + .thenApply(selectedNeighbors -> { + debug(l -> { + l.debug("selected neighbors={}", + selectedNeighbors.stream() + .map(selectedNeighbor -> + "(primaryKey=" + selectedNeighbor.getNodeReferenceWithDistance().getPrimaryKey() + + ",distance=" + selectedNeighbor.getNodeReferenceWithDistance().getDistance() + ")") + .collect(Collectors.joining(","))); + }); + return selectedNeighbors; + }); } private CompletableFuture> extendCandidatesIfNecessary(@Nonnull final StorageAdapter storageAdapter, @@ -1015,6 +1051,7 @@ private void writeLonelyNodeOnLayer(@Nonnull final Sto storageAdapter.getNodeFactory() .create(primaryKey, vector, ImmutableList.of()), layer, new BaseNeighborsChangeSet<>(ImmutableList.of())); + debug(l -> l.debug("written lonely node at key={} on layer={}", primaryKey, layer)); } @Nonnull @@ -1029,4 +1066,10 @@ private int insertionLayer(@Nonnull final Random random) { double u = 1.0 - random.nextDouble(); // Avoid log(0) return (int) Math.floor(-Math.log(u) * lambda); } + + private void debug(@Nonnull final Consumer loggerConsumer) { + if (logger.isDebugEnabled()) { + loggerConsumer.accept(logger); + } + } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeHelpers.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSWHelpers.java similarity index 76% rename from fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeHelpers.java rename to fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSWHelpers.java index 1b81c9f847..28d66df5fa 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeHelpers.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSWHelpers.java @@ -1,5 +1,5 @@ /* - * NodeHelpers.java + * HNSWHelpers.java * * This source file is part of the FoundationDB open source project * @@ -20,15 +20,18 @@ package com.apple.foundationdb.async.hnsw; +import com.christianheina.langx.half4j.Half; + import javax.annotation.Nonnull; /** * Some helper methods for {@link Node}s. */ -public class NodeHelpers { +@SuppressWarnings("checkstyle:AbbreviationAsWordInName") +public class HNSWHelpers { private static final char[] hexArray = "0123456789ABCDEF".toCharArray(); - private NodeHelpers() { + private HNSWHelpers() { // nothing } @@ -47,4 +50,14 @@ static String bytesToHex(byte[] bytes) { } return "0x" + new String(hexChars).replaceFirst("^0+(?!$)", ""); } + + @Nonnull + public static Half halfValueOf(final double d) { + return Half.shortBitsToHalf(Half.halfToShortBits(Half.valueOf(d))); + } + + @Nonnull + public static Half halfValueOf(final float f) { + return Half.shortBitsToHalf(Half.halfToShortBits(Half.valueOf(f))); + } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java index 9967db60a8..dce0f18f24 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java @@ -83,4 +83,10 @@ public InliningNode asInliningNode() { public static NodeFactory factory() { return FACTORY; } + + @Override + public String toString() { + return "I[primaryKey=" + getPrimaryKey() + + ";neighbors=" + getNeighbors() + "]"; + } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java index 380eb48ef9..57a8d9b289 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java @@ -74,7 +74,7 @@ protected CompletableFuture> fetchNodeInternal(@No final byte[] key = keyValue.getKey(); final byte[] value = keyValue.getValue(); onReadListener.onKeyValueRead(key, value); - final Tuple neighborKeyTuple = Tuple.fromBytes(key); + final Tuple neighborKeyTuple = getDataSubspace().unpack(key); final Tuple neighborValueTuple = Tuple.fromBytes(value); final Tuple neighborPrimaryKey = neighborKeyTuple.getNestedTuple(2); // neighbor primary key @@ -90,8 +90,8 @@ protected CompletableFuture> fetchNodeInternal(@No } @Override - public void writeNode(@Nonnull final Transaction transaction, @Nonnull final Node node, - final int layer, @Nonnull final NeighborsChangeSet neighborsChangeSet) { + public void writeNodeInternal(@Nonnull final Transaction transaction, @Nonnull final Node node, + final int layer, @Nonnull final NeighborsChangeSet neighborsChangeSet) { final InliningNode inliningNode = node.asInliningNode(); neighborsChangeSet.writeDelta(this, transaction, layer, inliningNode, t -> true); diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java index f5473fd29e..06a7490690 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java @@ -24,6 +24,8 @@ import com.apple.foundationdb.tuple.Tuple; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Iterables; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; import java.util.List; @@ -34,6 +36,9 @@ * TODO. */ class InsertNeighborsChangeSet implements NeighborsChangeSet { + @Nonnull + private static final Logger logger = LoggerFactory.getLogger(InsertNeighborsChangeSet.class); + @Nonnull private final NeighborsChangeSet parent; @@ -72,6 +77,10 @@ public void writeDelta(@Nonnull final InliningStorageAdapter storageAdapter, @No if (tuplePredicate.test(primaryKey)) { storageAdapter.writeNeighbor(transaction, layer, node.asInliningNode(), entry.getValue().asNodeReferenceWithVector()); + if (logger.isDebugEnabled()) { + logger.debug("inserted neighbor of primaryKey={} targeting primaryKey={}", node.getPrimaryKey(), + primaryKey); + } } } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReference.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReference.java index f057b199c5..59b831d04d 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReference.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReference.java @@ -58,11 +58,15 @@ public int hashCode() { return Objects.hashCode(primaryKey); } + @Override + public String toString() { + return "NR[primaryKey=" + primaryKey + "]"; + } + @Nonnull public static Iterable primaryKeys(@Nonnull Iterable neighbors) { return () -> Streams.stream(neighbors) .map(NodeReference::getPrimaryKey) .iterator(); } - } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithVector.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithVector.java index c3223aa200..b13443b926 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithVector.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithVector.java @@ -49,4 +49,11 @@ public Vector getDoubleVector() { public NodeReferenceWithVector asNodeReferenceWithVector() { return this; } + + @Override + public String toString() { + return "NRV[primaryKey=" + getPrimaryKey() + + ";vector=" + vector.toString(3) + + "]"; + } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java index 5308ea3271..6a2e5fd01e 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java @@ -1,5 +1,5 @@ /* - * NodeHelpers.java + * HNSWHelpers.java * * This source file is part of the FoundationDB open source project * @@ -27,6 +27,7 @@ import java.util.Arrays; import java.util.Objects; import java.util.function.Supplier; +import java.util.stream.Collectors; /** * TODO. @@ -81,6 +82,23 @@ private int computeHashCode() { return Arrays.hashCode(data); } + @Override + public String toString() { + return toString(3); + } + + public String toString(final int limitDimensions) { + if (limitDimensions < data.length) { + return "[" + Arrays.stream(Arrays.copyOfRange(data, 0, limitDimensions)) + .map(String::valueOf) + .collect(Collectors.joining(",")) + ", ...]"; + } else { + return "[" + Arrays.stream(data) + .map(String::valueOf) + .collect(Collectors.joining(",")) + "]"; + } + } + public static class HalfVector extends Vector { @Nonnull private final Supplier toDoubleVectorSupplier; diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/rtree/NodeHelpers.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/rtree/NodeHelpers.java index db4e4cf636..a11ac8b462 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/rtree/NodeHelpers.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/rtree/NodeHelpers.java @@ -1,5 +1,5 @@ /* - * NodeHelpers.java + * HNSWHelpers.java * * This source file is part of the FoundationDB open source project * diff --git a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java new file mode 100644 index 0000000000..6eb7d4a1c9 --- /dev/null +++ b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java @@ -0,0 +1,222 @@ +/* + * RTreeModificationTest.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import com.apple.foundationdb.Database; +import com.apple.foundationdb.Transaction; +import com.apple.foundationdb.async.rtree.RTree; +import com.apple.foundationdb.test.TestDatabaseExtension; +import com.apple.foundationdb.test.TestExecutors; +import com.apple.foundationdb.test.TestSubspaceExtension; +import com.apple.foundationdb.tuple.Tuple; +import com.apple.test.Tags; +import com.christianheina.langx.half4j.Half; +import com.google.common.collect.ImmutableList; +import org.assertj.core.util.Lists; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.api.parallel.Execution; +import org.junit.jupiter.api.parallel.ExecutionMode; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nonnull; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.Random; + +/** + * Tests testing insert/update/deletes of data into/in/from {@link RTree}s. + */ +@Execution(ExecutionMode.CONCURRENT) +@SuppressWarnings("checkstyle:AbbreviationAsWordInName") +@Tag(Tags.RequiresFDB) +@Tag(Tags.Slow) +public class HNSWModificationTest { + private static final Logger logger = LoggerFactory.getLogger(HNSWModificationTest.class); + private static final int NUM_TEST_RUNS = 5; + private static final int NUM_SAMPLES = 10_000; + + @RegisterExtension + static final TestDatabaseExtension dbExtension = new TestDatabaseExtension(); + @RegisterExtension + TestSubspaceExtension rtSubspace = new TestSubspaceExtension(dbExtension); + @RegisterExtension + TestSubspaceExtension rtSecondarySubspace = new TestSubspaceExtension(dbExtension); + + private Database db; + + @BeforeEach + public void setUpDb() { + db = dbExtension.getDatabase(); + } + + @Test + public void testCompactSerialization() { + final Random random = new Random(0); + final CompactStorageAdapter storageAdapter = + new CompactStorageAdapter(HNSW.DEFAULT_CONFIG, CompactNode.factory(), rtSubspace.getSubspace(), + OnWriteListener.NOOP, OnReadListener.NOOP); + final Node originalNode = + db.run(tr -> { + final NodeFactory nodeFactory = storageAdapter.getNodeFactory(); + + final Node randomCompactNode = + createRandomCompactNode(random, nodeFactory, 768, 16); + + writeNode(tr, storageAdapter, randomCompactNode, 0); + return randomCompactNode; + }); + + db.run(tr -> storageAdapter.fetchNode(tr, 0, originalNode.getPrimaryKey()) + .thenAccept(node -> { + Assertions.assertAll( + () -> Assertions.assertInstanceOf(CompactNode.class, node), + () -> Assertions.assertEquals(NodeKind.COMPACT, node.getKind()), + () -> Assertions.assertEquals(node.getPrimaryKey(), originalNode.getPrimaryKey()), + () -> Assertions.assertEquals(node.asCompactNode().getVector(), + originalNode.asCompactNode().getVector()), + () -> { + final ArrayList neighbors = + Lists.newArrayList(node.getNeighbors()); + neighbors.sort(Comparator.comparing(NodeReference::getPrimaryKey)); + final ArrayList originalNeighbors = + Lists.newArrayList(originalNode.getNeighbors()); + originalNeighbors.sort(Comparator.comparing(NodeReference::getPrimaryKey)); + Assertions.assertEquals(neighbors, originalNeighbors); + } + ); + }).join()); + } + + @Test + public void testInliningSerialization() { + final Random random = new Random(0); + final InliningStorageAdapter storageAdapter = + new InliningStorageAdapter(HNSW.DEFAULT_CONFIG, InliningNode.factory(), rtSubspace.getSubspace(), + OnWriteListener.NOOP, OnReadListener.NOOP); + final Node originalNode = + db.run(tr -> { + final NodeFactory nodeFactory = storageAdapter.getNodeFactory(); + + final Node randomInliningNode = + createRandomInliningNode(random, nodeFactory, 768, 16); + + writeNode(tr, storageAdapter, randomInliningNode, 0); + return randomInliningNode; + }); + + db.run(tr -> storageAdapter.fetchNode(tr, 0, originalNode.getPrimaryKey()) + .thenAccept(node -> Assertions.assertAll( + () -> Assertions.assertInstanceOf(InliningNode.class, node), + () -> Assertions.assertEquals(NodeKind.INLINING, node.getKind()), + () -> Assertions.assertEquals(node.getPrimaryKey(), originalNode.getPrimaryKey()), + () -> { + final ArrayList neighbors = + Lists.newArrayList(node.getNeighbors()); + neighbors.sort(Comparator.comparing(NodeReference::getPrimaryKey)); // should not be necessary the way it is stored + final ArrayList originalNeighbors = + Lists.newArrayList(originalNode.getNeighbors()); + originalNeighbors.sort(Comparator.comparing(NodeReference::getPrimaryKey)); + Assertions.assertEquals(neighbors, originalNeighbors); + } + )).join()); + } + + @Test + public void testBasicInsert() { + final Random random = new Random(0); + final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool()); + + db.run(tr -> { + for (int i = 0; i < 10; i ++) { + hnsw.insert(tr, createRandomPrimaryKey(random), createRandomVector(random, 728)).join(); + } + return null; + }); + } + + private void writeNode(@Nonnull final Transaction transaction, + @Nonnull final StorageAdapter storageAdapter, + @Nonnull final Node node, + final int layer) { + final NeighborsChangeSet insertChangeSet = + new InsertNeighborsChangeSet<>(new BaseNeighborsChangeSet<>(ImmutableList.of()), + node.getNeighbors()); + storageAdapter.writeNode(transaction, node, layer, insertChangeSet); + } + + @Nonnull + private Node createRandomCompactNode(@Nonnull final Random random, + @Nonnull final NodeFactory nodeFactory, + final int dimensionality, + final int numberOfNeighbors) { + final Tuple primaryKey = createRandomPrimaryKey(random); + final ImmutableList.Builder neighborsBuilder = ImmutableList.builder(); + for (int i = 0; i < numberOfNeighbors; i ++) { + neighborsBuilder.add(createRandomNodeReference(random)); + } + + return nodeFactory.create(primaryKey, createRandomVector(random, dimensionality), neighborsBuilder.build()); + } + + @Nonnull + private Node createRandomInliningNode(@Nonnull final Random random, + @Nonnull final NodeFactory nodeFactory, + final int dimensionality, + final int numberOfNeighbors) { + final Tuple primaryKey = createRandomPrimaryKey(random); + final ImmutableList.Builder neighborsBuilder = ImmutableList.builder(); + for (int i = 0; i < numberOfNeighbors; i ++) { + neighborsBuilder.add(createRandomNodeReferenceWithVector(random, dimensionality)); + } + + return nodeFactory.create(primaryKey, createRandomVector(random, dimensionality), neighborsBuilder.build()); + } + + @Nonnull + private NodeReference createRandomNodeReference(@Nonnull final Random random) { + return new NodeReference(createRandomPrimaryKey(random)); + } + + @Nonnull + private NodeReferenceWithVector createRandomNodeReferenceWithVector(@Nonnull final Random random, final int dimensionality) { + return new NodeReferenceWithVector(createRandomPrimaryKey(random), createRandomVector(random, dimensionality)); + } + + @Nonnull + private static Tuple createRandomPrimaryKey(final @Nonnull Random random) { + return Tuple.from(random.nextLong()); + } + + @Nonnull + private Vector.HalfVector createRandomVector(@Nonnull final Random random, final int dimensionality) { + final Half[] components = new Half[dimensionality]; + for (int d = 0; d < dimensionality; d ++) { + // don't ask + components[d] = HNSWHelpers.halfValueOf(random.nextDouble()); + } + return new Vector.HalfVector(components); + } +} diff --git a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java index 9a968e161d..f05e2c9c7c 100644 --- a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java +++ b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java @@ -31,6 +31,7 @@ import com.apple.test.Tags; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; +import com.google.protobuf.ByteString; import com.google.protobuf.Message; import org.assertj.core.util.Streams; import org.junit.jupiter.api.Assertions; @@ -75,12 +76,13 @@ protected void openRecordStore(final FDBRecordContext context, final RecordMetaD static Function getRecordGenerator(@Nonnull final Random random) { return recNo -> { - final Iterable vector = - () -> IntStream.range(0, 4000).mapToDouble(index -> random.nextDouble()).iterator(); + final byte[] vector = new byte[768 * 2]; + random.nextBytes(vector); + //logRecord(recNo, vector); return TestRecordsVectorsProto.NaiveVectorRecord.newBuilder() .setRecNo(recNo) - .addAllVectorData(vector) + .setVectorData(ByteString.copyFrom(vector)) .build(); }; } diff --git a/fdb-record-layer-core/src/test/proto/test_records_vector.proto b/fdb-record-layer-core/src/test/proto/test_records_vector.proto index bafe614416..4522e85f6f 100644 --- a/fdb-record-layer-core/src/test/proto/test_records_vector.proto +++ b/fdb-record-layer-core/src/test/proto/test_records_vector.proto @@ -30,7 +30,7 @@ option (schema).store_record_versions = true; message NaiveVectorRecord { optional int64 rec_no = 1 [(field).primary_key = true]; - repeated double vector_data = 2; + optional bytes vector_data = 2; } message RecordTypeUnion { From 687f562085604755aed16647c6d5a3f7e8be523b Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Fri, 1 Aug 2025 12:26:34 +0200 Subject: [PATCH 15/34] added some testing --- .../async/hnsw/CompactStorageAdapter.java | 50 ++++++++-- .../apple/foundationdb/async/hnsw/HNSW.java | 23 +++++ .../async/hnsw/InliningStorageAdapter.java | 93 ++++++++++++++----- .../async/hnsw/StorageAdapter.java | 4 + .../async/hnsw/HNSWModificationTest.java | 17 ++++ 5 files changed, 156 insertions(+), 31 deletions(-) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java index 2ce21f7091..dcca69c1a7 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java @@ -20,9 +20,15 @@ package com.apple.foundationdb.async.hnsw; +import com.apple.foundationdb.KeyValue; +import com.apple.foundationdb.Range; import com.apple.foundationdb.ReadTransaction; +import com.apple.foundationdb.StreamingMode; import com.apple.foundationdb.Transaction; +import com.apple.foundationdb.async.AsyncIterable; +import com.apple.foundationdb.async.AsyncUtil; import com.apple.foundationdb.subspace.Subspace; +import com.apple.foundationdb.tuple.ByteArrayUtil; import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; import com.google.common.base.Verify; @@ -31,6 +37,7 @@ import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import java.util.List; import java.util.concurrent.CompletableFuture; @@ -65,23 +72,28 @@ public StorageAdapter asInliningStorageAdapter() { protected CompletableFuture> fetchNodeInternal(@Nonnull final ReadTransaction readTransaction, final int layer, @Nonnull final Tuple primaryKey) { - final byte[] key = getDataSubspace().pack(Tuple.from(layer, primaryKey)); + final byte[] keyBytes = getDataSubspace().pack(Tuple.from(layer, primaryKey)); - return readTransaction.get(key) + return readTransaction.get(keyBytes) .thenApply(valueBytes -> { if (valueBytes == null) { throw new IllegalStateException("cannot fetch node"); } - - final Tuple nodeTuple = Tuple.fromBytes(valueBytes); - final Node node = nodeFromTuples(primaryKey, nodeTuple); - final OnReadListener onReadListener = getOnReadListener(); - onReadListener.onNodeRead(node); - onReadListener.onKeyValueRead(key, valueBytes); - return node; + return nodeFromRaw(primaryKey, keyBytes, valueBytes); }); } + @Nonnull + private Node nodeFromRaw(final @Nonnull Tuple primaryKey, @Nonnull final byte[] keyBytes, + @Nonnull final byte[] valueBytes) { + final Tuple nodeTuple = Tuple.fromBytes(valueBytes); + final Node node = nodeFromTuples(primaryKey, nodeTuple); + final OnReadListener onReadListener = getOnReadListener(); + onReadListener.onNodeRead(node); + onReadListener.onKeyValueRead(keyBytes, valueBytes); + return node; + } + @Nonnull private Node nodeFromTuples(@Nonnull final Tuple primaryKey, @Nonnull final Tuple valueTuple) { @@ -111,7 +123,6 @@ private Node compactNodeFromTuples(@Nonnull final Tuple primaryKe return getNodeFactory().create(primaryKey, vector, nodeReferences); } - @Override public void writeNodeInternal(@Nonnull final Transaction transaction, @Nonnull final Node node, final int layer, @Nonnull final NeighborsChangeSet neighborsChangeSet) { @@ -139,4 +150,23 @@ public void writeNodeInternal(@Nonnull final Transaction transaction, @Nonnull f transaction.set(key, nodeTuple.pack()); getOnWriteListener().onNodeWritten(layer, node); } + + public Iterable> scanLayer(@Nonnull final ReadTransaction readTransaction, int layer, + @Nullable final Tuple lastPrimaryKey, int maxNumRead) { + final byte[] layerPrefix = getDataSubspace().pack(Tuple.from(layer)); + final Range range = + lastPrimaryKey == null + ? Range.startsWith(layerPrefix) + : new Range(ByteArrayUtil.strinc(getDataSubspace().pack(Tuple.from(layer, lastPrimaryKey))), + ByteArrayUtil.strinc(layerPrefix)); + final AsyncIterable itemsIterable = + readTransaction.getRange(range, maxNumRead, false, StreamingMode.ITERATOR); + + return AsyncUtil.mapIterable(itemsIterable, keyValue -> { + final byte[] key = keyValue.getKey(); + final byte[] value = keyValue.getValue(); + final Tuple primaryKey = getDataSubspace().unpack(key); + return nodeFromRaw(primaryKey, key, value); + }); + } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index b7cb4c3fd0..14f9522c5c 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -20,6 +20,7 @@ package com.apple.foundationdb.async.hnsw; +import com.apple.foundationdb.Database; import com.apple.foundationdb.ReadTransaction; import com.apple.foundationdb.Transaction; import com.apple.foundationdb.annotation.API; @@ -34,6 +35,7 @@ import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.google.common.collect.Sets; +import com.google.common.collect.Streams; import com.google.errorprone.annotations.CanIgnoreReturnValue; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -1054,6 +1056,27 @@ private void writeLonelyNodeOnLayer(@Nonnull final Sto debug(l -> l.debug("written lonely node at key={} on layer={}", primaryKey, layer)); } + public void scanLayer(@Nonnull final Database db, + final int layer, + final int batchSize, + @Nonnull final Consumer> nodeConsumer) { + final StorageAdapter storageAdapter = getStorageAdapterForLayer(layer); + final AtomicReference lastPrimaryKeyAtomic = new AtomicReference<>(); + Tuple newPrimaryKey; + do { + final Tuple lastPrimaryKey = lastPrimaryKeyAtomic.get(); + lastPrimaryKeyAtomic.set(null); + newPrimaryKey = db.run(tr -> { + Streams.stream(storageAdapter.scanLayer(tr, layer, lastPrimaryKey, batchSize)) + .forEach(node -> { + nodeConsumer.accept(node); + lastPrimaryKeyAtomic.set(node.getPrimaryKey()); + }); + return lastPrimaryKeyAtomic.get(); + }, executor); + } while (newPrimaryKey != null); + } + @Nonnull private StorageAdapter getStorageAdapterForLayer(final int layer) { return layer > 0 diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java index 57a8d9b289..b9b4eacc98 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java @@ -25,13 +25,17 @@ import com.apple.foundationdb.ReadTransaction; import com.apple.foundationdb.StreamingMode; import com.apple.foundationdb.Transaction; +import com.apple.foundationdb.async.AsyncIterable; import com.apple.foundationdb.async.AsyncUtil; import com.apple.foundationdb.subspace.Subspace; +import com.apple.foundationdb.tuple.ByteArrayUtil; import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; import com.google.common.collect.ImmutableList; import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.util.List; import java.util.concurrent.CompletableFuture; /** @@ -66,27 +70,35 @@ protected CompletableFuture> fetchNodeInternal(@No return AsyncUtil.collect(readTransaction.getRange(Range.startsWith(rangeKey), ReadTransaction.ROW_LIMIT_UNLIMITED, false, StreamingMode.WANT_ALL), readTransaction.getExecutor()) - .thenApply(keyValues -> { - final OnReadListener onReadListener = getOnReadListener(); - - final ImmutableList.Builder nodeReferencesWithVectorBuilder = ImmutableList.builder(); - for (final KeyValue keyValue : keyValues) { - final byte[] key = keyValue.getKey(); - final byte[] value = keyValue.getValue(); - onReadListener.onKeyValueRead(key, value); - final Tuple neighborKeyTuple = getDataSubspace().unpack(key); - final Tuple neighborValueTuple = Tuple.fromBytes(value); - - final Tuple neighborPrimaryKey = neighborKeyTuple.getNestedTuple(2); // neighbor primary key - final Vector neighborVector = StorageAdapter.vectorFromTuple(neighborValueTuple); // the entire value is the vector - nodeReferencesWithVectorBuilder.add(new NodeReferenceWithVector(neighborPrimaryKey, neighborVector)); - } - - final Node node = - getNodeFactory().create(primaryKey, null, nodeReferencesWithVectorBuilder.build()); - onReadListener.onNodeRead(node); - return node; - }); + .thenApply(keyValues -> nodeFromRaw(primaryKey, keyValues)); + } + + @Nonnull + private Node nodeFromRaw(final @Nonnull Tuple primaryKey, final List keyValues) { + final OnReadListener onReadListener = getOnReadListener(); + + final ImmutableList.Builder nodeReferencesWithVectorBuilder = ImmutableList.builder(); + for (final KeyValue keyValue : keyValues) { + nodeReferencesWithVectorBuilder.add(neighborFromRaw(keyValue.getKey(), keyValue.getValue())); + } + + final Node node = + getNodeFactory().create(primaryKey, null, nodeReferencesWithVectorBuilder.build()); + onReadListener.onNodeRead(node); + return node; + } + + @Nonnull + private NodeReferenceWithVector neighborFromRaw(final @Nonnull byte[] key, final byte[] value) { + final OnReadListener onReadListener = getOnReadListener(); + + onReadListener.onKeyValueRead(key, value); + final Tuple neighborKeyTuple = getDataSubspace().unpack(key); + final Tuple neighborValueTuple = Tuple.fromBytes(value); + + final Tuple neighborPrimaryKey = neighborKeyTuple.getNestedTuple(2); // neighbor primary key + final Vector neighborVector = StorageAdapter.vectorFromTuple(neighborValueTuple); // the entire value is the vector + return new NodeReferenceWithVector(neighborPrimaryKey, neighborVector); } @Override @@ -122,4 +134,43 @@ private byte[] getNeighborKey(final int layer, @Nonnull final Tuple neighborPrimaryKey) { return getDataSubspace().pack(Tuple.from(layer, node.getPrimaryKey(), neighborPrimaryKey)); } + + @Override + public Iterable> scanLayer(@Nonnull final ReadTransaction readTransaction, int layer, + @Nullable final Tuple lastPrimaryKey, int maxNumRead) { + final byte[] layerPrefix = getDataSubspace().pack(Tuple.from(layer)); + final Range range = + lastPrimaryKey == null + ? Range.startsWith(layerPrefix) + : new Range(ByteArrayUtil.strinc(getDataSubspace().pack(Tuple.from(layer, lastPrimaryKey))), + ByteArrayUtil.strinc(layerPrefix)); + final AsyncIterable itemsIterable = + readTransaction.getRange(range, + maxNumRead, false, StreamingMode.ITERATOR); + int numRead = 0; + Tuple nodePrimaryKey = null; + ImmutableList.Builder> nodeBuilder = ImmutableList.builder(); + ImmutableList.Builder neighborsBuilder = ImmutableList.builder(); + for (final KeyValue item: itemsIterable) { + final NodeReferenceWithVector neighbor = + neighborFromRaw(item.getKey(), item.getValue()); + final Tuple primaryKeyFromNodeReference = neighbor.getPrimaryKey(); + if (nodePrimaryKey == null) { + nodePrimaryKey = primaryKeyFromNodeReference; + } else { + if (!nodePrimaryKey.equals(primaryKeyFromNodeReference)) { + nodeBuilder.add(getNodeFactory().create(nodePrimaryKey, null, neighborsBuilder.build())); + } + } + neighborsBuilder.add(neighbor); + numRead ++; + } + + // there may be a rest + if (numRead > 0 && numRead < maxNumRead) { + nodeBuilder.add(getNodeFactory().create(nodePrimaryKey, null, neighborsBuilder.build())); + } + + return nodeBuilder.build(); + } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java index 53e1a5750a..1bdbeef00a 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java @@ -28,6 +28,7 @@ import com.google.common.base.Verify; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import java.util.concurrent.CompletableFuture; /** @@ -91,6 +92,9 @@ CompletableFuture> fetchNode(@Nonnull ReadTransaction readTransaction, void writeNode(@Nonnull Transaction transaction, @Nonnull Node node, int layer, @Nonnull NeighborsChangeSet changeSet); + Iterable> scanLayer(@Nonnull ReadTransaction readTransaction, int layer, @Nullable Tuple lastPrimaryKey, + int maxNumRead); + @Nonnull static CompletableFuture fetchEntryNodeReference(@Nonnull final ReadTransaction readTransaction, @Nonnull final Subspace subspace, diff --git a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java index 6eb7d4a1c9..0b112fd2f0 100644 --- a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java +++ b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java @@ -157,6 +157,23 @@ public void testBasicInsert() { }); } + @Test + public void testBasicInsertAndScanLayer() { + final Random random = new Random(0); + final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool()); + + db.run(tr -> { + for (int i = 0; i < 20; i ++) { + hnsw.insert(tr, createRandomPrimaryKey(random), createRandomVector(random, 728)).join(); + } + return null; + }); + + hnsw.scanLayer(db, 0, 100, node -> { + System.out.println(node); + }); + } + private void writeNode(@Nonnull final Transaction transaction, @Nonnull final StorageAdapter storageAdapter, @Nonnull final Node node, From f1cda1029954aac9dfd2b619fcfd545c75978039 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Fri, 1 Aug 2025 18:04:04 +0200 Subject: [PATCH 16/34] inserts and searches work; logic to dump layers works --- .../async/hnsw/CompactStorageAdapter.java | 2 +- .../apple/foundationdb/async/hnsw/HNSW.java | 7 ++- .../async/hnsw/HNSWModificationTest.java | 62 ++++++++++++++++--- 3 files changed, 60 insertions(+), 11 deletions(-) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java index dcca69c1a7..bc67f63a8b 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java @@ -165,7 +165,7 @@ public Iterable> scanLayer(@Nonnull final ReadTransaction re return AsyncUtil.mapIterable(itemsIterable, keyValue -> { final byte[] key = keyValue.getKey(); final byte[] value = keyValue.getValue(); - final Tuple primaryKey = getDataSubspace().unpack(key); + final Tuple primaryKey = getDataSubspace().unpack(key).getNestedTuple(1); return nodeFromRaw(primaryKey, key, value); }); } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index 14f9522c5c..8d7db00100 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -796,6 +796,8 @@ private CompletableFuture(new BaseNeighborsChangeSet<>(ImmutableList.of()), newNode.getNeighbors()); + storageAdapter.writeNode(transaction, newNode, layer, newNodeChangeSet); + // create change sets for each selected neighbor and insert new node into them final Map> neighborChangeSetMap = Maps.newLinkedHashMap(); @@ -824,7 +826,6 @@ private CompletableFuture { - storageAdapter.writeNode(transaction, newNode, layer, newNodeChangeSet); for (int i = 0; i < selectedNeighbors.size(); i++) { final NodeReferenceAndNode selectedNeighbor = selectedNeighbors.get(i); final NeighborsChangeSet changeSet = changeSets.get(i); @@ -895,6 +896,8 @@ private CompletableFuture if (selectedNeighborNode.getNeighbors().size() < mMax) { return CompletableFuture.completedFuture(null); } else { + debug(l -> l.debug("pruning neighborhood of key={} which has numNeighbors={} out of mMax={}", + selectedNeighborNode.getPrimaryKey(), selectedNeighborNode.getNeighbors().size(), mMax)); return fetchNeighborhood(storageAdapter, transaction, layer, neighborChangeSet.merge(), nodeCache) .thenCompose(nodeReferenceWithVectors -> { final ImmutableList.Builder nodeReferencesWithDistancesBuilder = @@ -1079,7 +1082,7 @@ public void scanLayer(@Nonnull final Database db, @Nonnull private StorageAdapter getStorageAdapterForLayer(final int layer) { - return layer > 0 + return false && layer > 0 ? new InliningStorageAdapter(getConfig(), InliningNode.factory(), getSubspace(), getOnWriteListener(), getOnReadListener()) : new CompactStorageAdapter(getConfig(), CompactNode.factory(), getSubspace(), getOnWriteListener(), getOnReadListener()); } diff --git a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java index 0b112fd2f0..5a2b1c7d62 100644 --- a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java +++ b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java @@ -42,9 +42,13 @@ import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; +import java.io.BufferedWriter; +import java.io.FileWriter; +import java.io.IOException; import java.util.ArrayList; import java.util.Comparator; import java.util.Random; +import java.util.concurrent.atomic.AtomicLong; /** * Tests testing insert/update/deletes of data into/in/from {@link RTree}s. @@ -147,31 +151,68 @@ public void testInliningSerialization() { @Test public void testBasicInsert() { final Random random = new Random(0); + final AtomicLong nextNodeId = new AtomicLong(0L); final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool()); db.run(tr -> { for (int i = 0; i < 10; i ++) { - hnsw.insert(tr, createRandomPrimaryKey(random), createRandomVector(random, 728)).join(); + hnsw.insert(tr, createNextPrimaryKey(nextNodeId), createRandomVector(random, 728)).join(); } return null; }); } @Test - public void testBasicInsertAndScanLayer() { + public void testBasicInsertAndScanLayer() throws Exception { final Random random = new Random(0); - final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool()); + final AtomicLong nextNodeId = new AtomicLong(0L); + final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool(), + HNSW.DEFAULT_CONFIG.toBuilder().setM(4).setMMax(4).setMMax0(10).build(), + OnWriteListener.NOOP, OnReadListener.NOOP); db.run(tr -> { - for (int i = 0; i < 20; i ++) { - hnsw.insert(tr, createRandomPrimaryKey(random), createRandomVector(random, 728)).join(); + for (int i = 0; i < 100; i ++) { + hnsw.insert(tr, createNextPrimaryKey(nextNodeId), createRandomVector(random, 2)).join(); } return null; }); - hnsw.scanLayer(db, 0, 100, node -> { - System.out.println(node); - }); + int layer = 0; + while (true) { + if (!dumpLayer(hnsw, layer++)) { + break; + } + } + } + + private boolean dumpLayer(final HNSW hnsw, final int layer) throws IOException { + final String verticesFileName = "/Users/nseemann/Downloads/vertices-" + layer + ".csv"; + final String edgesFileName = "/Users/nseemann/Downloads/edges-" + layer + ".csv"; + + final AtomicLong numReadAtomic = new AtomicLong(0L); + try (final BufferedWriter verticesWriter = new BufferedWriter(new FileWriter(verticesFileName)); + final BufferedWriter edgesWriter = new BufferedWriter(new FileWriter(edgesFileName))) { + hnsw.scanLayer(db, layer, 100, node -> { + final CompactNode compactNode = node.asCompactNode(); + final Vector vector = compactNode.getVector(); + try { + verticesWriter.write(compactNode.getPrimaryKey().getLong(0) + "," + + vector.getComponent(0) + "," + + vector.getComponent(1)); + verticesWriter.newLine(); + + for (final var neighbor : compactNode.getNeighbors()) { + edgesWriter.write(compactNode.getPrimaryKey().getLong(0) + "," + + neighbor.getPrimaryKey().getLong(0)); + edgesWriter.newLine(); + } + numReadAtomic.getAndIncrement(); + } catch (final IOException e) { + throw new RuntimeException("unable to write to file", e); + } + }); + } + return numReadAtomic.get() != 0; } private void writeNode(@Nonnull final Transaction transaction, @@ -227,6 +268,11 @@ private static Tuple createRandomPrimaryKey(final @Nonnull Random random) { return Tuple.from(random.nextLong()); } + @Nonnull + private static Tuple createNextPrimaryKey(@Nonnull final AtomicLong nextIdAtomic) { + return Tuple.from(nextIdAtomic.getAndIncrement()); + } + @Nonnull private Vector.HalfVector createRandomVector(@Nonnull final Random random, final int dimensionality) { final Half[] components = new Half[dimensionality]; From 82322709cdf1a7942f5a654ad8d07a6bf9854312 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Sat, 2 Aug 2025 23:08:37 +0200 Subject: [PATCH 17/34] doing some performance tests --- .../async/hnsw/CompactStorageAdapter.java | 23 ++-- .../apple/foundationdb/async/hnsw/HNSW.java | 38 +++++-- .../async/hnsw/InliningStorageAdapter.java | 14 +-- .../async/hnsw/OnReadListener.java | 9 +- .../async/hnsw/StorageAdapter.java | 2 +- .../async/hnsw/HNSWModificationTest.java | 104 +++++++++++++++++- gradle/scripts/log4j-test.properties | 2 +- 7 files changed, 150 insertions(+), 42 deletions(-) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java index bc67f63a8b..f88b679d60 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java @@ -79,18 +79,18 @@ protected CompletableFuture> fetchNodeInternal(@Nonnull fina if (valueBytes == null) { throw new IllegalStateException("cannot fetch node"); } - return nodeFromRaw(primaryKey, keyBytes, valueBytes); + return nodeFromRaw(layer, primaryKey, keyBytes, valueBytes); }); } @Nonnull - private Node nodeFromRaw(final @Nonnull Tuple primaryKey, @Nonnull final byte[] keyBytes, - @Nonnull final byte[] valueBytes) { + private Node nodeFromRaw(final int layer, final @Nonnull Tuple primaryKey, + @Nonnull final byte[] keyBytes, @Nonnull final byte[] valueBytes) { final Tuple nodeTuple = Tuple.fromBytes(valueBytes); final Node node = nodeFromTuples(primaryKey, nodeTuple); final OnReadListener onReadListener = getOnReadListener(); - onReadListener.onNodeRead(node); - onReadListener.onKeyValueRead(keyBytes, valueBytes); + onReadListener.onNodeRead(layer, node); + onReadListener.onKeyValueRead(layer, keyBytes, valueBytes); return node; } @@ -139,16 +139,17 @@ public void writeNodeInternal(@Nonnull final Transaction transaction, @Nonnull f for (final NodeReference neighborReference : neighbors) { neighborItems.add(neighborReference.getPrimaryKey()); } - if (logger.isDebugEnabled()) { - logger.debug("written neighbors of primaryKey={}, oldSize={}, newSize={}", node.getPrimaryKey(), - node.getNeighbors().size(), neighborItems.size()); - } - nodeItems.add(Tuple.fromList(neighborItems)); + final Tuple nodeTuple = Tuple.fromList(nodeItems); transaction.set(key, nodeTuple.pack()); getOnWriteListener().onNodeWritten(layer, node); + + if (logger.isDebugEnabled()) { + logger.debug("written neighbors of primaryKey={}, oldSize={}, newSize={}", node.getPrimaryKey(), + node.getNeighbors().size(), neighborItems.size()); + } } public Iterable> scanLayer(@Nonnull final ReadTransaction readTransaction, int layer, @@ -166,7 +167,7 @@ public Iterable> scanLayer(@Nonnull final ReadTransaction re final byte[] key = keyValue.getKey(); final byte[] value = keyValue.getValue(); final Tuple primaryKey = getDataSubspace().unpack(key).getNestedTuple(1); - return nodeFromRaw(primaryKey, key, value); + return nodeFromRaw(layer, primaryKey, key, value); }); } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index 8d7db00100..07504aa52b 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -41,7 +41,9 @@ import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; +import java.util.ArrayList; import java.util.Collection; +import java.util.Collections; import java.util.Comparator; import java.util.List; import java.util.Map; @@ -74,7 +76,7 @@ public class HNSW { public static final int DEFAULT_M = 16; public static final int DEFAULT_M_MAX = DEFAULT_M; public static final int DEFAULT_M_MAX_0 = 2 * DEFAULT_M; - public static final int DEFAULT_EF_SEARCH = 64; + public static final int DEFAULT_EF_SEARCH = 100; public static final int DEFAULT_EF_CONSTRUCTION = 200; public static final boolean DEFAULT_EXTEND_CANDIDATES = false; public static final boolean DEFAULT_KEEP_PRUNED_CONNECTIONS = false; @@ -403,6 +405,7 @@ public OnReadListener getOnReadListener() { @SuppressWarnings("checkstyle:MethodName") // method name introduced by paper @Nonnull public CompletableFuture>> kNearestNeighborsSearch(@Nonnull final ReadTransaction readTransaction, + final int k, final int efSearch, @Nonnull final Vector queryVector) { return StorageAdapter.fetchEntryNodeReference(readTransaction, getSubspace(), getOnReadListener()) @@ -450,7 +453,17 @@ public CompletableFuture { + // reverse the original deque + final int size = searchResult.size(); + final int start = Math.max(0, size - k); + + final ArrayList> topKReversed = + Lists.newArrayList(searchResult.subList(start, size)); + Collections.reverse(topKReversed); + return topKReversed; + }); }); } @@ -579,14 +592,13 @@ private CompletableFuture }).thenCompose(ignored -> fetchSomeNodesIfNotCached(storageAdapter, readTransaction, layer, nearestNeighbors, nodeCache)) .thenApply(searchResult -> { - debug(l -> { - l.debug("searched layer={} for efSearch={} with result=={}", layer, efSearch, - searchResult.stream() - .map(nodeReferenceAndNode -> - "(primaryKey=" + nodeReferenceAndNode.getNodeReferenceWithDistance().getPrimaryKey() + - ",distance=" + nodeReferenceAndNode.getNodeReferenceWithDistance().getDistance() + ")") - .collect(Collectors.joining(","))); - }); + debug(l -> + l.debug("searched layer={} for efSearch={} with result=={}", layer, efSearch, + searchResult.stream() + .map(nodeReferenceAndNode -> + "(primaryKey=" + nodeReferenceAndNode.getNodeReferenceWithDistance().getPrimaryKey() + + ",distance=" + nodeReferenceAndNode.getNodeReferenceWithDistance().getDistance() + ")") + .collect(Collectors.joining(",")))); return searchResult; }); } @@ -1093,6 +1105,12 @@ private int insertionLayer(@Nonnull final Random random) { return (int) Math.floor(-Math.log(u) * lambda); } + private void info(@Nonnull final Consumer loggerConsumer) { + if (logger.isInfoEnabled()) { + loggerConsumer.accept(logger); + } + } + private void debug(@Nonnull final Consumer loggerConsumer) { if (logger.isDebugEnabled()) { loggerConsumer.accept(logger); diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java index b9b4eacc98..d025116a28 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java @@ -70,29 +70,29 @@ protected CompletableFuture> fetchNodeInternal(@No return AsyncUtil.collect(readTransaction.getRange(Range.startsWith(rangeKey), ReadTransaction.ROW_LIMIT_UNLIMITED, false, StreamingMode.WANT_ALL), readTransaction.getExecutor()) - .thenApply(keyValues -> nodeFromRaw(primaryKey, keyValues)); + .thenApply(keyValues -> nodeFromRaw(layer, primaryKey, keyValues)); } @Nonnull - private Node nodeFromRaw(final @Nonnull Tuple primaryKey, final List keyValues) { + private Node nodeFromRaw(final int layer, final @Nonnull Tuple primaryKey, final List keyValues) { final OnReadListener onReadListener = getOnReadListener(); final ImmutableList.Builder nodeReferencesWithVectorBuilder = ImmutableList.builder(); for (final KeyValue keyValue : keyValues) { - nodeReferencesWithVectorBuilder.add(neighborFromRaw(keyValue.getKey(), keyValue.getValue())); + nodeReferencesWithVectorBuilder.add(neighborFromRaw(layer, keyValue.getKey(), keyValue.getValue())); } final Node node = getNodeFactory().create(primaryKey, null, nodeReferencesWithVectorBuilder.build()); - onReadListener.onNodeRead(node); + onReadListener.onNodeRead(layer, node); return node; } @Nonnull - private NodeReferenceWithVector neighborFromRaw(final @Nonnull byte[] key, final byte[] value) { + private NodeReferenceWithVector neighborFromRaw(final int layer, final @Nonnull byte[] key, final byte[] value) { final OnReadListener onReadListener = getOnReadListener(); - onReadListener.onKeyValueRead(key, value); + onReadListener.onKeyValueRead(layer, key, value); final Tuple neighborKeyTuple = getDataSubspace().unpack(key); final Tuple neighborValueTuple = Tuple.fromBytes(value); @@ -153,7 +153,7 @@ public Iterable> scanLayer(@Nonnull final ReadTran ImmutableList.Builder neighborsBuilder = ImmutableList.builder(); for (final KeyValue item: itemsIterable) { final NodeReferenceWithVector neighbor = - neighborFromRaw(item.getKey(), item.getValue()); + neighborFromRaw(layer, item.getKey(), item.getValue()); final Tuple primaryKeyFromNodeReference = neighbor.getPrimaryKey(); if (nodePrimaryKey == null) { nodePrimaryKey = primaryKeyFromNodeReference; diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java index beb8530cd6..753648cf77 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnReadListener.java @@ -30,19 +30,16 @@ public interface OnReadListener { OnReadListener NOOP = new OnReadListener() { }; - default void onSlotIndexEntryRead(@Nonnull final byte[] key) { - // nothing - } - default CompletableFuture> onAsyncRead(@Nonnull CompletableFuture> future) { return future; } - default void onNodeRead(@Nonnull Node node) { + default void onNodeRead(int layer, @Nonnull Node node) { // nothing } - default void onKeyValueRead(@Nonnull byte[] key, + default void onKeyValueRead(int layer, + @Nonnull byte[] key, @Nonnull byte[] value) { // nothing } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java index 1bdbeef00a..5909e4ff7b 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/StorageAdapter.java @@ -107,7 +107,7 @@ static CompletableFuture fetchEntryNodeReference(@Nonnull fi if (valueBytes == null) { return null; // not a single node in the index } - onReadListener.onKeyValueRead(key, valueBytes); + onReadListener.onKeyValueRead(-1, key, valueBytes); final Tuple entryTuple = Tuple.fromBytes(valueBytes); final int lMax = (int)entryTuple.getLong(0); diff --git a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java index 5a2b1c7d62..98e2035ce5 100644 --- a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java +++ b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java @@ -30,6 +30,7 @@ import com.apple.test.Tags; import com.christianheina.langx.half4j.Half; import com.google.common.collect.ImmutableList; +import com.google.common.collect.Maps; import org.assertj.core.util.Lists; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; @@ -47,7 +48,10 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Comparator; +import java.util.List; +import java.util.Map; import java.util.Random; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; /** @@ -151,14 +155,48 @@ public void testInliningSerialization() { @Test public void testBasicInsert() { final Random random = new Random(0); - final AtomicLong nextNodeId = new AtomicLong(0L); - final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool()); + final AtomicLong nextNodeIdAtomic = new AtomicLong(0L); - db.run(tr -> { - for (int i = 0; i < 10; i ++) { - hnsw.insert(tr, createNextPrimaryKey(nextNodeId), createRandomVector(random, 728)).join(); + final TestOnReadListener onReadListener = new TestOnReadListener(); + + final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool(), + HNSW.DEFAULT_CONFIG.toBuilder().setMetric(Metric.COSINE_METRIC).setEfConstruction(34).setM(16).setMMax(16).setMMax0(32).build(), + OnWriteListener.NOOP, onReadListener); + + for (int i = 0; i < 10000;) { + i += basicInsertBatch(hnsw, random, 100, nextNodeIdAtomic, onReadListener); + } + + onReadListener.reset(); + final long beginTs = System.nanoTime(); + final List> result = + db.run(tr -> hnsw.kNearestNeighborsSearch(tr, 10, 20, createRandomVector(random, 768)).join()); + final long endTs = System.nanoTime(); + + for (NodeReferenceAndNode nodeReferenceAndNode : result) { + final NodeReferenceWithDistance nodeReferenceWithDistance = nodeReferenceAndNode.getNodeReferenceWithDistance(); + logger.info("nodeId ={} at distance={}", nodeReferenceWithDistance.getPrimaryKey().getLong(0), + nodeReferenceWithDistance.getDistance()); + } + System.out.println(onReadListener.getNodeCountByLayer()); + System.out.println(onReadListener.getBytesReadByLayer()); + + logger.info("search transaction took elapsedTime={}ms", TimeUnit.NANOSECONDS.toMillis(endTs - beginTs)); + } + + private int basicInsertBatch(@Nonnull final HNSW hnsw, @Nonnull final Random random, final int batchSize, + @Nonnull final AtomicLong nextNodeIdAtomic, @Nonnull final TestOnReadListener onReadListener) { + return db.run(tr -> { + onReadListener.reset(); + final long nextNodeId = nextNodeIdAtomic.get(); + final long beginTs = System.nanoTime(); + for (int i = 0; i < batchSize; i ++) { + hnsw.insert(tr, createNextPrimaryKey(nextNodeIdAtomic), createRandomVector(random, 768)).join(); } - return null; + final long endTs = System.nanoTime(); + logger.info("inserted batchSize={} records starting at nodeId={} took elapsedTime={}ms, readCounts={}, MSums={}", batchSize, nextNodeId, + TimeUnit.NANOSECONDS.toMillis(endTs - beginTs), onReadListener.getNodeCountByLayer(), onReadListener.getSumMByLayer()); + return batchSize; }); } @@ -185,6 +223,18 @@ public void testBasicInsertAndScanLayer() throws Exception { } } + @Test + public void testManyVectors() { + final Random random = new Random(); + for (long l = 0L; l < 3000000; l ++) { + final Vector.HalfVector randomVector = createRandomVector(random, 768); + final Tuple vectorTuple = StorageAdapter.tupleFromVector(randomVector); + final Vector roundTripVector = StorageAdapter.vectorFromTuple(vectorTuple); + Vector.comparativeDistance(Metric.EuclideanMetric.EUCLIDEAN_METRIC, randomVector, roundTripVector); + Assertions.assertEquals(randomVector, roundTripVector); + } + } + private boolean dumpLayer(final HNSW hnsw, final int layer) throws IOException { final String verticesFileName = "/Users/nseemann/Downloads/vertices-" + layer + ".csv"; final String edgesFileName = "/Users/nseemann/Downloads/edges-" + layer + ".csv"; @@ -282,4 +332,46 @@ private Vector.HalfVector createRandomVector(@Nonnull final Random random, final } return new Vector.HalfVector(components); } + + private static class TestOnReadListener implements OnReadListener { + final Map nodeCountByLayer; + final Map sumMByLayer; + final Map bytesReadByLayer; + + public TestOnReadListener() { + this.nodeCountByLayer = Maps.newConcurrentMap(); + this.sumMByLayer = Maps.newConcurrentMap(); + this.bytesReadByLayer = Maps.newConcurrentMap(); + } + + public Map getNodeCountByLayer() { + return nodeCountByLayer; + } + + public Map getBytesReadByLayer() { + return bytesReadByLayer; + } + + public Map getSumMByLayer() { + return sumMByLayer; + } + + public void reset() { + nodeCountByLayer.clear(); + bytesReadByLayer.clear(); + sumMByLayer.clear(); + } + + @Override + public void onNodeRead(final int layer, @Nonnull final Node node) { + nodeCountByLayer.compute(layer, (l, oldValue) -> (oldValue == null ? 0 : oldValue) + 1L); + sumMByLayer.compute(layer, (l, oldValue) -> (oldValue == null ? 0 : oldValue) + node.getNeighbors().size()); + } + + @Override + public void onKeyValueRead(final int layer, @Nonnull final byte[] key, @Nonnull final byte[] value) { + bytesReadByLayer.compute(layer, (l, oldValue) -> (oldValue == null ? 0 : oldValue) + + key.length + value.length); + } + } } diff --git a/gradle/scripts/log4j-test.properties b/gradle/scripts/log4j-test.properties index 447ee2f55a..1ae7583751 100644 --- a/gradle/scripts/log4j-test.properties +++ b/gradle/scripts/log4j-test.properties @@ -26,7 +26,7 @@ appender.console.name = STDOUT appender.console.layout.type = PatternLayout appender.console.layout.pattern = %d [%level] %logger{1.} - %m %X%n%ex{full} -rootLogger.level = debug +rootLogger.level = info rootLogger.appenderRefs = stdout rootLogger.appenderRef.stdout.ref = STDOUT From 2d88cc90352371fecc59a2540ecadb3a1b0b8fab Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Tue, 5 Aug 2025 09:48:49 +0200 Subject: [PATCH 18/34] better test helpers --- .../apple/foundationdb/async/hnsw/HNSW.java | 41 ++-- .../async/hnsw/HNSWModificationTest.java | 198 +++++++++++++++++- 2 files changed, 211 insertions(+), 28 deletions(-) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index 07504aa52b..ba5365759f 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -41,9 +41,7 @@ import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; -import java.util.ArrayList; import java.util.Collection; -import java.util.Collections; import java.util.Comparator; import java.util.List; import java.util.Map; @@ -51,6 +49,7 @@ import java.util.Queue; import java.util.Random; import java.util.Set; +import java.util.TreeSet; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; import java.util.concurrent.PriorityBlockingQueue; @@ -455,14 +454,23 @@ public CompletableFuture { - // reverse the original deque - final int size = searchResult.size(); - final int start = Math.max(0, size - k); - - final ArrayList> topKReversed = - Lists.newArrayList(searchResult.subList(start, size)); - Collections.reverse(topKReversed); - return topKReversed; + // reverse the original queue + final TreeSet> sortedTopK = + new TreeSet<>( + Comparator.comparing(nodeReferenceAndNode -> + nodeReferenceAndNode.getNodeReferenceWithDistance().getDistance())); + + for (final NodeReferenceAndNode nodeReferenceAndNode : searchResult) { + if (sortedTopK.size() < k || sortedTopK.last().getNodeReferenceWithDistance().getDistance() > + nodeReferenceAndNode.getNodeReferenceWithDistance().getDistance()) { + sortedTopK.add(nodeReferenceAndNode); + } + + if (sortedTopK.size() > k) { + sortedTopK.remove(sortedTopK.last()); + } + } + return ImmutableList.copyOf(sortedTopK); }); }); } @@ -592,13 +600,12 @@ private CompletableFuture }).thenCompose(ignored -> fetchSomeNodesIfNotCached(storageAdapter, readTransaction, layer, nearestNeighbors, nodeCache)) .thenApply(searchResult -> { - debug(l -> - l.debug("searched layer={} for efSearch={} with result=={}", layer, efSearch, - searchResult.stream() - .map(nodeReferenceAndNode -> - "(primaryKey=" + nodeReferenceAndNode.getNodeReferenceWithDistance().getPrimaryKey() + - ",distance=" + nodeReferenceAndNode.getNodeReferenceWithDistance().getDistance() + ")") - .collect(Collectors.joining(",")))); + debug(l -> l.debug("searched layer={} for efSearch={} with result=={}", layer, efSearch, + searchResult.stream() + .map(nodeReferenceAndNode -> + "(primaryKey=" + nodeReferenceAndNode.getNodeReferenceWithDistance().getPrimaryKey() + + ",distance=" + nodeReferenceAndNode.getNodeReferenceWithDistance().getDistance() + ")") + .collect(Collectors.joining(",")))); return searchResult; }); } diff --git a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java index 98e2035ce5..1a63dc1ebe 100644 --- a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java +++ b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java @@ -22,6 +22,8 @@ import com.apple.foundationdb.Database; import com.apple.foundationdb.Transaction; +import com.apple.foundationdb.async.AsyncUtil; +import com.apple.foundationdb.async.hnsw.Vector.HalfVector; import com.apple.foundationdb.async.rtree.RTree; import com.apple.foundationdb.test.TestDatabaseExtension; import com.apple.foundationdb.test.TestExecutors; @@ -36,23 +38,34 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.api.parallel.Execution; import org.junit.jupiter.api.parallel.ExecutionMode; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; +import java.io.BufferedReader; import java.io.BufferedWriter; +import java.io.FileReader; import java.io.FileWriter; import java.io.IOException; import java.util.ArrayList; import java.util.Comparator; import java.util.List; import java.util.Map; +import java.util.NavigableSet; +import java.util.Objects; import java.util.Random; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentSkipListSet; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; /** * Tests testing insert/update/deletes of data into/in/from {@link RTree}s. @@ -159,18 +172,20 @@ public void testBasicInsert() { final TestOnReadListener onReadListener = new TestOnReadListener(); + final int dimensions = 128; final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool(), - HNSW.DEFAULT_CONFIG.toBuilder().setMetric(Metric.COSINE_METRIC).setEfConstruction(34).setM(16).setMMax(16).setMMax0(32).build(), + HNSW.DEFAULT_CONFIG.toBuilder().setMetric(Metric.EUCLIDEAN_METRIC).setM(32).setMMax(32).setMMax0(64).build(), OnWriteListener.NOOP, onReadListener); - for (int i = 0; i < 10000;) { - i += basicInsertBatch(hnsw, random, 100, nextNodeIdAtomic, onReadListener); + for (int i = 0; i < 1000;) { + i += basicInsertBatch(100, nextNodeIdAtomic, onReadListener, + tr -> hnsw.insert(tr, createNextPrimaryKey(nextNodeIdAtomic), createRandomVector(random, dimensions))); } onReadListener.reset(); final long beginTs = System.nanoTime(); final List> result = - db.run(tr -> hnsw.kNearestNeighborsSearch(tr, 10, 20, createRandomVector(random, 768)).join()); + db.run(tr -> hnsw.kNearestNeighborsSearch(tr, 10, 100, createRandomVector(random, dimensions)).join()); final long endTs = System.nanoTime(); for (NodeReferenceAndNode nodeReferenceAndNode : result) { @@ -184,14 +199,15 @@ public void testBasicInsert() { logger.info("search transaction took elapsedTime={}ms", TimeUnit.NANOSECONDS.toMillis(endTs - beginTs)); } - private int basicInsertBatch(@Nonnull final HNSW hnsw, @Nonnull final Random random, final int batchSize, - @Nonnull final AtomicLong nextNodeIdAtomic, @Nonnull final TestOnReadListener onReadListener) { + private int basicInsertBatch(final int batchSize, + @Nonnull final AtomicLong nextNodeIdAtomic, @Nonnull final TestOnReadListener onReadListener, + @Nonnull final Function> insertFunction) { return db.run(tr -> { onReadListener.reset(); final long nextNodeId = nextNodeIdAtomic.get(); final long beginTs = System.nanoTime(); for (int i = 0; i < batchSize; i ++) { - hnsw.insert(tr, createNextPrimaryKey(nextNodeIdAtomic), createRandomVector(random, 768)).join(); + insertFunction.apply(tr).join(); } final long endTs = System.nanoTime(); logger.info("inserted batchSize={} records starting at nodeId={} took elapsedTime={}ms, readCounts={}, MSums={}", batchSize, nextNodeId, @@ -200,6 +216,91 @@ private int basicInsertBatch(@Nonnull final HNSW hnsw, @Nonnull final Random ran }); } + @Test + @Timeout(value = 150, unit = TimeUnit.MINUTES) + public void testSIFTInsert10k() throws Exception { + final Metric metric = Metric.EUCLIDEAN_METRIC; + final int k = 10; + final AtomicLong nextNodeIdAtomic = new AtomicLong(0L); + + final TestOnReadListener onReadListener = new TestOnReadListener(); + + final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool(), + HNSW.DEFAULT_CONFIG.toBuilder().setMetric(metric).setM(32).setMMax(32).setMMax0(64).build(), + OnWriteListener.NOOP, onReadListener); + + final String tsvFile = "/Users/nseemann/Downloads/train-100k.tsv"; + final int dimensions = 128; + + final AtomicReference queryVectorAtomic = new AtomicReference<>(); + final NavigableSet trueResults = new ConcurrentSkipListSet<>( + Comparator.comparing(NodeReferenceWithDistance::getDistance)); + + try (BufferedReader br = new BufferedReader(new FileReader(tsvFile))) { + for (int i = 0; i < 10000;) { + i += basicInsertBatch(100, nextNodeIdAtomic, onReadListener, + tr -> { + final String line; + try { + line = br.readLine(); + } catch (IOException e) { + throw new RuntimeException(e); + } + + final String[] values = Objects.requireNonNull(line).split("\t"); + Assertions.assertEquals(dimensions, values.length); + final Half[] halfs = new Half[dimensions]; + + for (int c = 0; c < values.length; c++) { + final String value = values[c]; + halfs[c] = HNSWHelpers.halfValueOf(Double.parseDouble(value)); + } + final Tuple currentPrimaryKey = createNextPrimaryKey(nextNodeIdAtomic); + final HalfVector currentVector = new HalfVector(halfs); + final HalfVector queryVector = queryVectorAtomic.get(); + if (queryVector == null) { + queryVectorAtomic.set(currentVector); + return AsyncUtil.DONE; + } else { + final double currentDistance = + Vector.comparativeDistance(metric, currentVector, queryVector); + if (trueResults.size() < k || trueResults.last().getDistance() > currentDistance) { + trueResults.add( + new NodeReferenceWithDistance(currentPrimaryKey, currentVector, + Vector.comparativeDistance(metric, currentVector, queryVector))); + } + if (trueResults.size() > k) { + trueResults.remove(trueResults.last()); + } + return hnsw.insert(tr, currentPrimaryKey, currentVector); + } + }); + } + } + + onReadListener.reset(); + final long beginTs = System.nanoTime(); + final List> results = + db.run(tr -> hnsw.kNearestNeighborsSearch(tr, k, 100, queryVectorAtomic.get()).join()); + final long endTs = System.nanoTime(); + + for (NodeReferenceAndNode nodeReferenceAndNode : results) { + final NodeReferenceWithDistance nodeReferenceWithDistance = nodeReferenceAndNode.getNodeReferenceWithDistance(); + logger.info("retrieved result nodeId = {} at distance= {}", nodeReferenceWithDistance.getPrimaryKey().getLong(0), + nodeReferenceWithDistance.getDistance()); + } + + for (final NodeReferenceWithDistance nodeReferenceWithDistance : trueResults) { + logger.info("true result nodeId ={} at distance={}", nodeReferenceWithDistance.getPrimaryKey().getLong(0), + nodeReferenceWithDistance.getDistance()); + } + + System.out.println(onReadListener.getNodeCountByLayer()); + System.out.println(onReadListener.getBytesReadByLayer()); + + logger.info("search transaction took elapsedTime={}ms", TimeUnit.NANOSECONDS.toMillis(endTs - beginTs)); + } + @Test public void testBasicInsertAndScanLayer() throws Exception { final Random random = new Random(0); @@ -224,10 +325,10 @@ public void testBasicInsertAndScanLayer() throws Exception { } @Test - public void testManyVectors() { + public void testManyRandomVectors() { final Random random = new Random(); for (long l = 0L; l < 3000000; l ++) { - final Vector.HalfVector randomVector = createRandomVector(random, 768); + final HalfVector randomVector = createRandomVector(random, 768); final Tuple vectorTuple = StorageAdapter.tupleFromVector(randomVector); final Vector roundTripVector = StorageAdapter.vectorFromTuple(vectorTuple); Vector.comparativeDistance(Metric.EuclideanMetric.EUCLIDEAN_METRIC, randomVector, roundTripVector); @@ -235,6 +336,81 @@ public void testManyVectors() { } } + @Test + @Timeout(value = 150, unit = TimeUnit.MINUTES) + public void testSIFTVectors() throws Exception { + final AtomicLong nextNodeIdAtomic = new AtomicLong(0L); + + final TestOnReadListener onReadListener = new TestOnReadListener(); + + final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool(), + HNSW.DEFAULT_CONFIG.toBuilder().setMetric(Metric.EUCLIDEAN_METRIC).setM(32).setMMax(32).setMMax0(64).build(), + OnWriteListener.NOOP, onReadListener); + + + final String tsvFile = "/Users/nseemann/Downloads/train-100k.tsv"; + final int dimensions = 128; + final var referenceVector = createRandomVector(new Random(0), dimensions); + long count = 0L; + double mean = 0.0d; + double mean2 = 0.0d; + + try (BufferedReader br = new BufferedReader(new FileReader(tsvFile))) { + for (int i = 0; i < 100_000; i ++) { + final String line; + try { + line = br.readLine(); + } catch (IOException e) { + throw new RuntimeException(e); + } + + final String[] values = Objects.requireNonNull(line).split("\t"); + Assertions.assertEquals(dimensions, values.length); + final Half[] halfs = new Half[dimensions]; + for (int c = 0; c < values.length; c++) { + final String value = values[c]; + halfs[c] = HNSWHelpers.halfValueOf(Double.parseDouble(value)); + } + final HalfVector newVector = new HalfVector(halfs); + final double distance = Vector.comparativeDistance(Metric.EUCLIDEAN_METRIC, referenceVector, newVector); + count++; + final double delta = distance - mean; + mean += delta / count; + final double delta2 = distance - mean; + mean2 += delta * delta2; + } + } + final double sampleVariance = mean2 / (count - 1); + final double standardDeviation = Math.sqrt(sampleVariance); + logger.info("mean={}, sample_variance={}, stddeviation={}, cv={}", mean, sampleVariance, standardDeviation, + standardDeviation / mean); + } + + + @ParameterizedTest + @ValueSource(ints = {2, 3, 10, 100, 768}) + public void testManyVectorsStandardDeviation(final int dimensionality) { + final Random random = new Random(); + final Metric metric = Metric.EuclideanMetric.EUCLIDEAN_METRIC; + long count = 0L; + double mean = 0.0d; + double mean2 = 0.0d; + for (long i = 0L; i < 100000; i ++) { + final HalfVector vector1 = createRandomVector(random, dimensionality); + final HalfVector vector2 = createRandomVector(random, dimensionality); + final double distance = Vector.comparativeDistance(metric, vector1, vector2); + count = i + 1; + final double delta = distance - mean; + mean += delta / count; + final double delta2 = distance - mean; + mean2 += delta * delta2; + } + final double sampleVariance = mean2 / (count - 1); + final double standardDeviation = Math.sqrt(sampleVariance); + logger.info("mean={}, sample_variance={}, stddeviation={}, cv={}", mean, sampleVariance, standardDeviation, + standardDeviation / mean); + } + private boolean dumpLayer(final HNSW hnsw, final int layer) throws IOException { final String verticesFileName = "/Users/nseemann/Downloads/vertices-" + layer + ".csv"; final String edgesFileName = "/Users/nseemann/Downloads/edges-" + layer + ".csv"; @@ -324,13 +500,13 @@ private static Tuple createNextPrimaryKey(@Nonnull final AtomicLong nextIdAtomic } @Nonnull - private Vector.HalfVector createRandomVector(@Nonnull final Random random, final int dimensionality) { + private HalfVector createRandomVector(@Nonnull final Random random, final int dimensionality) { final Half[] components = new Half[dimensionality]; for (int d = 0; d < dimensionality; d ++) { // don't ask components[d] = HNSWHelpers.halfValueOf(random.nextDouble()); } - return new Vector.HalfVector(components); + return new HalfVector(components); } private static class TestOnReadListener implements OnReadListener { From 71c583b91be2f0241c329449ed52d8e0305bebc3 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Tue, 5 Aug 2025 10:02:04 +0200 Subject: [PATCH 19/34] fix for bad bug that didn't wait for a future --- .../foundationdb/async/MoreAsyncUtil.java | 19 ++-- .../apple/foundationdb/async/hnsw/HNSW.java | 93 ++++++++----------- .../async/hnsw/HNSWModificationTest.java | 2 +- 3 files changed, 52 insertions(+), 62 deletions(-) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java index 85ea6ad045..a49e1fb7a0 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java @@ -46,10 +46,10 @@ import java.util.concurrent.ThreadFactory; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiConsumer; import java.util.function.BiFunction; import java.util.function.Function; -import java.util.function.IntFunction; import java.util.function.IntPredicate; import java.util.function.IntUnaryOperator; import java.util.function.Predicate; @@ -1058,22 +1058,25 @@ public static CompletableFuture swallowException(@Nonnull CompletableFutur return result; } - public static CompletableFuture forLoop(final int startI, @Nonnull final IntPredicate conditionPredicate, - @Nonnull final IntUnaryOperator stepFunction, - @Nonnull final IntFunction> body, - @Nonnull final Executor executor) { + public static CompletableFuture forLoop(final int startI, @Nullable final U startU, + @Nonnull final IntPredicate conditionPredicate, + @Nonnull final IntUnaryOperator stepFunction, + @Nonnull final BiFunction> body, + @Nonnull final Executor executor) { final AtomicInteger loopVariableAtomic = new AtomicInteger(startI); + final AtomicReference lastResultAtomic = new AtomicReference<>(startU); return AsyncUtil.whileTrue(() -> { final int loopVariable = loopVariableAtomic.get(); if (!conditionPredicate.test(loopVariable)) { return AsyncUtil.READY_FALSE; } - return body.apply(loopVariable) - .thenApply(ignored -> { + return body.apply(loopVariable, lastResultAtomic.get()) + .thenApply(result -> { loopVariableAtomic.set(stepFunction.applyAsInt(loopVariable)); + lastResultAtomic.set(result); return true; }); - }, executor); + }, executor).thenApply(ignored -> lastResultAtomic.get()); } @SuppressWarnings("unchecked") diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index ba5365759f..4859c7f2e9 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -426,23 +426,14 @@ public CompletableFuture nodeReferenceAtomic = - new AtomicReference<>(entryState); - - return MoreAsyncUtil.forLoop(entryLayer, - layer -> layer > 0, - layer -> layer - 1, - layer -> { - final var storageAdapter = getStorageAdapterForLayer(layer); - final var greedyIn = nodeReferenceAtomic.get(); - return greedySearchLayer(storageAdapter, readTransaction, greedyIn, layer, - queryVector) - .thenApply(greedyState -> { - nodeReferenceAtomic.set(greedyState); - return null; - }); - }, executor) - .thenApply(ignored -> nodeReferenceAtomic.get()); + return MoreAsyncUtil.forLoop(entryLayer, entryState, + layer -> layer > 0, + layer -> layer - 1, + (layer, previousNodeReference) -> { + final var storageAdapter = getStorageAdapterForLayer(layer); + return greedySearchLayer(storageAdapter, readTransaction, previousNodeReference, + layer, queryVector); + }, executor); }).thenCompose(nodeReference -> { if (nodeReference == null) { return CompletableFuture.completedFuture(null); @@ -747,46 +738,42 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N debug(l -> l.debug("entry node with key {} at layer {}", entryNodeReference.getPrimaryKey(), lMax)); - final AtomicReference nodeReferenceAtomic = - new AtomicReference<>(new NodeReferenceWithDistance(entryNodeReference.getPrimaryKey(), + final NodeReferenceWithDistance initialNodeReference = + new NodeReferenceWithDistance(entryNodeReference.getPrimaryKey(), entryNodeReference.getVector(), - Vector.comparativeDistance(metric, entryNodeReference.getVector(), newVector))); - MoreAsyncUtil.forLoop(lMax, - layer -> layer > insertionLayer, - layer -> layer - 1, - layer -> { - final StorageAdapter storageAdapter = getStorageAdapterForLayer(layer); - return greedySearchLayer(storageAdapter, transaction, - nodeReferenceAtomic.get(), layer, newVector) - .thenApply(nodeReference -> { - nodeReferenceAtomic.set(nodeReference); - return null; - }); - }, executor); - - debug(l -> { - final NodeReference nodeReference = nodeReferenceAtomic.get(); - l.debug("nearest entry point at lMax={} is at key={}", lMax, nodeReference.getPrimaryKey()); - }); - - final AtomicReference> nearestNeighborsAtomic = - new AtomicReference<>(ImmutableList.of(nodeReferenceAtomic.get())); - - return MoreAsyncUtil.forLoop(Math.min(lMax, insertionLayer), - layer -> layer >= 0, - layer -> layer - 1, - layer -> { - final StorageAdapter storageAdapter = getStorageAdapterForLayer(layer); - return insertIntoLayer(storageAdapter, transaction, - nearestNeighborsAtomic.get(), layer, newPrimaryKey, newVector) - .thenCompose(nearestNeighbors -> { - nearestNeighborsAtomic.set(nearestNeighbors); - return AsyncUtil.DONE; - }); - }, executor); + Vector.comparativeDistance(metric, entryNodeReference.getVector(), newVector)); + return MoreAsyncUtil.forLoop(lMax, initialNodeReference, + layer -> layer > insertionLayer, + layer -> layer - 1, + (layer, previousNodeReference) -> { + final StorageAdapter storageAdapter = getStorageAdapterForLayer(layer); + return greedySearchLayer(storageAdapter, transaction, + previousNodeReference, layer, newVector); + }, executor) + .thenCompose(nodeReference -> + insertIntoLayers(transaction, newPrimaryKey, newVector, nodeReference, + lMax, insertionLayer)); }).thenCompose(ignored -> AsyncUtil.DONE); } + @Nonnull + private CompletableFuture insertIntoLayers(final @Nonnull Transaction transaction, + final @Nonnull Tuple newPrimaryKey, + final @Nonnull Vector newVector, + final NodeReferenceWithDistance nodeReference, final int lMax, final int insertionLayer) { + debug(l -> { + l.debug("nearest entry point at lMax={} is at key={}", lMax, nodeReference.getPrimaryKey()); + }); + return MoreAsyncUtil.>forLoop(Math.min(lMax, insertionLayer), ImmutableList.of(nodeReference), + layer -> layer >= 0, + layer -> layer - 1, + (layer, previousNodeReferences) -> { + final StorageAdapter storageAdapter = getStorageAdapterForLayer(layer); + return insertIntoLayer(storageAdapter, transaction, + previousNodeReferences, layer, newPrimaryKey, newVector); + }, executor).thenCompose(ignored -> AsyncUtil.DONE); + } + @Nonnull private CompletableFuture> insertIntoLayer(@Nonnull final StorageAdapter storageAdapter, @Nonnull final Transaction transaction, diff --git a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java index 1a63dc1ebe..fd6609bf20 100644 --- a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java +++ b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java @@ -237,7 +237,7 @@ public void testSIFTInsert10k() throws Exception { Comparator.comparing(NodeReferenceWithDistance::getDistance)); try (BufferedReader br = new BufferedReader(new FileReader(tsvFile))) { - for (int i = 0; i < 10000;) { + for (int i = 0; i < 1000;) { i += basicInsertBatch(100, nextNodeIdAtomic, onReadListener, tr -> { final String line; From bf6a02a73347d13ea5ada26d96704ab4109ebb29 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Wed, 6 Aug 2025 11:42:12 +0200 Subject: [PATCH 20/34] batch insert --- .../foundationdb/async/MoreAsyncUtil.java | 16 +- .../apple/foundationdb/async/hnsw/HNSW.java | 163 ++++++++++++++++-- .../async/hnsw/HNSWModificationTest.java | 133 ++++++++++++-- 3 files changed, 275 insertions(+), 37 deletions(-) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java index a49e1fb7a0..0278a36751 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java @@ -23,7 +23,6 @@ import com.apple.foundationdb.annotation.API; import com.apple.foundationdb.util.LoggableException; import com.google.common.base.Suppliers; -import com.google.common.collect.ImmutableList; import com.google.common.collect.Lists; import com.google.common.util.concurrent.ThreadFactoryBuilder; @@ -31,11 +30,11 @@ import javax.annotation.Nullable; import java.util.ArrayDeque; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.NoSuchElementException; -import java.util.Objects; import java.util.Queue; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; @@ -1105,23 +1104,14 @@ public static CompletableFuture> forEach(@Nonnull final Iterable< final int index = indexAtomic.getAndIncrement(); working.add(body.apply(currentItem) - .thenAccept(resultNode -> { - Objects.requireNonNull(resultNode); - resultArray[index] = resultNode; - })); + .thenAccept(result -> resultArray[index] = result)); } if (working.isEmpty()) { return AsyncUtil.READY_FALSE; } return AsyncUtil.whenAny(working).thenApply(ignored -> true); - }, executor).thenApply(ignored -> { - final ImmutableList.Builder resultBuilder = ImmutableList.builder(); - for (final Object o : resultArray) { - resultBuilder.add((U)o); - } - return resultBuilder.build(); - }); + }, executor).thenApply(ignored -> Arrays.asList((U[])resultArray)); } /** diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index 4859c7f2e9..31cb055fcc 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -59,6 +59,9 @@ import java.util.function.Function; import java.util.stream.Collectors; +import static com.apple.foundationdb.async.MoreAsyncUtil.forEach; +import static com.apple.foundationdb.async.MoreAsyncUtil.forLoop; + /** * TODO. */ @@ -70,6 +73,7 @@ public class HNSW { public static final int MAX_CONCURRENT_NODE_READS = 16; public static final int MAX_CONCURRENT_NEIGHBOR_FETCHES = 3; + public static final int MAX_CONCURRENT_SEARCHES = 10; @Nonnull public static final Random DEFAULT_RANDOM = new Random(0L); @Nonnull public static final Metric DEFAULT_METRIC = new Metric.EuclideanMetric(); public static final int DEFAULT_M = 16; @@ -697,12 +701,17 @@ private CompletableFuture< @Nonnull final Iterable nodeReferences, @Nonnull final Function fetchBypassFunction, @Nonnull final BiFunction, U> biMapFunction) { - return MoreAsyncUtil.forEach(nodeReferences, + return forEach(nodeReferences, currentNeighborReference -> fetchNodeIfNecessaryAndApply(storageAdapter, readTransaction, layer, currentNeighborReference, fetchBypassFunction, biMapFunction), MAX_CONCURRENT_NODE_READS, getExecutor()); } + @Nonnull + public CompletableFuture insert(@Nonnull final Transaction transaction, @Nonnull final NodeReferenceWithVector nodeReferenceWithVector) { + return insert(transaction, nodeReferenceWithVector.getPrimaryKey(), nodeReferenceWithVector.getVector()); + } + @Nonnull public CompletableFuture insert(@Nonnull final Transaction transaction, @Nonnull final Tuple newPrimaryKey, @Nonnull final Vector newVector) { @@ -720,9 +729,9 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N new EntryNodeReference(newPrimaryKey, newVector, insertionLayer), getOnWriteListener()); debug(l -> l.debug("written entry node reference with key={} on layer={}", newPrimaryKey, insertionLayer)); } else { - final int entryNodeLayer = entryNodeReference.getLayer(); - if (insertionLayer > entryNodeLayer) { - writeLonelyNodes(transaction, newPrimaryKey, newVector, insertionLayer, entryNodeLayer); + final int lMax = entryNodeReference.getLayer(); + if (insertionLayer > lMax) { + writeLonelyNodes(transaction, newPrimaryKey, newVector, insertionLayer, lMax); StorageAdapter.writeEntryNodeReference(transaction, getSubspace(), new EntryNodeReference(newPrimaryKey, newVector, insertionLayer), getOnWriteListener()); debug(l -> l.debug("written entry node reference with key={} on layer={}", newPrimaryKey, insertionLayer)); @@ -757,13 +766,104 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N } @Nonnull - private CompletableFuture insertIntoLayers(final @Nonnull Transaction transaction, - final @Nonnull Tuple newPrimaryKey, - final @Nonnull Vector newVector, - final NodeReferenceWithDistance nodeReference, final int lMax, final int insertionLayer) { - debug(l -> { - l.debug("nearest entry point at lMax={} is at key={}", lMax, nodeReference.getPrimaryKey()); - }); + public CompletableFuture insertBatch(@Nonnull final Transaction transaction, + @Nonnull List batch) { + final Metric metric = getConfig().getMetric(); + + // determine the layer each item should be inserted at + final Random random = getConfig().getRandom(); + final List batchWithLayers = Lists.newArrayListWithCapacity(batch.size()); + for (final NodeReferenceWithVector current : batch) { + batchWithLayers.add(new NodeReferenceWithLayer(current.getPrimaryKey(), current.getVector(), + insertionLayer(random))); + } + // sort the layers in reverse order + batchWithLayers.sort(Comparator.comparing(NodeReferenceWithLayer::getL).reversed()); + + return StorageAdapter.fetchEntryNodeReference(transaction, getSubspace(), getOnReadListener()) + .thenCompose(entryNodeReference -> { + final int lMax = entryNodeReference == null ? -1 : entryNodeReference.getLayer(); + + return forEach(batchWithLayers, + item -> { + if (lMax == -1) { + return CompletableFuture.completedFuture(null); + } + + final Vector itemVector = item.getVector(); + final int itemL = item.getL(); + + final NodeReferenceWithDistance initialNodeReference = + new NodeReferenceWithDistance(entryNodeReference.getPrimaryKey(), + entryNodeReference.getVector(), + Vector.comparativeDistance(metric, entryNodeReference.getVector(), itemVector)); + + return MoreAsyncUtil.forLoop(lMax, initialNodeReference, + layer -> layer > itemL, + layer -> layer - 1, + (layer, previousNodeReference) -> { + final StorageAdapter storageAdapter = getStorageAdapterForLayer(layer); + return greedySearchLayer(storageAdapter, transaction, + previousNodeReference, layer, itemVector); + }, executor); + }, MAX_CONCURRENT_SEARCHES, getExecutor()) + .thenCompose(searchEntryReferences -> + forLoop(0, entryNodeReference, + index -> index < batchWithLayers.size(), + index -> index + 1, + (index, currentEntryNodeReference) -> { + final NodeReferenceWithLayer item = batchWithLayers.get(index); + final Tuple itemPrimaryKey = item.getPrimaryKey(); + final Vector itemVector = item.getVector(); + final int itemL = item.getL(); + + final EntryNodeReference newEntryNodeReference; + final int currentLMax; + + if (entryNodeReference == null) { + // this is the first node + writeLonelyNodes(transaction, itemPrimaryKey, itemVector, itemL, -1); + newEntryNodeReference = + new EntryNodeReference(itemPrimaryKey, itemVector, itemL); + StorageAdapter.writeEntryNodeReference(transaction, getSubspace(), + newEntryNodeReference, getOnWriteListener()); + debug(l -> l.debug("written entry node reference with key={} on layer={}", itemPrimaryKey, itemL)); + + return CompletableFuture.completedFuture(newEntryNodeReference); + } else { + currentLMax = currentEntryNodeReference.getLayer(); + if (itemL > currentLMax) { + writeLonelyNodes(transaction, itemPrimaryKey, itemVector, itemL, lMax); + newEntryNodeReference = + new EntryNodeReference(itemPrimaryKey, itemVector, itemL); + StorageAdapter.writeEntryNodeReference(transaction, getSubspace(), + newEntryNodeReference, getOnWriteListener()); + debug(l -> l.debug("written entry node reference with key={} on layer={}", itemPrimaryKey, itemL)); + } else { + newEntryNodeReference = entryNodeReference; + } + } + + debug(l -> l.debug("entry node with key {} at layer {}", + currentEntryNodeReference.getPrimaryKey(), currentLMax)); + + final var currentSearchEntry = + searchEntryReferences.get(index); + + return insertIntoLayers(transaction, itemPrimaryKey, itemVector, currentSearchEntry, + lMax, itemL).thenApply(ignored -> newEntryNodeReference); + }, getExecutor())); + }).thenCompose(ignored -> AsyncUtil.DONE); + } + + @Nonnull + private CompletableFuture insertIntoLayers(@Nonnull final Transaction transaction, + @Nonnull final Tuple newPrimaryKey, + @Nonnull final Vector newVector, + @Nonnull final NodeReferenceWithDistance nodeReference, + final int lMax, + final int insertionLayer) { + debug(l -> l.debug("nearest entry point at lMax={} is at key={}", lMax, nodeReference.getPrimaryKey())); return MoreAsyncUtil.>forLoop(Math.min(lMax, insertionLayer), ImmutableList.of(nodeReference), layer -> layer >= 0, layer -> layer - 1, @@ -817,7 +917,7 @@ private CompletableFuture { final Node selectedNeighborNode = selectedNeighbor.getNode(); final NeighborsChangeSet changeSet = @@ -1110,4 +1210,43 @@ private void debug(@Nonnull final Consumer loggerConsumer) { loggerConsumer.accept(logger); } } + + private static class NodeReferenceWithLayer extends NodeReferenceWithVector { + @SuppressWarnings("checkstyle:MemberName") + private final int l; + + public NodeReferenceWithLayer(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, + final int l) { + super(primaryKey, vector); + this.l = l; + } + + public int getL() { + return l; + } + } + + private static class NodeReferenceWithSearchEntry extends NodeReferenceWithVector { + @SuppressWarnings("checkstyle:MemberName") + private final int l; + @Nonnull + private final NodeReferenceWithDistance nodeReferenceWithDistance; + + public NodeReferenceWithSearchEntry(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, + final int l, + @Nonnull final NodeReferenceWithDistance nodeReferenceWithDistance) { + super(primaryKey, vector); + this.l = l; + this.nodeReferenceWithDistance = nodeReferenceWithDistance; + } + + public int getL() { + return l; + } + + @Nonnull + public NodeReferenceWithDistance getNodeReferenceWithDistance() { + return nodeReferenceWithDistance; + } + } } diff --git a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java index fd6609bf20..c60cc4b51c 100644 --- a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java +++ b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java @@ -22,7 +22,6 @@ import com.apple.foundationdb.Database; import com.apple.foundationdb.Transaction; -import com.apple.foundationdb.async.AsyncUtil; import com.apple.foundationdb.async.hnsw.Vector.HalfVector; import com.apple.foundationdb.async.rtree.RTree; import com.apple.foundationdb.test.TestDatabaseExtension; @@ -60,7 +59,6 @@ import java.util.NavigableSet; import java.util.Objects; import java.util.Random; -import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentSkipListSet; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; @@ -178,8 +176,8 @@ public void testBasicInsert() { OnWriteListener.NOOP, onReadListener); for (int i = 0; i < 1000;) { - i += basicInsertBatch(100, nextNodeIdAtomic, onReadListener, - tr -> hnsw.insert(tr, createNextPrimaryKey(nextNodeIdAtomic), createRandomVector(random, dimensions))); + i += basicInsertBatch(hnsw, 100, nextNodeIdAtomic, onReadListener, + tr -> new NodeReferenceWithVector(createNextPrimaryKey(nextNodeIdAtomic), createRandomVector(random, dimensions))); } onReadListener.reset(); @@ -199,15 +197,18 @@ public void testBasicInsert() { logger.info("search transaction took elapsedTime={}ms", TimeUnit.NANOSECONDS.toMillis(endTs - beginTs)); } - private int basicInsertBatch(final int batchSize, + private int basicInsertBatch(final HNSW hnsw, final int batchSize, @Nonnull final AtomicLong nextNodeIdAtomic, @Nonnull final TestOnReadListener onReadListener, - @Nonnull final Function> insertFunction) { + @Nonnull final Function insertFunction) { return db.run(tr -> { onReadListener.reset(); final long nextNodeId = nextNodeIdAtomic.get(); final long beginTs = System.nanoTime(); for (int i = 0; i < batchSize; i ++) { - insertFunction.apply(tr).join(); + final var newNodeReference = insertFunction.apply(tr); + if (newNodeReference != null) { + hnsw.insert(tr, newNodeReference).join(); + } } final long endTs = System.nanoTime(); logger.info("inserted batchSize={} records starting at nodeId={} took elapsedTime={}ms, readCounts={}, MSums={}", batchSize, nextNodeId, @@ -216,6 +217,29 @@ private int basicInsertBatch(final int batchSize, }); } + private int insertBatch(final HNSW hnsw, final int batchSize, + @Nonnull final AtomicLong nextNodeIdAtomic, @Nonnull final TestOnReadListener onReadListener, + @Nonnull final Function insertFunction) { + return db.run(tr -> { + onReadListener.reset(); + final long nextNodeId = nextNodeIdAtomic.get(); + final long beginTs = System.nanoTime(); + final ImmutableList.Builder nodeReferenceWithVectorBuilder = + ImmutableList.builder(); + for (int i = 0; i < batchSize; i ++) { + final var newNodeReference = insertFunction.apply(tr); + if (newNodeReference != null) { + nodeReferenceWithVectorBuilder.add(newNodeReference); + } + } + hnsw.insertBatch(tr, nodeReferenceWithVectorBuilder.build()).join(); + final long endTs = System.nanoTime(); + logger.info("inserted batch batchSize={} records starting at nodeId={} took elapsedTime={}ms, readCounts={}, MSums={}", batchSize, nextNodeId, + TimeUnit.NANOSECONDS.toMillis(endTs - beginTs), onReadListener.getNodeCountByLayer(), onReadListener.getSumMByLayer()); + return batchSize; + }); + } + @Test @Timeout(value = 150, unit = TimeUnit.MINUTES) public void testSIFTInsert10k() throws Exception { @@ -237,8 +261,93 @@ public void testSIFTInsert10k() throws Exception { Comparator.comparing(NodeReferenceWithDistance::getDistance)); try (BufferedReader br = new BufferedReader(new FileReader(tsvFile))) { - for (int i = 0; i < 1000;) { - i += basicInsertBatch(100, nextNodeIdAtomic, onReadListener, + for (int i = 0; i < 10000;) { + i += basicInsertBatch(hnsw, 100, nextNodeIdAtomic, onReadListener, + tr -> { + final String line; + try { + line = br.readLine(); + } catch (IOException e) { + throw new RuntimeException(e); + } + + final String[] values = Objects.requireNonNull(line).split("\t"); + Assertions.assertEquals(dimensions, values.length); + final Half[] halfs = new Half[dimensions]; + + for (int c = 0; c < values.length; c++) { + final String value = values[c]; + halfs[c] = HNSWHelpers.halfValueOf(Double.parseDouble(value)); + } + final Tuple currentPrimaryKey = createNextPrimaryKey(nextNodeIdAtomic); + final HalfVector currentVector = new HalfVector(halfs); + final HalfVector queryVector = queryVectorAtomic.get(); + if (queryVector == null) { + queryVectorAtomic.set(currentVector); + return null; + } else { + final double currentDistance = + Vector.comparativeDistance(metric, currentVector, queryVector); + if (trueResults.size() < k || trueResults.last().getDistance() > currentDistance) { + trueResults.add( + new NodeReferenceWithDistance(currentPrimaryKey, currentVector, + Vector.comparativeDistance(metric, currentVector, queryVector))); + } + if (trueResults.size() > k) { + trueResults.remove(trueResults.last()); + } + return new NodeReferenceWithVector(currentPrimaryKey, currentVector); + } + }); + } + } + + onReadListener.reset(); + final long beginTs = System.nanoTime(); + final List> results = + db.run(tr -> hnsw.kNearestNeighborsSearch(tr, k, 100, queryVectorAtomic.get()).join()); + final long endTs = System.nanoTime(); + + for (NodeReferenceAndNode nodeReferenceAndNode : results) { + final NodeReferenceWithDistance nodeReferenceWithDistance = nodeReferenceAndNode.getNodeReferenceWithDistance(); + logger.info("retrieved result nodeId = {} at distance= {}", nodeReferenceWithDistance.getPrimaryKey().getLong(0), + nodeReferenceWithDistance.getDistance()); + } + + for (final NodeReferenceWithDistance nodeReferenceWithDistance : trueResults) { + logger.info("true result nodeId ={} at distance={}", nodeReferenceWithDistance.getPrimaryKey().getLong(0), + nodeReferenceWithDistance.getDistance()); + } + + System.out.println(onReadListener.getNodeCountByLayer()); + System.out.println(onReadListener.getBytesReadByLayer()); + + logger.info("search transaction took elapsedTime={}ms", TimeUnit.NANOSECONDS.toMillis(endTs - beginTs)); + } + + @Test + @Timeout(value = 150, unit = TimeUnit.MINUTES) + public void testSIFTInsert10kWithBatchInsert() throws Exception { + final Metric metric = Metric.EUCLIDEAN_METRIC; + final int k = 10; + final AtomicLong nextNodeIdAtomic = new AtomicLong(0L); + + final TestOnReadListener onReadListener = new TestOnReadListener(); + + final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool(), + HNSW.DEFAULT_CONFIG.toBuilder().setMetric(metric).setM(32).setMMax(32).setMMax0(64).build(), + OnWriteListener.NOOP, onReadListener); + + final String tsvFile = "/Users/nseemann/Downloads/train-100k.tsv"; + final int dimensions = 128; + + final AtomicReference queryVectorAtomic = new AtomicReference<>(); + final NavigableSet trueResults = new ConcurrentSkipListSet<>( + Comparator.comparing(NodeReferenceWithDistance::getDistance)); + + try (BufferedReader br = new BufferedReader(new FileReader(tsvFile))) { + for (int i = 0; i < 10000;) { + i += insertBatch(hnsw, 100, nextNodeIdAtomic, onReadListener, tr -> { final String line; try { @@ -260,7 +369,7 @@ public void testSIFTInsert10k() throws Exception { final HalfVector queryVector = queryVectorAtomic.get(); if (queryVector == null) { queryVectorAtomic.set(currentVector); - return AsyncUtil.DONE; + return null; } else { final double currentDistance = Vector.comparativeDistance(metric, currentVector, queryVector); @@ -272,7 +381,7 @@ public void testSIFTInsert10k() throws Exception { if (trueResults.size() > k) { trueResults.remove(trueResults.last()); } - return hnsw.insert(tr, currentPrimaryKey, currentVector); + return new NodeReferenceWithVector(currentPrimaryKey, currentVector); } }); } @@ -306,7 +415,7 @@ public void testBasicInsertAndScanLayer() throws Exception { final Random random = new Random(0); final AtomicLong nextNodeId = new AtomicLong(0L); final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool(), - HNSW.DEFAULT_CONFIG.toBuilder().setM(4).setMMax(4).setMMax0(10).build(), + HNSW.DEFAULT_CONFIG.toBuilder().setM(4).setMMax(4).setMMax0(4).build(), OnWriteListener.NOOP, OnReadListener.NOOP); db.run(tr -> { From 5d61de35cbdd2d5a102859a99178a7d1cd66c740 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Tue, 12 Aug 2025 10:18:30 +0200 Subject: [PATCH 21/34] adding new distance rank comparison and some index helpers --- .../apple/foundationdb/async/hnsw/Metric.java | 31 - .../foundationdb/async/hnsw/Metrics.java | 43 ++ .../async/hnsw/HNSWModificationTest.java | 19 +- .../record/metadata/IndexOptions.java | 9 + .../record/metadata/IndexTypes.java | 5 + .../indexes/VectorIndexHelper.java | 114 ++++ .../indexes/VectorIndexMaintainer.java | 588 ++++++++++++++++++ .../indexes/VectorIndexMaintainerFactory.java | 158 +++++ .../record/query/expressions/Comparisons.java | 219 ++++++- .../src/main/proto/record_query_plan.proto | 6 + 10 files changed, 1145 insertions(+), 47 deletions(-) create mode 100644 fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metrics.java create mode 100644 fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexHelper.java create mode 100644 fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java create mode 100644 fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainerFactory.java diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java index d3fc11c082..6e236a5d10 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metric.java @@ -23,12 +23,6 @@ import javax.annotation.Nonnull; public interface Metric { - ManhattanMetric MANHATTAN_METRIC = new ManhattanMetric(); - EuclideanMetric EUCLIDEAN_METRIC = new EuclideanMetric(); - EuclideanSquareMetric EUCLIDEAN_SQUARE_METRIC = new EuclideanSquareMetric(); - CosineMetric COSINE_METRIC = new CosineMetric(); - DotProductMetric DOT_PRODUCT_METRIC = new DotProductMetric(); - double distance(Double[] vector1, Double[] vector2); default double comparativeDistance(Double[] vector1, Double[] vector2) { @@ -54,31 +48,6 @@ private static void validate(Double[] vector1, Double[] vector2) { } } - @Nonnull - static ManhattanMetric manhattanMetric() { - return Metric.MANHATTAN_METRIC; - } - - @Nonnull - static EuclideanMetric euclideanMetric() { - return Metric.EUCLIDEAN_METRIC; - } - - @Nonnull - static EuclideanSquareMetric euclideanSquareMetric() { - return Metric.EUCLIDEAN_SQUARE_METRIC; - } - - @Nonnull - static CosineMetric cosineMetric() { - return Metric.COSINE_METRIC; - } - - @Nonnull - static DotProductMetric dotProductMetric() { - return Metric.DOT_PRODUCT_METRIC; - } - class ManhattanMetric implements Metric { @Override public double distance(final Double[] vector1, final Double[] vector2) { diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metrics.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metrics.java new file mode 100644 index 0000000000..8c30faf852 --- /dev/null +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Metrics.java @@ -0,0 +1,43 @@ +/* + * Metric.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.async.hnsw; + +import javax.annotation.Nonnull; + +public enum Metrics { + MANHATTAN_METRIC(new Metric.ManhattanMetric()), + EUCLIDEAN_METRIC(new Metric.EuclideanMetric()), + EUCLIDEAN_SQUARE_METRIC(new Metric.EuclideanSquareMetric()), + COSINE_METRIC(new Metric.CosineMetric()), + DOT_PRODUCT_METRIC(new Metric.DotProductMetric()); + + @Nonnull + private final Metric metric; + + Metrics(@Nonnull final Metric metric) { + this.metric = metric; + } + + @Nonnull + public Metric getMetric() { + return metric; + } +} diff --git a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java index c60cc4b51c..85aea0e2ab 100644 --- a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java +++ b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java @@ -1,5 +1,5 @@ /* - * RTreeModificationTest.java + * HNSWModificationTest.java * * This source file is part of the FoundationDB open source project * @@ -172,7 +172,8 @@ public void testBasicInsert() { final int dimensions = 128; final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool(), - HNSW.DEFAULT_CONFIG.toBuilder().setMetric(Metric.EUCLIDEAN_METRIC).setM(32).setMMax(32).setMMax0(64).build(), + HNSW.DEFAULT_CONFIG.toBuilder().setMetric(Metrics.EUCLIDEAN_METRIC.getMetric()) + .setM(32).setMMax(32).setMMax0(64).build(), OnWriteListener.NOOP, onReadListener); for (int i = 0; i < 1000;) { @@ -243,7 +244,7 @@ private int insertBatch(final HNSW hnsw, final int batchSize, @Test @Timeout(value = 150, unit = TimeUnit.MINUTES) public void testSIFTInsert10k() throws Exception { - final Metric metric = Metric.EUCLIDEAN_METRIC; + final Metric metric = Metrics.EUCLIDEAN_METRIC.getMetric(); final int k = 10; final AtomicLong nextNodeIdAtomic = new AtomicLong(0L); @@ -328,7 +329,7 @@ public void testSIFTInsert10k() throws Exception { @Test @Timeout(value = 150, unit = TimeUnit.MINUTES) public void testSIFTInsert10kWithBatchInsert() throws Exception { - final Metric metric = Metric.EUCLIDEAN_METRIC; + final Metric metric = Metrics.EUCLIDEAN_METRIC.getMetric(); final int k = 10; final AtomicLong nextNodeIdAtomic = new AtomicLong(0L); @@ -440,7 +441,7 @@ public void testManyRandomVectors() { final HalfVector randomVector = createRandomVector(random, 768); final Tuple vectorTuple = StorageAdapter.tupleFromVector(randomVector); final Vector roundTripVector = StorageAdapter.vectorFromTuple(vectorTuple); - Vector.comparativeDistance(Metric.EuclideanMetric.EUCLIDEAN_METRIC, randomVector, roundTripVector); + Vector.comparativeDistance(Metrics.EUCLIDEAN_METRIC.getMetric(), randomVector, roundTripVector); Assertions.assertEquals(randomVector, roundTripVector); } } @@ -453,7 +454,8 @@ public void testSIFTVectors() throws Exception { final TestOnReadListener onReadListener = new TestOnReadListener(); final HNSW hnsw = new HNSW(rtSubspace.getSubspace(), TestExecutors.defaultThreadPool(), - HNSW.DEFAULT_CONFIG.toBuilder().setMetric(Metric.EUCLIDEAN_METRIC).setM(32).setMMax(32).setMMax0(64).build(), + HNSW.DEFAULT_CONFIG.toBuilder().setMetric(Metrics.EUCLIDEAN_METRIC.getMetric()) + .setM(32).setMMax(32).setMMax0(64).build(), OnWriteListener.NOOP, onReadListener); @@ -481,7 +483,8 @@ public void testSIFTVectors() throws Exception { halfs[c] = HNSWHelpers.halfValueOf(Double.parseDouble(value)); } final HalfVector newVector = new HalfVector(halfs); - final double distance = Vector.comparativeDistance(Metric.EUCLIDEAN_METRIC, referenceVector, newVector); + final double distance = Vector.comparativeDistance(Metrics.EUCLIDEAN_METRIC.getMetric(), + referenceVector, newVector); count++; final double delta = distance - mean; mean += delta / count; @@ -500,7 +503,7 @@ public void testSIFTVectors() throws Exception { @ValueSource(ints = {2, 3, 10, 100, 768}) public void testManyVectorsStandardDeviation(final int dimensionality) { final Random random = new Random(); - final Metric metric = Metric.EuclideanMetric.EUCLIDEAN_METRIC; + final Metric metric = Metrics.EUCLIDEAN_METRIC.getMetric(); long count = 0L; double mean = 0.0d; double mean2 = 0.0d; diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/IndexOptions.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/IndexOptions.java index 2b66805b2f..21dc8f23b4 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/IndexOptions.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/IndexOptions.java @@ -223,6 +223,15 @@ public class IndexOptions { */ public static final String RTREE_USE_NODE_SLOT_INDEX = "rtreeUseNodeSlotIndex"; + public static final String HNSW_METRIC = "hnswMetric"; + public static final String HNSW_M = "hnswM"; + public static final String HNSW_M_MAX = "hnswMax"; + public static final String HNSW_M_MAX_0 = "hnswMax0"; + public static final String HNSW_EF_SEARCH = "hnswEfSearch"; + public static final String HNSW_EF_CONSTRUCTION = "hnswEfConstruction"; + public static final String HNSW_EXTEND_CANDIDATES = "hnswExtendCandidates"; + public static final String HNSW_KEEP_PRUNED_CONNECTIONS = "hnswKeepPrunedConnections"; + private IndexOptions() { } } diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/IndexTypes.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/IndexTypes.java index 1d19171093..8d10f26d9e 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/IndexTypes.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/IndexTypes.java @@ -164,6 +164,11 @@ public class IndexTypes { */ public static final String MULTIDIMENSIONAL = "multidimensional"; + /** + * An index using an HNSW structure. + */ + public static final String VECTOR = "vector"; + private IndexTypes() { } } diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexHelper.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexHelper.java new file mode 100644 index 0000000000..9c7e87628b --- /dev/null +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexHelper.java @@ -0,0 +1,114 @@ +/* + * VectorIndexHelper.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.record.provider.foundationdb.indexes; + +import com.apple.foundationdb.annotation.API; +import com.apple.foundationdb.async.hnsw.HNSW; +import com.apple.foundationdb.async.hnsw.Metrics; +import com.apple.foundationdb.record.metadata.Index; +import com.apple.foundationdb.record.metadata.IndexOptions; +import com.apple.foundationdb.record.provider.common.StoreTimer; + +import javax.annotation.Nonnull; + +/** + * Helper functions for index maintainers that use a {@link HNSW}. + */ +@API(API.Status.EXPERIMENTAL) +public class VectorIndexHelper { + private VectorIndexHelper() { + } + + /** + * Parse standard options into {@link HNSW.Config}. + * @param index the index definition to get options from + * @return parsed config options + */ + public static HNSW.Config getConfig(@Nonnull final Index index) { + final HNSW.ConfigBuilder builder = HNSW.newConfigBuilder(); + final String hnswMetricOption = index.getOption(IndexOptions.HNSW_METRIC); + if (hnswMetricOption != null) { + builder.setMetric(Metrics.valueOf(hnswMetricOption).getMetric()); + } + final String hnswMOption = index.getOption(IndexOptions.HNSW_M); + if (hnswMOption != null) { + builder.setM(Integer.parseInt(hnswMOption)); + } + final String hnswMMaxOption = index.getOption(IndexOptions.HNSW_M_MAX); + if (hnswMMaxOption != null) { + builder.setMMax(Integer.parseInt(hnswMMaxOption)); + } + final String hnswMMax0Option = index.getOption(IndexOptions.HNSW_M_MAX_0); + if (hnswMMax0Option != null) { + builder.setMMax0(Integer.parseInt(hnswMMax0Option)); + } + final String hnswEfSearchOption = index.getOption(IndexOptions.HNSW_EF_SEARCH); + if (hnswEfSearchOption != null) { + builder.setEfSearch(Integer.parseInt(hnswEfSearchOption)); + } + final String hnswEfConstructionOption = index.getOption(IndexOptions.HNSW_EF_CONSTRUCTION); + if (hnswEfConstructionOption != null) { + builder.setEfConstruction(Integer.parseInt(hnswEfConstructionOption)); + } + final String hnswExtendCandidatesOption = index.getOption(IndexOptions.HNSW_EXTEND_CANDIDATES); + if (hnswExtendCandidatesOption != null) { + builder.setExtendCandidates(Boolean.parseBoolean(hnswExtendCandidatesOption)); + } + final String hnswKeepPrunedConnectionsOption = index.getOption(IndexOptions.HNSW_KEEP_PRUNED_CONNECTIONS); + if (hnswKeepPrunedConnectionsOption != null) { + builder.setKeepPrunedConnections(Boolean.parseBoolean(hnswKeepPrunedConnectionsOption)); + } + + return builder.build(); + } + + /** + * Instrumentation events specific to R-tree index maintenance. + */ + public enum Events implements StoreTimer.DetailEvent { + VECTOR_SCAN("scanning the HNSW of a vector index"), + VECTOR_SKIP_SCAN("skip scan the prefix tuples of a vector index scan"), + VECTOR_MODIFICATION("modifying the HNSW of a vector index"); + + private final String title; + private final String logKey; + + Events(String title, String logKey) { + this.title = title; + this.logKey = (logKey != null) ? logKey : StoreTimer.DetailEvent.super.logKey(); + } + + Events(String title) { + this(title, null); + } + + @Override + public String title() { + return title; + } + + @Override + @Nonnull + public String logKey() { + return this.logKey; + } + } +} diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java new file mode 100644 index 0000000000..0815e84ccb --- /dev/null +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java @@ -0,0 +1,588 @@ +/* + * VectorIndexMaintainer.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.record.provider.foundationdb.indexes; + +import com.apple.foundationdb.KeyValue; +import com.apple.foundationdb.Range; +import com.apple.foundationdb.ReadTransaction; +import com.apple.foundationdb.Transaction; +import com.apple.foundationdb.annotation.API; +import com.apple.foundationdb.async.AsyncIterator; +import com.apple.foundationdb.async.AsyncUtil; +import com.apple.foundationdb.async.rtree.ChildSlot; +import com.apple.foundationdb.async.rtree.ItemSlot; +import com.apple.foundationdb.async.rtree.Node; +import com.apple.foundationdb.async.rtree.NodeHelpers; +import com.apple.foundationdb.async.rtree.OnReadListener; +import com.apple.foundationdb.async.rtree.OnWriteListener; +import com.apple.foundationdb.async.rtree.RTree; +import com.apple.foundationdb.async.rtree.RTreeHilbertCurveHelpers; +import com.apple.foundationdb.record.CursorStreamingMode; +import com.apple.foundationdb.record.EndpointType; +import com.apple.foundationdb.record.ExecuteProperties; +import com.apple.foundationdb.record.IndexEntry; +import com.apple.foundationdb.record.IndexScanType; +import com.apple.foundationdb.record.PipelineOperation; +import com.apple.foundationdb.record.RecordCoreException; +import com.apple.foundationdb.record.RecordCursor; +import com.apple.foundationdb.record.RecordCursorContinuation; +import com.apple.foundationdb.record.RecordCursorProto; +import com.apple.foundationdb.record.RecordCursorResult; +import com.apple.foundationdb.record.ScanProperties; +import com.apple.foundationdb.record.TupleRange; +import com.apple.foundationdb.record.cursors.AsyncIteratorCursor; +import com.apple.foundationdb.record.cursors.AsyncLockCursor; +import com.apple.foundationdb.record.cursors.ChainedCursor; +import com.apple.foundationdb.record.cursors.CursorLimitManager; +import com.apple.foundationdb.record.cursors.LazyCursor; +import com.apple.foundationdb.record.locking.LockIdentifier; +import com.apple.foundationdb.record.metadata.Key; +import com.apple.foundationdb.record.metadata.expressions.DimensionsKeyExpression; +import com.apple.foundationdb.record.metadata.expressions.KeyExpression; +import com.apple.foundationdb.record.metadata.expressions.KeyWithValueExpression; +import com.apple.foundationdb.record.metadata.expressions.ThenKeyExpression; +import com.apple.foundationdb.record.provider.common.StoreTimer; +import com.apple.foundationdb.record.provider.foundationdb.FDBIndexableRecord; +import com.apple.foundationdb.record.provider.foundationdb.FDBStoreTimer; +import com.apple.foundationdb.record.provider.foundationdb.IndexMaintainerState; +import com.apple.foundationdb.record.provider.foundationdb.IndexScanBounds; +import com.apple.foundationdb.record.provider.foundationdb.KeyValueCursor; +import com.apple.foundationdb.record.provider.foundationdb.MultidimensionalIndexScanBounds; +import com.apple.foundationdb.record.query.QueryToKeyMatcher; +import com.apple.foundationdb.subspace.Subspace; +import com.apple.foundationdb.tuple.ByteArrayUtil; +import com.apple.foundationdb.tuple.ByteArrayUtil2; +import com.apple.foundationdb.tuple.Tuple; +import com.apple.foundationdb.tuple.TupleHelpers; +import com.google.common.base.Preconditions; +import com.google.common.base.Verify; +import com.google.common.collect.Lists; +import com.google.protobuf.ByteString; +import com.google.protobuf.InvalidProtocolBufferException; +import com.google.protobuf.Message; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.math.BigInteger; +import java.util.List; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Executor; +import java.util.function.Function; +import java.util.stream.Collectors; + +/** + * An index maintainer for keeping a {@link RTree}. + */ +@API(API.Status.EXPERIMENTAL) +public class VectorIndexMaintainer extends StandardIndexMaintainer { + private static final byte nodeSlotIndexSubspaceIndicator = 0x00; + @Nonnull + private final RTree.Config config; + + public VectorIndexMaintainer(IndexMaintainerState state) { + super(state); + this.config = MultiDimensionalIndexHelper.getConfig(state.index); + } + + @SuppressWarnings("resource") + @Nonnull + @Override + public RecordCursor scan(@Nonnull final IndexScanBounds scanBounds, @Nullable final byte[] continuation, + @Nonnull final ScanProperties scanProperties) { + if (!scanBounds.getScanType().equals(IndexScanType.BY_VALUE)) { + throw new RecordCoreException("Can only scan multidimensional index by value."); + } + if (!(scanBounds instanceof MultidimensionalIndexScanBounds)) { + throw new RecordCoreException("Need proper multidimensional index scan bounds."); + } + final MultidimensionalIndexScanBounds mDScanBounds = (MultidimensionalIndexScanBounds)scanBounds; + + final DimensionsKeyExpression dimensionsKeyExpression = getDimensionsKeyExpression(state.index.getRootExpression()); + final int prefixSize = dimensionsKeyExpression.getPrefixSize(); + + final ExecuteProperties executeProperties = scanProperties.getExecuteProperties(); + final ScanProperties innerScanProperties = scanProperties.with(ExecuteProperties::clearSkipAndLimit); + final CursorLimitManager cursorLimitManager = new CursorLimitManager(state.context, innerScanProperties); + final Subspace indexSubspace = getIndexSubspace(); + final Subspace nodeSlotIndexSubspace = getNodeSlotIndexSubspace(); + final FDBStoreTimer timer = Objects.requireNonNull(state.context.getTimer()); + + // + // Skip-scan through the prefixes in a way that we only consider each distinct prefix. That skip scan + // forms the outer of a join with an inner that searches the R-tree for that prefix using the + // spatial predicates of the scan bounds. + // + return RecordCursor.flatMapPipelined(prefixSkipScan(prefixSize, timer, mDScanBounds, innerScanProperties), + (prefixTuple, innerContinuation) -> { + final Subspace rtSubspace; + final Subspace rtNodeSlotIndexSubspace; + if (prefixTuple != null) { + Verify.verify(prefixTuple.size() == prefixSize); + rtSubspace = indexSubspace.subspace(prefixTuple); + rtNodeSlotIndexSubspace = nodeSlotIndexSubspace.subspace(prefixTuple); + } else { + rtSubspace = indexSubspace; + rtNodeSlotIndexSubspace = nodeSlotIndexSubspace; + } + + final Continuation parsedContinuation = Continuation.fromBytes(innerContinuation); + final BigInteger lastHilbertValue = + parsedContinuation == null ? null : parsedContinuation.getLastHilbertValue(); + final Tuple lastKey = parsedContinuation == null ? null : parsedContinuation.getLastKey(); + + final RTree rTree = new RTree(rtSubspace, rtNodeSlotIndexSubspace, getExecutor(), config, + RTreeHilbertCurveHelpers::hilbertValue, NodeHelpers::newRandomNodeId, + OnWriteListener.NOOP, new OnRead(cursorLimitManager, timer)); + final ReadTransaction transaction = state.context.readTransaction(true); + return new LazyCursor<>(state.context.acquireReadLock(new LockIdentifier(rtSubspace)) + .thenApply(lock -> new AsyncLockCursor<>(lock, new ItemSlotCursor(getExecutor(), + rTree.scan(transaction, lastHilbertValue, lastKey, + mDScanBounds::overlapsMbrApproximately, + (low, high) -> mDScanBounds.getSuffixRange().overlaps(low, high)), + cursorLimitManager, timer))), state.context.getExecutor()) + .filter(itemSlot -> lastHilbertValue == null || lastKey == null || + itemSlot.compareHilbertValueAndKey(lastHilbertValue, lastKey) > 0) + .filter(itemSlot -> mDScanBounds.containsPosition(itemSlot.getPosition())) + .filter(itemSlot -> mDScanBounds.getSuffixRange().contains(itemSlot.getKeySuffix())) + .map(itemSlot -> { + final List keyItems = Lists.newArrayList(); + if (prefixTuple != null) { + keyItems.addAll(prefixTuple.getItems()); + } + keyItems.addAll(itemSlot.getPosition().getCoordinates().getItems()); + keyItems.addAll(itemSlot.getKeySuffix().getItems()); + return new IndexEntry(state.index, Tuple.fromList(keyItems), itemSlot.getValue()); + }); + }, + continuation, + state.store.getPipelineSize(PipelineOperation.INDEX_TO_RECORD)) + .skipThenLimit(executeProperties.getSkip(), executeProperties.getReturnedRowLimit()); + } + + @Nonnull + @Override + public RecordCursor scan(@Nonnull final IndexScanType scanType, @Nonnull final TupleRange range, + @Nullable final byte[] continuation, @Nonnull final ScanProperties scanProperties) { + throw new RecordCoreException("index maintainer does not support this scan api"); + } + + @Nonnull + private Function> prefixSkipScan(final int prefixSize, + @Nonnull final StoreTimer timer, + @Nonnull final MultidimensionalIndexScanBounds mDScanBounds, + @Nonnull final ScanProperties innerScanProperties) { + final Function> outerFunction; + if (prefixSize > 0) { + outerFunction = outerContinuation -> timer.instrument(MultiDimensionalIndexHelper.Events.MULTIDIMENSIONAL_SKIP_SCAN, + new ChainedCursor<>(state.context, + lastKeyOptional -> nextPrefixTuple(mDScanBounds.getPrefixRange(), + prefixSize, lastKeyOptional.orElse(null), innerScanProperties), + Tuple::pack, + Tuple::fromBytes, + outerContinuation, + innerScanProperties)); + } else { + outerFunction = outerContinuation -> RecordCursor.fromFuture(CompletableFuture.completedFuture(null)); + } + return outerFunction; + } + + @SuppressWarnings({"resource", "PMD.CloseResource"}) + private CompletableFuture> nextPrefixTuple(@Nonnull final TupleRange prefixRange, + final int prefixSize, + @Nullable final Tuple lastPrefixTuple, + @Nonnull final ScanProperties scanProperties) { + final Subspace indexSubspace = getIndexSubspace(); + final KeyValueCursor cursor; + if (lastPrefixTuple == null) { + cursor = KeyValueCursor.Builder.withSubspace(indexSubspace) + .setContext(state.context) + .setRange(prefixRange) + .setContinuation(null) + .setScanProperties(scanProperties.setStreamingMode(CursorStreamingMode.ITERATOR) + .with(innerExecuteProperties -> innerExecuteProperties.setReturnedRowLimit(1))) + .build(); + } else { + KeyValueCursor.Builder builder = KeyValueCursor.Builder.withSubspace(indexSubspace) + .setContext(state.context) + .setContinuation(null) + .setScanProperties(scanProperties) + .setScanProperties(scanProperties.setStreamingMode(CursorStreamingMode.ITERATOR) + .with(innerExecuteProperties -> innerExecuteProperties.setReturnedRowLimit(1))); + + cursor = builder.setLow(indexSubspace.pack(lastPrefixTuple), EndpointType.RANGE_EXCLUSIVE) + .setHigh(prefixRange.getHigh(), prefixRange.getHighEndpoint()) + .build(); + } + + return cursor.onNext().thenApply(next -> { + cursor.close(); + if (next.hasNext()) { + final KeyValue kv = Objects.requireNonNull(next.get()); + return Optional.of(TupleHelpers.subTuple(indexSubspace.unpack(kv.getKey()), 0, prefixSize)); + } + return Optional.empty(); + }); + } + + @Override + protected CompletableFuture updateIndexKeys(@Nonnull final FDBIndexableRecord savedRecord, + final boolean remove, + @Nonnull final List indexEntries) { + final DimensionsKeyExpression dimensionsKeyExpression = getDimensionsKeyExpression(state.index.getRootExpression()); + final int prefixSize = dimensionsKeyExpression.getPrefixSize(); + final int dimensionsSize = dimensionsKeyExpression.getDimensionsSize(); + final Subspace indexSubspace = getIndexSubspace(); + final Subspace nodeSlotIndexSubspace = getNodeSlotIndexSubspace(); + final var futures = indexEntries.stream().map(indexEntry -> { + final var indexKeyItems = indexEntry.getKey().getItems(); + final Tuple prefixKey = Tuple.fromList(indexKeyItems.subList(0, prefixSize)); + + final Subspace rtSubspace; + final Subspace rtNodeSlotIndexSubspace; + if (prefixSize > 0) { + rtSubspace = indexSubspace.subspace(prefixKey); + rtNodeSlotIndexSubspace = nodeSlotIndexSubspace.subspace(prefixKey); + } else { + rtSubspace = indexSubspace; + rtNodeSlotIndexSubspace = nodeSlotIndexSubspace; + } + return state.context.doWithWriteLock(new LockIdentifier(rtSubspace), () -> { + final RTree.Point point = + validatePoint(new RTree.Point(Tuple.fromList(indexKeyItems.subList(prefixSize, prefixSize + dimensionsSize)))); + + final List primaryKeyParts = Lists.newArrayList(savedRecord.getPrimaryKey().getItems()); + state.index.trimPrimaryKey(primaryKeyParts); + final List keySuffixParts = + Lists.newArrayList(indexKeyItems.subList(prefixSize + dimensionsSize, indexKeyItems.size())); + keySuffixParts.addAll(primaryKeyParts); + final Tuple keySuffix = Tuple.fromList(keySuffixParts); + final FDBStoreTimer timer = Objects.requireNonNull(getTimer()); + final RTree rTree = new RTree(rtSubspace, rtNodeSlotIndexSubspace, getExecutor(), config, + RTreeHilbertCurveHelpers::hilbertValue, NodeHelpers::newRandomNodeId, new OnWrite(timer), + OnReadListener.NOOP); + if (remove) { + return rTree.delete(state.transaction, point, keySuffix); + } else { + return rTree.insertOrUpdate(state.transaction, + point, + keySuffix, + indexEntry.getValue()); + } + }); + }).collect(Collectors.toList()); + return AsyncUtil.whenAll(futures); + } + + @Override + public boolean canDeleteWhere(@Nonnull final QueryToKeyMatcher matcher, @Nonnull final Key.Evaluated evaluated) { + if (!super.canDeleteWhere(matcher, evaluated)) { + return false; + } + return evaluated.size() <= getDimensionsKeyExpression(state.index.getRootExpression()).getPrefixSize(); + } + + @Override + public CompletableFuture deleteWhere(@Nonnull final Transaction tr, @Nonnull final Tuple prefix) { + Verify.verify(getDimensionsKeyExpression(state.index.getRootExpression()).getPrefixSize() >= prefix.size()); + return super.deleteWhere(tr, prefix).thenApply(v -> { + // NOTE: Range.startsWith(), Subspace.range() and so on cover keys *strictly* within the range, but we sometimes + // store data at the prefix key itself. + final Subspace nodeSlotIndexSubspace = getNodeSlotIndexSubspace(); + final byte[] key = nodeSlotIndexSubspace.pack(prefix); + state.context.clear(new Range(key, ByteArrayUtil.strinc(key))); + return v; + }); + } + + @Nonnull + private Subspace getNodeSlotIndexSubspace() { + return getSecondarySubspace().subspace(Tuple.from(nodeSlotIndexSubspaceIndicator)); + } + + /** + * Traverse from the root of a key expression of a multidimensional index to the {@link DimensionsKeyExpression}. + * @param root the root {@link KeyExpression} of the index definition + * @return a {@link DimensionsKeyExpression} + */ + @Nonnull + public static DimensionsKeyExpression getDimensionsKeyExpression(@Nonnull final KeyExpression root) { + if (root instanceof KeyWithValueExpression) { + KeyExpression innerKey = ((KeyWithValueExpression)root).getInnerKey(); + while (innerKey instanceof ThenKeyExpression) { + innerKey = ((ThenKeyExpression)innerKey).getChildren().get(0); + } + if (innerKey instanceof DimensionsKeyExpression) { + return (DimensionsKeyExpression)innerKey; + } + throw new RecordCoreException("structure of multidimensional index is not supported"); + } + return (DimensionsKeyExpression)root; + } + + @Nonnull + private static RTree.Point validatePoint(@Nonnull RTree.Point point) { + for (int d = 0; d < point.getNumDimensions(); d ++) { + Object coordinate = point.getCoordinate(d); + Preconditions.checkArgument(coordinate == null || coordinate instanceof Long, + "dimension coordinates must be of type long"); + } + return point; + } + + static class OnRead implements OnReadListener { + @Nonnull + private final CursorLimitManager cursorLimitManager; + @Nonnull + private final FDBStoreTimer timer; + + public OnRead(@Nonnull final CursorLimitManager cursorLimitManager, + @Nonnull final FDBStoreTimer timer) { + this.cursorLimitManager = cursorLimitManager; + this.timer = timer; + } + + @Override + public CompletableFuture onAsyncRead(@Nonnull final CompletableFuture future) { + return timer.instrument(MultiDimensionalIndexHelper.Events.MULTIDIMENSIONAL_SCAN, future); + } + + @Override + public void onNodeRead(@Nonnull final Node node) { + switch (node.getKind()) { + case LEAF: + timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_LEAF_NODE_READS); + break; + case INTERMEDIATE: + timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_INTERMEDIATE_NODE_READS); + break; + default: + throw new RecordCoreException("unsupported kind of node"); + } + } + + @Override + public void onKeyValueRead(@Nonnull final Node node, @Nullable final byte[] key, @Nullable final byte[] value) { + final int keyLength = key == null ? 0 : key.length; + final int valueLength = value == null ? 0 : value.length; + + final int totalLength = keyLength + valueLength; + cursorLimitManager.reportScannedBytes(totalLength); + cursorLimitManager.tryRecordScan(); + timer.increment(FDBStoreTimer.Counts.LOAD_INDEX_KEY); + timer.increment(FDBStoreTimer.Counts.LOAD_INDEX_KEY_BYTES, keyLength); + timer.increment(FDBStoreTimer.Counts.LOAD_INDEX_VALUE_BYTES, valueLength); + + switch (node.getKind()) { + case LEAF: + timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_LEAF_NODE_READ_BYTES, totalLength); + break; + case INTERMEDIATE: + timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_INTERMEDIATE_NODE_READ_BYTES, totalLength); + break; + default: + throw new RecordCoreException("unsupported kind of node"); + } + } + + @Override + public void onChildNodeDiscard(@Nonnull final ChildSlot childSlot) { + timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_CHILD_NODE_DISCARDS); + } + } + + static class OnWrite implements OnWriteListener { + @Nonnull + private final FDBStoreTimer timer; + + public OnWrite(@Nonnull final FDBStoreTimer timer) { + this.timer = timer; + } + + @Override + public CompletableFuture onAsyncReadForWrite(@Nonnull final CompletableFuture future) { + return timer.instrument(MultiDimensionalIndexHelper.Events.MULTIDIMENSIONAL_MODIFICATION, future); + } + + @Override + public void onNodeWritten(@Nonnull final Node node) { + switch (node.getKind()) { + case LEAF: + timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_LEAF_NODE_WRITES); + break; + case INTERMEDIATE: + timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_INTERMEDIATE_NODE_WRITES); + break; + default: + throw new RecordCoreException("unsupported kind of node"); + } + } + + @Override + public void onKeyValueWritten(@Nonnull final Node node, @Nullable final byte[] key, @Nullable final byte[] value) { + final int keyLength = key == null ? 0 : key.length; + final int valueLength = value == null ? 0 : value.length; + + final int totalLength = keyLength + valueLength; + timer.increment(FDBStoreTimer.Counts.SAVE_INDEX_KEY); + timer.increment(FDBStoreTimer.Counts.SAVE_INDEX_KEY_BYTES, keyLength); + timer.increment(FDBStoreTimer.Counts.SAVE_INDEX_VALUE_BYTES, valueLength); + + switch (node.getKind()) { + case LEAF: + timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_LEAF_NODE_WRITE_BYTES, totalLength); + break; + case INTERMEDIATE: + timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_INTERMEDIATE_NODE_WRITE_BYTES, totalLength); + break; + default: + throw new RecordCoreException("unsupported kind of node"); + } + } + } + + static class ItemSlotCursor extends AsyncIteratorCursor { + @Nonnull + private final CursorLimitManager cursorLimitManager; + @Nonnull + private final FDBStoreTimer timer; + + public ItemSlotCursor(@Nonnull final Executor executor, @Nonnull final AsyncIterator iterator, + @Nonnull final CursorLimitManager cursorLimitManager, @Nonnull final FDBStoreTimer timer) { + super(executor, iterator); + this.cursorLimitManager = cursorLimitManager; + this.timer = timer; + } + + @Nonnull + @Override + public CompletableFuture> onNext() { + if (nextResult != null && !nextResult.hasNext()) { + // This guard is needed to guarantee that if onNext is called multiple times after the cursor has + // returned a result without a value, then the same NoNextReason is returned each time. Without this guard, + // one might return SCAN_LIMIT_REACHED (for example) after returning a result with SOURCE_EXHAUSTED because + // of the tryRecordScan check. + return CompletableFuture.completedFuture(nextResult); + } else if (cursorLimitManager.tryRecordScan()) { + return iterator.onHasNext().thenApply(hasNext -> { + if (hasNext) { + final ItemSlot itemSlot = iterator.next(); + timer.increment(FDBStoreTimer.Counts.LOAD_SCAN_ENTRY); + timer.increment(FDBStoreTimer.Counts.LOAD_KEY_VALUE); + valuesSeen++; + nextResult = RecordCursorResult.withNextValue(itemSlot, new Continuation(itemSlot.getHilbertValue(), itemSlot.getKey())); + } else { + // Source iterator is exhausted. + nextResult = RecordCursorResult.exhausted(); + } + return nextResult; + }); + } else { // a limit must have been exceeded + final Optional stoppedReason = cursorLimitManager.getStoppedReason(); + if (stoppedReason.isEmpty()) { + throw new RecordCoreException("limit manager stopped cursor but did not report a reason"); + } + Verify.verifyNotNull(nextResult, "should have seen at least one record"); + nextResult = RecordCursorResult.withoutNextValue(nextResult.getContinuation(), stoppedReason.get()); + return CompletableFuture.completedFuture(nextResult); + } + } + } + + private static class Continuation implements RecordCursorContinuation { + @Nullable + final BigInteger lastHilbertValue; + @Nullable + final Tuple lastKey; + @Nullable + private ByteString cachedByteString; + @Nullable + private byte[] cachedBytes; + + private Continuation(@Nullable final BigInteger lastHilbertValue, @Nullable final Tuple lastKey) { + this.lastHilbertValue = lastHilbertValue; + this.lastKey = lastKey; + } + + @Nullable + public BigInteger getLastHilbertValue() { + return lastHilbertValue; + } + + @Nullable + public Tuple getLastKey() { + return lastKey; + } + + @Nonnull + @Override + public ByteString toByteString() { + if (isEnd()) { + return ByteString.EMPTY; + } + + if (cachedByteString == null) { + cachedByteString = RecordCursorProto.MultidimensionalIndexScanContinuation.newBuilder() + .setLastHilbertValue(ByteString.copyFrom(Objects.requireNonNull(lastHilbertValue).toByteArray())) + .setLastKey(ByteString.copyFrom(Objects.requireNonNull(lastKey).pack())) + .build() + .toByteString(); + } + return cachedByteString; + } + + @Nullable + @Override + public byte[] toBytes() { + if (isEnd()) { + return null; + } + if (cachedBytes == null) { + cachedBytes = toByteString().toByteArray(); + } + return cachedBytes; + } + + @Override + public boolean isEnd() { + return lastHilbertValue == null || lastKey == null; + } + + @Nullable + private static Continuation fromBytes(@Nullable byte[] continuationBytes) { + if (continuationBytes != null) { + final RecordCursorProto.MultidimensionalIndexScanContinuation parsed; + try { + parsed = RecordCursorProto.MultidimensionalIndexScanContinuation.parseFrom(continuationBytes); + } catch (InvalidProtocolBufferException ex) { + throw new RecordCoreException("error parsing continuation", ex) + .addLogInfo("raw_bytes", ByteArrayUtil2.loggable(continuationBytes)); + } + return new Continuation(new BigInteger(parsed.getLastHilbertValue().toByteArray()), + Tuple.fromBytes(parsed.getLastKey().toByteArray())); + } else { + return null; + } + } + } +} diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainerFactory.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainerFactory.java new file mode 100644 index 0000000000..6b6ba33db6 --- /dev/null +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainerFactory.java @@ -0,0 +1,158 @@ +/* + * VectorIndexMaintainerFactory.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2023 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.record.provider.foundationdb.indexes; + +import com.apple.foundationdb.annotation.API; +import com.apple.foundationdb.async.hnsw.HNSW.Config; +import com.apple.foundationdb.record.logging.LogMessageKeys; +import com.apple.foundationdb.record.metadata.Index; +import com.apple.foundationdb.record.metadata.IndexOptions; +import com.apple.foundationdb.record.metadata.IndexTypes; +import com.apple.foundationdb.record.metadata.IndexValidator; +import com.apple.foundationdb.record.metadata.MetaDataException; +import com.apple.foundationdb.record.metadata.MetaDataValidator; +import com.apple.foundationdb.record.provider.foundationdb.IndexMaintainer; +import com.apple.foundationdb.record.provider.foundationdb.IndexMaintainerFactory; +import com.apple.foundationdb.record.provider.foundationdb.IndexMaintainerState; +import com.google.auto.service.AutoService; + +import javax.annotation.Nonnull; +import java.util.Arrays; +import java.util.Set; + +/** + * A factory for {@link VectorIndexMaintainer} indexes. + */ +@AutoService(IndexMaintainerFactory.class) +@API(API.Status.EXPERIMENTAL) +public class VectorIndexMaintainerFactory implements IndexMaintainerFactory { + static final String[] TYPES = { IndexTypes.VECTOR}; + + @Override + @Nonnull + public Iterable getIndexTypes() { + return Arrays.asList(TYPES); + } + + @Override + @Nonnull + public IndexValidator getIndexValidator(Index index) { + return new IndexValidator(index) { + @Override + public void validate(@Nonnull MetaDataValidator metaDataValidator) { + super.validate(metaDataValidator); + validateNotVersion(); + validateStructure(); + } + + /** + * TODO. + */ + private void validateStructure() { + // + // There is no structural constraint on the key expression of the index. We just happen to interpret + // things in specific ways: + // + // - without GroupingKeyExpression: + // - one HNSW for the entire table (ungrouped HNSW) + // - first column of the expression gives us access to the field containing the vector + // - with GroupingKeyExpression: + // - one HNSW for each grouping prefix + // - first column in the grouped expression gives us access to the field containing the vector + // + // In any case, the vector is always a half-precision-encoded vector of dimensionality + // blob.length / 2 (for now). + // + // TODO We do not support extraneous columns to support advanced covering index scans for now. That + // Will probably encoded by a KeyWithValueExpression in the root position (but not now) + // + } + + @Override + public void validateChangedOptions(@Nonnull final Index oldIndex, + @Nonnull final Set changedOptions) { + if (!changedOptions.isEmpty()) { + // Allow changing from unspecified to the default (or vice versa), but not otherwise. + final Config oldOptions = VectorIndexHelper.getConfig(oldIndex); + final Config newOptions = VectorIndexHelper.getConfig(index); + if (changedOptions.contains(IndexOptions.HNSW_METRIC)) { + if (oldOptions.getMetric() != newOptions.getMetric()) { + throw new MetaDataException("HNSW metric changed", + LogMessageKeys.INDEX_NAME, index.getName()); + } + changedOptions.remove(IndexOptions.HNSW_METRIC); + } + if (changedOptions.contains(IndexOptions.HNSW_M)) { + if (oldOptions.getM() != newOptions.getM()) { + throw new MetaDataException("HNSW M changed", + LogMessageKeys.INDEX_NAME, index.getName()); + } + changedOptions.remove(IndexOptions.HNSW_M); + } + if (changedOptions.contains(IndexOptions.HNSW_M_MAX)) { + if (oldOptions.getMMax() != newOptions.getMMax()) { + throw new MetaDataException("HNSW mMax changed", + LogMessageKeys.INDEX_NAME, index.getName()); + } + changedOptions.remove(IndexOptions.HNSW_M_MAX); + } + if (changedOptions.contains(IndexOptions.HNSW_M_MAX_0)) { + if (oldOptions.getMMax0() != newOptions.getMMax0()) { + throw new MetaDataException("HNSW mMax0 changed", + LogMessageKeys.INDEX_NAME, index.getName()); + } + changedOptions.remove(IndexOptions.HNSW_M_MAX_0); + } + // efSearch can be overridden in every scenario + changedOptions.remove(IndexOptions.HNSW_EF_SEARCH); + if (changedOptions.contains(IndexOptions.HNSW_EF_CONSTRUCTION)) { + if (oldOptions.getEfConstruction() != newOptions.getEfConstruction()) { + throw new MetaDataException("HNSW efConstruction changed", + LogMessageKeys.INDEX_NAME, index.getName()); + } + changedOptions.remove(IndexOptions.HNSW_EF_CONSTRUCTION); + } + if (changedOptions.contains(IndexOptions.HNSW_EXTEND_CANDIDATES)) { + if (oldOptions.isExtendCandidates() != newOptions.isExtendCandidates()) { + throw new MetaDataException("HNSW extendCandidates changed", + LogMessageKeys.INDEX_NAME, index.getName()); + } + changedOptions.remove(IndexOptions.HNSW_EXTEND_CANDIDATES); + } + if (changedOptions.contains(IndexOptions.HNSW_KEEP_PRUNED_CONNECTIONS)) { + if (oldOptions.isKeepPrunedConnections() != newOptions.isKeepPrunedConnections()) { + throw new MetaDataException("HNSW keepPrunedConnections changed", + LogMessageKeys.INDEX_NAME, index.getName()); + } + changedOptions.remove(IndexOptions.HNSW_KEEP_PRUNED_CONNECTIONS); + } + } + super.validateChangedOptions(oldIndex, changedOptions); + } + }; + } + + @Override + @Nonnull + public IndexMaintainer getIndexMaintainer(@Nonnull final IndexMaintainerState state) { + return new VectorIndexMaintainer(state); + } +} diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java index 1305ab01c4..5b3cd5eefc 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java @@ -38,6 +38,7 @@ import com.apple.foundationdb.record.metadata.expressions.TupleFieldsHelper; import com.apple.foundationdb.record.planprotos.PComparison; import com.apple.foundationdb.record.planprotos.PComparison.PComparisonType; +import com.apple.foundationdb.record.planprotos.PDistanceRankValueComparison; import com.apple.foundationdb.record.planprotos.PInvertedFunctionComparison; import com.apple.foundationdb.record.planprotos.PListComparison; import com.apple.foundationdb.record.planprotos.PMultiColumnComparison; @@ -56,9 +57,6 @@ import com.apple.foundationdb.record.query.plan.cascades.ConstrainedBoolean; import com.apple.foundationdb.record.query.plan.cascades.Correlated; import com.apple.foundationdb.record.query.plan.cascades.CorrelationIdentifier; -import com.apple.foundationdb.record.query.plan.explain.DefaultExplainFormatter; -import com.apple.foundationdb.record.query.plan.explain.ExplainTokens; -import com.apple.foundationdb.record.query.plan.explain.ExplainTokensWithPrecedence; import com.apple.foundationdb.record.query.plan.cascades.UsesValueEquivalence; import com.apple.foundationdb.record.query.plan.cascades.ValueEquivalence; import com.apple.foundationdb.record.query.plan.cascades.WithValue; @@ -68,6 +66,9 @@ import com.apple.foundationdb.record.query.plan.cascades.values.QuantifiedObjectValue; import com.apple.foundationdb.record.query.plan.cascades.values.Value; import com.apple.foundationdb.record.query.plan.cascades.values.translation.TranslationMap; +import com.apple.foundationdb.record.query.plan.explain.DefaultExplainFormatter; +import com.apple.foundationdb.record.query.plan.explain.ExplainTokens; +import com.apple.foundationdb.record.query.plan.explain.ExplainTokensWithPrecedence; import com.apple.foundationdb.record.query.plan.plans.QueryResult; import com.apple.foundationdb.record.query.plan.serialization.PlanSerialization; import com.apple.foundationdb.record.util.ProtoUtils; @@ -1504,6 +1505,12 @@ public static class ValueComparison implements Comparison { @Nonnull private final Supplier hashCodeSupplier; + protected ValueComparison(@Nonnull final PlanSerializationContext serializationContext, + @Nonnull final PValueComparison valueComparisonProto) { + this(Type.fromProto(serializationContext, Objects.requireNonNull(valueComparisonProto.getType())), + Value.fromValueProto(serializationContext, Objects.requireNonNull(valueComparisonProto.getComparandValue()))); + } + public ValueComparison(@Nonnull final Type type, @Nonnull final Value comparandValue) { this(type, comparandValue, ParameterRelationshipGraph.unbound()); @@ -1660,7 +1667,7 @@ public int hashCode() { } public int computeHashCode() { - return Objects.hash(type.name(), relatedByEquality()); + return Objects.hash(type.name(), getComparandValue(), relatedByEquality()); } private Set relatedByEquality() { @@ -1687,7 +1694,12 @@ public Comparison withParameterRelationshipMap(@Nonnull final ParameterRelations @Nonnull @Override - public PValueComparison toProto(@Nonnull final PlanSerializationContext serializationContext) { + public Message toProto(@Nonnull final PlanSerializationContext serializationContext) { + return toValueComparisonProto(serializationContext); + } + + @Nonnull + public PValueComparison toValueComparisonProto(@Nonnull final PlanSerializationContext serializationContext) { return PValueComparison.newBuilder() .setType(type.toProto(serializationContext)) .setComparandValue(comparandValue.toValueProto(serializationContext)) @@ -1697,14 +1709,13 @@ public PValueComparison toProto(@Nonnull final PlanSerializationContext serializ @Nonnull @Override public PComparison toComparisonProto(@Nonnull final PlanSerializationContext serializationContext) { - return PComparison.newBuilder().setValueComparison(toProto(serializationContext)).build(); + return PComparison.newBuilder().setValueComparison(toValueComparisonProto(serializationContext)).build(); } @Nonnull public static ValueComparison fromProto(@Nonnull final PlanSerializationContext serializationContext, @Nonnull final PValueComparison valueComparisonProto) { - return new ValueComparison(Type.fromProto(serializationContext, Objects.requireNonNull(valueComparisonProto.getType())), - Value.fromValueProto(serializationContext, Objects.requireNonNull(valueComparisonProto.getComparandValue()))); + return new ValueComparison(serializationContext, valueComparisonProto); } /** @@ -1727,6 +1738,198 @@ public ValueComparison fromProto(@Nonnull final PlanSerializationContext seriali } } + public static class DistanceRankValueComparison extends ValueComparison { + private static final ObjectPlanHash BASE_HASH = new ObjectPlanHash("Distance-Rank-Value-Comparison"); + + @Nonnull + private final Value limitValue; + + protected DistanceRankValueComparison(@Nonnull PlanSerializationContext serializationContext, + @Nonnull final PDistanceRankValueComparison distanceRankValueComparisonProto) { + super(serializationContext, distanceRankValueComparisonProto.getSuper()); + this.limitValue = Value.fromValueProto(serializationContext, + Objects.requireNonNull(distanceRankValueComparisonProto.getLimitValue())); + } + + public DistanceRankValueComparison(@Nonnull final Type type, @Nonnull final Value comparandValue, + @Nonnull final Value limitValue) { + this(type, comparandValue, ParameterRelationshipGraph.unbound(), limitValue); + } + + public DistanceRankValueComparison(@Nonnull final Type type, @Nonnull final Value comparandValue, + @Nonnull final ParameterRelationshipGraph parameterRelationshipGraph, + @Nonnull final Value limitValue) { + super(type, comparandValue, parameterRelationshipGraph); + this.limitValue = limitValue; + } + + @Nonnull + public Value getLimitValue() { + return limitValue; + } + + @Nonnull + @Override + public Comparison withType(@Nonnull final Type newType) { + if (getType() == newType) { + return this; + } + return new ValueComparison(newType, getComparandValue(), parameterRelationshipGraph); + } + + @Nonnull + @Override + @SuppressWarnings("PMD.CompareObjectsWithEquals") + public ValueComparison withValue(@Nonnull final Value value) { + if (getComparandValue() == value) { + return this; + } + return new ValueComparison(getType(), value); + } + + @Nonnull + @Override + @SuppressWarnings("PMD.CompareObjectsWithEquals") + public Optional replaceValuesMaybe(@Nonnull final Function> replacementFunction) { + return replacementFunction.apply(getComparandValue()) + .flatMap(replacedComparandValue -> + replacementFunction.apply(getLimitValue()).map(replacedLimitValue -> { + if (replacedComparandValue == getComparandValue() && + replacedLimitValue == getLimitValue()) { + return this; + } + return new DistanceRankValueComparison(getType(), replacedComparandValue, parameterRelationshipGraph, + replacedLimitValue); + })); + } + + @Nonnull + @Override + public Comparison translateCorrelations(@Nonnull final TranslationMap translationMap, final boolean shouldSimplifyValues) { + if (getComparandValue().getCorrelatedTo() + .stream() + .noneMatch(translationMap::containsSourceAlias) && + getLimitValue().getCorrelatedTo() + .stream() + .noneMatch(translationMap::containsSourceAlias)) { + return this; + } + + return new DistanceRankValueComparison(getType(), + getComparandValue().translateCorrelations(translationMap, shouldSimplifyValues), + parameterRelationshipGraph, + getLimitValue().translateCorrelations(translationMap, shouldSimplifyValues)); + } + + @Nonnull + @Override + public Set getCorrelatedTo() { + return ImmutableSet.builder() + .addAll(getComparandValue().getCorrelatedTo()) + .addAll(getLimitValue().getCorrelatedTo()) + .build(); + } + + @Nonnull + @Override + public ConstrainedBoolean semanticEqualsTyped(@Nonnull final Comparison other, @Nonnull final ValueEquivalence valueEquivalence) { + return super.semanticEqualsTyped(other, valueEquivalence) + .compose(ignored -> getLimitValue() + .semanticEquals(((DistanceRankValueComparison)other).getLimitValue(), + valueEquivalence)); + } + + @Nullable + @Override + @SuppressWarnings("PMD.CompareObjectsWithEquals") + public Boolean eval(@Nullable FDBRecordStoreBase store, @Nonnull EvaluationContext context, @Nullable Object v) { + throw new RecordCoreException("this comparison can only be evaluated using an index"); + } + + @Nonnull + @Override + public String typelessString() { + return getComparandValue() + ":" + getLimitValue(); + } + + @Override + public String toString() { + return explain().getExplainTokens().render(DefaultExplainFormatter.forDebugging()).toString(); + } + + @Nonnull + @Override + public ExplainTokensWithPrecedence explain() { + return ExplainTokensWithPrecedence.of(new ExplainTokens().addKeyword(getType().name()) + .addWhitespace().addNested(getComparandValue().explain().getExplainTokens()) + .addKeyword(":").addWhitespace() + .addNested(getLimitValue().explain().getExplainTokens())); + } + + public int computeHashCode() { + return Objects.hash(getType().name(), getComparandValue(), getLimitValue()); + } + + @Override + public int planHash(@Nonnull final PlanHashMode mode) { + switch (mode.getKind()) { + case LEGACY: + case FOR_CONTINUATION: + return PlanHashable.objectsPlanHash(mode, BASE_HASH, getType(), getComparandValue(), getLimitValue()); + default: + throw new UnsupportedOperationException("Hash Kind " + mode.name() + " is not supported"); + } + } + + @Nonnull + @Override + public Comparison withParameterRelationshipMap(@Nonnull final ParameterRelationshipGraph parameterRelationshipGraph) { + Verify.verify(this.parameterRelationshipGraph.isUnbound()); + return new DistanceRankValueComparison(getType(), getComparandValue(), parameterRelationshipGraph, + getLimitValue()); + } + + @Nonnull + @Override + public PDistanceRankValueComparison toProto(@Nonnull final PlanSerializationContext serializationContext) { + return PDistanceRankValueComparison.newBuilder() + .setSuper(super.toValueComparisonProto(serializationContext)) + .setLimitValue(getLimitValue().toValueProto(serializationContext)) + .build(); + } + + @Nonnull + @Override + public PComparison toComparisonProto(@Nonnull final PlanSerializationContext serializationContext) { + return PComparison.newBuilder().setDistanceRankValueComparison(toProto(serializationContext)).build(); + } + + @Nonnull + public static DistanceRankValueComparison fromProto(@Nonnull final PlanSerializationContext serializationContext, + @Nonnull final PDistanceRankValueComparison distanceRankValueComparisonProto) { + return new DistanceRankValueComparison(serializationContext, distanceRankValueComparisonProto); + } + + /** + * Deserializer. + */ + @AutoService(PlanDeserializer.class) + public static class Deserializer implements PlanDeserializer { + @Nonnull + @Override + public Class getProtoMessageClass() { + return PDistanceRankValueComparison.class; + } + + @Nonnull + @Override + public DistanceRankValueComparison fromProto(@Nonnull final PlanSerializationContext serializationContext, + @Nonnull final PDistanceRankValueComparison distanceRankValueComparisonProto) { + return DistanceRankValueComparison.fromProto(serializationContext, distanceRankValueComparisonProto); + } + } + } + /** * A comparison with a list of values. */ diff --git a/fdb-record-layer-core/src/main/proto/record_query_plan.proto b/fdb-record-layer-core/src/main/proto/record_query_plan.proto index e10c907404..c338c61a2a 100644 --- a/fdb-record-layer-core/src/main/proto/record_query_plan.proto +++ b/fdb-record-layer-core/src/main/proto/record_query_plan.proto @@ -1205,6 +1205,7 @@ message PComparison { PRecordTypeComparison record_type_comparison = 10; PConversionSimpleComparison conversion_simple_comparison = 11; PConversionParameterComparison conversion_parameter_comparison = 12; + PDistanceRankValueComparison distance_rank_value_comparison = 13; } } @@ -1271,6 +1272,11 @@ message PRecordTypeComparison { optional string record_type_name = 1; } +message PDistanceRankValueComparison { + optional PValueComparison super = 1; + optional PValue limitValue = 2; +} + // // Query Predicates // From b19042c41d5c8f66ac18a5ab84042f7b0f20f379 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Wed, 13 Aug 2025 14:55:18 +0200 Subject: [PATCH 22/34] index maintenance code complete --- .../foundationdb/VectorIndexScanBounds.java | 237 ++++++++++++ .../VectorIndexScanComparisons.java | 355 ++++++++++++++++++ 2 files changed, 592 insertions(+) create mode 100644 fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanBounds.java create mode 100644 fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanComparisons.java diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanBounds.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanBounds.java new file mode 100644 index 0000000000..f1437f9318 --- /dev/null +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanBounds.java @@ -0,0 +1,237 @@ +/* + * MultidimensionalIndexScanBounds.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2022 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.record.provider.foundationdb; + +import com.apple.foundationdb.annotation.API; +import com.apple.foundationdb.async.rtree.RTree; +import com.apple.foundationdb.record.IndexScanType; +import com.apple.foundationdb.record.TupleRange; +import com.apple.foundationdb.tuple.Tuple; +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import javax.annotation.Nonnull; +import java.util.List; +import java.util.stream.Collectors; + +/** + * {@link IndexScanBounds} for a multidimensional index scan. A multidimensional scan bounds object contains a + * {@link #prefixRange} and a {@link SpatialPredicate} which can be almost arbitrarily complex. The prefix range + * is a regular tuple range informing the index maintainer how to constrain the search over the non-multidimensional + * fields that can be viewed as a prefix whose data is stored in a regular one-dimensional index. The spatial predicate + * implements methods to quickly establish geometric overlap and containment. Spatial predicates do have internal + * structure as they can delegate to contained other spatial predicates thus allowing to form e.g. logical conjuncts + * or disjuncts. + */ +@API(API.Status.EXPERIMENTAL) +public class MultidimensionalIndexScanBounds implements IndexScanBounds { + @Nonnull + private final TupleRange prefixRange; + + @Nonnull + private final SpatialPredicate spatialPredicate; + + @Nonnull + private final TupleRange suffixRange; + + public MultidimensionalIndexScanBounds(@Nonnull final TupleRange prefixRange, + @Nonnull final SpatialPredicate spatialPredicate, + @Nonnull final TupleRange suffixRange) { + this.prefixRange = prefixRange; + this.spatialPredicate = spatialPredicate; + this.suffixRange = suffixRange; + } + + @Nonnull + @Override + public IndexScanType getScanType() { + return IndexScanType.BY_VALUE; + } + + @Nonnull + public TupleRange getPrefixRange() { + return prefixRange; + } + + @Nonnull + public SpatialPredicate getSpatialPredicate() { + return spatialPredicate; + } + + @Nonnull + public TupleRange getSuffixRange() { + return suffixRange; + } + + /** + * Method to compute if the rectangle handed in overlaps with this scan bounds object. This method is invoked when + * the R-tree data structure of a multidimensional index is searched. Note that this method can be implemented using + * a best-effort approach as it is permissible to indicate overlap between {@code mbr} and {@code this} when there + * is in fact no overlap. The rate of false-positives directly influences the search performance in the + * multidimensional index. + * @param mbr the minimum-bounding {@link RTree.Rectangle} + * @return {@code true} if {@code this} overlaps with {@code mbr} + */ + public boolean overlapsMbrApproximately(@Nonnull RTree.Rectangle mbr) { + return spatialPredicate.overlapsMbrApproximately(mbr); + } + + /** + * Method to compute if the point handed in is contained by this scan bounds object. + * @param position the {@link RTree.Point} + * @return {@code true} if {@code position} is contained by {@code this} + */ + public boolean containsPosition(@Nonnull RTree.Point position) { + return spatialPredicate.containsPosition(position); + } + + /** + * Spatial predicate. The implementing classes form a boolean algebra of sorts. Most notably {@link Hypercube} + * represents the logical variables, while {@link And} and {@link Or} can be used to build up more complex powerful + * bounds. + */ + public interface SpatialPredicate { + SpatialPredicate TAUTOLOGY = new SpatialPredicate() { + @Override + public boolean overlapsMbrApproximately(@Nonnull final RTree.Rectangle mbr) { + return true; + } + + @Override + public boolean containsPosition(@Nonnull final RTree.Point position) { + return true; + } + }; + + boolean overlapsMbrApproximately(@Nonnull RTree.Rectangle mbr); + + boolean containsPosition(@Nonnull RTree.Point position); + } + + /** + * Scan bounds that consists of other {@link SpatialPredicate}s to form a logical OR. + */ + public static class Or implements SpatialPredicate { + @Nonnull + private final List children; + + public Or(@Nonnull final List children) { + this.children = ImmutableList.copyOf(children); + } + + @Override + public boolean overlapsMbrApproximately(@Nonnull final RTree.Rectangle mbr) { + return children.stream() + .anyMatch(child -> child.overlapsMbrApproximately(mbr)); + } + + @Override + public boolean containsPosition(@Nonnull final RTree.Point position) { + return children.stream() + .anyMatch(child -> child.containsPosition(position)); + } + + @Override + public String toString() { + return children.stream().map(Object::toString).collect(Collectors.joining(" or ")); + } + } + + /** + * Scan bounds that consists of other {@link SpatialPredicate}s to form a logical AND. + */ + public static class And implements SpatialPredicate { + @Nonnull + private final List children; + + public And(@Nonnull final List children) { + this.children = ImmutableList.copyOf(children); + } + + @Override + public boolean overlapsMbrApproximately(@Nonnull final RTree.Rectangle mbr) { + return children.stream() + .allMatch(child -> child.overlapsMbrApproximately(mbr)); + } + + @Override + public boolean containsPosition(@Nonnull final RTree.Point position) { + return children.stream() + .allMatch(child -> child.containsPosition(position)); + } + + @Override + public String toString() { + return children.stream().map(Object::toString).collect(Collectors.joining(" and ")); + } + } + + /** + * Scan bounds describing an n-dimensional hypercube. + */ + public static class Hypercube implements SpatialPredicate { + @Nonnull + private final List dimensionRanges; + + public Hypercube(@Nonnull final List dimensionRanges) { + this.dimensionRanges = ImmutableList.copyOf(dimensionRanges); + } + + @Override + public boolean overlapsMbrApproximately(@Nonnull final RTree.Rectangle mbr) { + Preconditions.checkArgument(mbr.getNumDimensions() == dimensionRanges.size()); + + for (int d = 0; d < mbr.getNumDimensions(); d++) { + final Tuple lowTuple = Tuple.from(mbr.getLow(d)); + final Tuple highTuple = Tuple.from(mbr.getHigh(d)); + final TupleRange dimensionRange = dimensionRanges.get(d); + if (!dimensionRange.overlaps(lowTuple, highTuple)) { + return false; + } + } + return true; + } + + @Override + public boolean containsPosition(@Nonnull final RTree.Point position) { + Preconditions.checkArgument(position.getNumDimensions() == dimensionRanges.size()); + + for (int d = 0; d < position.getNumDimensions(); d++) { + final Tuple coordinate = Tuple.from(position.getCoordinate(d)); + final TupleRange dimensionRange = dimensionRanges.get(d); + if (!dimensionRange.contains(coordinate)) { + return false; + } + } + return true; + } + + @Nonnull + public List getDimensionRanges() { + return dimensionRanges; + } + + @Override + public String toString() { + return "HyperCube:[" + dimensionRanges + "]"; + } + } +} diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanComparisons.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanComparisons.java new file mode 100644 index 0000000000..5b6967e54c --- /dev/null +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanComparisons.java @@ -0,0 +1,355 @@ +/* + * MultidimensionalIndexScanComparisons.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2022 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.record.provider.foundationdb; + +import com.apple.foundationdb.annotation.API; +import com.apple.foundationdb.annotation.SpotBugsSuppressWarnings; +import com.apple.foundationdb.record.EvaluationContext; +import com.apple.foundationdb.record.IndexScanType; +import com.apple.foundationdb.record.PlanDeserializer; +import com.apple.foundationdb.record.PlanHashable; +import com.apple.foundationdb.record.PlanSerializationContext; +import com.apple.foundationdb.record.TupleRange; +import com.apple.foundationdb.record.metadata.Index; +import com.apple.foundationdb.record.planprotos.PIndexScanParameters; +import com.apple.foundationdb.record.planprotos.PMultidimensionalIndexScanComparisons; +import com.apple.foundationdb.record.provider.foundationdb.MultidimensionalIndexScanBounds.Hypercube; +import com.apple.foundationdb.record.query.plan.ScanComparisons; +import com.apple.foundationdb.record.query.plan.cascades.AliasMap; +import com.apple.foundationdb.record.query.plan.cascades.CorrelationIdentifier; +import com.apple.foundationdb.record.query.plan.explain.ExplainTokens; +import com.apple.foundationdb.record.query.plan.explain.ExplainTokensWithPrecedence; +import com.apple.foundationdb.record.query.plan.cascades.explain.Attribute; +import com.apple.foundationdb.record.query.plan.cascades.values.translation.TranslationMap; +import com.google.auto.service.AutoService; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.util.List; +import java.util.Objects; +import java.util.Set; + +/** + * {@link ScanComparisons} for use in a multidimensional index scan. + */ +@API(API.Status.UNSTABLE) +public class MultidimensionalIndexScanComparisons implements IndexScanParameters { + @Nonnull + private final ScanComparisons prefixScanComparisons; + @Nonnull + private final List dimensionsScanComparisons; + @Nonnull + private final ScanComparisons suffixScanComparisons; + + public MultidimensionalIndexScanComparisons(@Nonnull final ScanComparisons prefixScanComparisons, + @Nonnull final List dimensionsScanComparisons, + @Nonnull final ScanComparisons suffixKeyComparisonRanges) { + this.prefixScanComparisons = prefixScanComparisons; + this.dimensionsScanComparisons = dimensionsScanComparisons; + this.suffixScanComparisons = suffixKeyComparisonRanges; + } + + @Nonnull + @Override + public IndexScanType getScanType() { + return IndexScanType.BY_VALUE; + } + + @Nonnull + public ScanComparisons getPrefixScanComparisons() { + return prefixScanComparisons; + } + + @Nonnull + public List getDimensionsScanComparisons() { + return dimensionsScanComparisons; + } + + @Nonnull + public ScanComparisons getSuffixScanComparisons() { + return suffixScanComparisons; + } + + @Nonnull + @Override + public MultidimensionalIndexScanBounds bind(@Nonnull final FDBRecordStoreBase store, @Nonnull final Index index, + @Nonnull final EvaluationContext context) { + final ImmutableList.Builder dimensionsTupleRangeBuilder = ImmutableList.builder(); + for (final ScanComparisons dimensionScanComparison : dimensionsScanComparisons) { + dimensionsTupleRangeBuilder.add(dimensionScanComparison.toTupleRange(store, context)); + } + final Hypercube hypercube = new Hypercube(dimensionsTupleRangeBuilder.build()); + return new MultidimensionalIndexScanBounds(prefixScanComparisons.toTupleRange(store, context), + hypercube, suffixScanComparisons.toTupleRange(store, context)); + } + + @Override + public int planHash(@Nonnull PlanHashMode mode) { + return PlanHashable.objectsPlanHash(mode, prefixScanComparisons, dimensionsScanComparisons, + suffixScanComparisons); + } + + @Override + public boolean isUnique(@Nonnull Index index) { + return prefixScanComparisons.isEquality() && prefixScanComparisons.size() == index.getColumnSize(); + } + + @Nonnull + @Override + public ExplainTokensWithPrecedence explain() { + @Nullable var tupleRange = prefixScanComparisons.toTupleRangeWithoutContext(); + final var prefix = tupleRange == null + ? prefixScanComparisons.explain().getExplainTokens() + : new ExplainTokens().addToString(tupleRange); + + final var dimensions = + new ExplainTokens().addSequence(() -> new ExplainTokens().addCommaAndWhiteSpace(), + () -> dimensionsScanComparisons.stream() + .map(dimensionScanComparisons -> { + @Nullable var dimensionTupleRange = dimensionScanComparisons.toTupleRangeWithoutContext(); + return dimensionTupleRange == null + ? dimensionScanComparisons.explain().getExplainTokens() + : new ExplainTokens().addToString(dimensionTupleRange); + }).iterator()); + + + tupleRange = suffixScanComparisons.toTupleRangeWithoutContext(); + final var suffix = tupleRange == null + ? suffixScanComparisons.explain().getExplainTokens() + : new ExplainTokens().addToString(tupleRange); + + return ExplainTokensWithPrecedence.of(prefix.addOptionalWhitespace().addToString(":{").addOptionalWhitespace() + .addNested(dimensions).addOptionalWhitespace().addToString("}:").addOptionalWhitespace().addNested(suffix)); + } + + @Override + public void getPlannerGraphDetails(@Nonnull ImmutableList.Builder detailsBuilder, @Nonnull ImmutableMap.Builder attributeMapBuilder) { + @Nullable TupleRange tupleRange = prefixScanComparisons.toTupleRangeWithoutContext(); + if (tupleRange != null) { + detailsBuilder.add("prefix: " + tupleRange.getLowEndpoint().toString(false) + "{{plow}}, {{phigh}}" + tupleRange.getHighEndpoint().toString(true)); + attributeMapBuilder.put("plow", Attribute.gml(tupleRange.getLow() == null ? "-∞" : tupleRange.getLow().toString())); + attributeMapBuilder.put("phigh", Attribute.gml(tupleRange.getHigh() == null ? "∞" : tupleRange.getHigh().toString())); + } else { + detailsBuilder.add("prefix comparisons: {{pcomparisons}}"); + attributeMapBuilder.put("pcomparisons", Attribute.gml(prefixScanComparisons.toString())); + } + + for (int d = 0; d < dimensionsScanComparisons.size(); d++) { + final ScanComparisons dimensionScanComparisons = dimensionsScanComparisons.get(d); + tupleRange = dimensionScanComparisons.toTupleRangeWithoutContext(); + if (tupleRange != null) { + detailsBuilder.add("dim" + d + ": " + tupleRange.getLowEndpoint().toString(false) + "{{dlow" + d + "}}, {{dhigh" + d + "}}" + tupleRange.getHighEndpoint().toString(true)); + attributeMapBuilder.put("dlow" + d, Attribute.gml(tupleRange.getLow() == null ? "-∞" : tupleRange.getLow().toString())); + attributeMapBuilder.put("dhigh" + d, Attribute.gml(tupleRange.getHigh() == null ? "∞" : tupleRange.getHigh().toString())); + } else { + detailsBuilder.add("dim" + d + " comparisons: " + "{{dcomparisons" + d + "}}"); + attributeMapBuilder.put("dcomparisons" + d, Attribute.gml(dimensionScanComparisons.toString())); + } + } + + tupleRange = suffixScanComparisons.toTupleRangeWithoutContext(); + if (tupleRange != null) { + detailsBuilder.add("suffix: " + tupleRange.getLowEndpoint().toString(false) + "{{slow}}, {{shigh}}" + tupleRange.getHighEndpoint().toString(true)); + attributeMapBuilder.put("slow", Attribute.gml(tupleRange.getLow() == null ? "-∞" : tupleRange.getLow().toString())); + attributeMapBuilder.put("shigh", Attribute.gml(tupleRange.getHigh() == null ? "∞" : tupleRange.getHigh().toString())); + } else { + detailsBuilder.add("suffix comparisons: {{scomparisons}}"); + attributeMapBuilder.put("scomparisons", Attribute.gml(suffixScanComparisons.toString())); + } + } + + @Nonnull + @Override + public Set getCorrelatedTo() { + final ImmutableSet.Builder correlatedToBuilder = ImmutableSet.builder(); + correlatedToBuilder.addAll(prefixScanComparisons.getCorrelatedTo()); + correlatedToBuilder.addAll(dimensionsScanComparisons.stream() + .flatMap(dimensionScanComparison -> dimensionScanComparison.getCorrelatedTo().stream()).iterator()); + correlatedToBuilder.addAll(suffixScanComparisons.getCorrelatedTo()); + return correlatedToBuilder.build(); + } + + @Nonnull + @Override + public IndexScanParameters rebase(@Nonnull final AliasMap translationMap) { + return translateCorrelations(TranslationMap.rebaseWithAliasMap(translationMap), false); + } + + @Override + @SuppressWarnings("PMD.CompareObjectsWithEquals") + public boolean semanticEquals(@Nullable final Object other, @Nonnull final AliasMap aliasMap) { + if (this == other) { + return true; + } + if (other == null || getClass() != other.getClass()) { + return false; + } + + final MultidimensionalIndexScanComparisons that = (MultidimensionalIndexScanComparisons)other; + + if (!prefixScanComparisons.semanticEquals(that.prefixScanComparisons, aliasMap)) { + return false; + } + if (dimensionsScanComparisons.size() != that.dimensionsScanComparisons.size()) { + return false; + } + for (int i = 0; i < dimensionsScanComparisons.size(); i++) { + final ScanComparisons dimensionScanComparison = dimensionsScanComparisons.get(i); + final ScanComparisons otherDimensionScanComparison = that.dimensionsScanComparisons.get(i); + if (!dimensionScanComparison.semanticEquals(otherDimensionScanComparison, aliasMap)) { + return false; + } + } + return suffixScanComparisons.semanticEquals(that.suffixScanComparisons, aliasMap); + } + + @Override + public int semanticHashCode() { + int hashCode = prefixScanComparisons.semanticHashCode(); + for (final ScanComparisons dimensionScanComparison : dimensionsScanComparisons) { + hashCode = 31 * hashCode + dimensionScanComparison.semanticHashCode(); + } + return 31 * hashCode + suffixScanComparisons.semanticHashCode(); + } + + @Nonnull + @Override + @SuppressWarnings("PMD.CompareObjectsWithEquals") + public IndexScanParameters translateCorrelations(@Nonnull final TranslationMap translationMap, + final boolean shouldSimplifyValues) { + final ScanComparisons translatedPrefixScanComparisons = + prefixScanComparisons.translateCorrelations(translationMap, shouldSimplifyValues); + + final ImmutableList.Builder translatedDimensionScanComparisonBuilder = ImmutableList.builder(); + boolean isSameDimensionsScanComparisons = true; + for (final ScanComparisons dimensionScanComparisons : dimensionsScanComparisons) { + final ScanComparisons translatedDimensionScanComparison = + dimensionScanComparisons.translateCorrelations(translationMap, shouldSimplifyValues); + if (translatedDimensionScanComparison != dimensionScanComparisons) { + isSameDimensionsScanComparisons = false; + } + translatedDimensionScanComparisonBuilder.add(translatedDimensionScanComparison); + } + + final ScanComparisons translatedSuffixKeyScanComparisons = + suffixScanComparisons.translateCorrelations(translationMap, shouldSimplifyValues); + + if (translatedPrefixScanComparisons != prefixScanComparisons || !isSameDimensionsScanComparisons || + translatedSuffixKeyScanComparisons != suffixScanComparisons) { + return withComparisons(translatedPrefixScanComparisons, translatedDimensionScanComparisonBuilder.build(), + translatedSuffixKeyScanComparisons); + } + return this; + } + + @Nonnull + protected MultidimensionalIndexScanComparisons withComparisons(@Nonnull final ScanComparisons prefixScanComparisons, + @Nonnull final List dimensionsComparisonRanges, + @Nonnull final ScanComparisons suffixKeyScanComparisons) { + return new MultidimensionalIndexScanComparisons(prefixScanComparisons, dimensionsComparisonRanges, + suffixKeyScanComparisons); + } + + @Override + public String toString() { + return "BY_VALUE(MD):" + prefixScanComparisons + ":" + dimensionsScanComparisons + ":" + suffixScanComparisons; + } + + @Override + @SpotBugsSuppressWarnings("EQ_UNUSUAL") + @SuppressWarnings("EqualsWhichDoesntCheckParameterClass") + public boolean equals(final Object o) { + return semanticEquals(o, AliasMap.emptyMap()); + } + + @Override + public int hashCode() { + return semanticHashCode(); + } + + @Nonnull + @Override + public PMultidimensionalIndexScanComparisons toProto(@Nonnull final PlanSerializationContext serializationContext) { + final PMultidimensionalIndexScanComparisons.Builder builder = PMultidimensionalIndexScanComparisons.newBuilder(); + builder.setPrefixScanComparisons(prefixScanComparisons.toProto(serializationContext)); + for (final ScanComparisons dimensionsScanComparison : dimensionsScanComparisons) { + builder.addDimensionsScanComparisons(dimensionsScanComparison.toProto(serializationContext)); + } + builder.setSuffixScanComparisons(suffixScanComparisons.toProto(serializationContext)); + return builder.build(); + } + + @Nonnull + @Override + public PIndexScanParameters toIndexScanParametersProto(@Nonnull final PlanSerializationContext serializationContext) { + return PIndexScanParameters.newBuilder().setMultidimensionalIndexScanComparisons(toProto(serializationContext)).build(); + } + + @Nonnull + public static MultidimensionalIndexScanComparisons fromProto(@Nonnull final PlanSerializationContext serializationContext, + @Nonnull final PMultidimensionalIndexScanComparisons multidimensionalIndexScanComparisonsProto) { + final ImmutableList.Builder dimensionScanComparisonsBuilder = ImmutableList.builder(); + for (int i = 0; i < multidimensionalIndexScanComparisonsProto.getDimensionsScanComparisonsCount(); i ++) { + dimensionScanComparisonsBuilder.add(ScanComparisons.fromProto(serializationContext, + multidimensionalIndexScanComparisonsProto.getDimensionsScanComparisons(i))); + } + return new MultidimensionalIndexScanComparisons(ScanComparisons.fromProto(serializationContext, Objects.requireNonNull(multidimensionalIndexScanComparisonsProto.getPrefixScanComparisons())), + dimensionScanComparisonsBuilder.build(), + ScanComparisons.fromProto(serializationContext, Objects.requireNonNull(multidimensionalIndexScanComparisonsProto.getSuffixScanComparisons()))); + } + + @Nonnull + public static MultidimensionalIndexScanComparisons byValue(@Nullable ScanComparisons prefixScanComparisons, + @Nonnull final List dimensionsComparisonRanges, + @Nullable ScanComparisons suffixKeyScanComparisons) { + if (prefixScanComparisons == null) { + prefixScanComparisons = ScanComparisons.EMPTY; + } + + if (suffixKeyScanComparisons == null) { + suffixKeyScanComparisons = ScanComparisons.EMPTY; + } + + return new MultidimensionalIndexScanComparisons(prefixScanComparisons, dimensionsComparisonRanges, suffixKeyScanComparisons); + } + + /** + * Deserializer. + */ + @AutoService(PlanDeserializer.class) + public static class Deserializer implements PlanDeserializer { + @Nonnull + @Override + public Class getProtoMessageClass() { + return PMultidimensionalIndexScanComparisons.class; + } + + @Nonnull + @Override + public MultidimensionalIndexScanComparisons fromProto(@Nonnull final PlanSerializationContext serializationContext, + @Nonnull final PMultidimensionalIndexScanComparisons multidimensionalIndexScanComparisonsProto) { + return MultidimensionalIndexScanComparisons.fromProto(serializationContext, multidimensionalIndexScanComparisonsProto); + } + } +} From 4bf0083e0e595993e7b7c1817f622171643db7fe Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Wed, 13 Aug 2025 14:55:48 +0200 Subject: [PATCH 23/34] index maintenance code complete --- .../async/hnsw/CompactStorageAdapter.java | 4 +- .../foundationdb/async/hnsw/HNSWHelpers.java | 2 +- .../async/hnsw/InliningStorageAdapter.java | 7 +- .../async/hnsw/NodeReferenceWithDistance.java | 2 +- .../async/hnsw/OnWriteListener.java | 3 +- .../async/hnsw/StorageAdapter.java | 22 +- .../apple/foundationdb/async/hnsw/Vector.java | 29 ++ .../fdb-record-layer-core.gradle | 1 + .../provider/foundationdb/FDBStoreTimer.java | 8 + .../foundationdb/VectorIndexScanBounds.java | 214 ++------- .../VectorIndexScanComparisons.java | 181 ++++--- .../indexes/VectorIndexHelper.java | 3 +- .../indexes/VectorIndexMaintainer.java | 440 +++++++----------- .../record/query/expressions/Comparisons.java | 27 +- .../src/main/proto/record_cursor.proto | 9 + .../src/main/proto/record_query_plan.proto | 10 + 16 files changed, 407 insertions(+), 555 deletions(-) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java index f88b679d60..b590513019 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java @@ -143,8 +143,10 @@ public void writeNodeInternal(@Nonnull final Transaction transaction, @Nonnull f final Tuple nodeTuple = Tuple.fromList(nodeItems); - transaction.set(key, nodeTuple.pack()); + final byte[] value = nodeTuple.pack(); + transaction.set(key, value); getOnWriteListener().onNodeWritten(layer, node); + getOnWriteListener().onKeyValueWritten(layer, key, value); if (logger.isDebugEnabled()) { logger.debug("written neighbors of primaryKey={}, oldSize={}, newSize={}", node.getPrimaryKey(), diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSWHelpers.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSWHelpers.java index 28d66df5fa..322b4f85b0 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSWHelpers.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSWHelpers.java @@ -41,7 +41,7 @@ private HNSWHelpers() { * @return a {@link String} containing the hexadecimal representation of the byte array passed in */ @Nonnull - static String bytesToHex(byte[] bytes) { + public static String bytesToHex(byte[] bytes) { char[] hexChars = new char[bytes.length * 2]; for (int j = 0; j < bytes.length; j++) { int v = bytes[j] & 0xFF; diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java index d025116a28..db55b29597 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java @@ -117,9 +117,12 @@ private byte[] getNodeKey(final int layer, @Nonnull final Tuple primaryKey) { public void writeNeighbor(@Nonnull final Transaction transaction, final int layer, @Nonnull final Node node, @Nonnull final NodeReferenceWithVector neighbor) { - transaction.set(getNeighborKey(layer, node, neighbor.getPrimaryKey()), - StorageAdapter.tupleFromVector(neighbor.getVector()).pack()); + final byte[] neighborKey = getNeighborKey(layer, node, neighbor.getPrimaryKey()); + final byte[] value = StorageAdapter.tupleFromVector(neighbor.getVector()).pack(); + transaction.set(neighborKey, + value); getOnWriteListener().onNeighborWritten(layer, node, neighbor); + getOnWriteListener().onKeyValueWritten(layer, neighborKey, value); } public void deleteNeighbor(@Nonnull final Transaction transaction, final int layer, diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithDistance.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithDistance.java index 96a3a23720..bc9470735c 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithDistance.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithDistance.java @@ -26,7 +26,7 @@ import javax.annotation.Nonnull; import java.util.Objects; -class NodeReferenceWithDistance extends NodeReferenceWithVector { +public class NodeReferenceWithDistance extends NodeReferenceWithVector { private final double distance; public NodeReferenceWithDistance(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java index d57450d434..fd4a096208 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/OnWriteListener.java @@ -43,8 +43,7 @@ default void onNeighborDeleted(final int layer, @Nonnull final Node vectorFromTuple(final Tuple vectorTuple) { - final byte[] vectorAsBytes = vectorTuple.getBytes(0); - final int bytesLength = vectorAsBytes.length; + static Vector.HalfVector vectorFromTuple(final Tuple vectorTuple) { + return vectorFromBytes(vectorTuple.getBytes(0)); + } + + @Nonnull + static Vector.HalfVector vectorFromBytes(final byte[] vectorBytes) { + final int bytesLength = vectorBytes.length; Verify.verify(bytesLength % 2 == 0); final int componentSize = bytesLength >>> 1; final Half[] vectorHalfs = new Half[componentSize]; for (int i = 0; i < componentSize; i ++) { - vectorHalfs[i] = Half.shortBitsToHalf(shortFromBytes(vectorAsBytes, i << 1)); + vectorHalfs[i] = Half.shortBitsToHalf(shortFromBytes(vectorBytes, i << 1)); } return new Vector.HalfVector(vectorHalfs); } + @Nonnull @SuppressWarnings("PrimitiveArrayArgumentToVarargsMethod") static Tuple tupleFromVector(final Vector vector) { + return Tuple.from(bytesFromVector(vector)); + } + + @Nonnull + static byte[] bytesFromVector(final Vector vector) { final byte[] vectorBytes = new byte[2 * vector.size()]; for (int i = 0; i < vector.size(); i ++) { final byte[] componentBytes = bytesFromShort(Half.halfToShortBits(vector.getComponent(i))); @@ -154,7 +164,7 @@ static Tuple tupleFromVector(final Vector vector) { vectorBytes[indexTimesTwo] = componentBytes[0]; vectorBytes[indexTimesTwo + 1] = componentBytes[1]; } - return Tuple.from(vectorBytes); + return vectorBytes; } static short shortFromBytes(final byte[] bytes, final int offset) { diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java index 6a2e5fd01e..bfa179ea2b 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java @@ -58,6 +58,9 @@ public R[] getData() { return data; } + @Nonnull + public abstract byte[] getRawData(); + @Nonnull public abstract Vector toHalfVector(); @@ -102,10 +105,13 @@ public String toString(final int limitDimensions) { public static class HalfVector extends Vector { @Nonnull private final Supplier toDoubleVectorSupplier; + @Nonnull + private final Supplier toRawDataSupplier; public HalfVector(@Nonnull final Half[] data) { super(data); this.toDoubleVectorSupplier = Suppliers.memoize(this::computeDoubleVector); + this.toRawDataSupplier = Suppliers.memoize(this::computeRawData); } @Nonnull @@ -128,6 +134,22 @@ public DoubleVector computeDoubleVector() { } return new DoubleVector(result); } + + @Nonnull + @Override + public byte[] getRawData() { + return toRawDataSupplier.get(); + } + + @Nonnull + private byte[] computeRawData() { + return StorageAdapter.bytesFromVector(this); + } + + @Nonnull + public static HalfVector halfVectorFromBytes(@Nonnull final byte[] vectorBytes) { + return StorageAdapter.vectorFromBytes(vectorBytes); + } } public static class DoubleVector extends Vector { @@ -159,6 +181,13 @@ public HalfVector computeHalfVector() { public DoubleVector toDoubleVector() { return this; } + + @Nonnull + @Override + public byte[] getRawData() { + // TODO + throw new UnsupportedOperationException("not implemented yet"); + } } static double distance(@Nonnull Metric metric, diff --git a/fdb-record-layer-core/fdb-record-layer-core.gradle b/fdb-record-layer-core/fdb-record-layer-core.gradle index 41fcfd996a..e6a7416e47 100644 --- a/fdb-record-layer-core/fdb-record-layer-core.gradle +++ b/fdb-record-layer-core/fdb-record-layer-core.gradle @@ -31,6 +31,7 @@ dependencies { api(libs.protobuf) implementation(libs.slf4j.api) implementation(libs.guava) + implementation(libs.half4j) compileOnly(libs.jsr305) compileOnly(libs.autoService) annotationProcessor(libs.autoService) diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/FDBStoreTimer.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/FDBStoreTimer.java index 07d1af133d..a462cbea0a 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/FDBStoreTimer.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/FDBStoreTimer.java @@ -757,6 +757,14 @@ public enum Counts implements Count { LOCKS_ATTEMPTED("number of attempts to register a lock", false), /** Count of the locks released. */ LOCKS_RELEASED("number of locks released", false), + VECTOR_NODE_READS("intermediate nodes read", false), + VECTOR_NODE_READ_BYTES("intermediate node bytes read", true), + VECTOR_NODE0_READS("intermediate nodes read", false), + VECTOR_NODE0_READ_BYTES("intermediate node bytes read", true), + VECTOR_NODE_WRITES("intermediate nodes written", false), + VECTOR_NODE_WRITE_BYTES("intermediate node bytes written", true), + VECTOR_NODE0_WRITES("intermediate nodes written", false), + VECTOR_NODE0_WRITE_BYTES("intermediate node bytes written", true), ; private final String title; diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanBounds.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanBounds.java index f1437f9318..b3131bb0cb 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanBounds.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanBounds.java @@ -21,42 +21,41 @@ package com.apple.foundationdb.record.provider.foundationdb; import com.apple.foundationdb.annotation.API; -import com.apple.foundationdb.async.rtree.RTree; +import com.apple.foundationdb.async.hnsw.Vector; import com.apple.foundationdb.record.IndexScanType; +import com.apple.foundationdb.record.RecordCoreException; import com.apple.foundationdb.record.TupleRange; -import com.apple.foundationdb.tuple.Tuple; -import com.google.common.base.Preconditions; -import com.google.common.collect.ImmutableList; +import com.apple.foundationdb.record.query.expressions.Comparisons; import javax.annotation.Nonnull; -import java.util.List; -import java.util.stream.Collectors; +import javax.annotation.Nullable; /** - * {@link IndexScanBounds} for a multidimensional index scan. A multidimensional scan bounds object contains a - * {@link #prefixRange} and a {@link SpatialPredicate} which can be almost arbitrarily complex. The prefix range - * is a regular tuple range informing the index maintainer how to constrain the search over the non-multidimensional - * fields that can be viewed as a prefix whose data is stored in a regular one-dimensional index. The spatial predicate - * implements methods to quickly establish geometric overlap and containment. Spatial predicates do have internal - * structure as they can delegate to contained other spatial predicates thus allowing to form e.g. logical conjuncts - * or disjuncts. + * TODO. */ @API(API.Status.EXPERIMENTAL) -public class MultidimensionalIndexScanBounds implements IndexScanBounds { +public class VectorIndexScanBounds implements IndexScanBounds { @Nonnull private final TupleRange prefixRange; @Nonnull - private final SpatialPredicate spatialPredicate; + private final Comparisons.Type comparisonType; + @Nullable + private final Vector queryVector; + private final int limit; @Nonnull private final TupleRange suffixRange; - public MultidimensionalIndexScanBounds(@Nonnull final TupleRange prefixRange, - @Nonnull final SpatialPredicate spatialPredicate, - @Nonnull final TupleRange suffixRange) { + public VectorIndexScanBounds(@Nonnull final TupleRange prefixRange, + @Nonnull final Comparisons.Type comparisonType, + @Nullable final Vector queryVector, + final int limit, + @Nonnull final TupleRange suffixRange) { this.prefixRange = prefixRange; - this.spatialPredicate = spatialPredicate; + this.comparisonType = comparisonType; + this.queryVector = queryVector; + this.limit = limit; this.suffixRange = suffixRange; } @@ -72,166 +71,45 @@ public TupleRange getPrefixRange() { } @Nonnull - public SpatialPredicate getSpatialPredicate() { - return spatialPredicate; + public Comparisons.Type getComparisonType() { + return comparisonType; } - @Nonnull - public TupleRange getSuffixRange() { - return suffixRange; - } - - /** - * Method to compute if the rectangle handed in overlaps with this scan bounds object. This method is invoked when - * the R-tree data structure of a multidimensional index is searched. Note that this method can be implemented using - * a best-effort approach as it is permissible to indicate overlap between {@code mbr} and {@code this} when there - * is in fact no overlap. The rate of false-positives directly influences the search performance in the - * multidimensional index. - * @param mbr the minimum-bounding {@link RTree.Rectangle} - * @return {@code true} if {@code this} overlaps with {@code mbr} - */ - public boolean overlapsMbrApproximately(@Nonnull RTree.Rectangle mbr) { - return spatialPredicate.overlapsMbrApproximately(mbr); + @Nullable + public Vector getQueryVector() { + return queryVector; } - /** - * Method to compute if the point handed in is contained by this scan bounds object. - * @param position the {@link RTree.Point} - * @return {@code true} if {@code position} is contained by {@code this} - */ - public boolean containsPosition(@Nonnull RTree.Point position) { - return spatialPredicate.containsPosition(position); - } - - /** - * Spatial predicate. The implementing classes form a boolean algebra of sorts. Most notably {@link Hypercube} - * represents the logical variables, while {@link And} and {@link Or} can be used to build up more complex powerful - * bounds. - */ - public interface SpatialPredicate { - SpatialPredicate TAUTOLOGY = new SpatialPredicate() { - @Override - public boolean overlapsMbrApproximately(@Nonnull final RTree.Rectangle mbr) { - return true; - } - - @Override - public boolean containsPosition(@Nonnull final RTree.Point position) { - return true; - } - }; - - boolean overlapsMbrApproximately(@Nonnull RTree.Rectangle mbr); - - boolean containsPosition(@Nonnull RTree.Point position); + public int getLimit() { + return limit; } - /** - * Scan bounds that consists of other {@link SpatialPredicate}s to form a logical OR. - */ - public static class Or implements SpatialPredicate { - @Nonnull - private final List children; - - public Or(@Nonnull final List children) { - this.children = ImmutableList.copyOf(children); - } - - @Override - public boolean overlapsMbrApproximately(@Nonnull final RTree.Rectangle mbr) { - return children.stream() - .anyMatch(child -> child.overlapsMbrApproximately(mbr)); - } - - @Override - public boolean containsPosition(@Nonnull final RTree.Point position) { - return children.stream() - .anyMatch(child -> child.containsPosition(position)); - } - - @Override - public String toString() { - return children.stream().map(Object::toString).collect(Collectors.joining(" or ")); + public int getAdjustedLimit() { + switch (getComparisonType()) { + case DISTANCE_RANK_LESS_THAN: + return limit - 1; + case DISTANCE_RANK_LESS_THAN_OR_EQUAL: + return limit; + default: + throw new RecordCoreException("unsupported comparison"); } } - /** - * Scan bounds that consists of other {@link SpatialPredicate}s to form a logical AND. - */ - public static class And implements SpatialPredicate { - @Nonnull - private final List children; - - public And(@Nonnull final List children) { - this.children = ImmutableList.copyOf(children); - } - - @Override - public boolean overlapsMbrApproximately(@Nonnull final RTree.Rectangle mbr) { - return children.stream() - .allMatch(child -> child.overlapsMbrApproximately(mbr)); - } - - @Override - public boolean containsPosition(@Nonnull final RTree.Point position) { - return children.stream() - .allMatch(child -> child.containsPosition(position)); - } - - @Override - public String toString() { - return children.stream().map(Object::toString).collect(Collectors.joining(" and ")); - } + @Nonnull + public TupleRange getSuffixRange() { + return suffixRange; } - /** - * Scan bounds describing an n-dimensional hypercube. - */ - public static class Hypercube implements SpatialPredicate { - @Nonnull - private final List dimensionRanges; - - public Hypercube(@Nonnull final List dimensionRanges) { - this.dimensionRanges = ImmutableList.copyOf(dimensionRanges); - } - - @Override - public boolean overlapsMbrApproximately(@Nonnull final RTree.Rectangle mbr) { - Preconditions.checkArgument(mbr.getNumDimensions() == dimensionRanges.size()); - - for (int d = 0; d < mbr.getNumDimensions(); d++) { - final Tuple lowTuple = Tuple.from(mbr.getLow(d)); - final Tuple highTuple = Tuple.from(mbr.getHigh(d)); - final TupleRange dimensionRange = dimensionRanges.get(d); - if (!dimensionRange.overlaps(lowTuple, highTuple)) { - return false; - } - } - return true; - } - - @Override - public boolean containsPosition(@Nonnull final RTree.Point position) { - Preconditions.checkArgument(position.getNumDimensions() == dimensionRanges.size()); - - for (int d = 0; d < position.getNumDimensions(); d++) { - final Tuple coordinate = Tuple.from(position.getCoordinate(d)); - final TupleRange dimensionRange = dimensionRanges.get(d); - if (!dimensionRange.contains(coordinate)) { - return false; - } - } - return true; - } - - @Nonnull - public List getDimensionRanges() { - return dimensionRanges; - } - - @Override - public String toString() { - return "HyperCube:[" + dimensionRanges + "]"; + public boolean isWithinLimit(int rank) { + switch (getComparisonType()) { + case DISTANCE_RANK_EQUALS: + return rank == limit; + case DISTANCE_RANK_LESS_THAN: + return rank < limit; + case DISTANCE_RANK_LESS_THAN_OR_EQUAL: + return rank <= limit; + default: + throw new RecordCoreException("unsupported comparison"); } } } diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanComparisons.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanComparisons.java index 5b6967e54c..88e69b6748 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanComparisons.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/VectorIndexScanComparisons.java @@ -30,15 +30,16 @@ import com.apple.foundationdb.record.TupleRange; import com.apple.foundationdb.record.metadata.Index; import com.apple.foundationdb.record.planprotos.PIndexScanParameters; -import com.apple.foundationdb.record.planprotos.PMultidimensionalIndexScanComparisons; -import com.apple.foundationdb.record.provider.foundationdb.MultidimensionalIndexScanBounds.Hypercube; +import com.apple.foundationdb.record.planprotos.PVectorIndexScanComparisons; +import com.apple.foundationdb.record.query.expressions.Comparisons; +import com.apple.foundationdb.record.query.expressions.Comparisons.DistanceRankValueComparison; import com.apple.foundationdb.record.query.plan.ScanComparisons; import com.apple.foundationdb.record.query.plan.cascades.AliasMap; import com.apple.foundationdb.record.query.plan.cascades.CorrelationIdentifier; -import com.apple.foundationdb.record.query.plan.explain.ExplainTokens; -import com.apple.foundationdb.record.query.plan.explain.ExplainTokensWithPrecedence; import com.apple.foundationdb.record.query.plan.cascades.explain.Attribute; import com.apple.foundationdb.record.query.plan.cascades.values.translation.TranslationMap; +import com.apple.foundationdb.record.query.plan.explain.ExplainTokens; +import com.apple.foundationdb.record.query.plan.explain.ExplainTokensWithPrecedence; import com.google.auto.service.AutoService; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -46,7 +47,6 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; -import java.util.List; import java.util.Objects; import java.util.Set; @@ -54,19 +54,19 @@ * {@link ScanComparisons} for use in a multidimensional index scan. */ @API(API.Status.UNSTABLE) -public class MultidimensionalIndexScanComparisons implements IndexScanParameters { +public class VectorIndexScanComparisons implements IndexScanParameters { @Nonnull private final ScanComparisons prefixScanComparisons; @Nonnull - private final List dimensionsScanComparisons; + private final DistanceRankValueComparison distanceRankValueComparison; @Nonnull private final ScanComparisons suffixScanComparisons; - public MultidimensionalIndexScanComparisons(@Nonnull final ScanComparisons prefixScanComparisons, - @Nonnull final List dimensionsScanComparisons, - @Nonnull final ScanComparisons suffixKeyComparisonRanges) { + public VectorIndexScanComparisons(@Nonnull final ScanComparisons prefixScanComparisons, + @Nonnull final DistanceRankValueComparison distanceRankValueComparison, + @Nonnull final ScanComparisons suffixKeyComparisonRanges) { this.prefixScanComparisons = prefixScanComparisons; - this.dimensionsScanComparisons = dimensionsScanComparisons; + this.distanceRankValueComparison = distanceRankValueComparison; this.suffixScanComparisons = suffixKeyComparisonRanges; } @@ -82,8 +82,8 @@ public ScanComparisons getPrefixScanComparisons() { } @Nonnull - public List getDimensionsScanComparisons() { - return dimensionsScanComparisons; + public DistanceRankValueComparison getDistanceRankValueComparison() { + return distanceRankValueComparison; } @Nonnull @@ -93,20 +93,16 @@ public ScanComparisons getSuffixScanComparisons() { @Nonnull @Override - public MultidimensionalIndexScanBounds bind(@Nonnull final FDBRecordStoreBase store, @Nonnull final Index index, - @Nonnull final EvaluationContext context) { - final ImmutableList.Builder dimensionsTupleRangeBuilder = ImmutableList.builder(); - for (final ScanComparisons dimensionScanComparison : dimensionsScanComparisons) { - dimensionsTupleRangeBuilder.add(dimensionScanComparison.toTupleRange(store, context)); - } - final Hypercube hypercube = new Hypercube(dimensionsTupleRangeBuilder.build()); - return new MultidimensionalIndexScanBounds(prefixScanComparisons.toTupleRange(store, context), - hypercube, suffixScanComparisons.toTupleRange(store, context)); + public VectorIndexScanBounds bind(@Nonnull final FDBRecordStoreBase store, @Nonnull final Index index, + @Nonnull final EvaluationContext context) { + return new VectorIndexScanBounds(prefixScanComparisons.toTupleRange(store, context), + distanceRankValueComparison.getType(), distanceRankValueComparison.getVector(store, context), + distanceRankValueComparison.getLimit(store, context), suffixScanComparisons.toTupleRange(store, context)); } @Override public int planHash(@Nonnull PlanHashMode mode) { - return PlanHashable.objectsPlanHash(mode, prefixScanComparisons, dimensionsScanComparisons, + return PlanHashable.objectsPlanHash(mode, prefixScanComparisons, distanceRankValueComparison, suffixScanComparisons); } @@ -123,16 +119,19 @@ public ExplainTokensWithPrecedence explain() { ? prefixScanComparisons.explain().getExplainTokens() : new ExplainTokens().addToString(tupleRange); - final var dimensions = - new ExplainTokens().addSequence(() -> new ExplainTokens().addCommaAndWhiteSpace(), - () -> dimensionsScanComparisons.stream() - .map(dimensionScanComparisons -> { - @Nullable var dimensionTupleRange = dimensionScanComparisons.toTupleRangeWithoutContext(); - return dimensionTupleRange == null - ? dimensionScanComparisons.explain().getExplainTokens() - : new ExplainTokens().addToString(dimensionTupleRange); - }).iterator()); - + ExplainTokens distanceRank; + try { + @Nullable var vector = distanceRankValueComparison.getVector(null, null); + int limit = distanceRankValueComparison.getLimit(null, null); + distanceRank = + new ExplainTokens().addNested(vector == null + ? new ExplainTokens().addKeyword("null") + : new ExplainTokens().addToString(vector)); + distanceRank.addKeyword(distanceRankValueComparison.getType().name()).addWhitespace().addToString(limit); + } catch (final Comparisons.EvaluationContextRequiredException e) { + distanceRank = + new ExplainTokens().addNested(distanceRankValueComparison.explain().getExplainTokens()); + } tupleRange = suffixScanComparisons.toTupleRangeWithoutContext(); final var suffix = tupleRange == null @@ -140,9 +139,10 @@ public ExplainTokensWithPrecedence explain() { : new ExplainTokens().addToString(tupleRange); return ExplainTokensWithPrecedence.of(prefix.addOptionalWhitespace().addToString(":{").addOptionalWhitespace() - .addNested(dimensions).addOptionalWhitespace().addToString("}:").addOptionalWhitespace().addNested(suffix)); + .addNested(distanceRank).addOptionalWhitespace().addToString("}:").addOptionalWhitespace().addNested(suffix)); } + @SuppressWarnings("checkstyle:VariableDeclarationUsageDistance") @Override public void getPlannerGraphDetails(@Nonnull ImmutableList.Builder detailsBuilder, @Nonnull ImmutableMap.Builder attributeMapBuilder) { @Nullable TupleRange tupleRange = prefixScanComparisons.toTupleRangeWithoutContext(); @@ -155,17 +155,16 @@ public void getPlannerGraphDetails(@Nonnull ImmutableList.Builder detail attributeMapBuilder.put("pcomparisons", Attribute.gml(prefixScanComparisons.toString())); } - for (int d = 0; d < dimensionsScanComparisons.size(); d++) { - final ScanComparisons dimensionScanComparisons = dimensionsScanComparisons.get(d); - tupleRange = dimensionScanComparisons.toTupleRangeWithoutContext(); - if (tupleRange != null) { - detailsBuilder.add("dim" + d + ": " + tupleRange.getLowEndpoint().toString(false) + "{{dlow" + d + "}}, {{dhigh" + d + "}}" + tupleRange.getHighEndpoint().toString(true)); - attributeMapBuilder.put("dlow" + d, Attribute.gml(tupleRange.getLow() == null ? "-∞" : tupleRange.getLow().toString())); - attributeMapBuilder.put("dhigh" + d, Attribute.gml(tupleRange.getHigh() == null ? "∞" : tupleRange.getHigh().toString())); - } else { - detailsBuilder.add("dim" + d + " comparisons: " + "{{dcomparisons" + d + "}}"); - attributeMapBuilder.put("dcomparisons" + d, Attribute.gml(dimensionScanComparisons.toString())); - } + try { + @Nullable var vector = distanceRankValueComparison.getVector(null, null); + int limit = distanceRankValueComparison.getLimit(null, null); + detailsBuilder.add("distanceRank: {{vector}} {{type}} {{limit}}"); + attributeMapBuilder.put("vector", Attribute.gml(String.valueOf(vector))); + attributeMapBuilder.put("type", Attribute.gml(distanceRankValueComparison.getType())); + attributeMapBuilder.put("limit", Attribute.gml(limit)); + } catch (final Comparisons.EvaluationContextRequiredException e) { + detailsBuilder.add("distanceRank: {{comparison}}"); + attributeMapBuilder.put("comparison", Attribute.gml(distanceRankValueComparison)); } tupleRange = suffixScanComparisons.toTupleRangeWithoutContext(); @@ -184,8 +183,7 @@ public void getPlannerGraphDetails(@Nonnull ImmutableList.Builder detail public Set getCorrelatedTo() { final ImmutableSet.Builder correlatedToBuilder = ImmutableSet.builder(); correlatedToBuilder.addAll(prefixScanComparisons.getCorrelatedTo()); - correlatedToBuilder.addAll(dimensionsScanComparisons.stream() - .flatMap(dimensionScanComparison -> dimensionScanComparison.getCorrelatedTo().stream()).iterator()); + correlatedToBuilder.addAll(distanceRankValueComparison.getCorrelatedTo()); correlatedToBuilder.addAll(suffixScanComparisons.getCorrelatedTo()); return correlatedToBuilder.build(); } @@ -206,30 +204,22 @@ public boolean semanticEquals(@Nullable final Object other, @Nonnull final Alias return false; } - final MultidimensionalIndexScanComparisons that = (MultidimensionalIndexScanComparisons)other; + final VectorIndexScanComparisons that = (VectorIndexScanComparisons)other; if (!prefixScanComparisons.semanticEquals(that.prefixScanComparisons, aliasMap)) { return false; } - if (dimensionsScanComparisons.size() != that.dimensionsScanComparisons.size()) { + + if (!distanceRankValueComparison.semanticEquals(that.distanceRankValueComparison, aliasMap)) { return false; } - for (int i = 0; i < dimensionsScanComparisons.size(); i++) { - final ScanComparisons dimensionScanComparison = dimensionsScanComparisons.get(i); - final ScanComparisons otherDimensionScanComparison = that.dimensionsScanComparisons.get(i); - if (!dimensionScanComparison.semanticEquals(otherDimensionScanComparison, aliasMap)) { - return false; - } - } return suffixScanComparisons.semanticEquals(that.suffixScanComparisons, aliasMap); } @Override public int semanticHashCode() { int hashCode = prefixScanComparisons.semanticHashCode(); - for (final ScanComparisons dimensionScanComparison : dimensionsScanComparisons) { - hashCode = 31 * hashCode + dimensionScanComparison.semanticHashCode(); - } + hashCode = 31 * hashCode + distanceRankValueComparison.semanticHashCode(); return 31 * hashCode + suffixScanComparisons.semanticHashCode(); } @@ -241,39 +231,32 @@ public IndexScanParameters translateCorrelations(@Nonnull final TranslationMap t final ScanComparisons translatedPrefixScanComparisons = prefixScanComparisons.translateCorrelations(translationMap, shouldSimplifyValues); - final ImmutableList.Builder translatedDimensionScanComparisonBuilder = ImmutableList.builder(); - boolean isSameDimensionsScanComparisons = true; - for (final ScanComparisons dimensionScanComparisons : dimensionsScanComparisons) { - final ScanComparisons translatedDimensionScanComparison = - dimensionScanComparisons.translateCorrelations(translationMap, shouldSimplifyValues); - if (translatedDimensionScanComparison != dimensionScanComparisons) { - isSameDimensionsScanComparisons = false; - } - translatedDimensionScanComparisonBuilder.add(translatedDimensionScanComparison); - } + final DistanceRankValueComparison translatedDistanceRankValueComparison = + distanceRankValueComparison.translateCorrelations(translationMap, shouldSimplifyValues); final ScanComparisons translatedSuffixKeyScanComparisons = suffixScanComparisons.translateCorrelations(translationMap, shouldSimplifyValues); - if (translatedPrefixScanComparisons != prefixScanComparisons || !isSameDimensionsScanComparisons || + if (translatedPrefixScanComparisons != prefixScanComparisons || + translatedDistanceRankValueComparison != distanceRankValueComparison || translatedSuffixKeyScanComparisons != suffixScanComparisons) { - return withComparisons(translatedPrefixScanComparisons, translatedDimensionScanComparisonBuilder.build(), + return withComparisons(translatedPrefixScanComparisons, translatedDistanceRankValueComparison, translatedSuffixKeyScanComparisons); } return this; } @Nonnull - protected MultidimensionalIndexScanComparisons withComparisons(@Nonnull final ScanComparisons prefixScanComparisons, - @Nonnull final List dimensionsComparisonRanges, - @Nonnull final ScanComparisons suffixKeyScanComparisons) { - return new MultidimensionalIndexScanComparisons(prefixScanComparisons, dimensionsComparisonRanges, + protected VectorIndexScanComparisons withComparisons(@Nonnull final ScanComparisons prefixScanComparisons, + @Nonnull final DistanceRankValueComparison distanceRankValueComparison, + @Nonnull final ScanComparisons suffixKeyScanComparisons) { + return new VectorIndexScanComparisons(prefixScanComparisons, distanceRankValueComparison, suffixKeyScanComparisons); } @Override public String toString() { - return "BY_VALUE(MD):" + prefixScanComparisons + ":" + dimensionsScanComparisons + ":" + suffixScanComparisons; + return "BY_VALUE(VECTOR):" + prefixScanComparisons + ":" + distanceRankValueComparison + ":" + suffixScanComparisons; } @Override @@ -290,12 +273,10 @@ public int hashCode() { @Nonnull @Override - public PMultidimensionalIndexScanComparisons toProto(@Nonnull final PlanSerializationContext serializationContext) { - final PMultidimensionalIndexScanComparisons.Builder builder = PMultidimensionalIndexScanComparisons.newBuilder(); + public PVectorIndexScanComparisons toProto(@Nonnull final PlanSerializationContext serializationContext) { + final PVectorIndexScanComparisons.Builder builder = PVectorIndexScanComparisons.newBuilder(); builder.setPrefixScanComparisons(prefixScanComparisons.toProto(serializationContext)); - for (final ScanComparisons dimensionsScanComparison : dimensionsScanComparisons) { - builder.addDimensionsScanComparisons(dimensionsScanComparison.toProto(serializationContext)); - } + builder.setDistanceRankValueComparison(distanceRankValueComparison.toProto(serializationContext)); builder.setSuffixScanComparisons(suffixScanComparisons.toProto(serializationContext)); return builder.build(); } @@ -303,26 +284,22 @@ public PMultidimensionalIndexScanComparisons toProto(@Nonnull final PlanSerializ @Nonnull @Override public PIndexScanParameters toIndexScanParametersProto(@Nonnull final PlanSerializationContext serializationContext) { - return PIndexScanParameters.newBuilder().setMultidimensionalIndexScanComparisons(toProto(serializationContext)).build(); + return PIndexScanParameters.newBuilder().setVectorIndexScanComparisons(toProto(serializationContext)).build(); } @Nonnull - public static MultidimensionalIndexScanComparisons fromProto(@Nonnull final PlanSerializationContext serializationContext, - @Nonnull final PMultidimensionalIndexScanComparisons multidimensionalIndexScanComparisonsProto) { - final ImmutableList.Builder dimensionScanComparisonsBuilder = ImmutableList.builder(); - for (int i = 0; i < multidimensionalIndexScanComparisonsProto.getDimensionsScanComparisonsCount(); i ++) { - dimensionScanComparisonsBuilder.add(ScanComparisons.fromProto(serializationContext, - multidimensionalIndexScanComparisonsProto.getDimensionsScanComparisons(i))); - } - return new MultidimensionalIndexScanComparisons(ScanComparisons.fromProto(serializationContext, Objects.requireNonNull(multidimensionalIndexScanComparisonsProto.getPrefixScanComparisons())), - dimensionScanComparisonsBuilder.build(), - ScanComparisons.fromProto(serializationContext, Objects.requireNonNull(multidimensionalIndexScanComparisonsProto.getSuffixScanComparisons()))); + public static VectorIndexScanComparisons fromProto(@Nonnull final PlanSerializationContext serializationContext, + @Nonnull final PVectorIndexScanComparisons vectorIndexScanComparisonsProto) { + return new VectorIndexScanComparisons(ScanComparisons.fromProto(serializationContext, + Objects.requireNonNull(vectorIndexScanComparisonsProto.getPrefixScanComparisons())), + Objects.requireNonNull(DistanceRankValueComparison.fromProto(serializationContext, vectorIndexScanComparisonsProto.getDistanceRankValueComparison())), + ScanComparisons.fromProto(serializationContext, Objects.requireNonNull(vectorIndexScanComparisonsProto.getSuffixScanComparisons()))); } @Nonnull - public static MultidimensionalIndexScanComparisons byValue(@Nullable ScanComparisons prefixScanComparisons, - @Nonnull final List dimensionsComparisonRanges, - @Nullable ScanComparisons suffixKeyScanComparisons) { + public static VectorIndexScanComparisons byValue(@Nullable ScanComparisons prefixScanComparisons, + @Nonnull final DistanceRankValueComparison distanceRankValueComparison, + @Nullable ScanComparisons suffixKeyScanComparisons) { if (prefixScanComparisons == null) { prefixScanComparisons = ScanComparisons.EMPTY; } @@ -331,25 +308,25 @@ public static MultidimensionalIndexScanComparisons byValue(@Nullable ScanCompari suffixKeyScanComparisons = ScanComparisons.EMPTY; } - return new MultidimensionalIndexScanComparisons(prefixScanComparisons, dimensionsComparisonRanges, suffixKeyScanComparisons); + return new VectorIndexScanComparisons(prefixScanComparisons, distanceRankValueComparison, suffixKeyScanComparisons); } /** * Deserializer. */ @AutoService(PlanDeserializer.class) - public static class Deserializer implements PlanDeserializer { + public static class Deserializer implements PlanDeserializer { @Nonnull @Override - public Class getProtoMessageClass() { - return PMultidimensionalIndexScanComparisons.class; + public Class getProtoMessageClass() { + return PVectorIndexScanComparisons.class; } @Nonnull @Override - public MultidimensionalIndexScanComparisons fromProto(@Nonnull final PlanSerializationContext serializationContext, - @Nonnull final PMultidimensionalIndexScanComparisons multidimensionalIndexScanComparisonsProto) { - return MultidimensionalIndexScanComparisons.fromProto(serializationContext, multidimensionalIndexScanComparisonsProto); + public VectorIndexScanComparisons fromProto(@Nonnull final PlanSerializationContext serializationContext, + @Nonnull final PVectorIndexScanComparisons vectorIndexScanComparisonsProto) { + return VectorIndexScanComparisons.fromProto(serializationContext, vectorIndexScanComparisonsProto); } } } diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexHelper.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexHelper.java index 9c7e87628b..eab705e67d 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexHelper.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexHelper.java @@ -85,8 +85,7 @@ public static HNSW.Config getConfig(@Nonnull final Index index) { */ public enum Events implements StoreTimer.DetailEvent { VECTOR_SCAN("scanning the HNSW of a vector index"), - VECTOR_SKIP_SCAN("skip scan the prefix tuples of a vector index scan"), - VECTOR_MODIFICATION("modifying the HNSW of a vector index"); + VECTOR_SKIP_SCAN("skip scan the prefix tuples of a vector index scan"); private final String title; private final String logKey; diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java index 0815e84ccb..6e77618f1b 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java @@ -21,20 +21,19 @@ package com.apple.foundationdb.record.provider.foundationdb.indexes; import com.apple.foundationdb.KeyValue; -import com.apple.foundationdb.Range; import com.apple.foundationdb.ReadTransaction; import com.apple.foundationdb.Transaction; import com.apple.foundationdb.annotation.API; -import com.apple.foundationdb.async.AsyncIterator; import com.apple.foundationdb.async.AsyncUtil; -import com.apple.foundationdb.async.rtree.ChildSlot; -import com.apple.foundationdb.async.rtree.ItemSlot; -import com.apple.foundationdb.async.rtree.Node; -import com.apple.foundationdb.async.rtree.NodeHelpers; -import com.apple.foundationdb.async.rtree.OnReadListener; -import com.apple.foundationdb.async.rtree.OnWriteListener; -import com.apple.foundationdb.async.rtree.RTree; -import com.apple.foundationdb.async.rtree.RTreeHilbertCurveHelpers; +import com.apple.foundationdb.async.hnsw.HNSW; +import com.apple.foundationdb.async.hnsw.HNSW.Config; +import com.apple.foundationdb.async.hnsw.Node; +import com.apple.foundationdb.async.hnsw.NodeReference; +import com.apple.foundationdb.async.hnsw.NodeReferenceAndNode; +import com.apple.foundationdb.async.hnsw.NodeReferenceWithDistance; +import com.apple.foundationdb.async.hnsw.OnReadListener; +import com.apple.foundationdb.async.hnsw.OnWriteListener; +import com.apple.foundationdb.async.hnsw.Vector; import com.apple.foundationdb.record.CursorStreamingMode; import com.apple.foundationdb.record.EndpointType; import com.apple.foundationdb.record.ExecuteProperties; @@ -45,35 +44,30 @@ import com.apple.foundationdb.record.RecordCursor; import com.apple.foundationdb.record.RecordCursorContinuation; import com.apple.foundationdb.record.RecordCursorProto; -import com.apple.foundationdb.record.RecordCursorResult; import com.apple.foundationdb.record.ScanProperties; import com.apple.foundationdb.record.TupleRange; -import com.apple.foundationdb.record.cursors.AsyncIteratorCursor; import com.apple.foundationdb.record.cursors.AsyncLockCursor; import com.apple.foundationdb.record.cursors.ChainedCursor; -import com.apple.foundationdb.record.cursors.CursorLimitManager; import com.apple.foundationdb.record.cursors.LazyCursor; +import com.apple.foundationdb.record.cursors.ListCursor; import com.apple.foundationdb.record.locking.LockIdentifier; import com.apple.foundationdb.record.metadata.Key; -import com.apple.foundationdb.record.metadata.expressions.DimensionsKeyExpression; import com.apple.foundationdb.record.metadata.expressions.KeyExpression; import com.apple.foundationdb.record.metadata.expressions.KeyWithValueExpression; -import com.apple.foundationdb.record.metadata.expressions.ThenKeyExpression; import com.apple.foundationdb.record.provider.common.StoreTimer; import com.apple.foundationdb.record.provider.foundationdb.FDBIndexableRecord; import com.apple.foundationdb.record.provider.foundationdb.FDBStoreTimer; import com.apple.foundationdb.record.provider.foundationdb.IndexMaintainerState; import com.apple.foundationdb.record.provider.foundationdb.IndexScanBounds; import com.apple.foundationdb.record.provider.foundationdb.KeyValueCursor; -import com.apple.foundationdb.record.provider.foundationdb.MultidimensionalIndexScanBounds; +import com.apple.foundationdb.record.provider.foundationdb.VectorIndexScanBounds; import com.apple.foundationdb.record.query.QueryToKeyMatcher; import com.apple.foundationdb.subspace.Subspace; -import com.apple.foundationdb.tuple.ByteArrayUtil; import com.apple.foundationdb.tuple.ByteArrayUtil2; import com.apple.foundationdb.tuple.Tuple; import com.apple.foundationdb.tuple.TupleHelpers; -import com.google.common.base.Preconditions; import com.google.common.base.Verify; +import com.google.common.collect.ImmutableList; import com.google.common.collect.Lists; import com.google.protobuf.ByteString; import com.google.protobuf.InvalidProtocolBufferException; @@ -81,27 +75,29 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; -import java.math.BigInteger; import java.util.List; import java.util.Objects; import java.util.Optional; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.Executor; import java.util.function.Function; import java.util.stream.Collectors; /** - * An index maintainer for keeping a {@link RTree}. + * An index maintainer for keeping an {@link HNSW}. */ @API(API.Status.EXPERIMENTAL) public class VectorIndexMaintainer extends StandardIndexMaintainer { - private static final byte nodeSlotIndexSubspaceIndicator = 0x00; @Nonnull - private final RTree.Config config; + private final Config config; public VectorIndexMaintainer(IndexMaintainerState state) { super(state); - this.config = MultiDimensionalIndexHelper.getConfig(state.index); + this.config = VectorIndexHelper.getConfig(state.index); + } + + @Nonnull + public Config getConfig() { + return config; } @SuppressWarnings("resource") @@ -110,21 +106,19 @@ public VectorIndexMaintainer(IndexMaintainerState state) { public RecordCursor scan(@Nonnull final IndexScanBounds scanBounds, @Nullable final byte[] continuation, @Nonnull final ScanProperties scanProperties) { if (!scanBounds.getScanType().equals(IndexScanType.BY_VALUE)) { - throw new RecordCoreException("Can only scan multidimensional index by value."); + throw new RecordCoreException("Can only scan vector index by value."); } - if (!(scanBounds instanceof MultidimensionalIndexScanBounds)) { - throw new RecordCoreException("Need proper multidimensional index scan bounds."); + if (!(scanBounds instanceof VectorIndexScanBounds)) { + throw new RecordCoreException("Need proper vector index scan bounds."); } - final MultidimensionalIndexScanBounds mDScanBounds = (MultidimensionalIndexScanBounds)scanBounds; + final VectorIndexScanBounds vectorIndexScanBounds = (VectorIndexScanBounds)scanBounds; - final DimensionsKeyExpression dimensionsKeyExpression = getDimensionsKeyExpression(state.index.getRootExpression()); - final int prefixSize = dimensionsKeyExpression.getPrefixSize(); + final KeyWithValueExpression keyWithValueExpression = getKeyWithValueExpression(state.index.getRootExpression()); + final int prefixSize = keyWithValueExpression.getSplitPoint(); final ExecuteProperties executeProperties = scanProperties.getExecuteProperties(); final ScanProperties innerScanProperties = scanProperties.with(ExecuteProperties::clearSkipAndLimit); - final CursorLimitManager cursorLimitManager = new CursorLimitManager(state.context, innerScanProperties); final Subspace indexSubspace = getIndexSubspace(); - final Subspace nodeSlotIndexSubspace = getNodeSlotIndexSubspace(); final FDBStoreTimer timer = Objects.requireNonNull(state.context.getTimer()); // @@ -132,53 +126,91 @@ public RecordCursor scan(@Nonnull final IndexScanBounds scanBounds, // forms the outer of a join with an inner that searches the R-tree for that prefix using the // spatial predicates of the scan bounds. // - return RecordCursor.flatMapPipelined(prefixSkipScan(prefixSize, timer, mDScanBounds, innerScanProperties), + return RecordCursor.flatMapPipelined(prefixSkipScan(prefixSize, timer, vectorIndexScanBounds, innerScanProperties), (prefixTuple, innerContinuation) -> { - final Subspace rtSubspace; - final Subspace rtNodeSlotIndexSubspace; + final Subspace hnswSubspace; if (prefixTuple != null) { Verify.verify(prefixTuple.size() == prefixSize); - rtSubspace = indexSubspace.subspace(prefixTuple); - rtNodeSlotIndexSubspace = nodeSlotIndexSubspace.subspace(prefixTuple); + hnswSubspace = indexSubspace.subspace(prefixTuple); } else { - rtSubspace = indexSubspace; - rtNodeSlotIndexSubspace = nodeSlotIndexSubspace; + hnswSubspace = indexSubspace; } - final Continuation parsedContinuation = Continuation.fromBytes(innerContinuation); - final BigInteger lastHilbertValue = - parsedContinuation == null ? null : parsedContinuation.getLastHilbertValue(); - final Tuple lastKey = parsedContinuation == null ? null : parsedContinuation.getLastKey(); + if (innerContinuation != null) { + final RecordCursorProto.VectorIndexScanContinuation parsedContinuation = + Continuation.fromBytes(innerContinuation); + final ImmutableList.Builder indexEntriesBuilder = ImmutableList.builder(); + for (int i = 0; i < parsedContinuation.getIndexEntriesCount(); i ++) { + final RecordCursorProto.VectorIndexScanContinuation.IndexEntry indexEntryProto = + parsedContinuation.getIndexEntries(i); + indexEntriesBuilder.add(new IndexEntry(state.index, + Tuple.fromBytes(indexEntryProto.getKey().toByteArray()), + Tuple.fromBytes(indexEntryProto.getValue().toByteArray()))); + } + return new ListCursor<>(indexEntriesBuilder.build(), + parsedContinuation.getInnerContinuation().toByteArray()); + } - final RTree rTree = new RTree(rtSubspace, rtNodeSlotIndexSubspace, getExecutor(), config, - RTreeHilbertCurveHelpers::hilbertValue, NodeHelpers::newRandomNodeId, - OnWriteListener.NOOP, new OnRead(cursorLimitManager, timer)); + final HNSW hnsw = new HNSW(hnswSubspace, getExecutor(), getConfig(), + OnWriteListener.NOOP, new OnRead(timer)); final ReadTransaction transaction = state.context.readTransaction(true); - return new LazyCursor<>(state.context.acquireReadLock(new LockIdentifier(rtSubspace)) - .thenApply(lock -> new AsyncLockCursor<>(lock, new ItemSlotCursor(getExecutor(), - rTree.scan(transaction, lastHilbertValue, lastKey, - mDScanBounds::overlapsMbrApproximately, - (low, high) -> mDScanBounds.getSuffixRange().overlaps(low, high)), - cursorLimitManager, timer))), state.context.getExecutor()) - .filter(itemSlot -> lastHilbertValue == null || lastKey == null || - itemSlot.compareHilbertValueAndKey(lastHilbertValue, lastKey) > 0) - .filter(itemSlot -> mDScanBounds.containsPosition(itemSlot.getPosition())) - .filter(itemSlot -> mDScanBounds.getSuffixRange().contains(itemSlot.getKeySuffix())) - .map(itemSlot -> { - final List keyItems = Lists.newArrayList(); - if (prefixTuple != null) { - keyItems.addAll(prefixTuple.getItems()); - } - keyItems.addAll(itemSlot.getPosition().getCoordinates().getItems()); - keyItems.addAll(itemSlot.getKeySuffix().getItems()); - return new IndexEntry(state.index, Tuple.fromList(keyItems), itemSlot.getValue()); - }); + return new LazyCursor<>( + state.context.acquireReadLock(new LockIdentifier(hnswSubspace)) + .thenApply(lock -> + new AsyncLockCursor<>(lock, + new LazyCursor<>( + kNearestNeighborSearch(prefixTuple, hnsw, transaction, vectorIndexScanBounds), + getExecutor()))), + state.context.getExecutor()); }, continuation, state.store.getPipelineSize(PipelineOperation.INDEX_TO_RECORD)) .skipThenLimit(executeProperties.getSkip(), executeProperties.getReturnedRowLimit()); } + @SuppressWarnings({"resource", "checkstyle:MethodName"}) + @Nonnull + private CompletableFuture> kNearestNeighborSearch(@Nullable final Tuple prefixTuple, + @Nonnull final HNSW hnsw, + @Nonnull final ReadTransaction transaction, + @Nonnull final VectorIndexScanBounds vectorIndexScanBounds) { + return hnsw.kNearestNeighborsSearch(transaction, vectorIndexScanBounds.getAdjustedLimit(), 100, + Objects.requireNonNull(vectorIndexScanBounds.getQueryVector()).toHalfVector()) + .thenApply(nearestNeighbors -> { + final ImmutableList.Builder nearestNeighborEntriesBuilder = ImmutableList.builder(); + for (final NodeReferenceAndNode nearestNeighbor : nearestNeighbors) { + if (vectorIndexScanBounds.getSuffixRange().contains(nearestNeighbor.getNode().getPrimaryKey())) { + nearestNeighborEntriesBuilder.add(toIndexEntry(prefixTuple, nearestNeighbor)); + } + } + final ImmutableList nearestNeighborsEntries = nearestNeighborEntriesBuilder.build(); + return new ListCursor<>(getExecutor(), nearestNeighborsEntries, 0) + .mapResult(result -> { + final RecordCursorContinuation continuation = result.getContinuation(); + if (continuation.isEnd()) { + return result; + } + return result.withContinuation(new Continuation(nearestNeighborsEntries, continuation)); + }); + }); + } + + @Nonnull + private IndexEntry toIndexEntry(final Tuple prefixTuple, final NodeReferenceAndNode nearestNeighbor) { + final List keyItems = Lists.newArrayList(); + if (prefixTuple != null) { + keyItems.addAll(prefixTuple.getItems()); + } + final Node node = nearestNeighbor.getNode(); + final NodeReferenceWithDistance nodeReferenceWithDistance = + nearestNeighbor.getNodeReferenceWithDistance(); + keyItems.addAll(node.getPrimaryKey().getItems()); + final List valueItems = Lists.newArrayList(); + valueItems.add(nodeReferenceWithDistance.getVector().getRawData()); + return new IndexEntry(state.index, Tuple.fromList(keyItems), + Tuple.fromList(valueItems)); + } + @Nonnull @Override public RecordCursor scan(@Nonnull final IndexScanType scanType, @Nonnull final TupleRange range, @@ -189,13 +221,13 @@ public RecordCursor scan(@Nonnull final IndexScanType scanType, @Non @Nonnull private Function> prefixSkipScan(final int prefixSize, @Nonnull final StoreTimer timer, - @Nonnull final MultidimensionalIndexScanBounds mDScanBounds, + @Nonnull final VectorIndexScanBounds vectorIndexScanBounds, @Nonnull final ScanProperties innerScanProperties) { final Function> outerFunction; if (prefixSize > 0) { outerFunction = outerContinuation -> timer.instrument(MultiDimensionalIndexHelper.Events.MULTIDIMENSIONAL_SKIP_SCAN, new ChainedCursor<>(state.context, - lastKeyOptional -> nextPrefixTuple(mDScanBounds.getPrefixRange(), + lastKeyOptional -> nextPrefixTuple(vectorIndexScanBounds.getPrefixRange(), prefixSize, lastKeyOptional.orElse(null), innerScanProperties), Tuple::pack, Tuple::fromBytes, @@ -249,45 +281,30 @@ private CompletableFuture> nextPrefixTuple(@Nonnull final TupleR protected CompletableFuture updateIndexKeys(@Nonnull final FDBIndexableRecord savedRecord, final boolean remove, @Nonnull final List indexEntries) { - final DimensionsKeyExpression dimensionsKeyExpression = getDimensionsKeyExpression(state.index.getRootExpression()); - final int prefixSize = dimensionsKeyExpression.getPrefixSize(); - final int dimensionsSize = dimensionsKeyExpression.getDimensionsSize(); + final KeyWithValueExpression keyWithValueExpression = getKeyWithValueExpression(state.index.getRootExpression()); + final int prefixSize = keyWithValueExpression.getColumnSize(); final Subspace indexSubspace = getIndexSubspace(); - final Subspace nodeSlotIndexSubspace = getNodeSlotIndexSubspace(); final var futures = indexEntries.stream().map(indexEntry -> { final var indexKeyItems = indexEntry.getKey().getItems(); final Tuple prefixKey = Tuple.fromList(indexKeyItems.subList(0, prefixSize)); final Subspace rtSubspace; - final Subspace rtNodeSlotIndexSubspace; if (prefixSize > 0) { rtSubspace = indexSubspace.subspace(prefixKey); - rtNodeSlotIndexSubspace = nodeSlotIndexSubspace.subspace(prefixKey); } else { rtSubspace = indexSubspace; - rtNodeSlotIndexSubspace = nodeSlotIndexSubspace; } return state.context.doWithWriteLock(new LockIdentifier(rtSubspace), () -> { - final RTree.Point point = - validatePoint(new RTree.Point(Tuple.fromList(indexKeyItems.subList(prefixSize, prefixSize + dimensionsSize)))); - final List primaryKeyParts = Lists.newArrayList(savedRecord.getPrimaryKey().getItems()); state.index.trimPrimaryKey(primaryKeyParts); - final List keySuffixParts = - Lists.newArrayList(indexKeyItems.subList(prefixSize + dimensionsSize, indexKeyItems.size())); - keySuffixParts.addAll(primaryKeyParts); - final Tuple keySuffix = Tuple.fromList(keySuffixParts); + final Tuple trimmedPrimaryKey = Tuple.fromList(primaryKeyParts); final FDBStoreTimer timer = Objects.requireNonNull(getTimer()); - final RTree rTree = new RTree(rtSubspace, rtNodeSlotIndexSubspace, getExecutor(), config, - RTreeHilbertCurveHelpers::hilbertValue, NodeHelpers::newRandomNodeId, new OnWrite(timer), - OnReadListener.NOOP); + final HNSW hnsw = new HNSW(rtSubspace, getExecutor(), getConfig(), new OnWrite(timer), OnReadListener.NOOP); if (remove) { - return rTree.delete(state.transaction, point, keySuffix); + throw new UnsupportedOperationException("not implemented"); } else { - return rTree.insertOrUpdate(state.transaction, - point, - keySuffix, - indexEntry.getValue()); + return hnsw.insert(state.transaction, trimmedPrimaryKey, + Vector.HalfVector.halfVectorFromBytes(indexEntry.getValue().getBytes(0))); } }); }).collect(Collectors.toList()); @@ -299,116 +316,63 @@ public boolean canDeleteWhere(@Nonnull final QueryToKeyMatcher matcher, @Nonnull if (!super.canDeleteWhere(matcher, evaluated)) { return false; } - return evaluated.size() <= getDimensionsKeyExpression(state.index.getRootExpression()).getPrefixSize(); + return evaluated.size() <= getKeyWithValueExpression(state.index.getRootExpression()).getColumnSize(); } @Override public CompletableFuture deleteWhere(@Nonnull final Transaction tr, @Nonnull final Tuple prefix) { - Verify.verify(getDimensionsKeyExpression(state.index.getRootExpression()).getPrefixSize() >= prefix.size()); - return super.deleteWhere(tr, prefix).thenApply(v -> { - // NOTE: Range.startsWith(), Subspace.range() and so on cover keys *strictly* within the range, but we sometimes - // store data at the prefix key itself. - final Subspace nodeSlotIndexSubspace = getNodeSlotIndexSubspace(); - final byte[] key = nodeSlotIndexSubspace.pack(prefix); - state.context.clear(new Range(key, ByteArrayUtil.strinc(key))); - return v; - }); - } - - @Nonnull - private Subspace getNodeSlotIndexSubspace() { - return getSecondarySubspace().subspace(Tuple.from(nodeSlotIndexSubspaceIndicator)); + Verify.verify(getKeyWithValueExpression(state.index.getRootExpression()).getColumnSize() >= prefix.size()); + return super.deleteWhere(tr, prefix); } /** - * Traverse from the root of a key expression of a multidimensional index to the {@link DimensionsKeyExpression}. - * @param root the root {@link KeyExpression} of the index definition - * @return a {@link DimensionsKeyExpression} + * TODO. */ @Nonnull - public static DimensionsKeyExpression getDimensionsKeyExpression(@Nonnull final KeyExpression root) { + private static KeyWithValueExpression getKeyWithValueExpression(@Nonnull final KeyExpression root) { if (root instanceof KeyWithValueExpression) { - KeyExpression innerKey = ((KeyWithValueExpression)root).getInnerKey(); - while (innerKey instanceof ThenKeyExpression) { - innerKey = ((ThenKeyExpression)innerKey).getChildren().get(0); - } - if (innerKey instanceof DimensionsKeyExpression) { - return (DimensionsKeyExpression)innerKey; - } - throw new RecordCoreException("structure of multidimensional index is not supported"); - } - return (DimensionsKeyExpression)root; - } - - @Nonnull - private static RTree.Point validatePoint(@Nonnull RTree.Point point) { - for (int d = 0; d < point.getNumDimensions(); d ++) { - Object coordinate = point.getCoordinate(d); - Preconditions.checkArgument(coordinate == null || coordinate instanceof Long, - "dimension coordinates must be of type long"); + return (KeyWithValueExpression)root; } - return point; + throw new RecordCoreException("structure of vector index is not supported"); } static class OnRead implements OnReadListener { - @Nonnull - private final CursorLimitManager cursorLimitManager; @Nonnull private final FDBStoreTimer timer; - public OnRead(@Nonnull final CursorLimitManager cursorLimitManager, - @Nonnull final FDBStoreTimer timer) { - this.cursorLimitManager = cursorLimitManager; + public OnRead(@Nonnull final FDBStoreTimer timer) { this.timer = timer; } @Override - public CompletableFuture onAsyncRead(@Nonnull final CompletableFuture future) { - return timer.instrument(MultiDimensionalIndexHelper.Events.MULTIDIMENSIONAL_SCAN, future); + public CompletableFuture> onAsyncRead(@Nonnull final CompletableFuture> future) { + return timer.instrument(VectorIndexHelper.Events.VECTOR_SCAN, future); } @Override - public void onNodeRead(@Nonnull final Node node) { - switch (node.getKind()) { - case LEAF: - timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_LEAF_NODE_READS); - break; - case INTERMEDIATE: - timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_INTERMEDIATE_NODE_READS); - break; - default: - throw new RecordCoreException("unsupported kind of node"); + public void onNodeRead(final int layer, @Nonnull final Node node) { + if (layer == 0) { + timer.increment(FDBStoreTimer.Counts.VECTOR_NODE0_READS); + } else { + timer.increment(FDBStoreTimer.Counts.VECTOR_NODE_READS); } } @Override - public void onKeyValueRead(@Nonnull final Node node, @Nullable final byte[] key, @Nullable final byte[] value) { - final int keyLength = key == null ? 0 : key.length; - final int valueLength = value == null ? 0 : value.length; + public void onKeyValueRead(final int layer, @Nonnull final byte[] key, @Nonnull final byte[] value) { + final int keyLength = key.length; + final int valueLength = value.length; - final int totalLength = keyLength + valueLength; - cursorLimitManager.reportScannedBytes(totalLength); - cursorLimitManager.tryRecordScan(); timer.increment(FDBStoreTimer.Counts.LOAD_INDEX_KEY); timer.increment(FDBStoreTimer.Counts.LOAD_INDEX_KEY_BYTES, keyLength); timer.increment(FDBStoreTimer.Counts.LOAD_INDEX_VALUE_BYTES, valueLength); - switch (node.getKind()) { - case LEAF: - timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_LEAF_NODE_READ_BYTES, totalLength); - break; - case INTERMEDIATE: - timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_INTERMEDIATE_NODE_READ_BYTES, totalLength); - break; - default: - throw new RecordCoreException("unsupported kind of node"); + if (layer == 0) { + timer.increment(FDBStoreTimer.Counts.VECTOR_NODE0_READ_BYTES); + } else { + timer.increment(FDBStoreTimer.Counts.VECTOR_NODE_READ_BYTES); } } - - @Override - public void onChildNodeDiscard(@Nonnull final ChildSlot childSlot) { - timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_CHILD_NODE_DISCARDS); - } } static class OnWrite implements OnWriteListener { @@ -420,118 +384,57 @@ public OnWrite(@Nonnull final FDBStoreTimer timer) { } @Override - public CompletableFuture onAsyncReadForWrite(@Nonnull final CompletableFuture future) { - return timer.instrument(MultiDimensionalIndexHelper.Events.MULTIDIMENSIONAL_MODIFICATION, future); - } - - @Override - public void onNodeWritten(@Nonnull final Node node) { - switch (node.getKind()) { - case LEAF: - timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_LEAF_NODE_WRITES); - break; - case INTERMEDIATE: - timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_INTERMEDIATE_NODE_WRITES); - break; - default: - throw new RecordCoreException("unsupported kind of node"); + public void onNodeWritten(final int layer, @Nonnull final Node node) { + if (layer == 0) { + timer.increment(FDBStoreTimer.Counts.VECTOR_NODE0_WRITES); + } else { + timer.increment(FDBStoreTimer.Counts.VECTOR_NODE_WRITES); } } @Override - public void onKeyValueWritten(@Nonnull final Node node, @Nullable final byte[] key, @Nullable final byte[] value) { - final int keyLength = key == null ? 0 : key.length; - final int valueLength = value == null ? 0 : value.length; + public void onKeyValueWritten(final int layer, @Nonnull final byte[] key, @Nonnull final byte[] value) { + final int keyLength = key.length; + final int valueLength = value.length; final int totalLength = keyLength + valueLength; timer.increment(FDBStoreTimer.Counts.SAVE_INDEX_KEY); timer.increment(FDBStoreTimer.Counts.SAVE_INDEX_KEY_BYTES, keyLength); timer.increment(FDBStoreTimer.Counts.SAVE_INDEX_VALUE_BYTES, valueLength); - switch (node.getKind()) { - case LEAF: - timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_LEAF_NODE_WRITE_BYTES, totalLength); - break; - case INTERMEDIATE: - timer.increment(FDBStoreTimer.Counts.MULTIDIMENSIONAL_INTERMEDIATE_NODE_WRITE_BYTES, totalLength); - break; - default: - throw new RecordCoreException("unsupported kind of node"); + if (layer == 0) { + timer.increment(FDBStoreTimer.Counts.VECTOR_NODE0_WRITE_BYTES, totalLength); + } else { + timer.increment(FDBStoreTimer.Counts.VECTOR_NODE_WRITE_BYTES, totalLength); } } } - static class ItemSlotCursor extends AsyncIteratorCursor { - @Nonnull - private final CursorLimitManager cursorLimitManager; + private static class Continuation implements RecordCursorContinuation { @Nonnull - private final FDBStoreTimer timer; - - public ItemSlotCursor(@Nonnull final Executor executor, @Nonnull final AsyncIterator iterator, - @Nonnull final CursorLimitManager cursorLimitManager, @Nonnull final FDBStoreTimer timer) { - super(executor, iterator); - this.cursorLimitManager = cursorLimitManager; - this.timer = timer; - } - + private final List indexEntries; @Nonnull - @Override - public CompletableFuture> onNext() { - if (nextResult != null && !nextResult.hasNext()) { - // This guard is needed to guarantee that if onNext is called multiple times after the cursor has - // returned a result without a value, then the same NoNextReason is returned each time. Without this guard, - // one might return SCAN_LIMIT_REACHED (for example) after returning a result with SOURCE_EXHAUSTED because - // of the tryRecordScan check. - return CompletableFuture.completedFuture(nextResult); - } else if (cursorLimitManager.tryRecordScan()) { - return iterator.onHasNext().thenApply(hasNext -> { - if (hasNext) { - final ItemSlot itemSlot = iterator.next(); - timer.increment(FDBStoreTimer.Counts.LOAD_SCAN_ENTRY); - timer.increment(FDBStoreTimer.Counts.LOAD_KEY_VALUE); - valuesSeen++; - nextResult = RecordCursorResult.withNextValue(itemSlot, new Continuation(itemSlot.getHilbertValue(), itemSlot.getKey())); - } else { - // Source iterator is exhausted. - nextResult = RecordCursorResult.exhausted(); - } - return nextResult; - }); - } else { // a limit must have been exceeded - final Optional stoppedReason = cursorLimitManager.getStoppedReason(); - if (stoppedReason.isEmpty()) { - throw new RecordCoreException("limit manager stopped cursor but did not report a reason"); - } - Verify.verifyNotNull(nextResult, "should have seen at least one record"); - nextResult = RecordCursorResult.withoutNextValue(nextResult.getContinuation(), stoppedReason.get()); - return CompletableFuture.completedFuture(nextResult); - } - } - } + private final RecordCursorContinuation innerContinuation; - private static class Continuation implements RecordCursorContinuation { - @Nullable - final BigInteger lastHilbertValue; - @Nullable - final Tuple lastKey; @Nullable private ByteString cachedByteString; @Nullable private byte[] cachedBytes; - private Continuation(@Nullable final BigInteger lastHilbertValue, @Nullable final Tuple lastKey) { - this.lastHilbertValue = lastHilbertValue; - this.lastKey = lastKey; + private Continuation(@Nonnull final List indexEntries, + @Nonnull final RecordCursorContinuation innerContinuation) { + this.indexEntries = ImmutableList.copyOf(indexEntries); + this.innerContinuation = innerContinuation; } - @Nullable - public BigInteger getLastHilbertValue() { - return lastHilbertValue; + @Nonnull + public List getIndexEntries() { + return indexEntries; } - @Nullable - public Tuple getLastKey() { - return lastKey; + @Nonnull + public RecordCursorContinuation getInnerContinuation() { + return innerContinuation; } @Nonnull @@ -542,9 +445,17 @@ public ByteString toByteString() { } if (cachedByteString == null) { - cachedByteString = RecordCursorProto.MultidimensionalIndexScanContinuation.newBuilder() - .setLastHilbertValue(ByteString.copyFrom(Objects.requireNonNull(lastHilbertValue).toByteArray())) - .setLastKey(ByteString.copyFrom(Objects.requireNonNull(lastKey).pack())) + final RecordCursorProto.VectorIndexScanContinuation.Builder builder = + RecordCursorProto.VectorIndexScanContinuation.newBuilder(); + for (final var indexEntry : indexEntries) { + builder.addIndexEntries(RecordCursorProto.VectorIndexScanContinuation.IndexEntry.newBuilder() + .setKey(ByteString.copyFrom(indexEntry.getKey().pack())) + .setValue(ByteString.copyFrom(indexEntry.getKey().pack())) + .build()); + } + + cachedByteString = builder + .setInnerContinuation(Objects.requireNonNull(innerContinuation.toByteString())) .build() .toByteString(); } @@ -565,23 +476,16 @@ public byte[] toBytes() { @Override public boolean isEnd() { - return lastHilbertValue == null || lastKey == null; + return getInnerContinuation().isEnd(); } - @Nullable - private static Continuation fromBytes(@Nullable byte[] continuationBytes) { - if (continuationBytes != null) { - final RecordCursorProto.MultidimensionalIndexScanContinuation parsed; - try { - parsed = RecordCursorProto.MultidimensionalIndexScanContinuation.parseFrom(continuationBytes); - } catch (InvalidProtocolBufferException ex) { - throw new RecordCoreException("error parsing continuation", ex) - .addLogInfo("raw_bytes", ByteArrayUtil2.loggable(continuationBytes)); - } - return new Continuation(new BigInteger(parsed.getLastHilbertValue().toByteArray()), - Tuple.fromBytes(parsed.getLastKey().toByteArray())); - } else { - return null; + @Nonnull + private static RecordCursorProto.VectorIndexScanContinuation fromBytes(@Nonnull byte[] continuationBytes) { + try { + return RecordCursorProto.VectorIndexScanContinuation.parseFrom(continuationBytes); + } catch (InvalidProtocolBufferException ex) { + throw new RecordCoreException("error parsing continuation", ex) + .addLogInfo("raw_bytes", ByteArrayUtil2.loggable(continuationBytes)); } } } diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java index 5b3cd5eefc..b3719a4e62 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java @@ -22,6 +22,7 @@ import com.apple.foundationdb.annotation.API; import com.apple.foundationdb.annotation.SpotBugsSuppressWarnings; +import com.apple.foundationdb.async.hnsw.Vector; import com.apple.foundationdb.record.Bindings; import com.apple.foundationdb.record.EvaluationContext; import com.apple.foundationdb.record.ObjectPlanHash; @@ -633,7 +634,13 @@ public enum Type { @API(API.Status.EXPERIMENTAL) SORT(false), @API(API.Status.EXPERIMENTAL) - LIKE; + LIKE, + @API(API.Status.EXPERIMENTAL) + DISTANCE_RANK_EQUALS(true), + @API(API.Status.EXPERIMENTAL) + DISTANCE_RANK_LESS_THAN, + @API(API.Status.EXPERIMENTAL) + DISTANCE_RANK_LESS_THAN_OR_EQUAL; @Nonnull private static final Supplier> protoEnumBiMapSupplier = @@ -1760,6 +1767,9 @@ public DistanceRankValueComparison(@Nonnull final Type type, @Nonnull final Valu @Nonnull final ParameterRelationshipGraph parameterRelationshipGraph, @Nonnull final Value limitValue) { super(type, comparandValue, parameterRelationshipGraph); + Verify.verify(type == Type.DISTANCE_RANK_EQUALS || + type == Type.DISTANCE_RANK_LESS_THAN || + type == Type.DISTANCE_RANK_LESS_THAN_OR_EQUAL); this.limitValue = limitValue; } @@ -1805,7 +1815,8 @@ public Optional replaceValuesMaybe(@Nonnull final Function getVector(@Nullable final FDBRecordStoreBase store, final @Nullable EvaluationContext context) { + return (Vector)getComparand(store, context); + } + + public int getLimit(@Nullable final FDBRecordStoreBase store, final @Nullable EvaluationContext context) { + if (context == null) { + throw EvaluationContextRequiredException.instance(); + } + return (int)Objects.requireNonNull(getLimitValue().eval(store, context)); + } + @Nonnull public static DistanceRankValueComparison fromProto(@Nonnull final PlanSerializationContext serializationContext, @Nonnull final PDistanceRankValueComparison distanceRankValueComparisonProto) { diff --git a/fdb-record-layer-core/src/main/proto/record_cursor.proto b/fdb-record-layer-core/src/main/proto/record_cursor.proto index 04095bf779..ddeeff03e7 100644 --- a/fdb-record-layer-core/src/main/proto/record_cursor.proto +++ b/fdb-record-layer-core/src/main/proto/record_cursor.proto @@ -127,6 +127,15 @@ message MultidimensionalIndexScanContinuation { optional bytes lastKey = 2; } +message VectorIndexScanContinuation { + message IndexEntry { + optional bytes key = 1; + optional bytes value = 2; + } + repeated IndexEntry indexEntries = 1; + optional bytes inner_continuation = 2; +} + message TempTableInsertContinuation { optional bytes child_continuation = 1; optional planprotos.PTempTable tempTable = 2; diff --git a/fdb-record-layer-core/src/main/proto/record_query_plan.proto b/fdb-record-layer-core/src/main/proto/record_query_plan.proto index c338c61a2a..93a089e58e 100644 --- a/fdb-record-layer-core/src/main/proto/record_query_plan.proto +++ b/fdb-record-layer-core/src/main/proto/record_query_plan.proto @@ -1189,6 +1189,9 @@ message PComparison { TEXT_CONTAINS_ANY_PREFIX = 17; SORT = 18; LIKE = 19; + DISTANCE_RANK_EQUALS = 20; + DISTANCE_RANK_LESS_THAN = 21; + DISTANCE_RANK_LESS_THAN_OR_EQUAL = 22; } extensions 5000 to max; @@ -1620,6 +1623,7 @@ message PIndexScanParameters { PIndexScanComparisons index_scan_comparisons = 2; PMultidimensionalIndexScanComparisons multidimensional_index_scan_comparisons = 3; PTimeWindowScanComparisons time_window_scan_comparisons = 4; + PVectorIndexScanComparisons vector_index_scan_comparisons = 5; } } @@ -1655,6 +1659,12 @@ message PTimeWindowScanComparisons { optional PTimeWindowForFunction time_window = 2; } +message PVectorIndexScanComparisons { + optional PScanComparisons prefix_scan_comparisons = 1; + optional PDistanceRankValueComparison distance_rank_value_comparison = 2; + optional PScanComparisons suffix_scan_comparisons = 3; +} + enum PIndexFetchMethod { SCAN_AND_FETCH = 1; USE_REMOTE_FETCH = 2; From d4b8cb596915801b1e0b8720586a6bf740101f89 Mon Sep 17 00:00:00 2001 From: Youssef Hatem Date: Fri, 8 Aug 2025 19:23:20 +0100 Subject: [PATCH 24/34] - Support Vector datatype. --- .../query/plan/cascades/typing/Type.java | 197 ++++++++++++++++-- .../plan/cascades/typing/TypeRepository.java | 3 +- .../main/proto/record_metadata_options.proto | 7 +- .../src/main/proto/record_query_plan.proto | 8 + .../relational/api/metadata/DataType.java | 77 ++++++- .../src/main/antlr/RelationalLexer.g4 | 7 + .../src/main/antlr/RelationalParser.g4 | 12 +- ...RecordLayerStoreSchemaTemplateCatalog.java | 10 +- .../recordlayer/metadata/DataTypeUtils.java | 23 +- .../recordlayer/query/SemanticAnalyzer.java | 124 +++++++++++ .../query/visitors/BaseVisitor.java | 10 +- .../query/visitors/DdlVisitor.java | 35 ++-- .../query/visitors/DelegatingVisitor.java | 8 +- .../query/visitors/TypedVisitor.java | 4 - .../recordlayer/query/VectorTypeTest.java | 70 +++++++ .../src/test/java/YamlIntegrationTests.java | 9 +- 16 files changed, 557 insertions(+), 47 deletions(-) create mode 100644 fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java index e6cc85c390..9080b6130f 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java @@ -24,6 +24,7 @@ import com.apple.foundationdb.record.PlanSerializable; import com.apple.foundationdb.record.PlanSerializationContext; import com.apple.foundationdb.record.RecordCoreException; +import com.apple.foundationdb.record.RecordMetaDataOptionsProto; import com.apple.foundationdb.record.TupleFieldsProto; import com.apple.foundationdb.record.logging.LogMessageKeys; import com.apple.foundationdb.record.planprotos.PType; @@ -38,6 +39,7 @@ import com.apple.foundationdb.record.planprotos.PType.PRelationType; import com.apple.foundationdb.record.planprotos.PType.PTypeCode; import com.apple.foundationdb.record.planprotos.PType.PUuidType; +import com.apple.foundationdb.record.planprotos.PType.PVectorType; import com.apple.foundationdb.record.provider.foundationdb.FDBRecordVersion; import com.apple.foundationdb.record.query.plan.cascades.Narrowable; import com.apple.foundationdb.record.query.plan.cascades.NullableArrayTypeUtils; @@ -416,12 +418,18 @@ static List fromTyped(@Nonnull List typedList) { private static Type fromProtoType(@Nullable Descriptors.GenericDescriptor descriptor, @Nonnull Descriptors.FieldDescriptor.Type protoType, @Nonnull FieldDescriptorProto.Label protoLabel, + @Nullable DescriptorProtos.FieldOptions fieldOptions, boolean isNullable) { - final var typeCode = TypeCode.fromProtobufType(protoType); + final var typeCode = TypeCode.fromProtobufFieldDescriptor(protoType, fieldOptions); if (protoLabel == FieldDescriptorProto.Label.LABEL_REPEATED) { // collection type - return fromProtoTypeToArray(descriptor, protoType, typeCode, false); + return fromProtoTypeToArray(descriptor, protoType, typeCode, fieldOptions, false); } else if (typeCode.isPrimitive()) { + final var fieldOptionMaybe = Optional.ofNullable(fieldOptions).map(f -> f.getExtension(RecordMetaDataOptionsProto.field)); + if (fieldOptionMaybe.isPresent() && fieldOptionMaybe.get().hasVectorOptions()) { + final var vectorOptions = fieldOptionMaybe.get().getVectorOptions(); + return Type.Vector.of(isNullable, vectorOptions.getPrecision(), vectorOptions.getDimensions()); + } return primitiveType(typeCode, isNullable); } else if (typeCode == TypeCode.ENUM) { final var enumDescriptor = (Descriptors.EnumDescriptor)Objects.requireNonNull(descriptor); @@ -432,8 +440,10 @@ private static Type fromProtoType(@Nullable Descriptors.GenericDescriptor descri if (NullableArrayTypeUtils.describesWrappedArray(messageDescriptor)) { // find TypeCode of array elements final var elementField = messageDescriptor.findFieldByName(NullableArrayTypeUtils.getRepeatedFieldName()); - final var elementTypeCode = TypeCode.fromProtobufType(elementField.getType()); - return fromProtoTypeToArray(descriptor, protoType, elementTypeCode, true); + final var elementTypeCode = TypeCode.fromProtobufFieldDescriptor(elementField.getType(), elementField.getOptions()); + return fromProtoTypeToArray(descriptor, protoType, elementTypeCode, fieldOptions, true); + } else if (TupleFieldsProto.UUID.getDescriptor().equals(messageDescriptor)) { + return Type.uuidType(isNullable); } else { return Record.fromFieldDescriptorsMap(isNullable, Record.toFieldDescriptorMap(messageDescriptor.getFields())); } @@ -451,7 +461,9 @@ private static Type fromProtoType(@Nullable Descriptors.GenericDescriptor descri @Nonnull private static Array fromProtoTypeToArray(@Nullable Descriptors.GenericDescriptor descriptor, @Nonnull Descriptors.FieldDescriptor.Type protoType, - @Nonnull TypeCode typeCode, boolean isNullable) { + @Nonnull TypeCode typeCode, + @Nullable DescriptorProtos.FieldOptions fieldOptions, + boolean isNullable) { if (typeCode.isPrimitive()) { final var primitiveType = primitiveType(typeCode, false); return new Array(isNullable, primitiveType); @@ -463,10 +475,10 @@ private static Array fromProtoTypeToArray(@Nullable Descriptors.GenericDescripto if (isNullable) { Descriptors.Descriptor wrappedDescriptor = ((Descriptors.Descriptor)Objects.requireNonNull(descriptor)).findFieldByName(NullableArrayTypeUtils.getRepeatedFieldName()).getMessageType(); Objects.requireNonNull(wrappedDescriptor); - return new Array(true, fromProtoType(wrappedDescriptor, Descriptors.FieldDescriptor.Type.MESSAGE, FieldDescriptorProto.Label.LABEL_OPTIONAL, false)); + return new Array(true, fromProtoType(wrappedDescriptor, Descriptors.FieldDescriptor.Type.MESSAGE, FieldDescriptorProto.Label.LABEL_OPTIONAL, fieldOptions, false)); } else { // case 2: any arbitrary sub message we don't understand - return new Array(false, fromProtoType(descriptor, protoType, FieldDescriptorProto.Label.LABEL_OPTIONAL, false)); + return new Array(false, fromProtoType(descriptor, protoType, FieldDescriptorProto.Label.LABEL_OPTIONAL, fieldOptions, false)); } } } @@ -707,6 +719,7 @@ enum TypeCode { INT(Integer.class, FieldDescriptorProto.Type.TYPE_INT32, true, true), LONG(Long.class, FieldDescriptorProto.Type.TYPE_INT64, true, true), STRING(String.class, FieldDescriptorProto.Type.TYPE_STRING, true, false), + VECTOR(ByteString.class, FieldDescriptorProto.Type.TYPE_BYTES, true, false), VERSION(FDBRecordVersion.class, FieldDescriptorProto.Type.TYPE_BYTES, true, false), ENUM(Enum.class, FieldDescriptorProto.Type.TYPE_ENUM, false, false), RECORD(Message.class, null, false, false), @@ -813,11 +826,12 @@ private static BiMap, TypeCode> computeClassToTypeCodeMap() { /** * Generates a {@link TypeCode} that corresponds to the given protobuf * {@link com.google.protobuf.DescriptorProtos.FieldDescriptorProto.Type}. - * @param protobufType The protobuf type. + * @param protobufType The protobuf descriptor of the type. * @return A corresponding {@link TypeCode} instance. */ @Nonnull - public static TypeCode fromProtobufType(@Nonnull final Descriptors.FieldDescriptor.Type protobufType) { + public static TypeCode fromProtobufFieldDescriptor(@Nonnull final Descriptors.FieldDescriptor.Type protobufType, + @Nullable final DescriptorProtos.FieldOptions fieldOptions) { switch (protobufType) { case DOUBLE: return TypeCode.DOUBLE; @@ -845,7 +859,15 @@ public static TypeCode fromProtobufType(@Nonnull final Descriptors.FieldDescript case MESSAGE: return TypeCode.RECORD; case BYTES: + { + if (fieldOptions != null) { + final var recordTypeOptions = fieldOptions.getExtension(RecordMetaDataOptionsProto.field); + if (recordTypeOptions.hasVectorOptions()) { + return TypeCode.VECTOR; + } + } return TypeCode.BYTES; + } default: throw new IllegalArgumentException("unknown protobuf type " + protobufType); } @@ -953,6 +975,7 @@ class Primitive implements Type { @Nonnull private final TypeCode typeCode; + @Nonnull private final Supplier hashCodeSupplier = Suppliers.memoize(this::computeHashCode); private Primitive(final boolean isNullable, @Nonnull final TypeCode typeCode) { @@ -985,12 +1008,12 @@ public void addProtoField(@Nonnull final TypeRepository.Builder typeRepositoryBu @Nonnull final Optional ignored, @Nonnull final FieldDescriptorProto.Label label) { final var protoType = Objects.requireNonNull(getTypeCode().getProtoType()); - descriptorBuilder.addField(FieldDescriptorProto.newBuilder() + final var fieldDescriptorBuilder = FieldDescriptorProto.newBuilder() .setNumber(fieldNumber) .setName(fieldName) .setType(protoType) - .setLabel(label) - .build()); + .setLabel(label); + descriptorBuilder.addField(fieldDescriptorBuilder.build()); } @Override @@ -1174,6 +1197,154 @@ public boolean equals(final Object other) { } } + class Vector implements Type { + + private final boolean isNullable; + + private final int precision; + + private final int dimensions; + + private Vector(final boolean isNullable, final int precision, final int dimensions) { + this.isNullable = isNullable; + this.precision = precision; + this.dimensions = dimensions; + } + + @Nonnull + @SuppressWarnings("PMD.ReplaceVectorWithList") + public static Vector of(final boolean isNullable, final int precision, final int dimensions) { + return new Vector(isNullable, precision, dimensions); + } + + @Override + public TypeCode getTypeCode() { + return TypeCode.VECTOR; + } + + @Override + public boolean isPrimitive() { + return true; + } + + @Override + public boolean isNullable() { + return isNullable; + } + + @Nonnull + @Override + public Type withNullability(final boolean newIsNullable) { + if (isNullable == newIsNullable) { + return this; + } + return new Vector(newIsNullable, precision, dimensions); + } + + public int getPrecision() { + return precision; + } + + public int getDimensions() { + return dimensions; + } + + @Nonnull + @Override + public ExplainTokens describe() { + final var resultExplainTokens = new ExplainTokens(); + resultExplainTokens.addKeyword(getTypeCode().toString()); + return resultExplainTokens.addOptionalWhitespace().addOpeningParen().addOptionalWhitespace() + .addNested(new ExplainTokens().addToString(precision).addToString(", ").addToString(dimensions)).addOptionalWhitespace() + .addClosingParen(); + } + + @Override + public void addProtoField(@Nonnull final TypeRepository.Builder typeRepositoryBuilder, + @Nonnull final DescriptorProto.Builder descriptorBuilder, final int fieldNumber, + @Nonnull final String fieldName, @Nonnull final Optional typeNameOptional, + @Nonnull final FieldDescriptorProto.Label label) { + final var protoType = Objects.requireNonNull(getTypeCode().getProtoType()); + FieldDescriptorProto.Builder builder = FieldDescriptorProto.newBuilder() + .setNumber(fieldNumber) + .setName(fieldName) + .setType(protoType) + .setLabel(label); + final var fieldOptions = RecordMetaDataOptionsProto.FieldOptions.newBuilder() + .setVectorOptions( + RecordMetaDataOptionsProto.FieldOptions.VectorOptions + .newBuilder() + .setPrecision(precision) + .setDimensions(dimensions) + .build()) + .build(); + builder.getOptionsBuilder().setExtension(RecordMetaDataOptionsProto.field, fieldOptions); + typeNameOptional.ifPresent(builder::setTypeName); + descriptorBuilder.addField(builder); + } + + @Nonnull + @Override + public PType toTypeProto(@Nonnull final PlanSerializationContext serializationContext) { + return PType.newBuilder().setVectorType(toProto(serializationContext)).build(); + } + + @Nonnull + @Override + public PVectorType toProto(@Nonnull final PlanSerializationContext serializationContext) { + final PVectorType.Builder vectorTypeBuilder = PVectorType.newBuilder() + .setIsNullable(isNullable) + .setDimensions(dimensions) + .setPrecision(precision); + return vectorTypeBuilder.build(); + } + + @Nonnull + @SuppressWarnings("PMD.ReplaceVectorWithList") + public static Vector fromProto(@Nonnull final PlanSerializationContext serializationContext, + @Nonnull final PVectorType vectorTypeProto) { + Verify.verify(vectorTypeProto.hasIsNullable()); + return new Vector(vectorTypeProto.getIsNullable(), vectorTypeProto.getPrecision(), vectorTypeProto.getDimensions()); + } + + /** + * Deserializer. + */ + @AutoService(PlanDeserializer.class) + public static class Deserializer implements PlanDeserializer { + @Nonnull + @Override + public Class getProtoMessageClass() { + return PVectorType.class; + } + + @Nonnull + @Override + @SuppressWarnings("PMD.ReplaceVectorWithList") + public Vector fromProto(@Nonnull final PlanSerializationContext serializationContext, + @Nonnull final PVectorType vectorTypeProto) { + return Vector.fromProto(serializationContext, vectorTypeProto); + } + } + + @Override + public int hashCode() { + return Objects.hash(getTypeCode().name(), precision, dimensions); + } + + @Override + @SuppressWarnings("PMD.ReplaceVectorWithList") + public boolean equals(final Object o) { + if (o == null || getClass() != o.getClass()) { + return false; + } + final Vector vector = (Vector)o; + return isNullable == vector.isNullable + && precision == vector.precision + && dimensions == vector.dimensions; + } + } + /** * The none type is an unresolved type meaning that an entity returning a none type should resolve the * type to a regular type as the runtime does not support a none-typed data producer. Only the empty array constant @@ -2248,10 +2419,12 @@ public static Record fromFieldDescriptorsMap(final boolean isNullable, @Nonnull final var fieldsBuilder = ImmutableList.builder(); for (final var entry : Objects.requireNonNull(fieldDescriptorMap).entrySet()) { final var fieldDescriptor = entry.getValue(); + final var fieldOptions = fieldDescriptor.getOptions(); fieldsBuilder.add( new Field(fromProtoType(getTypeSpecificDescriptor(fieldDescriptor), fieldDescriptor.getType(), fieldDescriptor.toProto().getLabel(), + fieldOptions, !fieldDescriptor.isRequired()), Optional.of(entry.getKey()), Optional.of(fieldDescriptor.getNumber()))); diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/TypeRepository.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/TypeRepository.java index c8aa4d2d2e..30e6de5c3d 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/TypeRepository.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/TypeRepository.java @@ -20,6 +20,7 @@ package com.apple.foundationdb.record.query.plan.cascades.typing; +import com.apple.foundationdb.record.RecordMetaDataOptionsProto; import com.apple.foundationdb.record.TupleFieldsProto; import com.google.common.base.Preconditions; import com.google.common.base.Verify; @@ -67,7 +68,7 @@ public class TypeRepository { public static final TypeRepository EMPTY_SCHEMA = empty(); @Nonnull - public static final List DEPENDENCIES = List.of(TupleFieldsProto.getDescriptor()); + public static final List DEPENDENCIES = List.of(TupleFieldsProto.getDescriptor(), RecordMetaDataOptionsProto.getDescriptor()); @Nonnull private final FileDescriptorSet fileDescSet; diff --git a/fdb-record-layer-core/src/main/proto/record_metadata_options.proto b/fdb-record-layer-core/src/main/proto/record_metadata_options.proto index 26ecbeb261..34721c09c5 100644 --- a/fdb-record-layer-core/src/main/proto/record_metadata_options.proto +++ b/fdb-record-layer-core/src/main/proto/record_metadata_options.proto @@ -60,8 +60,13 @@ message FieldOptions { repeated Index.Option options = 3; // Note: there is no way to specify these in a .proto file. } optional IndexOption index = 3; + message VectorOptions { + optional int32 precision = 1 [default = 16]; + optional int32 dimensions = 2 [default = 768]; + } + optional VectorOptions vectorOptions = 4; } extend google.protobuf.FieldOptions { - optional FieldOptions field = 1233; + optional FieldOptions field = 1239; } diff --git a/fdb-record-layer-core/src/main/proto/record_query_plan.proto b/fdb-record-layer-core/src/main/proto/record_query_plan.proto index 93a089e58e..c2fe657786 100644 --- a/fdb-record-layer-core/src/main/proto/record_query_plan.proto +++ b/fdb-record-layer-core/src/main/proto/record_query_plan.proto @@ -52,6 +52,7 @@ message PType { RELATION = 15; NONE = 16; UUID = 17; + VECTOR = 18; } message PPrimitiveType { @@ -75,6 +76,12 @@ message PType { // nothing } + message PVectorType { + optional bool is_nullable = 1; + optional int32 precision = 2; + optional int32 dimensions = 3; + } + message PAnyRecordType { optional bool is_nullable = 1; } @@ -123,6 +130,7 @@ message PType { PArrayType array_type = 8; PAnyRecordType any_record_type = 9; PUuidType uuid_type = 10; + PVectorType vector_type = 11; } } diff --git a/fdb-relational-api/src/main/java/com/apple/foundationdb/relational/api/metadata/DataType.java b/fdb-relational-api/src/main/java/com/apple/foundationdb/relational/api/metadata/DataType.java index 56991fb11b..8872a71221 100644 --- a/fdb-relational-api/src/main/java/com/apple/foundationdb/relational/api/metadata/DataType.java +++ b/fdb-relational-api/src/main/java/com/apple/foundationdb/relational/api/metadata/DataType.java @@ -70,6 +70,7 @@ public abstract class DataType { typeCodeJdbcTypeMap.put(Code.ENUM, Types.OTHER); typeCodeJdbcTypeMap.put(Code.UUID, Types.OTHER); typeCodeJdbcTypeMap.put(Code.BYTES, Types.BINARY); + typeCodeJdbcTypeMap.put(Code.VECTOR, Types.OTHER); typeCodeJdbcTypeMap.put(Code.VERSION, Types.BINARY); typeCodeJdbcTypeMap.put(Code.STRUCT, Types.STRUCT); typeCodeJdbcTypeMap.put(Code.ARRAY, Types.ARRAY); @@ -729,6 +730,79 @@ public String toString() { } } + public static final class VectorType extends DataType { + private final int precision; + + private final int dimensions; + + @Nonnull + private final Supplier hashCodeSupplier = Suppliers.memoize(this::computeHashCode); + + private VectorType(final boolean isNullable, int precision, int dimensions) { + super(isNullable, true, Code.VECTOR); + this.precision = precision; + this.dimensions = dimensions; + } + + @Override + public boolean isResolved() { + return true; + } + + @Nonnull + @Override + public DataType withNullable(final boolean isNullable) { + if (isNullable == this.isNullable()) { + return this; + } + return new VectorType(isNullable, precision, dimensions); + } + + @Nonnull + @Override + public DataType resolve(@Nonnull final Map resolutionMap) { + return this; + } + + public int getPrecision() { + return precision; + } + + public int getDimensions() { + return dimensions; + } + + private int computeHashCode() { + return Objects.hash(getCode(), precision, dimensions, isNullable()); + } + + @Override + public int hashCode() { + return hashCodeSupplier.get(); + } + + @Override + public boolean equals(final Object o) { + if (o == null || getClass() != o.getClass()) { + return false; + } + final VectorType that = (VectorType)o; + return precision == that.precision + && dimensions == that.dimensions + && isNullable() == that.isNullable(); + } + + @Override + public String toString() { + return "vector(p=" + precision + ", d=" + dimensions + ")" + (isNullable() ? " ∪ ∅" : ""); + } + + @Nonnull + public static VectorType of(int precision, int dimensions, boolean isNullable) { + return new VectorType(isNullable, precision, dimensions); + } + } + public static final class VersionType extends DataType { @Nonnull private static final VersionType NOT_NULLABLE_INSTANCE = new VersionType(false); @@ -1504,7 +1578,8 @@ public enum Code { STRUCT, ARRAY, UNKNOWN, - NULL + NULL, + VECTOR, } @SuppressWarnings("PMD.AvoidFieldNameMatchingTypeName") diff --git a/fdb-relational-core/src/main/antlr/RelationalLexer.g4 b/fdb-relational-core/src/main/antlr/RelationalLexer.g4 index b729d6b5b5..aa2cd94702 100644 --- a/fdb-relational-core/src/main/antlr/RelationalLexer.g4 +++ b/fdb-relational-core/src/main/antlr/RelationalLexer.g4 @@ -730,6 +730,13 @@ USER_RESOURCES: 'USER_RESOURCES'; VALIDATION: 'VALIDATION'; VALUE: 'VALUE'; VARIABLES: 'VARIABLES'; +VECTOR: 'VECTOR'; +VECTOR16: 'VECTOR16'; +VECTOR32: 'VECTOR32'; +VECTOR64: 'VECTOR64'; +FLOATVECTOR: 'FLOATVECTOR'; +DOUBLEVECTOR: 'DOUBLEVECTOR'; +HALFVECTOR: 'HALFVECTOR'; VIEW: 'VIEW'; VIRTUAL: 'VIRTUAL'; VISIBLE: 'VISIBLE'; diff --git a/fdb-relational-core/src/main/antlr/RelationalParser.g4 b/fdb-relational-core/src/main/antlr/RelationalParser.g4 index 139de5984c..417df54640 100644 --- a/fdb-relational-core/src/main/antlr/RelationalParser.g4 +++ b/fdb-relational-core/src/main/antlr/RelationalParser.g4 @@ -138,7 +138,17 @@ columnType : primitiveType | customType=uid; primitiveType - : BOOLEAN | INTEGER | BIGINT | FLOAT | DOUBLE | STRING | BYTES; + : BOOLEAN | INTEGER | BIGINT | FLOAT | DOUBLE | STRING | BYTES | vectorType; + +vectorType + : VECTOR '(' length=DECIMAL_LITERAL ')' + | VECTOR16 '(' length=DECIMAL_LITERAL ')' + | VECTOR32 '(' length=DECIMAL_LITERAL ')' + | VECTOR64 '(' length=DECIMAL_LITERAL ')' + | HALFVECTOR '(' length=DECIMAL_LITERAL ')' + | FLOATVECTOR '(' length=DECIMAL_LITERAL ')' + | DOUBLEVECTOR '(' length=DECIMAL_LITERAL ')' + ; columnConstraint : nullNotnull #nullColumnConstraint diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/catalog/RecordLayerStoreSchemaTemplateCatalog.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/catalog/RecordLayerStoreSchemaTemplateCatalog.java index cea1c2128b..c755308d7e 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/catalog/RecordLayerStoreSchemaTemplateCatalog.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/catalog/RecordLayerStoreSchemaTemplateCatalog.java @@ -76,6 +76,14 @@ */ class RecordLayerStoreSchemaTemplateCatalog implements SchemaTemplateCatalog { + @Nonnull + private static final com.google.protobuf.ExtensionRegistry registry = com.google.protobuf.ExtensionRegistry.newInstance(); + + static { + registry.add(com.apple.foundationdb.record.RecordMetaDataOptionsProto.field); + registry.add(com.apple.foundationdb.record.RecordMetaDataOptionsProto.record); + } + @Nonnull private final RecordLayerSchema catalogSchema; @@ -210,7 +218,7 @@ private static SchemaTemplate toSchemaTemplate(@Nonnull final Message message) t // deserialization of the same message over and over again. final Descriptors.Descriptor descriptor = message.getDescriptorForType(); final ByteString bs = Assert.castUnchecked(message.getField(descriptor.findFieldByName(SchemaTemplateSystemTable.METADATA)), ByteString.class); - final RecordMetaData metaData = RecordMetaData.build(RecordMetaDataProto.MetaData.parseFrom(bs.toByteArray())); + final RecordMetaData metaData = RecordMetaData.newBuilder().setRecords(RecordMetaDataProto.MetaData.parseFrom(bs.toByteArray(), registry)).getRecordMetaData(); final String name = message.getField(descriptor.findFieldByName(SchemaTemplateSystemTable.TEMPLATE_NAME)).toString(); int templateVersion = (int) message.getField(descriptor.findFieldByName(SchemaTemplateSystemTable.TEMPLATE_VERSION)); return RecordLayerSchemaTemplate.fromRecordMetadata(metaData, name, templateVersion); diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/DataTypeUtils.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/DataTypeUtils.java index 2c3820be2c..d58ad4ed0c 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/DataTypeUtils.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/DataTypeUtils.java @@ -61,6 +61,12 @@ public static DataType toRelationalType(@Nonnull final Type type) { final var typeCode = type.getTypeCode(); + // if the type code is BYTES, + if (typeCode.equals(Type.TypeCode.VECTOR)) { + final var vectorType = (Type.Vector)type; + return DataType.VectorType.of(vectorType.getPrecision(), vectorType.getDimensions(), vectorType.isNullable()); + } + if (typeCode == Type.TypeCode.ANY || typeCode == Type.TypeCode.NONE || typeCode == Type.TypeCode.NULL || typeCode == Type.TypeCode.UNKNOWN) { return DataType.UnknownType.instance(); } @@ -112,10 +118,19 @@ public static Type toRecordLayerType(@Nonnull final DataType type) { return primitivesMap.get(type); } + // handle bytes with fixed size and precision. + if (type.getCode().equals(DataType.Code.VECTOR)) { + final var vectorType = (DataType.VectorType)type; + final var precision = vectorType.getPrecision(); + final var dimensions = vectorType.getDimensions(); + return Type.Vector.of(type.isNullable(), precision, dimensions); + } + switch (type.getCode()) { case STRUCT: final var struct = (DataType.StructType) type; - final var fields = struct.getFields().stream().map(field -> Type.Record.Field.of(DataTypeUtils.toRecordLayerType(field.getType()), Optional.of(field.getName()), Optional.of(field.getIndex()))).collect(Collectors.toList()); + final var fields = struct.getFields().stream().map(field -> Type.Record.Field.of(DataTypeUtils.toRecordLayerType(field.getType()), + Optional.of(field.getName()), Optional.of(field.getIndex()))).collect(Collectors.toList()); return Type.Record.fromFieldsWithName(struct.getName(), struct.isNullable(), fields); case ARRAY: final var asArray = (DataType.ArrayType) type; @@ -123,11 +138,13 @@ public static Type toRecordLayerType(@Nonnull final DataType type) { // but since in RL we store the elements as a 'repeated' field, there is not a way to tell if an element is explicitly 'null'. // The current RL behavior loses the nullability information even if the constituent of Type.Array is explicitly marked 'nullable'. Hence, // the check here avoids silently swallowing the requirement. - Assert.thatUnchecked(asArray.getElementType().getCode() == DataType.Code.NULL || !asArray.getElementType().isNullable(), ErrorCode.UNSUPPORTED_OPERATION, "No support for nullable array elements."); + Assert.thatUnchecked(asArray.getElementType().getCode() == DataType.Code.NULL || !asArray.getElementType().isNullable(), + ErrorCode.UNSUPPORTED_OPERATION, "No support for nullable array elements."); return new Type.Array(asArray.isNullable(), toRecordLayerType(asArray.getElementType())); case ENUM: final var asEnum = (DataType.EnumType) type; - final List enumValues = asEnum.getValues().stream().map(v -> new Type.Enum.EnumValue(v.getName(), v.getNumber())).collect(Collectors.toList()); + final List enumValues = asEnum.getValues().stream().map(v -> new Type.Enum.EnumValue(v.getName(), + v.getNumber())).collect(Collectors.toList()); return new Type.Enum(asEnum.isNullable(), enumValues, asEnum.getName()); case UNKNOWN: return new Type.Any(); diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/SemanticAnalyzer.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/SemanticAnalyzer.java index 648bbba39d..d391223ef3 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/SemanticAnalyzer.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/SemanticAnalyzer.java @@ -502,6 +502,130 @@ public Optional lookupNestedField(@Nonnull Identifier requestedIdent return Optional.of(nestedAttribute); } + public static final class ParsedTypeInfo { + @Nullable + private final RelationalParser.PrimitiveTypeContext primitiveTypeContext; + + @Nullable + private final Identifier customType; + + private final boolean isNullable; + + private final boolean isRepeated; + + private ParsedTypeInfo(@Nullable final RelationalParser.PrimitiveTypeContext primitiveTypeContext, + @Nullable final Identifier customType, final boolean isNullable, final boolean isRepeated) { + this.primitiveTypeContext = primitiveTypeContext; + this.customType = customType; + this.isNullable = isNullable; + this.isRepeated = isRepeated; + } + + public boolean hasPrimitiveType() { + return primitiveTypeContext != null; + } + + @Nullable + public RelationalParser.PrimitiveTypeContext getPrimitiveTypeContext() { + return primitiveTypeContext; + } + + public boolean hasCustomType() { + return customType != null; + } + + @Nullable + public Identifier getCustomType() { + return customType; + } + + public boolean isNullable() { + return isNullable; + } + + public boolean isRepeated() { + return isRepeated; + } + + @Nonnull + public static ParsedTypeInfo ofPrimitiveType(@Nonnull final RelationalParser.PrimitiveTypeContext primitiveTypeContext, + final boolean isNullable, final boolean isRepeated) { + return new ParsedTypeInfo(primitiveTypeContext, null, isNullable, isRepeated); + } + + @Nonnull + public static ParsedTypeInfo ofCustomType(@Nonnull final Identifier customType, + final boolean isNullable, final boolean isRepeated) { + return new ParsedTypeInfo(null, customType, isNullable, isRepeated); + } + } + + @Nonnull + public DataType lookupType(@Nonnull final ParsedTypeInfo parsedTypeInfo, + @Nonnull final Function> dataTypeProvider) { + DataType type; + final var isNullable = parsedTypeInfo.isNullable(); + if (parsedTypeInfo.hasCustomType()) { + final var typeName = Assert.notNullUnchecked(parsedTypeInfo.getCustomType()).getName(); + final var maybeFound = dataTypeProvider.apply(typeName); + // if we cannot find the type now, mark it, we will try to resolve it later on via a second pass. + type = maybeFound.orElseGet(() -> DataType.UnresolvedType.of(typeName, isNullable)); + } else { + final var primitiveType = Assert.notNullUnchecked(parsedTypeInfo.getPrimitiveTypeContext()); + if (primitiveType.vectorType() != null) { + final var ctx = primitiveType.vectorType(); + int precision = 16; + if (ctx.VECTOR32() != null || ctx.FLOATVECTOR() != null) { + precision = 32; + } else if (ctx.VECTOR64() != null || ctx.DOUBLEVECTOR() != null) { + precision = 64; + } + int length = Assert.castUnchecked(ParseHelpers.parseDecimal(ctx.length.getText()), Integer.class); + type = DataType.VectorType.of(precision, length, isNullable); + } else { + final var primitiveTypeName = parsedTypeInfo.getPrimitiveTypeContext().getText(); + + switch (primitiveTypeName.toUpperCase(Locale.ROOT)) { + case "STRING": + type = isNullable ? DataType.Primitives.NULLABLE_STRING.type() : DataType.Primitives.STRING.type(); + break; + case "INTEGER": + type = isNullable ? DataType.Primitives.NULLABLE_INTEGER.type() : DataType.Primitives.INTEGER.type(); + break; + case "BIGINT": + type = isNullable ? DataType.Primitives.NULLABLE_LONG.type() : DataType.Primitives.LONG.type(); + break; + case "DOUBLE": + type = isNullable ? DataType.Primitives.NULLABLE_DOUBLE.type() : DataType.Primitives.DOUBLE.type(); + break; + case "BOOLEAN": + type = isNullable ? DataType.Primitives.NULLABLE_BOOLEAN.type() : DataType.Primitives.BOOLEAN.type(); + break; + case "BYTES": + type = isNullable ? DataType.Primitives.NULLABLE_BYTES.type() : DataType.Primitives.BYTES.type(); + break; + case "FLOAT": + type = isNullable ? DataType.Primitives.NULLABLE_FLOAT.type() : DataType.Primitives.FLOAT.type(); + break; + default: + Assert.notNullUnchecked(metadataCatalog); + // assume it is a custom type, will fail in upper layers if the type can not be resolved. + // lookup the type (Struct, Table, or Enum) in the schema template metadata under construction. + final var maybeFound = dataTypeProvider.apply(primitiveTypeName); + // if we cannot find the type now, mark it, we will try to resolve it later on via a second pass. + type = maybeFound.orElseGet(() -> DataType.UnresolvedType.of(primitiveTypeName, isNullable)); + break; + } + } + } + + if (parsedTypeInfo.isRepeated()) { + return DataType.ArrayType.from(type.withNullable(false), isNullable); + } else { + return type; + } + } + @Nonnull public DataType lookupType(@Nonnull Identifier typeIdentifier, boolean isNullable, boolean isRepeated, @Nonnull Function> dataTypeProvider) { diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java index 394c91ec0b..9d8b4a23e0 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java @@ -370,10 +370,14 @@ public DataType visitColumnType(@Nonnull RelationalParser.ColumnTypeContext ctx) return ddlVisitor.visitColumnType(ctx); } - @Nonnull @Override - public DataType visitPrimitiveType(@Nonnull RelationalParser.PrimitiveTypeContext ctx) { - return ddlVisitor.visitPrimitiveType(ctx); + public Object visitPrimitiveType(final RelationalParser.PrimitiveTypeContext ctx) { + return visitChildren(ctx); + } + + @Override + public Object visitVectorType(final RelationalParser.VectorTypeContext ctx) { + return visitChildren(ctx); } @Nonnull diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java index 5575626681..9fd7482ace 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java @@ -56,7 +56,6 @@ import java.util.ArrayList; import java.util.List; import java.util.Locale; -import java.util.Optional; import java.util.function.Function; import java.util.stream.Collectors; @@ -99,11 +98,14 @@ public static DdlVisitor of(@Nonnull BaseVisitor delegate, @Override public DataType visitFunctionColumnType(@Nonnull final RelationalParser.FunctionColumnTypeContext ctx) { final var semanticAnalyzer = getDelegate().getSemanticAnalyzer(); + final SemanticAnalyzer.ParsedTypeInfo typeInfo; if (ctx.customType != null) { final var columnType = visitUid(ctx.customType); - return semanticAnalyzer.lookupType(columnType, true, false, metadataBuilder::findType); + typeInfo = SemanticAnalyzer.ParsedTypeInfo.ofCustomType(columnType, true, false); + } else { + typeInfo = SemanticAnalyzer.ParsedTypeInfo.ofPrimitiveType(ctx.primitiveType(), true, false); } - return visitPrimitiveType(ctx.primitiveType()).withNullable(true); + return semanticAnalyzer.lookupType(typeInfo, metadataBuilder::findType); } // TODO: remove @@ -111,20 +113,14 @@ public DataType visitFunctionColumnType(@Nonnull final RelationalParser.Function @Override public DataType visitColumnType(@Nonnull RelationalParser.ColumnTypeContext ctx) { final var semanticAnalyzer = getDelegate().getSemanticAnalyzer(); + final SemanticAnalyzer.ParsedTypeInfo typeInfo; if (ctx.customType != null) { final var columnType = visitUid(ctx.customType); - return semanticAnalyzer.lookupType(columnType, false, false, metadataBuilder::findType); + typeInfo = SemanticAnalyzer.ParsedTypeInfo.ofCustomType(columnType, true, false); + } else { + typeInfo = SemanticAnalyzer.ParsedTypeInfo.ofPrimitiveType(ctx.primitiveType(), true, false); } - return visitPrimitiveType(ctx.primitiveType()); - } - - // TODO: remove - @Nonnull - @Override - public DataType visitPrimitiveType(@Nonnull RelationalParser.PrimitiveTypeContext ctx) { - final var semanticAnalyzer = getDelegate().getSemanticAnalyzer(); - final var primitiveType = Identifier.of(ctx.getText()); - return semanticAnalyzer.lookupType(primitiveType, false, false, ignored -> Optional.empty()); + return semanticAnalyzer.lookupType(typeInfo, metadataBuilder::findType); } /** @@ -147,9 +143,16 @@ public RecordLayerColumn visitColumnDefinition(@Nonnull RelationalParser.ColumnD // but a way to represent it in RecordMetadata. Assert.thatUnchecked(isRepeated || isNullable, ErrorCode.UNSUPPORTED_OPERATION, "NOT NULL is only allowed for ARRAY column type"); containsNullableArray = containsNullableArray || (isRepeated && isNullable); - final var columnTypeId = ctx.columnType().customType != null ? visitUid(ctx.columnType().customType) : Identifier.of(ctx.columnType().getText()); + final var semanticAnalyzer = getDelegate().getSemanticAnalyzer(); - final var columnType = semanticAnalyzer.lookupType(columnTypeId, isNullable, isRepeated, metadataBuilder::findType); + final SemanticAnalyzer.ParsedTypeInfo typeInfo; + if (ctx.columnType().customType != null) { + final var columnType = visitUid(ctx.columnType().customType); + typeInfo = SemanticAnalyzer.ParsedTypeInfo.ofCustomType(columnType, true, false); + } else { + typeInfo = SemanticAnalyzer.ParsedTypeInfo.ofPrimitiveType(ctx.columnType().primitiveType(), true, false); + } + final var columnType = semanticAnalyzer.lookupType(typeInfo, metadataBuilder::findType); return RecordLayerColumn.newBuilder().setName(columnId.getName()).setDataType(columnType).build(); } diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java index 4f2663f3b9..ef0fbf67ba 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java @@ -200,12 +200,16 @@ public DataType visitColumnType(@Nonnull RelationalParser.ColumnTypeContext ctx) return getDelegate().visitColumnType(ctx); } - @Nonnull @Override - public DataType visitPrimitiveType(@Nonnull RelationalParser.PrimitiveTypeContext ctx) { + public Object visitPrimitiveType(final RelationalParser.PrimitiveTypeContext ctx) { return getDelegate().visitPrimitiveType(ctx); } + @Override + public Object visitVectorType(final RelationalParser.VectorTypeContext ctx) { + return getDelegate().visitVectorType(ctx); + } + @Nonnull @Override public Boolean visitNullColumnConstraint(@Nonnull RelationalParser.NullColumnConstraintContext ctx) { diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java index bf13ecc143..5d61f700e7 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java @@ -146,10 +146,6 @@ public interface TypedVisitor extends RelationalParserVisitor { @Override DataType visitColumnType(@Nonnull RelationalParser.ColumnTypeContext ctx); - @Nonnull - @Override - DataType visitPrimitiveType(@Nonnull RelationalParser.PrimitiveTypeContext ctx); - @Nonnull @Override Boolean visitNullColumnConstraint(@Nonnull RelationalParser.NullColumnConstraintContext ctx); diff --git a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java new file mode 100644 index 0000000000..ddff1180a0 --- /dev/null +++ b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java @@ -0,0 +1,70 @@ +/* + * VectorTypeTest.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.relational.recordlayer.query; + +import com.apple.foundationdb.relational.api.StructResultSetMetaData; +import com.apple.foundationdb.relational.api.metadata.DataType; +import com.apple.foundationdb.relational.recordlayer.EmbeddedRelationalExtension; +import com.apple.foundationdb.relational.utils.Ddl; +import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.Order; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +import javax.annotation.Nonnull; +import java.net.URI; +import java.util.stream.Stream; + +public class VectorTypeTest { + + @RegisterExtension + @Order(0) + public final EmbeddedRelationalExtension relationalExtension = new EmbeddedRelationalExtension(); + + + @Nonnull + public static Stream vectorArguments() { + return Stream.of( + Arguments.of("halfvector(512)", DataType.VectorType.of(16, 512, true)), + Arguments.of("vector16(512)", DataType.VectorType.of(16, 512, true)), + Arguments.of("doublevector(1024)", DataType.VectorType.of(64, 1024, true)), + Arguments.of("vector32(768)", DataType.VectorType.of(32, 768, true)), + Arguments.of("vector(256)", DataType.VectorType.of(16, 256, true))); + } + + @ParameterizedTest(name = "{0} evaluates to data type {1}") + @MethodSource("vectorArguments") + void vectorTest(@Nonnull final String ddlType, @Nonnull final DataType expectedType) throws Exception { + final String schemaTemplate = "create table t1(id bigint, col1 " + ddlType + ", primary key(id))"; + try (var ddl = Ddl.builder().database(URI.create("/TEST/QT")).relationalExtension(relationalExtension).schemaTemplate(schemaTemplate).build()) { + try (var statement = ddl.setSchemaAndGetConnection().createStatement()) { + statement.execute("select * from t1"); + final var metadata = statement.getResultSet().getMetaData(); + Assertions.assertThat(metadata).isInstanceOf(StructResultSetMetaData.class); + final var relationalMetadata = (StructResultSetMetaData)metadata; + final var type = relationalMetadata.getRelationalDataType().getFields().get(1).getType(); + Assertions.assertThat(type).isEqualTo(expectedType); + } + } + } +} diff --git a/yaml-tests/src/test/java/YamlIntegrationTests.java b/yaml-tests/src/test/java/YamlIntegrationTests.java index 138f9bcc75..8b7b83bf09 100644 --- a/yaml-tests/src/test/java/YamlIntegrationTests.java +++ b/yaml-tests/src/test/java/YamlIntegrationTests.java @@ -267,8 +267,8 @@ public void enumTest(YamlTest.Runner runner) throws Exception { } @TestTemplate - public void uuidProtoTest(YamlTest.Runner runner) throws Exception { - runner.runYamsql("uuid-proto.yamsql"); + public void uuidTest(YamlTest.Runner runner) throws Exception { + runner.runYamsql("uuid.yamsql"); } @TestTemplate @@ -305,4 +305,9 @@ public void literalExtractionTests(YamlTest.Runner runner) throws Exception { public void caseSensitivityTest(YamlTest.Runner runner) throws Exception { runner.runYamsql("case-sensitivity.yamsql"); } + + @TestTemplate + public void vectorTypeTests(YamlTest.Runner runner) throws Exception { + runner.runYamsql("vector-type.yamsql"); + } } From 108f88d00f7c2683631e8d15c6b13e4cf223106b Mon Sep 17 00:00:00 2001 From: Youssef Hatem Date: Wed, 13 Aug 2025 14:09:27 +0100 Subject: [PATCH 25/34] - ddl support for tables with HNSW. --- .../src/main/antlr/RelationalLexer.g4 | 4 + .../src/main/antlr/RelationalParser.g4 | 19 ++++- .../recordlayer/query/IndexGenerator.java | 49 +++++++++-- .../query/visitors/BaseVisitor.java | 24 ++++++ .../query/visitors/DdlVisitor.java | 85 ++++++++++++++++++- .../query/visitors/DelegatingVisitor.java | 25 ++++++ .../query/visitors/TypedVisitor.java | 20 +++++ .../recordlayer/query/VectorTypeTest.java | 16 ++++ 8 files changed, 232 insertions(+), 10 deletions(-) diff --git a/fdb-relational-core/src/main/antlr/RelationalLexer.g4 b/fdb-relational-core/src/main/antlr/RelationalLexer.g4 index aa2cd94702..7a3609b146 100644 --- a/fdb-relational-core/src/main/antlr/RelationalLexer.g4 +++ b/fdb-relational-core/src/main/antlr/RelationalLexer.g4 @@ -112,6 +112,7 @@ GET: 'GET'; GRANT: 'GRANT'; GROUP: 'GROUP'; HAVING: 'HAVING'; +HNSW: 'HNSW'; HIGH_PRIORITY: 'HIGH_PRIORITY'; HISTOGRAM: 'HISTOGRAM'; IF: 'IF'; @@ -160,6 +161,7 @@ OPTION: 'OPTION'; OPTIONAL: 'OPTIONAL'; OPTIONALLY: 'OPTIONALLY'; OR: 'OR'; +ORGANIZED: 'ORGANIZED'; ORDER: 'ORDER'; OUT: 'OUT'; OVER: 'OVER'; @@ -465,6 +467,7 @@ DO: 'DO'; DUMPFILE: 'DUMPFILE'; DUPLICATE: 'DUPLICATE'; DYNAMIC: 'DYNAMIC'; +EF_CONSTRUCTION: 'EF_CONSTRUCTION'; ENABLE: 'ENABLE'; ENCRYPTION: 'ENCRYPTION'; END: 'END'; @@ -537,6 +540,7 @@ LIST: 'LIST'; LOCAL: 'LOCAL'; LOGFILE: 'LOGFILE'; LOGS: 'LOGS'; +M: 'M'; MASTER: 'MASTER'; MASTER_AUTO_POSITION: 'MASTER_AUTO_POSITION'; MASTER_CONNECT_RETRY: 'MASTER_CONNECT_RETRY'; diff --git a/fdb-relational-core/src/main/antlr/RelationalParser.g4 b/fdb-relational-core/src/main/antlr/RelationalParser.g4 index 417df54640..68ab052877 100644 --- a/fdb-relational-core/src/main/antlr/RelationalParser.g4 +++ b/fdb-relational-core/src/main/antlr/RelationalParser.g4 @@ -122,9 +122,23 @@ structDefinition ; tableDefinition - : TABLE uid LEFT_ROUND_BRACKET columnDefinition (COMMA columnDefinition)* COMMA primaryKeyDefinition RIGHT_ROUND_BRACKET + : TABLE uid LEFT_ROUND_BRACKET columnDefinition (COMMA columnDefinition)* COMMA primaryKeyDefinition (COMMA organizedByClause)? RIGHT_ROUND_BRACKET ; +organizedByClause + : ORGANIZED BY HNSW '(' embeddingsCol=fullId partitionClause? ')' hnswConfigurations? + ; + +hnswConfigurations + : WITH '(' hnswConfiguration (COMMA hnswConfiguration)* ')' + ; + +hnswConfiguration + : M '=' mValue=DECIMAL_LITERAL + | EF_CONSTRUCTION '=' efConstructionValue=DECIMAL_LITERAL + ; + + columnDefinition : colName=uid columnType ARRAY? columnConstraint? ; @@ -1110,10 +1124,11 @@ frameRange | expression (PRECEDING | FOLLOWING) ; +*/ + partitionClause : PARTITION BY expression (',' expression)* ; -*/ scalarFunctionName : functionNameBase diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/IndexGenerator.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/IndexGenerator.java index 14c10f53c5..127ae476fe 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/IndexGenerator.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/IndexGenerator.java @@ -23,6 +23,7 @@ import com.apple.foundationdb.annotation.API; import com.apple.foundationdb.record.EvaluationContext; import com.apple.foundationdb.record.RecordCoreException; +import com.apple.foundationdb.record.expressions.RecordKeyExpressionProto; import com.apple.foundationdb.record.metadata.IndexOptions; import com.apple.foundationdb.record.metadata.IndexPredicate; import com.apple.foundationdb.record.metadata.IndexTypes; @@ -75,10 +76,12 @@ import com.apple.foundationdb.relational.util.NullableArrayUtils; import com.google.common.base.Verify; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.Iterables; import com.google.common.collect.Iterators; import com.google.common.collect.Lists; import com.google.common.collect.PeekingIterator; +import com.ibm.icu.impl.locale.XCldrStub; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -484,7 +487,7 @@ private KeyExpression removeBitmapBucketOffset(@Nonnull KeyExpression groupingEx } @Nonnull - private KeyExpression generate(@Nonnull List fields, @Nonnull Map orderingFunctions) { + private static KeyExpression generate(@Nonnull List fields, @Nonnull Map orderingFunctions) { if (fields.isEmpty()) { return EmptyKeyExpression.EMPTY; } else if (fields.size() == 1) { @@ -514,7 +517,7 @@ private KeyExpression generate(@Nonnull List fields, @Nonnull Map orderingFunctions) { + private static KeyExpression toKeyExpression(Value value, Map orderingFunctions) { var expr = toKeyExpression(value); if (orderingFunctions.containsKey(value)) { return function(orderingFunctions.get(value), expr); @@ -524,7 +527,7 @@ private KeyExpression toKeyExpression(Value value, Map orderingFu } @Nonnull - private KeyExpression toKeyExpression(@Nonnull Value value) { + private static KeyExpression toKeyExpression(@Nonnull Value value) { if (value instanceof VersionValue) { return VersionKeyExpression.VERSION; } else if (value instanceof FieldValue) { @@ -739,12 +742,12 @@ private Value dereference(@Nonnull Value value) { } @Nonnull - private KeyExpression toKeyExpression(@Nonnull List> fields) { + private static KeyExpression toKeyExpression(@Nonnull List> fields) { return toKeyExpression(fields, 0); } @Nonnull - private KeyExpression toKeyExpression(@Nonnull List> fields, int index) { + private static KeyExpression toKeyExpression(@Nonnull List> fields, int index) { Assert.thatUnchecked(!fields.isEmpty()); final var field = fields.get(index); final var keyExpression = toKeyExpression(field.getLeft(), field.getRight()); @@ -779,4 +782,40 @@ private static FieldKeyExpression toKeyExpression(@Nonnull String name, @Nonnull public static IndexGenerator from(@Nonnull RelationalExpression relationalExpression, boolean useLongBasedExtremumEver) { return new IndexGenerator(relationalExpression, useLongBasedExtremumEver); } + + @Nonnull + public static RecordLayerIndex generateHnswIndex(@Nonnull final String tableName, + @Nonnull final Expression embedding, + @Nonnull final Expressions partitionExpressions, + @Nonnull final Map options) { + final var embeddingKeyExpression = toKeyExpression(embedding.getUnderlying()); + final var partitionKeyExpression = generate(ImmutableList.copyOf(partitionExpressions.underlying().iterator()), + Collections.emptyMap()); + final var keyExpression = keyWithValue(concat(embeddingKeyExpression, partitionKeyExpression), + embeddingKeyExpression.getColumnSize()); + + final var indexOptions = options.entrySet().stream().map( entry -> { + final var key = entry.getKey(); + final var value = entry.getValue(); + switch (key) { + case "M": + return Map.entry(IndexOptions.HNSW_M, value); + case "EF_CONSTRUCTION": + return Map.entry(IndexOptions.HNSW_EF_CONSTRUCTION, value); + default: + throw new RelationalException("unknown HNSW option: " + key, ErrorCode.SYNTAX_ERROR) + .toUncheckedWrappedException(); + } + }).collect(ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue)); + + + final var builder = RecordLayerIndex.newBuilder(); + final var index = builder.setIndexType(IndexTypes.VECTOR) + .setName(tableName + "$hnsw") + .setOptions(indexOptions) + .setTableName(tableName) + .setKeyExpression(keyExpression) + .build(); + return index; + } } diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java index 9d8b4a23e0..bb2772ee02 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java @@ -55,6 +55,7 @@ import javax.annotation.Nullable; import java.net.URI; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Set; @@ -352,6 +353,23 @@ public RecordLayerTable visitTableDefinition(@Nonnull RelationalParser.TableDefi return ddlVisitor.visitTableDefinition(ctx); } + @Override + public Object visitOrganizedByClause(final RelationalParser.OrganizedByClauseContext ctx) { + return ddlVisitor.visitOrganizedByClause(ctx); + } + + @Nonnull + @Override + public Map visitHnswConfigurations(final RelationalParser.HnswConfigurationsContext ctx) { + return ddlVisitor.visitHnswConfigurations(ctx); + } + + @Nonnull + @Override + public NonnullPair visitHnswConfiguration(final RelationalParser.HnswConfigurationContext ctx) { + return ddlVisitor.visitHnswConfiguration(ctx); + } + @Nonnull @Override public Object visitColumnDefinition(@Nonnull RelationalParser.ColumnDefinitionContext ctx) { @@ -1506,6 +1524,12 @@ public Object visitWindowName(@Nonnull RelationalParser.WindowNameContext ctx) { return visitChildren(ctx); } + @Nonnull + @Override + public Expressions visitPartitionClause(final RelationalParser.PartitionClauseContext ctx) { + return ddlVisitor.visitPartitionClause(ctx); + } + @Nonnull @Override public Object visitScalarFunctionName(@Nonnull RelationalParser.ScalarFunctionNameContext ctx) { diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java index 9fd7482ace..8403940ebf 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java @@ -21,13 +21,17 @@ package com.apple.foundationdb.relational.recordlayer.query.visitors; import com.apple.foundationdb.annotation.API; +import com.apple.foundationdb.record.metadata.IndexTypes; import com.apple.foundationdb.record.query.plan.cascades.expressions.LogicalSortExpression; import com.apple.foundationdb.record.query.plan.cascades.values.PromoteValue; import com.apple.foundationdb.record.query.plan.cascades.values.ThrowsValue; +import com.apple.foundationdb.record.util.pair.NonnullPair; import com.apple.foundationdb.relational.api.Options; import com.apple.foundationdb.relational.api.ddl.MetadataOperationsFactory; import com.apple.foundationdb.relational.api.exceptions.ErrorCode; +import com.apple.foundationdb.relational.api.exceptions.RelationalException; import com.apple.foundationdb.relational.api.metadata.DataType; +import com.apple.foundationdb.relational.api.metadata.Index; import com.apple.foundationdb.relational.api.metadata.InvokedRoutine; import com.apple.foundationdb.relational.generated.RelationalParser; import com.apple.foundationdb.relational.recordlayer.metadata.DataTypeUtils; @@ -41,6 +45,8 @@ import com.apple.foundationdb.relational.recordlayer.query.Identifier; import com.apple.foundationdb.relational.recordlayer.query.IndexGenerator; import com.apple.foundationdb.relational.recordlayer.query.LogicalOperator; +import com.apple.foundationdb.relational.recordlayer.query.LogicalOperators; +import com.apple.foundationdb.relational.recordlayer.query.LogicalPlanFragment; import com.apple.foundationdb.relational.recordlayer.query.PreparedParams; import com.apple.foundationdb.relational.recordlayer.query.ProceduralPlan; import com.apple.foundationdb.relational.recordlayer.query.QueryParser; @@ -48,14 +54,17 @@ import com.apple.foundationdb.relational.recordlayer.query.functions.CompiledSqlFunction; import com.apple.foundationdb.relational.util.Assert; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import org.antlr.v4.runtime.ParserRuleContext; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import java.net.URI; import java.util.ArrayList; import java.util.List; import java.util.Locale; +import java.util.Map; import java.util.function.Function; import java.util.stream.Collectors; @@ -173,6 +182,62 @@ public RecordLayerTable visitTableDefinition(@Nonnull RelationalParser.TableDefi return tableBuilder.build(); } + public RecordLayerIndex visitIntrinsicIndex(@Nonnull final RelationalParser.TableDefinitionContext ctx) { + Assert.thatUnchecked(ctx.organizedByClause() != null); + final var ddlCatalog = metadataBuilder.build(); + // parse the index SQL query using the newly constructed metadata. + getDelegate().replaceSchemaTemplate(ddlCatalog); + + // create a synthetic plan fragment comprising the table access only. This is important for resolving the + // components of the clause correctly, including the embedding column and the partitioning columns. + final var tableName = visitUid(ctx.uid()); + final var logicalOperator = LogicalOperator.generateTableAccess(tableName, ImmutableSet.of(), getDelegate().getSemanticAnalyzer()); + + final var organizedByClause = ctx.organizedByClause(); + final var embeddingColumnId = visitFullId(organizedByClause.embeddingsCol); + final var embeddingColumn = getDelegate().getSemanticAnalyzer().resolveIdentifier(embeddingColumnId, LogicalOperators.ofSingle(logicalOperator)); + + getDelegate().pushPlanFragment().setOperator(logicalOperator); + final var partitionExpressions = (organizedByClause.partitionClause() == null) + ? Expressions.empty() + : visitPartitionClause(organizedByClause.partitionClause()); + getDelegate().popPlanFragment(); + final var indexOptions = (organizedByClause.hnswConfigurations() == null) + ? ImmutableMap.of() + : visitHnswConfigurations(organizedByClause.hnswConfigurations()); + + return IndexGenerator.generateHnswIndex(tableName.getName(), embeddingColumn, partitionExpressions, indexOptions); + } + + @Nullable + @Override + public Object visitOrganizedByClause(final RelationalParser.OrganizedByClauseContext ctx) { + return null; // postpone processing, it should start exactly after the table is fully resolved. + } + + @Nonnull + @Override + public Map visitHnswConfigurations(final RelationalParser.HnswConfigurationsContext ctx) { + return ctx.hnswConfiguration().stream().map(this::visitHnswConfiguration) + .collect(ImmutableMap.toImmutableMap( + NonnullPair::getLeft, + NonnullPair::getRight, + (existing, replacement) -> { + throw new RelationalException("duplicate configuration '" + existing + "'", ErrorCode.SYNTAX_ERROR) + .toUncheckedWrappedException(); + })); + } + + @Nonnull + @Override + public NonnullPair visitHnswConfiguration(final RelationalParser.HnswConfigurationContext ctx) { + if (ctx.mValue != null) { + return NonnullPair.of("M", ctx.mValue.getText()); + } + Assert.thatUnchecked(ctx.efConstructionValue != null); + return NonnullPair.of("EF_CONSTRUCTION", ctx.efConstructionValue.getText()); + } + @Nonnull @Override public RecordLayerTable visitStructDefinition(@Nonnull RelationalParser.StructDefinitionContext ctx) { @@ -254,9 +319,16 @@ public ProceduralPlan visitCreateSchemaTemplateStatement(@Nonnull RelationalPars indexClauses.add(templateClause.indexDefinition()); } } + final var indexes = ImmutableList.builder(); structClauses.build().stream().map(this::visitStructDefinition).map(RecordLayerTable::getDatatype).forEach(metadataBuilder::addAuxiliaryType); - tableClauses.build().stream().map(this::visitTableDefinition).forEach(metadataBuilder::addTable); - final var indexes = indexClauses.build().stream().map(this::visitIndexDefinition).collect(ImmutableList.toImmutableList()); + for (final var tableClause : tableClauses.build()) { + metadataBuilder.addTable(visitTableDefinition(tableClause)); + if (tableClause.organizedByClause() != null) { + visitIntrinsicIndex(tableClause); + } + } + + indexClauses.build().stream().map(this::visitIndexDefinition).forEach(indexes::add); // TODO: this is currently relying on the lexical order of the function to resolve function dependencies which // is limited. functionClauses.build().forEach(functionClause -> { @@ -264,7 +336,7 @@ public ProceduralPlan visitCreateSchemaTemplateStatement(@Nonnull RelationalPars functionClause.routineBody(), metadataBuilder.build()); metadataBuilder.addInvokedRoutine(invokedRoutine); }); - for (final var index : indexes) { + for (final var index : indexes.build()) { final var table = metadataBuilder.extractTable(index.getTableName()); final var tableWithIndex = RecordLayerTable.Builder.from(table).addIndex(index).build(); metadataBuilder.addTable(tableWithIndex); @@ -508,4 +580,11 @@ public DataType visitReturnsType(@Nonnull RelationalParser.ReturnsTypeContext ct public Boolean visitNullColumnConstraint(@Nonnull RelationalParser.NullColumnConstraintContext ctx) { return ctx.nullNotnull().NOT() == null; } + + @Nonnull + @Override + public Expressions visitPartitionClause(final RelationalParser.PartitionClauseContext ctx) { + return Expressions.of(ctx.expression().stream().map(expContext -> + Assert.castUnchecked(visit(expContext), Expression.class)).collect(ImmutableList.toImmutableList())); + } } diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java index ef0fbf67ba..f3bb965f90 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java @@ -45,6 +45,7 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; import java.util.List; +import java.util.Map; import java.util.Set; @API(API.Status.EXPERIMENTAL) @@ -182,6 +183,24 @@ public RecordLayerTable visitTableDefinition(@Nonnull RelationalParser.TableDefi return getDelegate().visitTableDefinition(ctx); } + @Nullable + @Override + public Object visitOrganizedByClause(final RelationalParser.OrganizedByClauseContext ctx) { + return getDelegate().visitOrganizedByClause(ctx); + } + + @Nonnull + @Override + public Map visitHnswConfigurations(final RelationalParser.HnswConfigurationsContext ctx) { + return getDelegate().visitHnswConfigurations(ctx); + } + + @Nonnull + @Override + public NonnullPair visitHnswConfiguration(final RelationalParser.HnswConfigurationContext ctx) { + return getDelegate().visitHnswConfiguration(ctx); + } + @Nonnull @Override public Object visitColumnDefinition(@Nonnull RelationalParser.ColumnDefinitionContext ctx) { @@ -1345,6 +1364,12 @@ public Object visitWindowName(@Nonnull RelationalParser.WindowNameContext ctx) { return getDelegate().visitWindowName(ctx); } + @Nonnull + @Override + public Expressions visitPartitionClause(final RelationalParser.PartitionClauseContext ctx) { + return getDelegate().visitPartitionClause(ctx); + } + @Nonnull @Override public Object visitScalarFunctionName(@Nonnull RelationalParser.ScalarFunctionNameContext ctx) { diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java index 5d61f700e7..fa92225328 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java @@ -20,9 +20,12 @@ package com.apple.foundationdb.relational.recordlayer.query.visitors; +import com.apple.foundationdb.record.expressions.RecordKeyExpressionProto; +import com.apple.foundationdb.record.metadata.expressions.KeyExpression; import com.apple.foundationdb.record.query.plan.cascades.predicates.CompatibleTypeEvolutionPredicate; import com.apple.foundationdb.record.util.pair.NonnullPair; import com.apple.foundationdb.relational.api.metadata.DataType; +import com.apple.foundationdb.relational.api.metadata.Index; import com.apple.foundationdb.relational.generated.RelationalParser; import com.apple.foundationdb.relational.generated.RelationalParserVisitor; import com.apple.foundationdb.relational.recordlayer.metadata.RecordLayerIndex; @@ -41,6 +44,7 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; import java.util.List; +import java.util.Map; import java.util.Set; /** @@ -134,6 +138,18 @@ public interface TypedVisitor extends RelationalParserVisitor { @Override RecordLayerTable visitTableDefinition(@Nonnull RelationalParser.TableDefinitionContext ctx); + @Nullable + @Override + Object visitOrganizedByClause(RelationalParser.OrganizedByClauseContext ctx); + + @Nonnull + @Override + Map visitHnswConfigurations(RelationalParser.HnswConfigurationsContext ctx); + + @Nonnull + @Override + NonnullPair visitHnswConfiguration(RelationalParser.HnswConfigurationContext ctx); + @Nonnull @Override Object visitColumnDefinition(@Nonnull RelationalParser.ColumnDefinitionContext ctx); @@ -846,6 +862,10 @@ public interface TypedVisitor extends RelationalParserVisitor { @Override Object visitWindowName(@Nonnull RelationalParser.WindowNameContext ctx); + @Nonnull + @Override + Expressions visitPartitionClause(final RelationalParser.PartitionClauseContext ctx); + @Nonnull @Override Object visitScalarFunctionName(@Nonnull RelationalParser.ScalarFunctionNameContext ctx); diff --git a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java index ddff1180a0..ae95490edc 100644 --- a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java +++ b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java @@ -26,6 +26,7 @@ import com.apple.foundationdb.relational.utils.Ddl; import org.assertj.core.api.Assertions; import org.junit.jupiter.api.Order; +import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; @@ -67,4 +68,19 @@ void vectorTest(@Nonnull final String ddlType, @Nonnull final DataType expectedT } } } + + @Test + void hnswDDlSomeTest() throws Exception { + final String schemaTemplate = "create table photos(zone string, recordId string, " + + "embedding vector(768), primary key (zone, recordId), organized by hnsw(embedding partition by zone) with (m = 10, ef_construction = 5))"; + try (var ddl = Ddl.builder().database(URI.create("/TEST/QT")).relationalExtension(relationalExtension).schemaTemplate(schemaTemplate).build()) { + try (var statement = ddl.setSchemaAndGetConnection().createStatement()) { + statement.execute("select * from t1"); + final var metadata = statement.getResultSet().getMetaData(); + Assertions.assertThat(metadata).isInstanceOf(StructResultSetMetaData.class); + final var relationalMetadata = (StructResultSetMetaData)metadata; + final var type = relationalMetadata.getRelationalDataType().getFields().get(1).getType(); + } + } + } } From f31e00c5bf87c4bbe08378f640d3d8221e2a57ad Mon Sep 17 00:00:00 2001 From: Youssef Hatem Date: Wed, 13 Aug 2025 17:50:46 +0100 Subject: [PATCH 26/34] - fixes and cleanups. --- .../query/plan/cascades/typing/Type.java | 14 +++- .../src/main/antlr/RelationalLexer.g4 | 6 +- .../src/main/antlr/RelationalParser.g4 | 6 +- .../recordlayer/query/IndexGenerator.java | 14 ++-- .../query/visitors/DdlVisitor.java | 76 +++++++++++-------- .../relational/api/ddl/IndexTest.java | 58 +++++++++++++- .../recordlayer/query/VectorTypeTest.java | 15 ---- .../src/test/java/YamlIntegrationTests.java | 5 -- 8 files changed, 129 insertions(+), 65 deletions(-) diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java index 9080b6130f..b747534868 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java @@ -719,7 +719,7 @@ enum TypeCode { INT(Integer.class, FieldDescriptorProto.Type.TYPE_INT32, true, true), LONG(Long.class, FieldDescriptorProto.Type.TYPE_INT64, true, true), STRING(String.class, FieldDescriptorProto.Type.TYPE_STRING, true, false), - VECTOR(ByteString.class, FieldDescriptorProto.Type.TYPE_BYTES, true, false), + VECTOR(Vector.JavaVectorType.class, FieldDescriptorProto.Type.TYPE_BYTES, true, false), VERSION(FDBRecordVersion.class, FieldDescriptorProto.Type.TYPE_BYTES, true, false), ENUM(Enum.class, FieldDescriptorProto.Type.TYPE_ENUM, false, false), RECORD(Message.class, null, false, false), @@ -1199,6 +1199,18 @@ public boolean equals(final Object other) { class Vector implements Type { + static final class JavaVectorType { + private final ByteString underlying; + + JavaVectorType(@Nonnull final ByteString underlying) { + this.underlying = underlying; + } + + public ByteString getUnderlying() { + return underlying; + } + } + private final boolean isNullable; private final int precision; diff --git a/fdb-relational-core/src/main/antlr/RelationalLexer.g4 b/fdb-relational-core/src/main/antlr/RelationalLexer.g4 index 7a3609b146..43f841233c 100644 --- a/fdb-relational-core/src/main/antlr/RelationalLexer.g4 +++ b/fdb-relational-core/src/main/antlr/RelationalLexer.g4 @@ -467,7 +467,6 @@ DO: 'DO'; DUMPFILE: 'DUMPFILE'; DUPLICATE: 'DUPLICATE'; DYNAMIC: 'DYNAMIC'; -EF_CONSTRUCTION: 'EF_CONSTRUCTION'; ENABLE: 'ENABLE'; ENCRYPTION: 'ENCRYPTION'; END: 'END'; @@ -510,6 +509,10 @@ HASH: 'HASH'; HELP: 'HELP'; HOST: 'HOST'; HOSTS: 'HOSTS'; +HNSW_M: 'HNSW_M'; +HNSW_MMAX: 'HNSW_MMAX'; +HNSW_MMAX0: 'HNSW_MMAX0'; +HNSW_EF_CONSTRUCTION: 'HNSW_EF_CONSTRUCTION'; IDENTIFIED: 'IDENTIFIED'; IGNORE_SERVER_IDS: 'IGNORE_SERVER_IDS'; IMPORT: 'IMPORT'; @@ -540,7 +543,6 @@ LIST: 'LIST'; LOCAL: 'LOCAL'; LOGFILE: 'LOGFILE'; LOGS: 'LOGS'; -M: 'M'; MASTER: 'MASTER'; MASTER_AUTO_POSITION: 'MASTER_AUTO_POSITION'; MASTER_CONNECT_RETRY: 'MASTER_CONNECT_RETRY'; diff --git a/fdb-relational-core/src/main/antlr/RelationalParser.g4 b/fdb-relational-core/src/main/antlr/RelationalParser.g4 index 68ab052877..39ad26350e 100644 --- a/fdb-relational-core/src/main/antlr/RelationalParser.g4 +++ b/fdb-relational-core/src/main/antlr/RelationalParser.g4 @@ -134,8 +134,10 @@ hnswConfigurations ; hnswConfiguration - : M '=' mValue=DECIMAL_LITERAL - | EF_CONSTRUCTION '=' efConstructionValue=DECIMAL_LITERAL + : HNSW_M '=' mValue=DECIMAL_LITERAL + | HNSW_MMAX '=' mMaxValue=DECIMAL_LITERAL + | HNSW_MMAX0 '=' mMax0Value=DECIMAL_LITERAL + | HNSW_EF_CONSTRUCTION '=' efConstructionValue=DECIMAL_LITERAL ; diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/IndexGenerator.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/IndexGenerator.java index 127ae476fe..c8987ea06e 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/IndexGenerator.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/IndexGenerator.java @@ -23,7 +23,6 @@ import com.apple.foundationdb.annotation.API; import com.apple.foundationdb.record.EvaluationContext; import com.apple.foundationdb.record.RecordCoreException; -import com.apple.foundationdb.record.expressions.RecordKeyExpressionProto; import com.apple.foundationdb.record.metadata.IndexOptions; import com.apple.foundationdb.record.metadata.IndexPredicate; import com.apple.foundationdb.record.metadata.IndexTypes; @@ -81,7 +80,6 @@ import com.google.common.collect.Iterators; import com.google.common.collect.Lists; import com.google.common.collect.PeekingIterator; -import com.ibm.icu.impl.locale.XCldrStub; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -791,17 +789,21 @@ public static RecordLayerIndex generateHnswIndex(@Nonnull final String tableName final var embeddingKeyExpression = toKeyExpression(embedding.getUnderlying()); final var partitionKeyExpression = generate(ImmutableList.copyOf(partitionExpressions.underlying().iterator()), Collections.emptyMap()); - final var keyExpression = keyWithValue(concat(embeddingKeyExpression, partitionKeyExpression), - embeddingKeyExpression.getColumnSize()); + final var keyExpression = keyWithValue(concat(partitionKeyExpression, embeddingKeyExpression), + partitionKeyExpression.getColumnSize()); final var indexOptions = options.entrySet().stream().map( entry -> { final var key = entry.getKey(); final var value = entry.getValue(); switch (key) { - case "M": + case "HNSW_M": return Map.entry(IndexOptions.HNSW_M, value); - case "EF_CONSTRUCTION": + case "HNSW_EF_CONSTRUCTION": return Map.entry(IndexOptions.HNSW_EF_CONSTRUCTION, value); + case "HNSW_MMAX": + return Map.entry(IndexOptions.HNSW_M_MAX, value); + case "HNSW_MMAX0": + return Map.entry(IndexOptions.HNSW_M_MAX_0, value); default: throw new RelationalException("unknown HNSW option: " + key, ErrorCode.SYNTAX_ERROR) .toUncheckedWrappedException(); diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java index 8403940ebf..3dad473a8f 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java @@ -125,9 +125,9 @@ public DataType visitColumnType(@Nonnull RelationalParser.ColumnTypeContext ctx) final SemanticAnalyzer.ParsedTypeInfo typeInfo; if (ctx.customType != null) { final var columnType = visitUid(ctx.customType); - typeInfo = SemanticAnalyzer.ParsedTypeInfo.ofCustomType(columnType, true, false); + typeInfo = SemanticAnalyzer.ParsedTypeInfo.ofCustomType(columnType, false, false); } else { - typeInfo = SemanticAnalyzer.ParsedTypeInfo.ofPrimitiveType(ctx.primitiveType(), true, false); + typeInfo = SemanticAnalyzer.ParsedTypeInfo.ofPrimitiveType(ctx.primitiveType(), false, false); } return semanticAnalyzer.lookupType(typeInfo, metadataBuilder::findType); } @@ -157,9 +157,9 @@ public RecordLayerColumn visitColumnDefinition(@Nonnull RelationalParser.ColumnD final SemanticAnalyzer.ParsedTypeInfo typeInfo; if (ctx.columnType().customType != null) { final var columnType = visitUid(ctx.columnType().customType); - typeInfo = SemanticAnalyzer.ParsedTypeInfo.ofCustomType(columnType, true, false); + typeInfo = SemanticAnalyzer.ParsedTypeInfo.ofCustomType(columnType, true, isRepeated); } else { - typeInfo = SemanticAnalyzer.ParsedTypeInfo.ofPrimitiveType(ctx.columnType().primitiveType(), true, false); + typeInfo = SemanticAnalyzer.ParsedTypeInfo.ofPrimitiveType(ctx.columnType().primitiveType(), true, isRepeated); } final var columnType = semanticAnalyzer.lookupType(typeInfo, metadataBuilder::findType); return RecordLayerColumn.newBuilder().setName(columnId.getName()).setDataType(columnType).build(); @@ -183,30 +183,33 @@ public RecordLayerTable visitTableDefinition(@Nonnull RelationalParser.TableDefi } public RecordLayerIndex visitIntrinsicIndex(@Nonnull final RelationalParser.TableDefinitionContext ctx) { - Assert.thatUnchecked(ctx.organizedByClause() != null); - final var ddlCatalog = metadataBuilder.build(); - // parse the index SQL query using the newly constructed metadata. - getDelegate().replaceSchemaTemplate(ddlCatalog); - - // create a synthetic plan fragment comprising the table access only. This is important for resolving the - // components of the clause correctly, including the embedding column and the partitioning columns. - final var tableName = visitUid(ctx.uid()); - final var logicalOperator = LogicalOperator.generateTableAccess(tableName, ImmutableSet.of(), getDelegate().getSemanticAnalyzer()); - - final var organizedByClause = ctx.organizedByClause(); - final var embeddingColumnId = visitFullId(organizedByClause.embeddingsCol); - final var embeddingColumn = getDelegate().getSemanticAnalyzer().resolveIdentifier(embeddingColumnId, LogicalOperators.ofSingle(logicalOperator)); - - getDelegate().pushPlanFragment().setOperator(logicalOperator); - final var partitionExpressions = (organizedByClause.partitionClause() == null) - ? Expressions.empty() - : visitPartitionClause(organizedByClause.partitionClause()); - getDelegate().popPlanFragment(); - final var indexOptions = (organizedByClause.hnswConfigurations() == null) - ? ImmutableMap.of() - : visitHnswConfigurations(organizedByClause.hnswConfigurations()); - - return IndexGenerator.generateHnswIndex(tableName.getName(), embeddingColumn, partitionExpressions, indexOptions); + return getDelegate().getPlanGenerationContext().withDisabledLiteralProcessing(() -> { + Assert.thatUnchecked(ctx.organizedByClause() != null); + final var ddlCatalog = metadataBuilder.build(); + // parse the index SQL query using the newly constructed metadata. + getDelegate().replaceSchemaTemplate(ddlCatalog); + + // create a synthetic plan fragment comprising the table access only. This is important for resolving the + // components of the clause correctly, including the embedding column and the partitioning columns. + final var tableName = visitUid(ctx.uid()); + final var logicalOperator = LogicalOperator.generateTableAccess(tableName, ImmutableSet.of(), getDelegate().getSemanticAnalyzer()); + + final var organizedByClause = ctx.organizedByClause(); + + final var embeddingColumnId = visitFullId(organizedByClause.embeddingsCol); + final var embeddingColumn = getDelegate().getSemanticAnalyzer().resolveIdentifier(embeddingColumnId, LogicalOperators.ofSingle(logicalOperator)); + + getDelegate().pushPlanFragment().setOperator(logicalOperator); + final var partitionExpressions = (organizedByClause.partitionClause() == null) + ? Expressions.empty() + : visitPartitionClause(organizedByClause.partitionClause()); + getDelegate().popPlanFragment(); + final var indexOptions = (organizedByClause.hnswConfigurations() == null) + ? ImmutableMap.of() + : visitHnswConfigurations(organizedByClause.hnswConfigurations()); + + return IndexGenerator.generateHnswIndex(tableName.getName(), embeddingColumn, partitionExpressions, indexOptions); + }); } @Nullable @@ -232,10 +235,19 @@ public Map visitHnswConfigurations(final RelationalParser.HnswCo @Override public NonnullPair visitHnswConfiguration(final RelationalParser.HnswConfigurationContext ctx) { if (ctx.mValue != null) { - return NonnullPair.of("M", ctx.mValue.getText()); + return NonnullPair.of("HNSW_M", ctx.mValue.getText()); } - Assert.thatUnchecked(ctx.efConstructionValue != null); - return NonnullPair.of("EF_CONSTRUCTION", ctx.efConstructionValue.getText()); + if (ctx.efConstructionValue != null) { + return NonnullPair.of("HNSW_EF_CONSTRUCTION", ctx.efConstructionValue.getText()); + } + if (ctx.mMaxValue != null) { + return NonnullPair.of("HNSW_MMAX", ctx.mMaxValue.getText()); + } + if (ctx.mMax0Value != null) { + return NonnullPair.of("HNSW_MMAX0", ctx.mMax0Value.getText()); + } + Assert.failUnchecked(ErrorCode.SYNTAX_ERROR, "unknown hnsw configuration" + ctx); + return null; } @Nonnull @@ -324,7 +336,7 @@ public ProceduralPlan visitCreateSchemaTemplateStatement(@Nonnull RelationalPars for (final var tableClause : tableClauses.build()) { metadataBuilder.addTable(visitTableDefinition(tableClause)); if (tableClause.organizedByClause() != null) { - visitIntrinsicIndex(tableClause); + indexes.add(visitIntrinsicIndex(tableClause)); } } diff --git a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/api/ddl/IndexTest.java b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/api/ddl/IndexTest.java index c61ada2fc4..a5118071a3 100644 --- a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/api/ddl/IndexTest.java +++ b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/api/ddl/IndexTest.java @@ -20,7 +20,9 @@ package com.apple.foundationdb.relational.api.ddl; +import com.apple.foundationdb.record.metadata.IndexOptions; import com.apple.foundationdb.record.metadata.IndexTypes; +import com.apple.foundationdb.record.metadata.expressions.FunctionKeyExpression; import com.apple.foundationdb.record.metadata.expressions.GroupingKeyExpression; import com.apple.foundationdb.record.metadata.expressions.KeyExpression; import com.apple.foundationdb.record.metadata.expressions.KeyWithValueExpression; @@ -40,6 +42,7 @@ import com.apple.foundationdb.relational.util.NullableArrayUtils; import com.apple.foundationdb.relational.utils.SimpleDatabaseRule; import com.apple.foundationdb.relational.utils.TestSchemas; +import com.google.common.collect.ImmutableMap; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Disabled; @@ -111,6 +114,11 @@ private void indexIs(@Nonnull final String stmt, @Nonnull final KeyExpression ex private void indexIs(@Nonnull final String stmt, @Nonnull final KeyExpression expectedKey, @Nonnull final String indexType, @Nonnull final Consumer validator) throws Exception { + indexIs(stmt, expectedKey, indexType, "MV1", validator); + } + + private void indexIs(@Nonnull final String stmt, @Nonnull final KeyExpression expectedKey, @Nonnull final String indexType, + @Nonnull final String indexName, @Nonnull final Consumer validator) throws Exception { shouldWorkWithInjectedFactory(stmt, new AbstractMetadataOperationsFactory() { @Nonnull @Override @@ -123,7 +131,7 @@ public ConstantAction getSaveSchemaTemplateConstantAction(@Nonnull final SchemaT Assertions.assertEquals(1, table.getIndexes().size(), "Incorrect number of indexes!"); final Index index = Assert.optionalUnchecked(table.getIndexes().stream().findFirst()); Assertions.assertInstanceOf(RecordLayerIndex.class, index); - Assertions.assertEquals("MV1", index.getName(), "Incorrect index name!"); + Assertions.assertEquals(indexName, index.getName(), "Incorrect index name!"); Assertions.assertEquals(indexType, index.getIndexType()); final KeyExpression actualKey = KeyExpression.fromProto(((RecordLayerIndex) index).getKeyExpression().toKeyExpression()); Assertions.assertEquals(expectedKey, actualKey); @@ -896,7 +904,7 @@ void createIndexWithOrderByInFromSelect() throws Exception { final String stmt = "CREATE SCHEMA TEMPLATE test_template " + "CREATE TYPE AS STRUCT A(x bigint) " + "CREATE TABLE T(p bigint, a A array, primary key(p))" + - "CREATE INDEX mv1 AS SELECT SQ.x from T AS t, (select M.x from t.a AS M order by M.x) SQ"; + "CREATE INDEX mv1 AS SELECT SQ.x from T AS t, (select Q.x from t.a AS Q order by Q.x) SQ"; shouldFailWith(stmt, ErrorCode.UNSUPPORTED_OPERATION, "order by is not supported in subquery"); } @@ -927,4 +935,50 @@ void createIndexWithOrderByMixedDirection() throws Exception { concat(field("COL1"), function("order_desc_nulls_last", field("COL2")), function("order_asc_nulls_last", field("COL3"))), IndexTypes.VALUE); } + + @Test + void createHnswIndex() throws Exception { + final String stmt = "CREATE SCHEMA TEMPLATE test_template " + + "create table photos(zone string, recordId string, " + + "embedding vector(768), primary key (zone, recordId), organized by hnsw(embedding partition by zone) " + + "with (hnsw_m = 10, hnsw_ef_construction = 5))"; + indexIs(stmt, + new KeyWithValueExpression(concat(field("ZONE"), field("EMBEDDING")), 1), + IndexTypes.VECTOR, "PHOTOS$hnsw", idx -> { + Assertions.assertInstanceOf(RecordLayerIndex.class, idx); + Assertions.assertEquals(ImmutableMap.of(IndexOptions.HNSW_M, "10", IndexOptions.HNSW_EF_CONSTRUCTION, "5"), + ((RecordLayerIndex)idx).getOptions()); + }); + } + + @Test + void createHnswIndexMultiplePartitions() throws Exception { + final String stmt = "CREATE SCHEMA TEMPLATE test_template " + + "create table photos(zone string, recordId string, name string," + + "embedding vector(768), primary key (zone, recordId), organized by hnsw(embedding partition by zone, name) " + + "with (hnsw_m = 10, hnsw_ef_construction = 5))"; + indexIs(stmt, + new KeyWithValueExpression(concat(field("ZONE"), field("NAME"), field("EMBEDDING")), 2), + IndexTypes.VECTOR, "PHOTOS$hnsw", idx -> { + Assertions.assertInstanceOf(RecordLayerIndex.class, idx); + Assertions.assertEquals(ImmutableMap.of(IndexOptions.HNSW_M, "10", IndexOptions.HNSW_EF_CONSTRUCTION, "5"), + ((RecordLayerIndex)idx).getOptions()); + }); + } + + @Test + void createHnswIndexPartitionArithmeticExpression() throws Exception { + final String stmt = "CREATE SCHEMA TEMPLATE test_template " + + "create table photos(zone string, recordId string, name string," + + "embedding vector(768), primary key (zone, recordId), organized by hnsw(embedding partition by zone + 3, name) " + + "with (hnsw_m = 10, hnsw_ef_construction = 5))"; + indexIs(stmt, + new KeyWithValueExpression(concat(FunctionKeyExpression.create("add", concat(field("ZONE"), value(3))) , + field("NAME"), field("EMBEDDING")), 2), + IndexTypes.VECTOR, "PHOTOS$hnsw", idx -> { + Assertions.assertInstanceOf(RecordLayerIndex.class, idx); + Assertions.assertEquals(ImmutableMap.of(IndexOptions.HNSW_M, "10", IndexOptions.HNSW_EF_CONSTRUCTION, "5"), + ((RecordLayerIndex)idx).getOptions()); + }); + } } diff --git a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java index ae95490edc..967b66aa8f 100644 --- a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java +++ b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java @@ -68,19 +68,4 @@ void vectorTest(@Nonnull final String ddlType, @Nonnull final DataType expectedT } } } - - @Test - void hnswDDlSomeTest() throws Exception { - final String schemaTemplate = "create table photos(zone string, recordId string, " + - "embedding vector(768), primary key (zone, recordId), organized by hnsw(embedding partition by zone) with (m = 10, ef_construction = 5))"; - try (var ddl = Ddl.builder().database(URI.create("/TEST/QT")).relationalExtension(relationalExtension).schemaTemplate(schemaTemplate).build()) { - try (var statement = ddl.setSchemaAndGetConnection().createStatement()) { - statement.execute("select * from t1"); - final var metadata = statement.getResultSet().getMetaData(); - Assertions.assertThat(metadata).isInstanceOf(StructResultSetMetaData.class); - final var relationalMetadata = (StructResultSetMetaData)metadata; - final var type = relationalMetadata.getRelationalDataType().getFields().get(1).getType(); - } - } - } } diff --git a/yaml-tests/src/test/java/YamlIntegrationTests.java b/yaml-tests/src/test/java/YamlIntegrationTests.java index 8b7b83bf09..e877d72973 100644 --- a/yaml-tests/src/test/java/YamlIntegrationTests.java +++ b/yaml-tests/src/test/java/YamlIntegrationTests.java @@ -305,9 +305,4 @@ public void literalExtractionTests(YamlTest.Runner runner) throws Exception { public void caseSensitivityTest(YamlTest.Runner runner) throws Exception { runner.runYamsql("case-sensitivity.yamsql"); } - - @TestTemplate - public void vectorTypeTests(YamlTest.Runner runner) throws Exception { - runner.runYamsql("vector-type.yamsql"); - } } From 976abda64facecfdaf425db20e4712620cece7c8 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Wed, 13 Aug 2025 22:34:09 +0200 Subject: [PATCH 27/34] get the build to build --- .../foundationdb/async/MoreAsyncUtil.java | 7 +- .../async/hnsw/AbstractStorageAdapter.java | 1 + .../async/hnsw/BaseNeighborsChangeSet.java | 2 + .../foundationdb/async/hnsw/CompactNode.java | 2 + .../async/hnsw/CompactStorageAdapter.java | 2 + .../async/hnsw/DeleteNeighborsChangeSet.java | 2 + .../async/hnsw/EntryNodeReference.java | 17 +++++ .../apple/foundationdb/async/hnsw/HNSW.java | 71 ++++++++----------- .../foundationdb/async/hnsw/InliningNode.java | 2 + .../async/hnsw/InliningStorageAdapter.java | 2 + .../async/hnsw/InsertNeighborsChangeSet.java | 2 + .../apple/foundationdb/async/hnsw/Node.java | 2 +- .../async/hnsw/NodeReferenceWithVector.java | 17 +++++ .../async/hnsw/HNSWModificationTest.java | 2 + .../provider/foundationdb/FDBRecordStore.java | 7 ++ .../indexes/VectorIndexMaintainerFactory.java | 1 + .../record/query/expressions/Comparisons.java | 21 ++++++ .../query/plan/cascades/typing/Type.java | 27 ++++--- .../indexes/VectorIndexSimpleTest.java | 4 +- .../indexes/VectorIndexTestBase.java | 69 +++++++++--------- .../src/test/proto/test_records_vector.proto | 10 ++- .../recordlayer/query/IndexGenerator.java | 3 +- .../query/visitors/DdlVisitor.java | 3 - .../query/visitors/TypedVisitor.java | 5 +- .../recordlayer/query/VectorTypeTest.java | 2 - gradle/codequality/pmd-rules.xml | 1 + 26 files changed, 179 insertions(+), 105 deletions(-) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java index 0278a36751..64e6d6b732 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/MoreAsyncUtil.java @@ -1057,6 +1057,7 @@ public static CompletableFuture swallowException(@Nonnull CompletableFutur return result; } + @Nonnull public static CompletableFuture forLoop(final int startI, @Nullable final U startU, @Nonnull final IntPredicate conditionPredicate, @Nonnull final IntUnaryOperator stepFunction, @@ -1064,7 +1065,7 @@ public static CompletableFuture forLoop(final int startI, @Nullable final @Nonnull final Executor executor) { final AtomicInteger loopVariableAtomic = new AtomicInteger(startI); final AtomicReference lastResultAtomic = new AtomicReference<>(startU); - return AsyncUtil.whileTrue(() -> { + return whileTrue(() -> { final int loopVariable = loopVariableAtomic.get(); if (!conditionPredicate.test(loopVariable)) { return AsyncUtil.READY_FALSE; @@ -1093,7 +1094,7 @@ public static CompletableFuture> forEach(@Nonnull final Iterable< final AtomicInteger indexAtomic = new AtomicInteger(0); final Object[] resultArray = new Object[toBeProcessed.size()]; - return AsyncUtil.whileTrue(() -> { + return whileTrue(() -> { working.removeIf(CompletableFuture::isDone); while (working.size() <= parallelism) { @@ -1110,7 +1111,7 @@ public static CompletableFuture> forEach(@Nonnull final Iterable< if (working.isEmpty()) { return AsyncUtil.READY_FALSE; } - return AsyncUtil.whenAny(working).thenApply(ignored -> true); + return whenAny(working).thenApply(ignored -> true); }, executor).thenApply(ignored -> Arrays.asList((U[])resultArray)); } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java index 5b074279eb..e3d0c943fc 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/AbstractStorageAdapter.java @@ -129,6 +129,7 @@ private Node checkNode(@Nullable final Node node) { return node; } + @Override public void writeNode(@Nonnull Transaction transaction, @Nonnull Node node, int layer, @Nonnull NeighborsChangeSet changeSet) { writeNodeInternal(transaction, node, layer, changeSet); diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/BaseNeighborsChangeSet.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/BaseNeighborsChangeSet.java index 397e818ac1..bb8271af39 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/BaseNeighborsChangeSet.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/BaseNeighborsChangeSet.java @@ -41,11 +41,13 @@ public BaseNeighborsChangeSet(@Nonnull final List neighbors) { } @Nullable + @Override public BaseNeighborsChangeSet getParent() { return null; } @Nonnull + @Override public List merge() { return neighbors; } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java index 23f2667b63..a6a28e778d 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactNode.java @@ -20,6 +20,7 @@ package com.apple.foundationdb.async.hnsw; +import com.apple.foundationdb.annotation.SpotBugsSuppressWarnings; import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; @@ -37,6 +38,7 @@ public class CompactNode extends AbstractNode { @SuppressWarnings("unchecked") @Nonnull @Override + @SpotBugsSuppressWarnings("NP_PARAMETER_MUST_BE_NONNULL_BUT_MARKED_AS_NULLABLE") public Node create(@Nonnull final Tuple primaryKey, @Nullable final Vector vector, @Nonnull final List neighbors) { return new CompactNode(primaryKey, Objects.requireNonNull(vector), (List)neighbors); diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java index b590513019..c3a04f86a2 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/CompactStorageAdapter.java @@ -154,6 +154,8 @@ public void writeNodeInternal(@Nonnull final Transaction transaction, @Nonnull f } } + @Nonnull + @Override public Iterable> scanLayer(@Nonnull final ReadTransaction readTransaction, int layer, @Nullable final Tuple lastPrimaryKey, int maxNumRead) { final byte[] layerPrefix = getDataSubspace().pack(Tuple.from(layer)); diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java index 0d4694b2fe..e431561119 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/DeleteNeighborsChangeSet.java @@ -52,11 +52,13 @@ public DeleteNeighborsChangeSet(@Nonnull final NeighborsChangeSet parent, } @Nonnull + @Override public NeighborsChangeSet getParent() { return parent; } @Nonnull + @Override public Iterable merge() { return Iterables.filter(getParent().merge(), current -> !deletedNeighborsPrimaryKeys.contains(current.getPrimaryKey())); diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/EntryNodeReference.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/EntryNodeReference.java index 3cbbb3c2dd..db81252e17 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/EntryNodeReference.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/EntryNodeReference.java @@ -24,6 +24,7 @@ import com.christianheina.langx.half4j.Half; import javax.annotation.Nonnull; +import java.util.Objects; class EntryNodeReference extends NodeReferenceWithVector { private final int layer; @@ -36,4 +37,20 @@ public EntryNodeReference(@Nonnull final Tuple primaryKey, @Nonnull final Vector public int getLayer() { return layer; } + + @Override + public boolean equals(final Object o) { + if (!(o instanceof EntryNodeReference)) { + return false; + } + if (!super.equals(o)) { + return false; + } + return layer == ((EntryNodeReference)o).layer; + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), layer); + } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index 31cb055fcc..42ec45b280 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -430,7 +430,7 @@ public CompletableFuture layer > 0, layer -> layer - 1, (layer, previousNodeReference) -> { @@ -751,7 +751,7 @@ public CompletableFuture insert(@Nonnull final Transaction transaction, @N new NodeReferenceWithDistance(entryNodeReference.getPrimaryKey(), entryNodeReference.getVector(), Vector.comparativeDistance(metric, entryNodeReference.getVector(), newVector)); - return MoreAsyncUtil.forLoop(lMax, initialNodeReference, + return forLoop(lMax, initialNodeReference, layer -> layer > insertionLayer, layer -> layer - 1, (layer, previousNodeReference) -> { @@ -778,7 +778,7 @@ public CompletableFuture insertBatch(@Nonnull final Transaction transactio insertionLayer(random))); } // sort the layers in reverse order - batchWithLayers.sort(Comparator.comparing(NodeReferenceWithLayer::getL).reversed()); + batchWithLayers.sort(Comparator.comparing(NodeReferenceWithLayer::getLayer).reversed()); return StorageAdapter.fetchEntryNodeReference(transaction, getSubspace(), getOnReadListener()) .thenCompose(entryNodeReference -> { @@ -791,14 +791,14 @@ public CompletableFuture insertBatch(@Nonnull final Transaction transactio } final Vector itemVector = item.getVector(); - final int itemL = item.getL(); + final int itemL = item.getLayer(); final NodeReferenceWithDistance initialNodeReference = new NodeReferenceWithDistance(entryNodeReference.getPrimaryKey(), entryNodeReference.getVector(), Vector.comparativeDistance(metric, entryNodeReference.getVector(), itemVector)); - return MoreAsyncUtil.forLoop(lMax, initialNodeReference, + return forLoop(lMax, initialNodeReference, layer -> layer > itemL, layer -> layer - 1, (layer, previousNodeReference) -> { @@ -815,7 +815,7 @@ public CompletableFuture insertBatch(@Nonnull final Transaction transactio final NodeReferenceWithLayer item = batchWithLayers.get(index); final Tuple itemPrimaryKey = item.getPrimaryKey(); final Vector itemVector = item.getVector(); - final int itemL = item.getL(); + final int itemL = item.getLayer(); final EntryNodeReference newEntryNodeReference; final int currentLMax; @@ -1078,14 +1078,13 @@ private CompletableFuture }).thenCompose(selectedNeighbors -> fetchSomeNodesIfNotCached(storageAdapter, readTransaction, layer, selectedNeighbors, nodeCache)) .thenApply(selectedNeighbors -> { - debug(l -> { - l.debug("selected neighbors={}", - selectedNeighbors.stream() - .map(selectedNeighbor -> - "(primaryKey=" + selectedNeighbor.getNodeReferenceWithDistance().getPrimaryKey() + - ",distance=" + selectedNeighbor.getNodeReferenceWithDistance().getDistance() + ")") - .collect(Collectors.joining(","))); - }); + debug(l -> + l.debug("selected neighbors={}", + selectedNeighbors.stream() + .map(selectedNeighbor -> + "(primaryKey=" + selectedNeighbor.getNodeReferenceWithDistance().getPrimaryKey() + + ",distance=" + selectedNeighbor.getNodeReferenceWithDistance().getDistance() + ")") + .collect(Collectors.joining(",")))); return selectedNeighbors; }); } @@ -1199,6 +1198,7 @@ private int insertionLayer(@Nonnull final Random random) { return (int) Math.floor(-Math.log(u) * lambda); } + @SuppressWarnings("PMD.UnusedPrivateMethod") private void info(@Nonnull final Consumer loggerConsumer) { if (logger.isInfoEnabled()) { loggerConsumer.accept(logger); @@ -1212,41 +1212,32 @@ private void debug(@Nonnull final Consumer loggerConsumer) { } private static class NodeReferenceWithLayer extends NodeReferenceWithVector { - @SuppressWarnings("checkstyle:MemberName") - private final int l; + private final int layer; public NodeReferenceWithLayer(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, - final int l) { + final int layer) { super(primaryKey, vector); - this.l = l; - } - - public int getL() { - return l; + this.layer = layer; } - } - - private static class NodeReferenceWithSearchEntry extends NodeReferenceWithVector { - @SuppressWarnings("checkstyle:MemberName") - private final int l; - @Nonnull - private final NodeReferenceWithDistance nodeReferenceWithDistance; - public NodeReferenceWithSearchEntry(@Nonnull final Tuple primaryKey, @Nonnull final Vector vector, - final int l, - @Nonnull final NodeReferenceWithDistance nodeReferenceWithDistance) { - super(primaryKey, vector); - this.l = l; - this.nodeReferenceWithDistance = nodeReferenceWithDistance; + public int getLayer() { + return layer; } - public int getL() { - return l; + @Override + public boolean equals(final Object o) { + if (!(o instanceof NodeReferenceWithLayer)) { + return false; + } + if (!super.equals(o)) { + return false; + } + return layer == ((NodeReferenceWithLayer)o).layer; } - @Nonnull - public NodeReferenceWithDistance getNodeReferenceWithDistance() { - return nodeReferenceWithDistance; + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), layer); } } } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java index dce0f18f24..48e2398950 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningNode.java @@ -20,6 +20,7 @@ package com.apple.foundationdb.async.hnsw; +import com.apple.foundationdb.annotation.SpotBugsSuppressWarnings; import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; @@ -57,6 +58,7 @@ public InliningNode(@Nonnull final Tuple primaryKey, @Nonnull @Override + @SpotBugsSuppressWarnings("NP_PARAMETER_MUST_BE_NONNULL_BUT_MARKED_AS_NULLABLE") public NodeReferenceWithVector getSelfReference(@Nullable final Vector vector) { return new NodeReferenceWithVector(getPrimaryKey(), Objects.requireNonNull(vector)); } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java index db55b29597..ebbfd4d698 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InliningStorageAdapter.java @@ -63,6 +63,7 @@ public StorageAdapter asInliningStorageAdapter() { } @Nonnull + @Override protected CompletableFuture> fetchNodeInternal(@Nonnull final ReadTransaction readTransaction, final int layer, @Nonnull final Tuple primaryKey) { @@ -138,6 +139,7 @@ private byte[] getNeighborKey(final int layer, return getDataSubspace().pack(Tuple.from(layer, node.getPrimaryKey(), neighborPrimaryKey)); } + @Nonnull @Override public Iterable> scanLayer(@Nonnull final ReadTransaction readTransaction, int layer, @Nullable final Tuple lastPrimaryKey, int maxNumRead) { diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java index 06a7490690..d68d3ae933 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/InsertNeighborsChangeSet.java @@ -57,11 +57,13 @@ public InsertNeighborsChangeSet(@Nonnull final NeighborsChangeSet parent, } @Nonnull + @Override public NeighborsChangeSet getParent() { return parent; } @Nonnull + @Override public Iterable merge() { return Iterables.concat(getParent().merge(), insertedNeighborsMap.values()); } diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java index 4e7b08c301..f2c623f882 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Node.java @@ -36,7 +36,7 @@ public interface Node { Tuple getPrimaryKey(); @Nonnull - N getSelfReference(@Nullable final Vector vector); + N getSelfReference(@Nullable Vector vector); @Nonnull List getNeighbors(); diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithVector.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithVector.java index b13443b926..e21b221622 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithVector.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeReferenceWithVector.java @@ -22,6 +22,7 @@ import com.apple.foundationdb.tuple.Tuple; import com.christianheina.langx.half4j.Half; +import com.google.common.base.Objects; import javax.annotation.Nonnull; @@ -50,6 +51,22 @@ public NodeReferenceWithVector asNodeReferenceWithVector() { return this; } + @Override + public boolean equals(final Object o) { + if (!(o instanceof NodeReferenceWithVector)) { + return false; + } + if (!super.equals(o)) { + return false; + } + return Objects.equal(vector, ((NodeReferenceWithVector)o).vector); + } + + @Override + public int hashCode() { + return Objects.hashCode(super.hashCode(), vector); + } + @Override public String toString() { return "NRV[primaryKey=" + getPrimaryKey() + diff --git a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java index 85aea0e2ab..a0238fd4fe 100644 --- a/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java +++ b/fdb-extensions/src/test/java/com/apple/foundationdb/async/hnsw/HNSWModificationTest.java @@ -35,6 +35,7 @@ import org.assertj.core.util.Lists; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; @@ -72,6 +73,7 @@ @SuppressWarnings("checkstyle:AbbreviationAsWordInName") @Tag(Tags.RequiresFDB) @Tag(Tags.Slow) +@Disabled public class HNSWModificationTest { private static final Logger logger = LoggerFactory.getLogger(HNSWModificationTest.class); private static final int NUM_TEST_RUNS = 5; diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/FDBRecordStore.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/FDBRecordStore.java index 23dd1f17f6..4537529749 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/FDBRecordStore.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/FDBRecordStore.java @@ -644,11 +644,18 @@ private FDBStoredRecord dryRunSetSizeInfo(@Nonnull Record return recordBuilder.build(); } + @SuppressWarnings("unchecked") @Nonnull private FDBStoredRecord serializeAndSaveRecord(@Nonnull RecordSerializer typedSerializer, @Nonnull final FDBStoredRecordBuilder recordBuilder, @Nonnull final RecordMetaData metaData, @Nullable FDBStoredSizes oldSizeInfo) { final Tuple primaryKey = recordBuilder.getPrimaryKey(); final FDBRecordVersion version = recordBuilder.getVersion(); + + // final M record = recordBuilder.getRecord(); + // M cleansed_rec = (M)record.toBuilder() + // .clearField(record.getDescriptorForType().findFieldByName("vector_data")) + // .build(); + final byte[] serialized = typedSerializer.serialize(metaData, recordBuilder.getRecordType(), recordBuilder.getRecord(), getTimer()); final FDBRecordVersion splitVersion = useOldVersionFormat() ? null : version; final SplitHelper.SizeInfo sizeInfo = new SplitHelper.SizeInfo(); diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainerFactory.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainerFactory.java index 6b6ba33db6..d8b2444eef 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainerFactory.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainerFactory.java @@ -87,6 +87,7 @@ private void validateStructure() { } @Override + @SuppressWarnings("PMD.CompareObjectsWithEquals") public void validateChangedOptions(@Nonnull final Index oldIndex, @Nonnull final Set changedOptions) { if (!changedOptions.isEmpty()) { diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java index b3719a4e62..d5daa2b1f0 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java @@ -1863,6 +1863,26 @@ public String typelessString() { return getComparandValue() + ":" + getLimitValue(); } + @Override + public final boolean equals(final Object o) { + if (!(o instanceof DistanceRankValueComparison)) { + return false; + } + final DistanceRankValueComparison that = (DistanceRankValueComparison)o; + if (!super.equals(o)) { + return false; + } + + return limitValue.equals(that.limitValue); + } + + @Override + public int hashCode() { + int result = super.hashCode(); + result = 31 * result + limitValue.hashCode(); + return result; + } + @Override public String toString() { return explain().getExplainTokens().render(DefaultExplainFormatter.forDebugging()).toString(); @@ -1877,6 +1897,7 @@ public ExplainTokensWithPrecedence explain() { .addNested(getLimitValue().explain().getExplainTokens())); } + @Override public int computeHashCode() { return Objects.hash(getType().name(), getComparandValue(), getLimitValue()); } diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java index b747534868..e72a4ea65e 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java @@ -1198,23 +1198,8 @@ public boolean equals(final Object other) { } class Vector implements Type { - - static final class JavaVectorType { - private final ByteString underlying; - - JavaVectorType(@Nonnull final ByteString underlying) { - this.underlying = underlying; - } - - public ByteString getUnderlying() { - return underlying; - } - } - private final boolean isNullable; - private final int precision; - private final int dimensions; private Vector(final boolean isNullable, final int precision, final int dimensions) { @@ -1319,6 +1304,18 @@ public static Vector fromProto(@Nonnull final PlanSerializationContext serializa return new Vector(vectorTypeProto.getIsNullable(), vectorTypeProto.getPrecision(), vectorTypeProto.getDimensions()); } + static final class JavaVectorType { + private final ByteString underlying; + + JavaVectorType(@Nonnull final ByteString underlying) { + this.underlying = underlying; + } + + public ByteString getUnderlying() { + return underlying; + } + } + /** * Deserializer. */ diff --git a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java index 5b069a802e..b7d3fda646 100644 --- a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java +++ b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java @@ -21,6 +21,7 @@ package com.apple.foundationdb.record.provider.foundationdb.indexes; import com.apple.test.Tags; +import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; import org.slf4j.Logger; @@ -30,12 +31,13 @@ * Tests for multidimensional type indexes. */ @Tag(Tags.RequiresFDB) +@Disabled public class VectorIndexSimpleTest extends VectorIndexTestBase { private static final Logger logger = LoggerFactory.getLogger(VectorIndexSimpleTest.class); @Test void basicReadTest() throws Exception { - super.basicReadTest(false); + super.basicReadTest(); } @Test diff --git a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java index f05e2c9c7c..a2c6712841 100644 --- a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java +++ b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java @@ -21,8 +21,14 @@ package com.apple.foundationdb.record.provider.foundationdb.indexes; import com.apple.foundationdb.async.AsyncUtil; +import com.apple.foundationdb.async.hnsw.Metrics; +import com.apple.foundationdb.async.hnsw.Vector; import com.apple.foundationdb.record.RecordMetaData; import com.apple.foundationdb.record.RecordMetaDataBuilder; +import com.apple.foundationdb.record.metadata.Index; +import com.apple.foundationdb.record.metadata.IndexOptions; +import com.apple.foundationdb.record.metadata.IndexTypes; +import com.apple.foundationdb.record.metadata.expressions.KeyWithValueExpression; import com.apple.foundationdb.record.provider.foundationdb.FDBRecordContext; import com.apple.foundationdb.record.provider.foundationdb.FDBStoredRecord; import com.apple.foundationdb.record.provider.foundationdb.query.FDBRecordStoreQueryTestBase; @@ -31,9 +37,10 @@ import com.apple.test.Tags; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.protobuf.ByteString; import com.google.protobuf.Message; -import org.assertj.core.util.Streams; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Tag; import org.slf4j.Logger; @@ -63,31 +70,40 @@ public abstract class VectorIndexTestBase extends FDBRecordStoreQueryTestBase { private static final SimpleDateFormat timeFormat = new SimpleDateFormat("MM/dd/yyyy HH:mm:ss"); + @CanIgnoreReturnValue + RecordMetaDataBuilder addVectorIndex(@Nonnull final RecordMetaDataBuilder metaDataBuilder) { + metaDataBuilder.addIndex("UngroupedVectorRecord", + new Index("UngroupedVectorIndex", new KeyWithValueExpression(field("vector_data"), 0), + IndexTypes.VECTOR, + ImmutableMap.of(IndexOptions.HNSW_METRIC, Metrics.EUCLIDEAN_METRIC.toString()))); + return metaDataBuilder; + } + protected void openRecordStore(FDBRecordContext context) throws Exception { openRecordStore(context, NO_HOOK); } protected void openRecordStore(final FDBRecordContext context, final RecordMetaDataHook hook) throws Exception { RecordMetaDataBuilder metaDataBuilder = RecordMetaData.newBuilder().setRecords(TestRecordsVectorsProto.getDescriptor()); - metaDataBuilder.getRecordType("NaiveVectorRecord").setPrimaryKey(field("rec_no")); + metaDataBuilder.getRecordType("UngroupedVectorRecord").setPrimaryKey(field("rec_no")); hook.apply(metaDataBuilder); createOrOpenRecordStore(context, metaDataBuilder.getRecordMetaData()); } static Function getRecordGenerator(@Nonnull final Random random) { return recNo -> { - final byte[] vector = new byte[768 * 2]; + final byte[] vector = new byte[128 * 2]; random.nextBytes(vector); //logRecord(recNo, vector); - return TestRecordsVectorsProto.NaiveVectorRecord.newBuilder() + return TestRecordsVectorsProto.UngroupedVectorRecord.newBuilder() .setRecNo(recNo) .setVectorData(ByteString.copyFrom(vector)) .build(); }; } - public void loadRecords(final boolean useAsync, @Nonnull final RecordMetaDataHook hook, final long seed, + public void saveRecords(final boolean useAsync, @Nonnull final RecordMetaDataHook hook, final long seed, final int numSamples) { final Random random = new Random(seed); final var recordGenerator = getRecordGenerator(random); @@ -154,59 +170,48 @@ private long batchAsync(final RecordMetaDataHook hook, final int numRecords, return numRecordsCommitted; } - private static void logRecord(final long recNo, @Nonnull final Iterable vector) { + private static void logRecord(final long recNo, @Nonnull final ByteString vectorData) { if (logger.isInfoEnabled()) { logger.info("recNo: {}; vectorData: [{})", - recNo, Streams.stream(vector).map(String::valueOf).collect(Collectors.joining(","))); + recNo, Vector.HalfVector.halfVectorFromBytes(vectorData.toByteArray())); } } - void basicReadTest(final boolean useAsync) throws Exception { - loadRecords(useAsync, NO_HOOK, 0, 5000); + void basicReadTest() throws Exception { + saveRecords(false, this::addVectorIndex, 0, 1000); try (FDBRecordContext context = openContext()) { - openRecordStore(context); - for (long l = 0; l < 5000; l ++) { + openRecordStore(context, this::addVectorIndex); + for (long l = 0; l < 1000; l ++) { FDBStoredRecord rec = recordStore.loadRecord(Tuple.from(l)); - //Thread.sleep(10); + assertNotNull(rec); - TestRecordsVectorsProto.NaiveVectorRecord.Builder recordBuilder = - TestRecordsVectorsProto.NaiveVectorRecord.newBuilder(); + TestRecordsVectorsProto.UngroupedVectorRecord.Builder recordBuilder = + TestRecordsVectorsProto.UngroupedVectorRecord.newBuilder(); recordBuilder.mergeFrom(rec.getRecord()); final var record = recordBuilder.build(); - //logRecord(record.getRecNo(), record.getVectorDataList()); + logRecord(record.getRecNo(), record.getVectorData()); } commit(context); } } void basicConcurrentReadTest(final boolean useAsync) throws Exception { - loadRecords(useAsync, NO_HOOK, 0, 5000); + saveRecords(useAsync, NO_HOOK, 0, 5000); for (int i = 0; i < 1; i ++) { try (FDBRecordContext context = openContext()) { openRecordStore(context); for (long l = 0; l < 5000; l += 20) { - final var batch = fetchBatchConcurrently(l, 20); - //batch.forEach(record -> logRecord(record.getRecNo(), record.getVectorDataList())); + fetchBatchConcurrently(l, 20); } commit(context); } } } - private List fetchBatchConcurrently(long startRecNo, int batchSize) throws Exception { + private List fetchBatchConcurrently(long startRecNo, int batchSize) throws Exception { final ImmutableList.Builder>> futureBatchBuilder = ImmutableList.builder(); for (int i = 0; i < batchSize; i ++) { - final long task = startRecNo + i; - //System.out.println("task " + task + " scheduled"); - final CompletableFuture> recordFuture = recordStore.loadRecordAsync(Tuple.from(startRecNo + i)).thenApply(r -> { -// try { -// Thread.sleep(10); -// } catch (InterruptedException e) { -// throw new RuntimeException(e); -// } - //System.out.println("task " + task + " done, thread: " + Thread.currentThread()); - return r; - }); + final CompletableFuture> recordFuture = recordStore.loadRecordAsync(Tuple.from(startRecNo + i)); futureBatchBuilder.add(recordFuture); } final var batchFuture = AsyncUtil.getAll(futureBatchBuilder.build()); @@ -215,8 +220,8 @@ private List fetchBatchConcurrently(l return batch.stream() .map(rec -> { assertNotNull(rec); - TestRecordsVectorsProto.NaiveVectorRecord.Builder recordBuilder = - TestRecordsVectorsProto.NaiveVectorRecord.newBuilder(); + TestRecordsVectorsProto.UngroupedVectorRecord.Builder recordBuilder = + TestRecordsVectorsProto.UngroupedVectorRecord.newBuilder(); recordBuilder.mergeFrom(rec.getRecord()); return recordBuilder.build(); }) diff --git a/fdb-record-layer-core/src/test/proto/test_records_vector.proto b/fdb-record-layer-core/src/test/proto/test_records_vector.proto index 4522e85f6f..925ebbfc32 100644 --- a/fdb-record-layer-core/src/test/proto/test_records_vector.proto +++ b/fdb-record-layer-core/src/test/proto/test_records_vector.proto @@ -28,11 +28,17 @@ import "record_metadata_options.proto"; option (schema).store_record_versions = true; -message NaiveVectorRecord { +message UngroupedVectorRecord { optional int64 rec_no = 1 [(field).primary_key = true]; optional bytes vector_data = 2; } +message GroupedVectorRecord { + optional int64 rec_no = 1 [(field).primary_key = true]; + optional int32 groupId = 2; + optional bytes vector_data = 3; +} + message RecordTypeUnion { - optional NaiveVectorRecord _NaiveVectorRecord = 1; + optional UngroupedVectorRecord _NaiveVectorRecord = 1; } diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/IndexGenerator.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/IndexGenerator.java index c8987ea06e..ad5156d75a 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/IndexGenerator.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/IndexGenerator.java @@ -812,12 +812,11 @@ public static RecordLayerIndex generateHnswIndex(@Nonnull final String tableName final var builder = RecordLayerIndex.newBuilder(); - final var index = builder.setIndexType(IndexTypes.VECTOR) + return builder.setIndexType(IndexTypes.VECTOR) .setName(tableName + "$hnsw") .setOptions(indexOptions) .setTableName(tableName) .setKeyExpression(keyExpression) .build(); - return index; } } diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java index 3dad473a8f..07476789c6 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java @@ -21,7 +21,6 @@ package com.apple.foundationdb.relational.recordlayer.query.visitors; import com.apple.foundationdb.annotation.API; -import com.apple.foundationdb.record.metadata.IndexTypes; import com.apple.foundationdb.record.query.plan.cascades.expressions.LogicalSortExpression; import com.apple.foundationdb.record.query.plan.cascades.values.PromoteValue; import com.apple.foundationdb.record.query.plan.cascades.values.ThrowsValue; @@ -31,7 +30,6 @@ import com.apple.foundationdb.relational.api.exceptions.ErrorCode; import com.apple.foundationdb.relational.api.exceptions.RelationalException; import com.apple.foundationdb.relational.api.metadata.DataType; -import com.apple.foundationdb.relational.api.metadata.Index; import com.apple.foundationdb.relational.api.metadata.InvokedRoutine; import com.apple.foundationdb.relational.generated.RelationalParser; import com.apple.foundationdb.relational.recordlayer.metadata.DataTypeUtils; @@ -46,7 +44,6 @@ import com.apple.foundationdb.relational.recordlayer.query.IndexGenerator; import com.apple.foundationdb.relational.recordlayer.query.LogicalOperator; import com.apple.foundationdb.relational.recordlayer.query.LogicalOperators; -import com.apple.foundationdb.relational.recordlayer.query.LogicalPlanFragment; import com.apple.foundationdb.relational.recordlayer.query.PreparedParams; import com.apple.foundationdb.relational.recordlayer.query.ProceduralPlan; import com.apple.foundationdb.relational.recordlayer.query.QueryParser; diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java index fa92225328..6ebd5c1bb9 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java @@ -20,12 +20,9 @@ package com.apple.foundationdb.relational.recordlayer.query.visitors; -import com.apple.foundationdb.record.expressions.RecordKeyExpressionProto; -import com.apple.foundationdb.record.metadata.expressions.KeyExpression; import com.apple.foundationdb.record.query.plan.cascades.predicates.CompatibleTypeEvolutionPredicate; import com.apple.foundationdb.record.util.pair.NonnullPair; import com.apple.foundationdb.relational.api.metadata.DataType; -import com.apple.foundationdb.relational.api.metadata.Index; import com.apple.foundationdb.relational.generated.RelationalParser; import com.apple.foundationdb.relational.generated.RelationalParserVisitor; import com.apple.foundationdb.relational.recordlayer.metadata.RecordLayerIndex; @@ -864,7 +861,7 @@ public interface TypedVisitor extends RelationalParserVisitor { @Nonnull @Override - Expressions visitPartitionClause(final RelationalParser.PartitionClauseContext ctx); + Expressions visitPartitionClause(RelationalParser.PartitionClauseContext ctx); @Nonnull @Override diff --git a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java index 967b66aa8f..b0e8d55d36 100644 --- a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java +++ b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java @@ -26,7 +26,6 @@ import com.apple.foundationdb.relational.utils.Ddl; import org.assertj.core.api.Assertions; import org.junit.jupiter.api.Order; -import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; @@ -37,7 +36,6 @@ import java.util.stream.Stream; public class VectorTypeTest { - @RegisterExtension @Order(0) public final EmbeddedRelationalExtension relationalExtension = new EmbeddedRelationalExtension(); diff --git a/gradle/codequality/pmd-rules.xml b/gradle/codequality/pmd-rules.xml index 500ef17c69..4d8745d875 100644 --- a/gradle/codequality/pmd-rules.xml +++ b/gradle/codequality/pmd-rules.xml @@ -16,6 +16,7 @@ + From ef5dc3b2c142496a793b9bde5147f35b6c862c51 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Thu, 14 Aug 2025 13:54:39 +0200 Subject: [PATCH 28/34] adding some test cases --- .../apple/foundationdb/async/hnsw/HNSW.java | 21 ++-- .../indexes/VectorIndexMaintainer.java | 3 +- .../record/query/expressions/Comparisons.java | 3 +- .../indexes/VectorIndexSimpleTest.java | 11 +- .../indexes/VectorIndexTestBase.java | 110 +++++++++++------- 5 files changed, 92 insertions(+), 56 deletions(-) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java index 42ec45b280..fb177c9d77 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java @@ -36,6 +36,7 @@ import com.google.common.collect.Maps; import com.google.common.collect.Sets; import com.google.common.collect.Streams; +import com.google.common.collect.TreeMultimap; import com.google.errorprone.annotations.CanIgnoreReturnValue; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -49,7 +50,6 @@ import java.util.Queue; import java.util.Random; import java.util.Set; -import java.util.TreeSet; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; import java.util.concurrent.PriorityBlockingQueue; @@ -450,22 +450,25 @@ public CompletableFuture { // reverse the original queue - final TreeSet> sortedTopK = - new TreeSet<>( - Comparator.comparing(nodeReferenceAndNode -> - nodeReferenceAndNode.getNodeReferenceWithDistance().getDistance())); + final TreeMultimap> sortedTopK = + TreeMultimap.create(Comparator.naturalOrder(), + Comparator.comparing(nodeReferenceAndNode -> nodeReferenceAndNode.getNode().getPrimaryKey())); for (final NodeReferenceAndNode nodeReferenceAndNode : searchResult) { - if (sortedTopK.size() < k || sortedTopK.last().getNodeReferenceWithDistance().getDistance() > + if (sortedTopK.size() < k || sortedTopK.keySet().last() > nodeReferenceAndNode.getNodeReferenceWithDistance().getDistance()) { - sortedTopK.add(nodeReferenceAndNode); + sortedTopK.put(nodeReferenceAndNode.getNodeReferenceWithDistance().getDistance(), + nodeReferenceAndNode); } if (sortedTopK.size() > k) { - sortedTopK.remove(sortedTopK.last()); + final Double lastKey = sortedTopK.keySet().last(); + final NodeReferenceAndNode lastNode = sortedTopK.get(lastKey).last(); + sortedTopK.remove(lastKey, lastNode); } } - return ImmutableList.copyOf(sortedTopK); + + return ImmutableList.copyOf(sortedTopK.values()); }); }); } diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java index 6e77618f1b..1a063211b4 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java @@ -299,7 +299,8 @@ protected CompletableFuture updateIndexKeys(@Nonnull f state.index.trimPrimaryKey(primaryKeyParts); final Tuple trimmedPrimaryKey = Tuple.fromList(primaryKeyParts); final FDBStoreTimer timer = Objects.requireNonNull(getTimer()); - final HNSW hnsw = new HNSW(rtSubspace, getExecutor(), getConfig(), new OnWrite(timer), OnReadListener.NOOP); + final HNSW hnsw = + new HNSW(rtSubspace, getExecutor(), getConfig(), new OnWrite(timer), OnReadListener.NOOP); if (remove) { throw new UnsupportedOperationException("not implemented"); } else { diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java index d5daa2b1f0..731f499f09 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java @@ -1938,7 +1938,8 @@ public PComparison toComparisonProto(@Nonnull final PlanSerializationContext ser @Nullable public Vector getVector(@Nullable final FDBRecordStoreBase store, final @Nullable EvaluationContext context) { - return (Vector)getComparand(store, context); + final Object comparand = getComparand(store, context); + return comparand == null ? null : Vector.HalfVector.halfVectorFromBytes((byte[])comparand); } public int getLimit(@Nullable final FDBRecordStoreBase store, final @Nullable EvaluationContext context) { diff --git a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java index b7d3fda646..8a8b4e8607 100644 --- a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java +++ b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java @@ -21,7 +21,6 @@ package com.apple.foundationdb.record.provider.foundationdb.indexes; import com.apple.test.Tags; -import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; import org.slf4j.Logger; @@ -31,17 +30,17 @@ * Tests for multidimensional type indexes. */ @Tag(Tags.RequiresFDB) -@Disabled public class VectorIndexSimpleTest extends VectorIndexTestBase { private static final Logger logger = LoggerFactory.getLogger(VectorIndexSimpleTest.class); + @Override @Test - void basicReadTest() throws Exception { - super.basicReadTest(); + void basicWriteReadTest() throws Exception { + super.basicWriteReadTest(); } @Test - void basicConcurrentReadTest() throws Exception { - super.basicConcurrentReadTest(false); + void basicWriteIndexReadTest() throws Exception { + super.basicWriteIndexReadTest(); } } diff --git a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java index a2c6712841..299c6ba775 100644 --- a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java +++ b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java @@ -23,19 +23,32 @@ import com.apple.foundationdb.async.AsyncUtil; import com.apple.foundationdb.async.hnsw.Metrics; import com.apple.foundationdb.async.hnsw.Vector; +import com.apple.foundationdb.record.IndexFetchMethod; +import com.apple.foundationdb.record.RecordCursorIterator; import com.apple.foundationdb.record.RecordMetaData; import com.apple.foundationdb.record.RecordMetaDataBuilder; import com.apple.foundationdb.record.metadata.Index; import com.apple.foundationdb.record.metadata.IndexOptions; import com.apple.foundationdb.record.metadata.IndexTypes; import com.apple.foundationdb.record.metadata.expressions.KeyWithValueExpression; +import com.apple.foundationdb.record.provider.foundationdb.FDBQueriedRecord; import com.apple.foundationdb.record.provider.foundationdb.FDBRecordContext; import com.apple.foundationdb.record.provider.foundationdb.FDBStoredRecord; +import com.apple.foundationdb.record.provider.foundationdb.VectorIndexScanComparisons; import com.apple.foundationdb.record.provider.foundationdb.query.FDBRecordStoreQueryTestBase; +import com.apple.foundationdb.record.query.expressions.Comparisons; +import com.apple.foundationdb.record.query.expressions.Comparisons.DistanceRankValueComparison; +import com.apple.foundationdb.record.query.plan.QueryPlanConstraint; +import com.apple.foundationdb.record.query.plan.ScanComparisons; +import com.apple.foundationdb.record.query.plan.cascades.typing.Type; +import com.apple.foundationdb.record.query.plan.cascades.values.LiteralValue; +import com.apple.foundationdb.record.query.plan.plans.RecordQueryFetchFromPartialRecordPlan; +import com.apple.foundationdb.record.query.plan.plans.RecordQueryIndexPlan; import com.apple.foundationdb.record.vector.TestRecordsVectorsProto; +import com.apple.foundationdb.record.vector.TestRecordsVectorsProto.UngroupedVectorRecord; import com.apple.foundationdb.tuple.Tuple; import com.apple.test.Tags; -import com.google.common.base.Preconditions; +import com.christianheina.langx.half4j.Half; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.errorprone.annotations.CanIgnoreReturnValue; @@ -49,14 +62,13 @@ import javax.annotation.Nonnull; import java.text.SimpleDateFormat; import java.util.ArrayList; -import java.util.Collections; import java.util.List; +import java.util.Objects; +import java.util.Optional; import java.util.Random; import java.util.concurrent.CompletableFuture; import java.util.function.Consumer; import java.util.function.Function; -import java.util.stream.Collectors; -import java.util.stream.IntStream; import static com.apple.foundationdb.record.metadata.Key.Expressions.field; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -92,20 +104,30 @@ protected void openRecordStore(final FDBRecordContext context, final RecordMetaD static Function getRecordGenerator(@Nonnull final Random random) { return recNo -> { - final byte[] vector = new byte[128 * 2]; + final byte[] vector = randomVectorData(random, 128); random.nextBytes(vector); //logRecord(recNo, vector); - return TestRecordsVectorsProto.UngroupedVectorRecord.newBuilder() + return UngroupedVectorRecord.newBuilder() .setRecNo(recNo) .setVectorData(ByteString.copyFrom(vector)) .build(); }; } - public void saveRecords(final boolean useAsync, @Nonnull final RecordMetaDataHook hook, final long seed, + static byte[] randomVectorData(final Random random, final int dimensions) { + // we do this in this convoluted way to make sure we won't get NaNs and other special surprises + final Half[] componentData = new Half[dimensions]; + for (int i = 0; i < componentData.length; i++) { + componentData[i] = Half.valueOf(random.nextFloat()); + } + + Vector.HalfVector vector = new Vector.HalfVector(componentData); + return vector.getRawData(); + } + + public void saveRecords(final boolean useAsync, @Nonnull final RecordMetaDataHook hook, @Nonnull final Random random, final int numSamples) { - final Random random = new Random(seed); final var recordGenerator = getRecordGenerator(random); if (useAsync) { Assertions.assertDoesNotThrow(() -> batchAsync(hook, numSamples, 100, recNo -> recordStore.saveRecordAsync(recordGenerator.apply(recNo)))); @@ -114,22 +136,6 @@ public void saveRecords(final boolean useAsync, @Nonnull final RecordMetaDataHoo } } - public void deleteRecords(final boolean useAsync, @Nonnull final RecordMetaDataHook hook, final long seed, final int numRecords, - final int numDeletes) throws Exception { - Preconditions.checkArgument(numDeletes <= numRecords); - final Random random = new Random(seed); - final List recNos = IntStream.range(0, numRecords) - .boxed() - .collect(Collectors.toList()); - Collections.shuffle(recNos, random); - final List recNosToBeDeleted = recNos.subList(0, numDeletes); - if (useAsync) { - batchAsync(hook, recNosToBeDeleted.size(), 500, recNo -> recordStore.deleteRecordAsync(Tuple.from(recNo))); - } else { - batch(hook, recNosToBeDeleted.size(), 500, recNo -> recordStore.deleteRecord(Tuple.from(recNo))); - } - } - private long batch(final RecordMetaDataHook hook, final int numRecords, final int batchSize, Consumer recordConsumer) throws Exception { long numRecordsCommitted = 0; while (numRecordsCommitted < numRecords) { @@ -177,16 +183,17 @@ private static void logRecord(final long recNo, @Nonnull final ByteString vector } } - void basicReadTest() throws Exception { - saveRecords(false, this::addVectorIndex, 0, 1000); + void basicWriteReadTest() throws Exception { + final Random random = new Random(); + saveRecords(false, this::addVectorIndex, random, 1000); try (FDBRecordContext context = openContext()) { openRecordStore(context, this::addVectorIndex); for (long l = 0; l < 1000; l ++) { FDBStoredRecord rec = recordStore.loadRecord(Tuple.from(l)); assertNotNull(rec); - TestRecordsVectorsProto.UngroupedVectorRecord.Builder recordBuilder = - TestRecordsVectorsProto.UngroupedVectorRecord.newBuilder(); + UngroupedVectorRecord.Builder recordBuilder = + UngroupedVectorRecord.newBuilder(); recordBuilder.mergeFrom(rec.getRecord()); final var record = recordBuilder.build(); logRecord(record.getRecNo(), record.getVectorData()); @@ -195,20 +202,45 @@ final var record = recordBuilder.build(); } } - void basicConcurrentReadTest(final boolean useAsync) throws Exception { - saveRecords(useAsync, NO_HOOK, 0, 5000); - for (int i = 0; i < 1; i ++) { - try (FDBRecordContext context = openContext()) { - openRecordStore(context); - for (long l = 0; l < 5000; l += 20) { - fetchBatchConcurrently(l, 20); + void basicWriteIndexReadTest() throws Exception { + final Random random = new Random(0); + saveRecords(false, this::addVectorIndex, random, 1000); + + final DistanceRankValueComparison distanceRankComparison = + new DistanceRankValueComparison(Comparisons.Type.DISTANCE_RANK_LESS_THAN_OR_EQUAL, + new LiteralValue<>(Type.Vector.of(false, 16, 128), + randomVectorData(random, 128)), + new LiteralValue<>(10)); + + final VectorIndexScanComparisons vectorIndexScanComparisons = + new VectorIndexScanComparisons(ScanComparisons.EMPTY, distanceRankComparison, ScanComparisons.EMPTY); + + final var baseRecordType = + Type.Record.fromFieldDescriptorsMap( + Type.Record.toFieldDescriptorMap(UngroupedVectorRecord.getDescriptor().getFields())); + + final var indexPlan = new RecordQueryIndexPlan("UngroupedVectorIndex", field("recNo"), + vectorIndexScanComparisons, IndexFetchMethod.SCAN_AND_FETCH, + RecordQueryFetchFromPartialRecordPlan.FetchIndexRecords.PRIMARY_KEY, false, false, + Optional.empty(), baseRecordType, QueryPlanConstraint.noConstraint()); + + try (FDBRecordContext context = openContext()) { + openRecordStore(context, this::addVectorIndex); + + try (RecordCursorIterator> cursor = executeQuery(indexPlan)) { + while (cursor.hasNext()) { + FDBQueriedRecord rec = cursor.next(); + UngroupedVectorRecord.Builder myrec = UngroupedVectorRecord.newBuilder(); + myrec.mergeFrom(Objects.requireNonNull(rec).getRecord()); + System.out.println(myrec); } - commit(context); } + + //commit(context); } } - private List fetchBatchConcurrently(long startRecNo, int batchSize) throws Exception { + private List fetchBatchConcurrently(long startRecNo, int batchSize) throws Exception { final ImmutableList.Builder>> futureBatchBuilder = ImmutableList.builder(); for (int i = 0; i < batchSize; i ++) { final CompletableFuture> recordFuture = recordStore.loadRecordAsync(Tuple.from(startRecNo + i)); @@ -220,8 +252,8 @@ private List fetchBatchConcurrent return batch.stream() .map(rec -> { assertNotNull(rec); - TestRecordsVectorsProto.UngroupedVectorRecord.Builder recordBuilder = - TestRecordsVectorsProto.UngroupedVectorRecord.newBuilder(); + UngroupedVectorRecord.Builder recordBuilder = + UngroupedVectorRecord.newBuilder(); recordBuilder.mergeFrom(rec.getRecord()); return recordBuilder.build(); }) From 27e58cdd3646b9df9522eaf09e58a1008317ace4 Mon Sep 17 00:00:00 2001 From: Normen Seemann Date: Thu, 14 Aug 2025 15:09:07 +0200 Subject: [PATCH 29/34] adding a scan with prefix --- .../indexes/VectorIndexMaintainer.java | 6 +- .../indexes/VectorIndexSimpleTest.java | 5 ++ .../indexes/VectorIndexTestBase.java | 88 ++++++++++++------- .../src/test/proto/test_records_vector.proto | 11 +-- 4 files changed, 67 insertions(+), 43 deletions(-) diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java index 1a063211b4..b48abdedd5 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java @@ -147,8 +147,10 @@ public RecordCursor scan(@Nonnull final IndexScanBounds scanBounds, Tuple.fromBytes(indexEntryProto.getKey().toByteArray()), Tuple.fromBytes(indexEntryProto.getValue().toByteArray()))); } - return new ListCursor<>(indexEntriesBuilder.build(), - parsedContinuation.getInnerContinuation().toByteArray()); + final ImmutableList indexEntries = indexEntriesBuilder.build(); + return new ListCursor<>(indexEntries, parsedContinuation.getInnerContinuation().toByteArray()) + .mapResult(result -> + result.withContinuation(new Continuation(indexEntries, result.getContinuation()))); } final HNSW hnsw = new HNSW(hnswSubspace, getExecutor(), getConfig(), diff --git a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java index 8a8b4e8607..6124371121 100644 --- a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java +++ b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexSimpleTest.java @@ -43,4 +43,9 @@ void basicWriteReadTest() throws Exception { void basicWriteIndexReadTest() throws Exception { super.basicWriteIndexReadTest(); } + + @Test + void basicWriteIndexReadGroupedTest() throws Exception { + super.basicWriteIndexReadGroupedTest(); + } } diff --git a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java index 299c6ba775..c03f3691f3 100644 --- a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java +++ b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java @@ -45,11 +45,10 @@ import com.apple.foundationdb.record.query.plan.plans.RecordQueryFetchFromPartialRecordPlan; import com.apple.foundationdb.record.query.plan.plans.RecordQueryIndexPlan; import com.apple.foundationdb.record.vector.TestRecordsVectorsProto; -import com.apple.foundationdb.record.vector.TestRecordsVectorsProto.UngroupedVectorRecord; +import com.apple.foundationdb.record.vector.TestRecordsVectorsProto.VectorRecord; import com.apple.foundationdb.tuple.Tuple; import com.apple.test.Tags; import com.christianheina.langx.half4j.Half; -import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.protobuf.ByteString; @@ -62,7 +61,6 @@ import javax.annotation.Nonnull; import java.text.SimpleDateFormat; import java.util.ArrayList; -import java.util.List; import java.util.Objects; import java.util.Optional; import java.util.Random; @@ -70,6 +68,7 @@ import java.util.function.Consumer; import java.util.function.Function; +import static com.apple.foundationdb.record.metadata.Key.Expressions.concat; import static com.apple.foundationdb.record.metadata.Key.Expressions.field; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -83,11 +82,15 @@ public abstract class VectorIndexTestBase extends FDBRecordStoreQueryTestBase { private static final SimpleDateFormat timeFormat = new SimpleDateFormat("MM/dd/yyyy HH:mm:ss"); @CanIgnoreReturnValue - RecordMetaDataBuilder addVectorIndex(@Nonnull final RecordMetaDataBuilder metaDataBuilder) { - metaDataBuilder.addIndex("UngroupedVectorRecord", + RecordMetaDataBuilder addVectorIndexes(@Nonnull final RecordMetaDataBuilder metaDataBuilder) { + metaDataBuilder.addIndex("VectorRecord", new Index("UngroupedVectorIndex", new KeyWithValueExpression(field("vector_data"), 0), IndexTypes.VECTOR, ImmutableMap.of(IndexOptions.HNSW_METRIC, Metrics.EUCLIDEAN_METRIC.toString()))); + metaDataBuilder.addIndex("VectorRecord", + new Index("GroupedVectorIndex", new KeyWithValueExpression(concat(field("group_id"), field("vector_data")), 1), + IndexTypes.VECTOR, + ImmutableMap.of(IndexOptions.HNSW_METRIC, Metrics.EUCLIDEAN_METRIC.toString()))); return metaDataBuilder; } @@ -97,7 +100,7 @@ protected void openRecordStore(FDBRecordContext context) throws Exception { protected void openRecordStore(final FDBRecordContext context, final RecordMetaDataHook hook) throws Exception { RecordMetaDataBuilder metaDataBuilder = RecordMetaData.newBuilder().setRecords(TestRecordsVectorsProto.getDescriptor()); - metaDataBuilder.getRecordType("UngroupedVectorRecord").setPrimaryKey(field("rec_no")); + metaDataBuilder.getRecordType("VectorRecord").setPrimaryKey(field("rec_no")); hook.apply(metaDataBuilder); createOrOpenRecordStore(context, metaDataBuilder.getRecordMetaData()); } @@ -107,14 +110,15 @@ static Function getRecordGenerator(@Nonnull final Random random) final byte[] vector = randomVectorData(random, 128); random.nextBytes(vector); - //logRecord(recNo, vector); - return UngroupedVectorRecord.newBuilder() + return VectorRecord.newBuilder() .setRecNo(recNo) .setVectorData(ByteString.copyFrom(vector)) + .setGroupId(recNo.intValue() % 2) .build(); }; } + @Nonnull static byte[] randomVectorData(final Random random, final int dimensions) { // we do this in this convoluted way to make sure we won't get NaNs and other special surprises final Half[] componentData = new Half[dimensions]; @@ -185,15 +189,15 @@ private static void logRecord(final long recNo, @Nonnull final ByteString vector void basicWriteReadTest() throws Exception { final Random random = new Random(); - saveRecords(false, this::addVectorIndex, random, 1000); + saveRecords(false, this::addVectorIndexes, random, 1000); try (FDBRecordContext context = openContext()) { - openRecordStore(context, this::addVectorIndex); + openRecordStore(context, this::addVectorIndexes); for (long l = 0; l < 1000; l ++) { FDBStoredRecord rec = recordStore.loadRecord(Tuple.from(l)); assertNotNull(rec); - UngroupedVectorRecord.Builder recordBuilder = - UngroupedVectorRecord.newBuilder(); + VectorRecord.Builder recordBuilder = + VectorRecord.newBuilder(); recordBuilder.mergeFrom(rec.getRecord()); final var record = recordBuilder.build(); logRecord(record.getRecNo(), record.getVectorData()); @@ -204,7 +208,7 @@ final var record = recordBuilder.build(); void basicWriteIndexReadTest() throws Exception { final Random random = new Random(0); - saveRecords(false, this::addVectorIndex, random, 1000); + saveRecords(false, this::addVectorIndexes, random, 1000); final DistanceRankValueComparison distanceRankComparison = new DistanceRankValueComparison(Comparisons.Type.DISTANCE_RANK_LESS_THAN_OR_EQUAL, @@ -217,7 +221,7 @@ void basicWriteIndexReadTest() throws Exception { final var baseRecordType = Type.Record.fromFieldDescriptorsMap( - Type.Record.toFieldDescriptorMap(UngroupedVectorRecord.getDescriptor().getFields())); + Type.Record.toFieldDescriptorMap(VectorRecord.getDescriptor().getFields())); final var indexPlan = new RecordQueryIndexPlan("UngroupedVectorIndex", field("recNo"), vectorIndexScanComparisons, IndexFetchMethod.SCAN_AND_FETCH, @@ -225,12 +229,12 @@ void basicWriteIndexReadTest() throws Exception { Optional.empty(), baseRecordType, QueryPlanConstraint.noConstraint()); try (FDBRecordContext context = openContext()) { - openRecordStore(context, this::addVectorIndex); + openRecordStore(context, this::addVectorIndexes); try (RecordCursorIterator> cursor = executeQuery(indexPlan)) { while (cursor.hasNext()) { FDBQueriedRecord rec = cursor.next(); - UngroupedVectorRecord.Builder myrec = UngroupedVectorRecord.newBuilder(); + VectorRecord.Builder myrec = VectorRecord.newBuilder(); myrec.mergeFrom(Objects.requireNonNull(rec).getRecord()); System.out.println(myrec); } @@ -240,23 +244,41 @@ void basicWriteIndexReadTest() throws Exception { } } - private List fetchBatchConcurrently(long startRecNo, int batchSize) throws Exception { - final ImmutableList.Builder>> futureBatchBuilder = ImmutableList.builder(); - for (int i = 0; i < batchSize; i ++) { - final CompletableFuture> recordFuture = recordStore.loadRecordAsync(Tuple.from(startRecNo + i)); - futureBatchBuilder.add(recordFuture); + void basicWriteIndexReadGroupedTest() throws Exception { + final Random random = new Random(0); + saveRecords(false, this::addVectorIndexes, random, 1000); + + final DistanceRankValueComparison distanceRankComparison = + new DistanceRankValueComparison(Comparisons.Type.DISTANCE_RANK_LESS_THAN_OR_EQUAL, + new LiteralValue<>(Type.Vector.of(false, 16, 128), + randomVectorData(random, 128)), + new LiteralValue<>(10)); + + final VectorIndexScanComparisons vectorIndexScanComparisons = + new VectorIndexScanComparisons(ScanComparisons.EMPTY, distanceRankComparison, ScanComparisons.EMPTY); + + final var baseRecordType = + Type.Record.fromFieldDescriptorsMap( + Type.Record.toFieldDescriptorMap(VectorRecord.getDescriptor().getFields())); + + final var indexPlan = new RecordQueryIndexPlan("GroupedVectorIndex", field("recNo"), + vectorIndexScanComparisons, IndexFetchMethod.SCAN_AND_FETCH, + RecordQueryFetchFromPartialRecordPlan.FetchIndexRecords.PRIMARY_KEY, false, false, + Optional.empty(), baseRecordType, QueryPlanConstraint.noConstraint()); + + try (FDBRecordContext context = openContext()) { + openRecordStore(context, this::addVectorIndexes); + + try (RecordCursorIterator> cursor = executeQuery(indexPlan)) { + while (cursor.hasNext()) { + FDBQueriedRecord rec = cursor.next(); + VectorRecord.Builder myrec = VectorRecord.newBuilder(); + myrec.mergeFrom(Objects.requireNonNull(rec).getRecord()); + System.out.println(myrec); + } + } + + //commit(context); } - final var batchFuture = AsyncUtil.getAll(futureBatchBuilder.build()); - final var batch = batchFuture.get(); - - return batch.stream() - .map(rec -> { - assertNotNull(rec); - UngroupedVectorRecord.Builder recordBuilder = - UngroupedVectorRecord.newBuilder(); - recordBuilder.mergeFrom(rec.getRecord()); - return recordBuilder.build(); - }) - .collect(ImmutableList.toImmutableList()); } } diff --git a/fdb-record-layer-core/src/test/proto/test_records_vector.proto b/fdb-record-layer-core/src/test/proto/test_records_vector.proto index 925ebbfc32..63ff019c9d 100644 --- a/fdb-record-layer-core/src/test/proto/test_records_vector.proto +++ b/fdb-record-layer-core/src/test/proto/test_records_vector.proto @@ -28,17 +28,12 @@ import "record_metadata_options.proto"; option (schema).store_record_versions = true; -message UngroupedVectorRecord { +message VectorRecord { optional int64 rec_no = 1 [(field).primary_key = true]; - optional bytes vector_data = 2; -} - -message GroupedVectorRecord { - optional int64 rec_no = 1 [(field).primary_key = true]; - optional int32 groupId = 2; + optional int32 group_id = 2; optional bytes vector_data = 3; } message RecordTypeUnion { - optional UngroupedVectorRecord _NaiveVectorRecord = 1; + optional VectorRecord _VectorRecord = 1; } From bde65da8e6c9dd124812fa343a95be9db4efe41e Mon Sep 17 00:00:00 2001 From: Youssef Hatem Date: Thu, 14 Aug 2025 16:59:48 +0100 Subject: [PATCH 30/34] - checkpoint. Generation of logical plan works correctly. --- .../apple/foundationdb/async/hnsw/Vector.java | 18 +++++- .../expressions/SelectExpression.java | 2 +- .../query/plan/cascades/typing/Type.java | 18 ++---- .../values/AbstractArrayConstructorValue.java | 6 ++ .../plan/cascades/values/ArithmeticValue.java | 47 ++++++++++++++- .../query/plan/cascades/values/RankValue.java | 25 ++++++++ .../fdb-relational-core.gradle | 1 + .../src/main/antlr/RelationalLexer.g4 | 6 +- .../src/main/antlr/RelationalParser.g4 | 29 +++++---- .../recordlayer/query/ParseHelpers.java | 4 ++ .../functions/SqlFunctionCatalogImpl.java | 3 + .../query/visitors/BaseVisitor.java | 29 +++++++-- .../query/visitors/DdlVisitor.java | 7 --- .../query/visitors/DelegatingVisitor.java | 23 +++++++- .../query/visitors/ExpressionVisitor.java | 59 +++++++++++++++++++ .../query/visitors/IdentifierVisitor.java | 6 ++ .../query/visitors/TypedVisitor.java | 17 +++++- .../recordlayer/query/VectorTypeTest.java | 21 +++++++ 18 files changed, 266 insertions(+), 55 deletions(-) diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java index bfa179ea2b..8a4ab62469 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java @@ -67,6 +67,8 @@ public R[] getData() { @Nonnull public abstract DoubleVector toDoubleVector(); + public abstract int precision(); + @Override public boolean equals(final Object o) { if (!(o instanceof Vector)) { @@ -126,6 +128,11 @@ public DoubleVector toDoubleVector() { return toDoubleVectorSupplier.get(); } + @Override + public int precision() { + return 16; + } + @Nonnull public DoubleVector computeDoubleVector() { Double[] result = new Double[data.length]; @@ -182,6 +189,11 @@ public DoubleVector toDoubleVector() { return this; } + @Override + public int precision() { + return 64; + } + @Nonnull @Override public byte[] getRawData() { @@ -190,9 +202,9 @@ public byte[] getRawData() { } } - static double distance(@Nonnull Metric metric, - @Nonnull final Vector vector1, - @Nonnull final Vector vector2) { + public static double distance(@Nonnull Metric metric, + @Nonnull final Vector vector1, + @Nonnull final Vector vector2) { return metric.distance(vector1.toDoubleVector().getData(), vector2.toDoubleVector().getData()); } diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/expressions/SelectExpression.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/expressions/SelectExpression.java index ad50724bcb..2dfcdbb3d5 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/expressions/SelectExpression.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/expressions/SelectExpression.java @@ -630,7 +630,7 @@ public PlannerGraph rewriteInternalPlannerGraph(@Nonnull final List 30 ? String.format("%02x", predicateString.hashCode()) : predicateString; + final var abbreviatedPredicateString = predicateString; return PlannerGraph.fromNodeAndChildGraphs( new PlannerGraph.LogicalOperatorNode(this, "SELECT " + resultValue, diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java index e72a4ea65e..b9a2ac21bb 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java @@ -659,6 +659,10 @@ static Type fromObject(@Nullable final Object object) { if (object instanceof DynamicMessage) { return Record.fromDescriptor(((DynamicMessage) object).getDescriptorForType()); } + if (object instanceof com.apple.foundationdb.async.hnsw.Vector) { + final var vector = (com.apple.foundationdb.async.hnsw.Vector)object; + return Type.Vector.of(false, vector.precision(), vector.size()); + } final var typeCode = typeCodeFromPrimitive(object); if (typeCode == TypeCode.NULL) { return Type.nullType(); @@ -719,7 +723,7 @@ enum TypeCode { INT(Integer.class, FieldDescriptorProto.Type.TYPE_INT32, true, true), LONG(Long.class, FieldDescriptorProto.Type.TYPE_INT64, true, true), STRING(String.class, FieldDescriptorProto.Type.TYPE_STRING, true, false), - VECTOR(Vector.JavaVectorType.class, FieldDescriptorProto.Type.TYPE_BYTES, true, false), + VECTOR(com.apple.foundationdb.async.hnsw.Vector.class, FieldDescriptorProto.Type.TYPE_BYTES, true, false), VERSION(FDBRecordVersion.class, FieldDescriptorProto.Type.TYPE_BYTES, true, false), ENUM(Enum.class, FieldDescriptorProto.Type.TYPE_ENUM, false, false), RECORD(Message.class, null, false, false), @@ -1304,18 +1308,6 @@ public static Vector fromProto(@Nonnull final PlanSerializationContext serializa return new Vector(vectorTypeProto.getIsNullable(), vectorTypeProto.getPrecision(), vectorTypeProto.getDimensions()); } - static final class JavaVectorType { - private final ByteString underlying; - - JavaVectorType(@Nonnull final ByteString underlying) { - this.underlying = underlying; - } - - public ByteString getUnderlying() { - return underlying; - } - } - /** * Deserializer. */ diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/AbstractArrayConstructorValue.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/AbstractArrayConstructorValue.java index cfca7bec5d..ea6a1077b8 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/AbstractArrayConstructorValue.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/AbstractArrayConstructorValue.java @@ -276,6 +276,12 @@ public static LightArrayConstructorValue of(@Nonnull final List return new LightArrayConstructorValue(children); } + @Nonnull + public static LightArrayConstructorValue of(@Nonnull final List children, @Nonnull final Type type) { + Verify.verify(!children.isEmpty()); + return new LightArrayConstructorValue(children, type); + } + @Nonnull public static LightArrayConstructorValue emptyArray(@Nonnull final Type elementType) { return new LightArrayConstructorValue(elementType); diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/ArithmeticValue.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/ArithmeticValue.java index f2b565ab1f..d7bb5d2ddd 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/ArithmeticValue.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/ArithmeticValue.java @@ -22,6 +22,8 @@ import com.apple.foundationdb.annotation.API; import com.apple.foundationdb.annotation.SpotBugsSuppressWarnings; +import com.apple.foundationdb.async.hnsw.Metric; +import com.apple.foundationdb.async.hnsw.Vector; import com.apple.foundationdb.record.EvaluationContext; import com.apple.foundationdb.record.ObjectPlanHash; import com.apple.foundationdb.record.PlanDeserializer; @@ -358,6 +360,28 @@ public BitmapBucketOffsetFn() { } } + /** + * Euclidean distance function. + */ + @AutoService(BuiltInFunction.class) + public static class EuclideanDistanceFn extends BuiltInFunction { + public EuclideanDistanceFn() { + super("euclidean_distance", + ImmutableList.of(Type.any(), Type.any()), ArithmeticValue::encapsulateInternal); + } + } + + /** + * Cosine distance function. + */ + @AutoService(BuiltInFunction.class) + public static class CosineDistanceFn extends BuiltInFunction { + public CosineDistanceFn() { + super("cosine_distance", + ImmutableList.of(Type.any(), Type.any()), ArithmeticValue::encapsulateInternal); + } + } + /** * Logical operator. */ @@ -372,7 +396,9 @@ public enum LogicalOperator { BITXOR("^", Precedence.BITWISE_XOR), BITMAP_BUCKET_NUMBER("bitmap_bucket_number", Precedence.NEVER_PARENS), BITMAP_BUCKET_OFFSET("bitmap_bucket_offset", Precedence.NEVER_PARENS), - BITMAP_BIT_POSITION("bitmap_bit_position", Precedence.NEVER_PARENS) + BITMAP_BIT_POSITION("bitmap_bit_position", Precedence.NEVER_PARENS), + EUCLIDEAN_DISTANCE("euclidean_distance", Precedence.NEVER_PARENS), + COSINE_DISTANCE("cosine_distance", Precedence.NEVER_PARENS) ; @Nonnull @@ -518,7 +544,24 @@ public enum PhysicalOperator { BITMAP_BIT_POSITION_LI(LogicalOperator.BITMAP_BIT_POSITION, TypeCode.LONG, TypeCode.INT, TypeCode.LONG, (l, r) -> Math.subtractExact((long)l, Math.multiplyExact(Math.floorDiv((long)l, (int)r), (int)r))), BITMAP_BIT_POSITION_II(LogicalOperator.BITMAP_BIT_POSITION, TypeCode.INT, TypeCode.INT, TypeCode.INT, (l, r) -> Math.subtractExact((int)l, Math.multiplyExact(Math.floorDiv((int)l, (int)r), (int)r))), - ; + + EUCLIDEAN_DISTANCE_VV(LogicalOperator.EUCLIDEAN_DISTANCE, TypeCode.VECTOR, TypeCode.VECTOR, TypeCode.DOUBLE, ((l, r) -> new Metric.EuclideanMetric().distance(((Vector)l).toDoubleVector().getData(), ((Vector)r).toDoubleVector().getData()))), +// EUCLIDEAN_DISTANCE_VA(LogicalOperator.EUCLIDEAN_DISTANCE, TypeCode.VECTOR, TypeCode.ARRAY, TypeCode.DOUBLE, ((l, r) -> { +// final List rhs = (List)r; +// final Vector lhs = (Vector)l; +// Verify.verify(rhs.size() == lhs.size()); +// if (rhs.isEmpty()) { +// return 0.0d; +// } +// final var doubleArray = new Double[rhs.size()]; +// for (int i = 0; i < rhs.size(); i++) { +// final var item = rhs.get(i); +// Verify.verify(item instanceof Number); +// doubleArray[i] = ((Number)item).doubleValue(); +// } +// return new Metric.CosineMetric().distance(((Vector)l).toDoubleVector().getData(), doubleArray); +// })), + COSINE_DISTANCE_VV(LogicalOperator.COSINE_DISTANCE, TypeCode.VECTOR, TypeCode.VECTOR, TypeCode.DOUBLE, ((l, r) -> new Metric.CosineMetric().distance(((Vector)l).toDoubleVector().getData(), ((Vector)r).toDoubleVector().getData()))); @Nonnull private static final Supplier> protoEnumBiMapSupplier = diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RankValue.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RankValue.java index 2935bfb318..78cac98e1d 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RankValue.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RankValue.java @@ -26,10 +26,15 @@ import com.apple.foundationdb.record.PlanSerializationContext; import com.apple.foundationdb.record.planprotos.PRankValue; import com.apple.foundationdb.record.planprotos.PValue; +import com.apple.foundationdb.record.query.plan.cascades.BuiltInFunction; import com.apple.foundationdb.record.query.plan.cascades.typing.Type; +import com.apple.foundationdb.record.query.plan.cascades.typing.Typed; import com.google.auto.service.AutoService; +import com.google.common.base.Verify; +import com.google.common.collect.ImmutableList; import javax.annotation.Nonnull; +import java.util.List; import java.util.Objects; /** @@ -111,4 +116,24 @@ public RankValue fromProto(@Nonnull final PlanSerializationContext serialization return RankValue.fromProto(serializationContext, rankValueProto); } } + + /** + * The {@code rank} window function. + */ + @AutoService(BuiltInFunction.class) + public static class RankFn extends BuiltInFunction { + + public RankFn() { + super("rank", ImmutableList.of(Type.any(), Type.any()), RankFn::encapsulateInternal); + } + + @Nonnull + private static RankValue encapsulateInternal(@Nonnull BuiltInFunction builtInFunction, + @Nonnull final List arguments) { + Verify.verify(arguments.size() == 2); // ordering expressions must be present, ok for now. + final var partitioningValuesList = (AbstractArrayConstructorValue)arguments.get(0); + final var argumentValuesList = (AbstractArrayConstructorValue)arguments.get(1); + return new RankValue(partitioningValuesList.getChildren(), argumentValuesList.getChildren()); + } + } } diff --git a/fdb-relational-core/fdb-relational-core.gradle b/fdb-relational-core/fdb-relational-core.gradle index 5c77feb8f4..cc57c45aad 100644 --- a/fdb-relational-core/fdb-relational-core.gradle +++ b/fdb-relational-core/fdb-relational-core.gradle @@ -40,6 +40,7 @@ dependencies { implementation(libs.log4j.core) implementation(libs.caffeine) implementation(libs.caffeine.guava) + implementation(libs.half4j) antlr(libs.antlr) diff --git a/fdb-relational-core/src/main/antlr/RelationalLexer.g4 b/fdb-relational-core/src/main/antlr/RelationalLexer.g4 index 43f841233c..b85df08791 100644 --- a/fdb-relational-core/src/main/antlr/RelationalLexer.g4 +++ b/fdb-relational-core/src/main/antlr/RelationalLexer.g4 @@ -969,6 +969,7 @@ DES_ENCRYPT: 'DES_ENCRYPT'; DIMENSION: 'DIMENSION'; DISJOINT: 'DISJOINT'; DRY: 'DRY'; +EUCLIDEAN_DISTANCE: 'EUCLIDEAN_DISTANCE'; ELT: 'ELT'; ENABLE_LONG_ROWS: 'ENABLE_LONG_ROWS'; ENCODE: 'ENCODE'; @@ -1289,7 +1290,7 @@ DECIMAL_LITERAL: DEC_DIGIT+ DECIMAL_TYPE_MODIFIER?; HEXADECIMAL_LITERAL: 'X' STRING_LITERAL; BASE64_LITERAL: 'B64' STRING_LITERAL; -REAL_LITERAL: (DEC_DIGIT+)? '.' DEC_DIGIT+ REAL_TYPE_MODIFIER? +REAL_LITERAL: MINUS? (DEC_DIGIT+)? '.' DEC_DIGIT+ REAL_TYPE_MODIFIER? | DEC_DIGIT+ '.' EXPONENT_NUM_PART REAL_TYPE_MODIFIER? | (DEC_DIGIT+)? '.' (DEC_DIGIT+ EXPONENT_NUM_PART) REAL_TYPE_MODIFIER? | DEC_DIGIT+ EXPONENT_NUM_PART REAL_TYPE_MODIFIER?; @@ -1352,9 +1353,10 @@ fragment SQUOTA_STRING: '\'' ('\'\'' | ~('\''))* '\''; fragment DEC_DIGIT: [0-9]; fragment BIT_STRING_L: 'B' '\'' [01]+ '\''; +fragment HALF_TYPE_MODIFIER: ('H'|'h'); fragment FLOAT_TYPE_MODIFIER: ('F'|'f'); fragment DOUBLE_TYPE_MODIFIER: ('D'|'d'); -fragment REAL_TYPE_MODIFIER: (FLOAT_TYPE_MODIFIER | DOUBLE_TYPE_MODIFIER); +fragment REAL_TYPE_MODIFIER: (HALF_TYPE_MODIFIER | FLOAT_TYPE_MODIFIER | DOUBLE_TYPE_MODIFIER); fragment INT_TYPE_MODIFIER: ('I' | 'i'); fragment LONG_TYPE_MODIFIER: ('L' | 'l'); fragment DECIMAL_TYPE_MODIFIER: (INT_TYPE_MODIFIER | LONG_TYPE_MODIFIER); diff --git a/fdb-relational-core/src/main/antlr/RelationalParser.g4 b/fdb-relational-core/src/main/antlr/RelationalParser.g4 index 39ad26350e..fb8ba98963 100644 --- a/fdb-relational-core/src/main/antlr/RelationalParser.g4 +++ b/fdb-relational-core/src/main/antlr/RelationalParser.g4 @@ -766,13 +766,14 @@ nullLiteral // done constant - : stringLiteral #stringConstant // done - | decimalLiteral #decimalConstant // done - | '-' decimalLiteral #negativeDecimalConstant // done - | bytesLiteral #bytesConstant // done - | booleanLiteral #booleanConstant // done - | BIT_STRING #bitStringConstant // done (unsupported) - | NOT? nullLiteral #nullConstant // done (unsupported) - if NOT exists. + : stringLiteral #stringConstant // done + | decimalLiteral #decimalConstant // done + | '-' decimalLiteral #negativeDecimalConstant // done + | bytesLiteral #bytesConstant // done + | booleanLiteral #booleanConstant // done + | BIT_STRING #bitStringConstant // done (unsupported) + | NOT? nullLiteral #nullConstant // done (unsupported) - if NOT exists. + | VECTOR '(' dimension=REAL_LITERAL (COMMA dimension=REAL_LITERAL)* ')' #vectorConstant // done ; @@ -951,8 +952,9 @@ ifNotExists functionCall : aggregateWindowedFunction #aggregateFunctionCall // done (supported) + | nonAggregateWindowedFunction #nonAggregateFunctionCall // done (supported) | specificFunction #specificFunctionCall // - | scalarFunctionName '(' functionArgs? ')' #scalarFunctionCall // done (unsupported) + | scalarFunctionName '(' functionArgs? ')' #scalarFunctionCall // done (supported) ; specificFunction @@ -1088,20 +1090,18 @@ nonAggregateWindowedFunction ; overClause - : OVER (/* '(' windowSpec? ')' |*/ windowName) + : OVER ( '(' windowSpec ')' | windowName) ; windowName : uid ; -//commented out until we want to support window functions -/* windowSpec - : windowName? partitionClause? orderByClause? frameClause? + : windowName? partitionClause? orderByClause? /*frameClause?*/ ; - +/* frameClause : frameUnits frameExtent ; @@ -1125,7 +1125,6 @@ frameRange | UNBOUNDED (PRECEDING | FOLLOWING) | expression (PRECEDING | FOLLOWING) ; - */ partitionClause @@ -1320,7 +1319,7 @@ functionNameBase | CREATE_DH_PARAMETERS | CREATE_DIGEST | CROSSES | CUME_DIST | DATABASE | DATE | DATEDIFF | DATE_FORMAT | DAY | DAYNAME | DAYOFMONTH | DAYOFWEEK | DAYOFYEAR | DECODE | DEGREES | DENSE_RANK | DES_DECRYPT - | DES_ENCRYPT | DIMENSION | DISJOINT | ELT | ENCODE + | DES_ENCRYPT | DIMENSION | DISJOINT | EUCLIDEAN_DISTANCE | ELT | ENCODE | ENCRYPT | ENDPOINT | ENVELOPE | EQUALS | EXP | EXPORT_SET | EXTERIORRING | EXTRACTVALUE | FIELD | FIND_IN_SET | FIRST_VALUE | FLOOR | FORMAT | FOUND_ROWS | FROM_BASE64 | FROM_DAYS diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/ParseHelpers.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/ParseHelpers.java index 686d3e7639..16910b1f40 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/ParseHelpers.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/ParseHelpers.java @@ -22,6 +22,7 @@ import com.apple.foundationdb.annotation.API; +import com.apple.foundationdb.async.hnsw.HNSWHelpers; import com.apple.foundationdb.record.query.plan.cascades.TreeLike; import com.apple.foundationdb.record.query.plan.cascades.typing.TypeRepository; import com.apple.foundationdb.record.query.plan.cascades.values.Value; @@ -69,6 +70,9 @@ public static Object parseDecimal(@Nonnull String valueAsString) { if (valueAsString.contains(".")) { final var lastCharacter = valueAsString.charAt(lastCharIdx); switch (lastCharacter) { + case 'h': + case 'H': // fallthrough + return HNSWHelpers.halfValueOf(Float.parseFloat(valueAsString.substring(0, lastCharIdx))); case 'f': // fallthrough case 'F': return Float.parseFloat(valueAsString.substring(0, lastCharIdx)); diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/functions/SqlFunctionCatalogImpl.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/functions/SqlFunctionCatalogImpl.java index 827fc863ee..0c32752161 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/functions/SqlFunctionCatalogImpl.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/functions/SqlFunctionCatalogImpl.java @@ -125,6 +125,8 @@ private static ImmutableMap BuiltInFunctionCatalog.resolve("bitmap_bit_position", 1 + argumentsCount)) .put("bitmap_bucket_offset", argumentsCount -> BuiltInFunctionCatalog.resolve("bitmap_bucket_offset", 1 + argumentsCount)) .put("bitmap_construct_agg", argumentsCount -> BuiltInFunctionCatalog.resolve("BITMAP_CONSTRUCT_AGG", argumentsCount)) + .put("euclidean_distance", argumentsCount -> BuiltInFunctionCatalog.resolve("euclidean_distance", argumentsCount)) + .put("cosine_distance", argumentsCount -> BuiltInFunctionCatalog.resolve("cosine_distance", argumentsCount)) .put("not", argumentsCount -> BuiltInFunctionCatalog.resolve("not", argumentsCount)) .put("and", argumentsCount -> BuiltInFunctionCatalog.resolve("and", argumentsCount)) .put("or", argumentsCount -> BuiltInFunctionCatalog.resolve("or", argumentsCount)) @@ -144,6 +146,7 @@ private static ImmutableMap BuiltInFunctionCatalog.resolve("isNull", argumentsCount)) .put("is not null", argumentsCount -> BuiltInFunctionCatalog.resolve("notNull", argumentsCount)) .put("range", argumentsCount -> BuiltInFunctionCatalog.resolve("range", argumentsCount)) + .put("rank", argumentsCount -> BuiltInFunctionCatalog.resolve("rank", argumentsCount)) .put("__pattern_for_like", argumentsCount -> BuiltInFunctionCatalog.resolve("patternForLike", argumentsCount)) .put("__internal_array", argumentsCount -> BuiltInFunctionCatalog.resolve("array", argumentsCount)) .put("__pick_value", argumentsCount -> BuiltInFunctionCatalog.resolve("pick", argumentsCount)) diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java index bb2772ee02..3c5d69f47e 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java @@ -1176,6 +1176,11 @@ public Expression visitNullConstant(@Nonnull RelationalParser.NullConstantContex return expressionVisitor.visitNullConstant(ctx); } + @Override + public Expression visitVectorConstant(final RelationalParser.VectorConstantContext ctx) { + return expressionVisitor.visitVectorConstant(ctx); + } + @Nonnull @Override public Object visitStringDataType(@Nonnull RelationalParser.StringDataTypeContext ctx) { @@ -1392,6 +1397,12 @@ public Expression visitAggregateFunctionCall(@Nonnull RelationalParser.Aggregate return expressionVisitor.visitAggregateFunctionCall(ctx); } + @Nonnull + @Override + public Expression visitNonAggregateFunctionCall(final RelationalParser.NonAggregateFunctionCallContext ctx) { + return expressionVisitor.visitNonAggregateFunctionCall(ctx) ; + } + @Nonnull @Override public Object visitSpecificFunctionCall(@Nonnull RelationalParser.SpecificFunctionCallContext ctx) { @@ -1508,26 +1519,32 @@ public Expression visitAggregateWindowedFunction(@Nonnull RelationalParser.Aggre @Nonnull @Override - public Object visitNonAggregateWindowedFunction(@Nonnull RelationalParser.NonAggregateWindowedFunctionContext ctx) { - return visitChildren(ctx); + public Expression visitNonAggregateWindowedFunction(@Nonnull RelationalParser.NonAggregateWindowedFunctionContext ctx) { + return expressionVisitor.visitNonAggregateWindowedFunction(ctx); } @Nonnull @Override - public Object visitOverClause(@Nonnull RelationalParser.OverClauseContext ctx) { - return visitChildren(ctx); + public NonnullPair> visitOverClause(@Nonnull RelationalParser.OverClauseContext ctx) { + return expressionVisitor.visitOverClause(ctx); + } + + @Nonnull + @Override + public Identifier visitWindowName(@Nonnull RelationalParser.WindowNameContext ctx) { + return identifierVisitor.visitWindowName(ctx); } @Nonnull @Override - public Object visitWindowName(@Nonnull RelationalParser.WindowNameContext ctx) { + public Object visitWindowSpec(final RelationalParser.WindowSpecContext ctx) { return visitChildren(ctx); } @Nonnull @Override public Expressions visitPartitionClause(final RelationalParser.PartitionClauseContext ctx) { - return ddlVisitor.visitPartitionClause(ctx); + return expressionVisitor.visitPartitionClause(ctx); } @Nonnull diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java index 07476789c6..5c337d4896 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java @@ -589,11 +589,4 @@ public DataType visitReturnsType(@Nonnull RelationalParser.ReturnsTypeContext ct public Boolean visitNullColumnConstraint(@Nonnull RelationalParser.NullColumnConstraintContext ctx) { return ctx.nullNotnull().NOT() == null; } - - @Nonnull - @Override - public Expressions visitPartitionClause(final RelationalParser.PartitionClauseContext ctx) { - return Expressions.of(ctx.expression().stream().map(expContext -> - Assert.castUnchecked(visit(expContext), Expression.class)).collect(ImmutableList.toImmutableList())); - } } diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java index f3bb965f90..c3a6474ee7 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java @@ -1016,6 +1016,11 @@ public Expression visitNullConstant(@Nonnull RelationalParser.NullConstantContex return getDelegate().visitNullConstant(ctx); } + @Override + public Expression visitVectorConstant(final RelationalParser.VectorConstantContext ctx) { + return getDelegate().visitVectorConstant(ctx); + } + @Nonnull @Override public Object visitStringDataType(@Nonnull RelationalParser.StringDataTypeContext ctx) { @@ -1232,6 +1237,12 @@ public Expression visitAggregateFunctionCall(@Nonnull RelationalParser.Aggregate return getDelegate().visitAggregateFunctionCall(ctx); } + @Nonnull + @Override + public Expression visitNonAggregateFunctionCall(final RelationalParser.NonAggregateFunctionCallContext ctx) { + return getDelegate().visitNonAggregateFunctionCall(ctx); + } + @Nonnull @Override public Object visitSpecificFunctionCall(@Nonnull RelationalParser.SpecificFunctionCallContext ctx) { @@ -1348,22 +1359,28 @@ public Expression visitAggregateWindowedFunction(@Nonnull RelationalParser.Aggre @Nonnull @Override - public Object visitNonAggregateWindowedFunction(@Nonnull RelationalParser.NonAggregateWindowedFunctionContext ctx) { + public Expression visitNonAggregateWindowedFunction(@Nonnull RelationalParser.NonAggregateWindowedFunctionContext ctx) { return getDelegate().visitNonAggregateWindowedFunction(ctx); } @Nonnull @Override - public Object visitOverClause(@Nonnull RelationalParser.OverClauseContext ctx) { + public NonnullPair> visitOverClause(@Nonnull RelationalParser.OverClauseContext ctx) { return getDelegate().visitOverClause(ctx); } @Nonnull @Override - public Object visitWindowName(@Nonnull RelationalParser.WindowNameContext ctx) { + public Identifier visitWindowName(@Nonnull RelationalParser.WindowNameContext ctx) { return getDelegate().visitWindowName(ctx); } + @Nonnull + @Override + public Object visitWindowSpec(final RelationalParser.WindowSpecContext ctx) { + return getDelegate().visitWindowSpec(ctx); + } + @Nonnull @Override public Expressions visitPartitionClause(final RelationalParser.PartitionClauseContext ctx) { diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/ExpressionVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/ExpressionVisitor.java index e6cd22899c..39d330c00b 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/ExpressionVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/ExpressionVisitor.java @@ -21,6 +21,7 @@ package com.apple.foundationdb.relational.recordlayer.query.visitors; import com.apple.foundationdb.annotation.API; +import com.apple.foundationdb.async.hnsw.Vector; import com.apple.foundationdb.record.query.plan.cascades.Quantifier; import com.apple.foundationdb.record.query.plan.cascades.predicates.CompatibleTypeEvolutionPredicate; import com.apple.foundationdb.record.query.plan.cascades.typing.Type; @@ -52,6 +53,7 @@ import com.apple.foundationdb.relational.recordlayer.query.TautologicalValue; import com.apple.foundationdb.relational.util.Assert; import com.apple.foundationdb.relational.util.ExcludeFromJacocoGeneratedReport; +import com.christianheina.langx.half4j.Half; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Streams; @@ -260,6 +262,32 @@ public Expression visitAggregateFunctionCall(@Nonnull RelationalParser.Aggregate return visitAggregateWindowedFunction(functionCon.aggregateWindowedFunction()); } + @Nonnull + @Override + public Expression visitNonAggregateFunctionCall(final RelationalParser.NonAggregateFunctionCallContext ctx) { + return visitNonAggregateWindowedFunction(ctx.nonAggregateWindowedFunction()); + } + + @Nonnull + @Override + public Expression visitNonAggregateWindowedFunction(@Nonnull final RelationalParser.NonAggregateWindowedFunctionContext ctx) { + Assert.notNullUnchecked(ctx.RANK(), ErrorCode.UNSUPPORTED_QUERY, "only RANK window function is currently supported"); + final var partitionAndOrderByValues = visitOverClause(ctx.overClause()); + + final var partitionExpressions = partitionAndOrderByValues.getLeft(); + final var partitionValues = Streams.stream(partitionExpressions.underlying()).collect(ImmutableList.toImmutableList()); + final var partitionArray = AbstractArrayConstructorValue.LightArrayConstructorValue.of(partitionValues, Type.any()); + + final var orderByExpressions = partitionAndOrderByValues.getRight(); + final var orderByValues = orderByExpressions.stream().map(r -> r.getExpression().getUnderlying()) + .collect(ImmutableList.toImmutableList()); + final var orderByArray = AbstractArrayConstructorValue.LightArrayConstructorValue.of(orderByValues, Type.any()); + + final var arguments = Expressions.of(ImmutableList.of(Expression.ofUnnamed(partitionArray), Expression.ofUnnamed(orderByArray))); + final var result = getDelegate().getSemanticAnalyzer().resolveScalarFunction("rank", arguments, false); + return result; + } + @Nonnull @Override public Expression visitAggregateWindowedFunction(@Nonnull RelationalParser.AggregateWindowedFunctionContext functionContext) { @@ -341,6 +369,22 @@ public Expression visitFunctionArg(@Nonnull RelationalParser.FunctionArgContext return Assert.castUnchecked(functionArgContext.expression().accept(this), Expression.class); } + @Nonnull + @Override + public NonnullPair> visitOverClause(@Nonnull RelationalParser.OverClauseContext ctx) { + Assert.isNullUnchecked(ctx.windowName(), ErrorCode.UNSUPPORTED_QUERY, "named window functions not supported"); + final var partitions = visitPartitionClause(ctx.windowSpec().partitionClause()); + final var orderByClause = visitOrderByClause(ctx.windowSpec().orderByClause()); + return NonnullPair.of(partitions, orderByClause); + } + + @Nonnull + @Override + public Expressions visitPartitionClause(final RelationalParser.PartitionClauseContext ctx) { + return Expressions.of(ctx.expression().stream().map(expContext -> + Assert.castUnchecked(visit(expContext), Expression.class)).collect(ImmutableList.toImmutableList())); + } + @Nonnull @Override public Expressions visitFunctionArgs(@Nonnull RelationalParser.FunctionArgsContext ctx) { @@ -651,6 +695,21 @@ public Expression visitNullLiteral(@Nonnull RelationalParser.NullLiteralContext return Expression.ofUnnamed(new NullValue(Type.nullType())); // do not strip nulls. } + @Override + public Expression visitVectorConstant(final RelationalParser.VectorConstantContext ctx) { + final var items = ctx.REAL_LITERAL().stream().map(l -> ParseHelpers.parseDecimal(l.getText())).collect(ImmutableList.toImmutableList()); + // todo: we should allow the user to mix and match literals and use max_type to calculate the array element type + // but ok for now. + Assert.thatUnchecked(items.stream().map(Object::getClass).distinct().count() == 1, ErrorCode.SYNTAX_ERROR, + "vector elements must be of the same type"); + final var firstItem = items.get(0); + if (firstItem instanceof Double) { + return Expression.ofUnnamed(LiteralValue.ofScalar(new Vector.DoubleVector(items.stream().map(i -> (Double)i).toArray(Double[]::new)))); + } + Assert.thatUnchecked(firstItem instanceof Half, ErrorCode.SYNTAX_ERROR, "vector element type " + firstItem.getClass().getSimpleName() + " not supported"); + return Expression.ofUnnamed(LiteralValue.ofScalar(new Vector.HalfVector(items.stream().map(i -> (Half)i).toArray(Half[]::new)))); + } + @Nonnull @Override public Expression visitStringConstant(@Nonnull RelationalParser.StringConstantContext ctx) { diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/IdentifierVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/IdentifierVisitor.java index 8ebe8e4899..e9c95acbf9 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/IdentifierVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/IdentifierVisitor.java @@ -122,4 +122,10 @@ public Identifier visitCollationName(@Nonnull RelationalParser.CollationNameCont public Identifier visitTableFunctionName(final RelationalParser.TableFunctionNameContext ctx) { return visitFullId(ctx.fullId()); } + + @Nonnull + @Override + public Identifier visitWindowName(@Nonnull final RelationalParser.WindowNameContext ctx) { + return visitUid(ctx.uid()); + } } diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java index 6ebd5c1bb9..2ab3136a34 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java @@ -628,6 +628,9 @@ public interface TypedVisitor extends RelationalParserVisitor { @Override Expression visitNullConstant(@Nonnull RelationalParser.NullConstantContext ctx); + @Override + Expression visitVectorConstant(RelationalParser.VectorConstantContext ctx); + @Nonnull @Override Object visitStringDataType(@Nonnull RelationalParser.StringDataTypeContext ctx); @@ -771,6 +774,10 @@ public interface TypedVisitor extends RelationalParserVisitor { @Override Expression visitAggregateFunctionCall(@Nonnull RelationalParser.AggregateFunctionCallContext ctx); + @Nonnull + @Override + Expression visitNonAggregateFunctionCall(final RelationalParser.NonAggregateFunctionCallContext ctx); + @Nonnull @Override Object visitSpecificFunctionCall(@Nonnull RelationalParser.SpecificFunctionCallContext ctx); @@ -849,15 +856,19 @@ public interface TypedVisitor extends RelationalParserVisitor { @Nonnull @Override - Object visitNonAggregateWindowedFunction(@Nonnull RelationalParser.NonAggregateWindowedFunctionContext ctx); + Expression visitNonAggregateWindowedFunction(@Nonnull RelationalParser.NonAggregateWindowedFunctionContext ctx); + + @Nonnull + @Override + NonnullPair> visitOverClause(@Nonnull RelationalParser.OverClauseContext ctx); @Nonnull @Override - Object visitOverClause(@Nonnull RelationalParser.OverClauseContext ctx); + Identifier visitWindowName(@Nonnull RelationalParser.WindowNameContext ctx); @Nonnull @Override - Object visitWindowName(@Nonnull RelationalParser.WindowNameContext ctx); + Object visitWindowSpec(final RelationalParser.WindowSpecContext ctx); @Nonnull @Override diff --git a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java index b0e8d55d36..6bee4c648a 100644 --- a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java +++ b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java @@ -23,9 +23,11 @@ import com.apple.foundationdb.relational.api.StructResultSetMetaData; import com.apple.foundationdb.relational.api.metadata.DataType; import com.apple.foundationdb.relational.recordlayer.EmbeddedRelationalExtension; +import com.apple.foundationdb.relational.recordlayer.Utils; import com.apple.foundationdb.relational.utils.Ddl; import org.assertj.core.api.Assertions; import org.junit.jupiter.api.Order; +import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; @@ -40,6 +42,9 @@ public class VectorTypeTest { @Order(0) public final EmbeddedRelationalExtension relationalExtension = new EmbeddedRelationalExtension(); + public VectorTypeTest() { + Utils.enableCascadesDebugger(); + } @Nonnull public static Stream vectorArguments() { @@ -66,4 +71,20 @@ void vectorTest(@Nonnull final String ddlType, @Nonnull final DataType expectedT } } } + + @Test + void selectFromHnsw() throws Exception { + final String schemaTemplate = "create table photos(zone string, recordId string, name string," + + "embedding vector(3), primary key (zone, recordId), organized by hnsw(embedding partition by zone + 3, name) " + + "with (hnsw_m = 10, hnsw_ef_construction = 5))"; + try (var ddl = Ddl.builder().database(URI.create("/TEST/QT")).relationalExtension(relationalExtension).schemaTemplate(schemaTemplate).build()) { + try (var statement = ddl.setSchemaAndGetConnection().createStatement()) { + final var result = statement.execute("SELECT * FROM photos WHERE RANK() OVER (PARTITION BY zone ORDER BY euclidean_distance(embedding, vector(1.2, -0.5, 3.14)) DESC) < 10"); + final var metadata = statement.getResultSet().getMetaData(); + Assertions.assertThat(metadata).isInstanceOf(StructResultSetMetaData.class); + final var relationalMetadata = (StructResultSetMetaData)metadata; + final var type = relationalMetadata.getRelationalDataType().getFields().get(1).getType(); + } + } + } } From 998d0b3f75a040506e95cb65d1a428d262388745 Mon Sep 17 00:00:00 2001 From: Youssef Hatem Date: Mon, 18 Aug 2025 11:20:45 +0100 Subject: [PATCH 31/34] Generating and planning HNSW SQL queries. - very ad-hoc, but "works". --- .../apple/foundationdb/async/hnsw/Vector.java | 8 + .../record/query/plan/ScanComparisons.java | 3 + .../query/plan/cascades/MatchCandidate.java | 29 ++ .../cascades/VectorIndexExpansionVisitor.java | 211 +++++++++++ .../VectorIndexScanMatchCandidate.java | 340 ++++++++++++++++++ .../expressions/SelectExpression.java | 2 +- .../cascades/predicates/RangeConstraints.java | 3 + .../rules/AbstractDataAccessRule.java | 12 + .../query/plan/cascades/typing/Type.java | 4 + .../plan/cascades/values/ArithmeticValue.java | 15 - .../values/EuclideanDistanceRankValue.java | 114 ++++++ .../plan/cascades/values/FieldValue.java | 5 + .../query/plan/cascades/values/RankValue.java | 60 ++++ .../values/RecordConstructorValue.java | 3 + .../plan/cascades/values/RelOpValue.java | 18 + .../query/plan/cascades/values/Value.java | 6 + .../plan/plans/RecordQueryIndexPlan.java | 15 +- .../src/main/proto/record_query_plan.proto | 5 + .../relational/recordlayer/MessageTuple.java | 14 +- .../query/visitors/BaseVisitor.java | 6 +- .../query/visitors/DelegatingVisitor.java | 6 +- .../query/visitors/ExpressionVisitor.java | 7 +- .../query/visitors/TypedVisitor.java | 6 +- .../recordlayer/query/VectorTypeTest.java | 23 +- 24 files changed, 876 insertions(+), 39 deletions(-) create mode 100644 fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/VectorIndexExpansionVisitor.java create mode 100644 fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/VectorIndexScanMatchCandidate.java create mode 100644 fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/EuclideanDistanceRankValue.java diff --git a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java index 8a4ab62469..e1c7e34e10 100644 --- a/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java +++ b/fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/Vector.java @@ -213,4 +213,12 @@ static double comparativeDistance(@Nonnull Metric metric, @Nonnull final Vector vector2) { return metric.comparativeDistance(vector1.toDoubleVector().getData(), vector2.toDoubleVector().getData()); } + + public static Vector fromBytes(@Nonnull final byte[] bytes, int precision) { + if (precision == 16) { + return HalfVector.halfVectorFromBytes(bytes); + } + // TODO + throw new UnsupportedOperationException("not implemented yet"); + } } diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/ScanComparisons.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/ScanComparisons.java index 4292dda86a..34903baa44 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/ScanComparisons.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/ScanComparisons.java @@ -151,6 +151,7 @@ public static ComparisonType getComparisonType(@Nonnull Comparisons.Comparison c switch (comparison.getType()) { case EQUALS: case IS_NULL: + case DISTANCE_RANK_EQUALS: return ComparisonType.EQUALITY; case LESS_THAN: case LESS_THAN_OR_EQUALS: @@ -159,6 +160,8 @@ public static ComparisonType getComparisonType(@Nonnull Comparisons.Comparison c case STARTS_WITH: case NOT_NULL: case SORT: + case DISTANCE_RANK_LESS_THAN: + case DISTANCE_RANK_LESS_THAN_OR_EQUAL: return ComparisonType.INEQUALITY; case NOT_EQUALS: default: diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/MatchCandidate.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/MatchCandidate.java index 8d7fbcef2e..653a816103 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/MatchCandidate.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/MatchCandidate.java @@ -361,6 +361,17 @@ static Iterable fromIndexDefinition(@Nonnull final RecordMetaDat isReverse ).ifPresent(resultBuilder::add); break; + case IndexTypes.VECTOR: + expandVectorIndexMatchCandidate( + index, + availableRecordTypeNames, + availableRecordTypes, + queriedRecordTypeNames, + queriedRecordTypes, + isReverse, + commonPrimaryKeyForIndex + ).ifPresent(resultBuilder::add); + break; default: break; } @@ -385,6 +396,24 @@ private static Optional expandValueIndexMatchCandidate(@Nonnull ); } + private static Optional expandVectorIndexMatchCandidate(@Nonnull final Index index, + @Nonnull final Set availableRecordTypeNames, + @Nonnull final Collection availableRecordTypes, + @Nonnull final Set queriedRecordTypeNames, + @Nonnull final Collection queriedRecordTypes, + final boolean isReverse, + @Nullable final KeyExpression commonPrimaryKeyForIndex) { + return expandIndexMatchCandidate(index, + availableRecordTypeNames, + availableRecordTypes, + queriedRecordTypeNames, + queriedRecordTypes, + isReverse, + commonPrimaryKeyForIndex, + new VectorIndexExpansionVisitor(index, queriedRecordTypes) + ); + } + private static Optional expandAggregateIndexMatchCandidate(@Nonnull final Index index, @Nonnull final Set availableRecordTypeNames, @Nonnull final Collection availableRecordTypes, diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/VectorIndexExpansionVisitor.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/VectorIndexExpansionVisitor.java new file mode 100644 index 0000000000..3f45564b09 --- /dev/null +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/VectorIndexExpansionVisitor.java @@ -0,0 +1,211 @@ +/* + * VectorIndexExpansionVisitor.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2021 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.record.query.plan.cascades; + +import com.apple.foundationdb.annotation.SpotBugsSuppressWarnings; +import com.apple.foundationdb.record.metadata.Index; +import com.apple.foundationdb.record.metadata.IndexTypes; +import com.apple.foundationdb.record.metadata.RecordType; +import com.apple.foundationdb.record.metadata.expressions.GroupingKeyExpression; +import com.apple.foundationdb.record.metadata.expressions.KeyExpression; +import com.apple.foundationdb.record.metadata.expressions.KeyWithValueExpression; +import com.apple.foundationdb.record.query.plan.cascades.debug.Debugger; +import com.apple.foundationdb.record.query.plan.cascades.expressions.MatchableSortExpression; +import com.apple.foundationdb.record.query.plan.cascades.predicates.Placeholder; +import com.apple.foundationdb.record.query.plan.cascades.predicates.PredicateWithValueAndRanges; +import com.apple.foundationdb.record.query.plan.cascades.values.EuclideanDistanceRankValue; +import com.apple.foundationdb.record.query.plan.cascades.values.Value; +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Objects; +import java.util.Set; +import java.util.function.Supplier; + +import static com.apple.foundationdb.record.metadata.Key.Expressions.concat; + +/** + * Class to expand vector index access into a candidate graph. The visitation methods are left unchanged from the super + * class {@link KeyExpressionExpansionVisitor}, this class merely provides a specific {@link #expand} method. + */ +public class VectorIndexExpansionVisitor extends KeyExpressionExpansionVisitor implements ExpansionVisitor { + @Nonnull + private static final Set SUPPORTED_INDEX_TYPES = Set.of( + IndexTypes.VECTOR + ); + + @Nonnull + private static final Set GROUPED_INDEX_TYPES = Set.of( + IndexTypes.RANK, + IndexTypes.PERMUTED_MAX, + IndexTypes.PERMUTED_MIN + ); + + @Nonnull + private final Index index; + @Nonnull + private final List queriedRecordTypes; + + public VectorIndexExpansionVisitor(@Nonnull Index index, @Nonnull Collection queriedRecordTypes) { + Preconditions.checkArgument(SUPPORTED_INDEX_TYPES.contains(index.getType())); + this.index = index; + this.queriedRecordTypes = ImmutableList.copyOf(queriedRecordTypes); + } + + @Nonnull + @Override + @SpotBugsSuppressWarnings("NP_PARAMETER_MUST_BE_NONNULL_BUT_MARKED_AS_NULLABLE") + public MatchCandidate expand(@Nonnull final Supplier baseQuantifierSupplier, + @Nullable final KeyExpression primaryKey, + final boolean isReverse) { + Debugger.updateIndex(PredicateWithValueAndRanges.class, old -> 0); + + final var baseQuantifier = baseQuantifierSupplier.get(); + final var allExpansionsBuilder = ImmutableList.builder(); + + // add the value for the flow of records + allExpansionsBuilder.add(GraphExpansion.ofQuantifier(baseQuantifier)); + + var rootExpression = index.getRootExpression(); + + if (rootExpression instanceof GroupingKeyExpression) { + if (GROUPED_INDEX_TYPES.contains(index.getType())) { + rootExpression = ((GroupingKeyExpression)rootExpression).getWholeKey(); + } else { + throw new UnsupportedOperationException("cannot create match candidate on grouping expression for unknown index type"); + } + } + + final int keyValueSplitPoint; + if (rootExpression instanceof KeyWithValueExpression) { + final KeyWithValueExpression keyWithValueExpression = (KeyWithValueExpression)rootExpression; + keyValueSplitPoint = keyWithValueExpression.getSplitPoint(); + rootExpression = keyWithValueExpression.getInnerKey(); + } else { + keyValueSplitPoint = -1; + } + + final var keyValues = Lists.newArrayList(); + final var valueValues = Lists.newArrayList(); + + final var initialState = + VisitorState.of(keyValues, + valueValues, + baseQuantifier, + ImmutableList.of(), + keyValueSplitPoint, + 0, + false, + true); + + final var keyValueExpansion = + pop(rootExpression.expand(push(initialState))) + .toBuilder() + .removeAllResultColumns() + .build(); + + final var distanceRankValuePlaceHolder = new EuclideanDistanceRankValue(ImmutableList.copyOf(keyValues), ImmutableList.copyOf(valueValues)) + .asPlaceholder(newParameterAlias()); + + allExpansionsBuilder.add(keyValueExpansion); + allExpansionsBuilder.add(GraphExpansion.ofPlaceholder(distanceRankValuePlaceHolder)); + + if (index.hasPredicate()) { + final var filteredIndexPredicate = Objects.requireNonNull(index.getPredicate()).toPredicate(baseQuantifier.getFlowedObjectValue()); + final var valueRangesMaybe = IndexPredicateExpansion.dnfPredicateToRanges(filteredIndexPredicate); + final var predicateExpansionBuilder = GraphExpansion.builder(); + if (valueRangesMaybe.isEmpty()) { // could not create DNF, store the predicate as-is. + allExpansionsBuilder.add(GraphExpansion.ofPredicate(filteredIndexPredicate)); + } else { + final var valueRanges = valueRangesMaybe.get(); + for (final var value : valueRanges.keySet()) { + // we check if the predicate value is a placeholder, if so, create a placeholder, otherwise, add it as a constraint. + final var maybePlaceholder = keyValueExpansion.getPlaceholders() + .stream() + .filter(existingPlaceholder -> existingPlaceholder.getValue().semanticEquals(value, AliasMap.emptyMap())) + .findFirst(); + if (maybePlaceholder.isEmpty()) { + predicateExpansionBuilder.addPredicate(PredicateWithValueAndRanges.ofRanges(value, ImmutableSet.copyOf(valueRanges.get(value)))); + } else { + predicateExpansionBuilder.addPlaceholder(maybePlaceholder.get().withExtraRanges(ImmutableSet.copyOf(valueRanges.get(value)))); + } + } + } + allExpansionsBuilder.add(predicateExpansionBuilder.build()); + } + + final var completeExpansion = GraphExpansion.ofOthers(allExpansionsBuilder.build()); + final var sealedExpansion = completeExpansion.seal(); + final var parameters = + sealedExpansion.getPlaceholders() + .stream() + .map(Placeholder::getParameterAlias) + .collect(ImmutableList.toImmutableList()); + final var matchableSortExpression = new MatchableSortExpression(parameters, isReverse, + sealedExpansion.buildSelectWithResultValue(baseQuantifier.getFlowedObjectValue())); + return new VectorIndexScanMatchCandidate(index, + queriedRecordTypes, + Traversal.withRoot(Reference.initialOf(matchableSortExpression)), + parameters, + baseQuantifier.getFlowedObjectType(), + baseQuantifier.getAlias(), + keyValues, + valueValues, + fullKey(index, primaryKey), + distanceRankValuePlaceHolder, + primaryKey); + } + + /** + * Compute the full key of an index (given that the index is a value index). + * + * @param index index to be expanded + * @param primaryKey primary key of the records the index ranges over. The primary key is used to determine + * parts in the index definition that already contain parts of the primary key. All primary key components + * that are not already part of the index key are appended to the index key. + * @return a {@link KeyExpression} describing the full index key as stored + */ + @Nonnull + public static KeyExpression fullKey(@Nonnull Index index, @Nullable final KeyExpression primaryKey) { + final KeyExpression rootExpression = index.getRootExpression() instanceof KeyWithValueExpression + ? ((KeyWithValueExpression)index.getRootExpression()).getKeyExpression() + : index.getRootExpression(); + if (primaryKey == null) { + return rootExpression; + } + final var trimmedPrimaryKeyComponents = new ArrayList<>(primaryKey.normalizeKeyForPositions()); + index.trimPrimaryKey(trimmedPrimaryKeyComponents); + if (trimmedPrimaryKeyComponents.isEmpty()) { + return rootExpression; + } + final var fullKeyListBuilder = ImmutableList.builder(); + fullKeyListBuilder.add(rootExpression); + fullKeyListBuilder.addAll(trimmedPrimaryKeyComponents); + return concat(fullKeyListBuilder.build()); + } +} diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/VectorIndexScanMatchCandidate.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/VectorIndexScanMatchCandidate.java new file mode 100644 index 0000000000..854ddf0e86 --- /dev/null +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/VectorIndexScanMatchCandidate.java @@ -0,0 +1,340 @@ +/* + * VectorIndexScanMatchCandidate.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2020 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.record.query.plan.cascades; + +import com.apple.foundationdb.record.RecordCoreException; +import com.apple.foundationdb.record.RecordMetaData; +import com.apple.foundationdb.record.metadata.Index; +import com.apple.foundationdb.record.metadata.RecordType; +import com.apple.foundationdb.record.metadata.expressions.KeyExpression; +import com.apple.foundationdb.record.provider.foundationdb.VectorIndexScanComparisons; +import com.apple.foundationdb.record.query.expressions.Comparisons; +import com.apple.foundationdb.record.query.plan.AvailableFields; +import com.apple.foundationdb.record.query.plan.ScanComparisons; +import com.apple.foundationdb.record.query.plan.cascades.predicates.Placeholder; +import com.apple.foundationdb.record.query.plan.cascades.typing.Type; +import com.apple.foundationdb.record.query.plan.cascades.values.Value; +import com.apple.foundationdb.record.query.plan.plans.RecordQueryCoveringIndexPlan; +import com.apple.foundationdb.record.query.plan.plans.RecordQueryFetchFromPartialRecordPlan; +import com.apple.foundationdb.record.query.plan.plans.RecordQueryIndexPlan; +import com.apple.foundationdb.record.query.plan.plans.RecordQueryPlan; +import com.apple.foundationdb.record.util.pair.NonnullPair; +import com.google.common.base.Suppliers; +import com.google.common.base.Verify; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.util.Collection; +import java.util.List; +import java.util.Optional; +import java.util.function.Supplier; + +/** + * Case class to represent a match candidate that is backed by an index. + */ +public class VectorIndexScanMatchCandidate implements ScanWithFetchMatchCandidate, ValueIndexLikeMatchCandidate { + /** + * Index metadata structure. + */ + @Nonnull + private final Index index; + + /** + * Record types this index is defined over. + */ + private final List queriedRecordTypes; + + /** + * Holds the parameter names for all necessary parameters that need to be bound during matching. + */ + @Nonnull + private final List parameters; + + /** + * Base type. + */ + @Nonnull + private final Type baseType; + + /** + * Base alias. + */ + @Nonnull + private final CorrelationIdentifier baseAlias; + + /** + * List of values that represent the key parts of the index represented by the candidate in the expanded graph. + */ + @Nonnull + private final List indexKeyValues; + + /** + * List of values that represent the value parts of the index represented by the candidate in the expanded graph. + */ + @Nonnull + private final List indexValueValues; + + /** + * Traversal object of the expanded index scan graph. + */ + @Nonnull + private final Traversal traversal; + + @Nonnull + private final KeyExpression fullKeyExpression; + + @Nullable + private final KeyExpression primaryKey; + + @Nonnull + private final Placeholder distancePredicatePlaceholder; + + @Nonnull + private final Supplier>> primaryKeyValuesOptionalSupplier; + + @Nonnull + private final Supplier> indexEntryToLogicalRecordOptionalSupplier; + + public VectorIndexScanMatchCandidate(@Nonnull final Index index, + @Nonnull final Collection queriedRecordTypes, + @Nonnull final Traversal traversal, + @Nonnull final List parameters, + @Nonnull final Type baseType, + @Nonnull final CorrelationIdentifier baseAlias, + @Nonnull final List indexKeyValues, + @Nonnull final List indexValueValues, + @Nonnull final KeyExpression fullKeyExpression, + @Nonnull final Placeholder distancePredicatePlaceholder, + @Nullable final KeyExpression primaryKey) { + this.index = index; + this.queriedRecordTypes = ImmutableList.copyOf(queriedRecordTypes); + this.traversal = traversal; + this.parameters = ImmutableList.copyOf(parameters); + this.baseType = baseType; + this.baseAlias = baseAlias; + this.indexKeyValues = ImmutableList.copyOf(indexKeyValues); + this.indexValueValues = ImmutableList.copyOf(indexValueValues); + this.fullKeyExpression = fullKeyExpression; + this.distancePredicatePlaceholder = distancePredicatePlaceholder; + this.primaryKey = primaryKey; + this.primaryKeyValuesOptionalSupplier = + Suppliers.memoize(() -> MatchCandidate.computePrimaryKeyValuesMaybe(primaryKey, baseType)); + this.indexEntryToLogicalRecordOptionalSupplier = + Suppliers.memoize(() -> ScanWithFetchMatchCandidate.computeIndexEntryToLogicalRecord(queriedRecordTypes, + baseAlias, baseType, indexKeyValues, indexValueValues)); + } + + @Override + public int getColumnSize() { + return index.getColumnSize(); + } + + @Override + public boolean isUnique() { + return index.isUnique(); + } + + @Nonnull + @Override + public String getName() { + return index.getName(); + } + + @Nonnull + @Override + public List getQueriedRecordTypes() { + return queriedRecordTypes; + } + + @Nonnull + @Override + public Traversal getTraversal() { + return traversal; + } + + @Nonnull + @Override + public List getSargableAliases() { + return parameters; + } + + @Nonnull + @Override + public List getOrderingAliases() { + return getSargableAliases(); + } + + @Nonnull + @Override + public Type getBaseType() { + return baseType; + } + + @Nonnull + public CorrelationIdentifier getBaseAlias() { + return baseAlias; + } + + @Nonnull + public List getIndexKeyValues() { + return indexKeyValues; + } + + @Nonnull + public List getIndexValueValues() { + return indexValueValues; + } + + @Nonnull + @Override + public KeyExpression getFullKeyExpression() { + return fullKeyExpression; + } + + @Override + public String toString() { + return "value[" + getName() + "]"; + } + + @Override + public boolean createsDuplicates() { + return index.getRootExpression().createsDuplicates(); + } + + @Nonnull + @Override + public Optional> getPrimaryKeyValuesMaybe() { + return primaryKeyValuesOptionalSupplier.get(); + } + + @Nonnull + private Optional getIndexEntryToLogicalRecordMaybe() { + return indexEntryToLogicalRecordOptionalSupplier.get(); + } + + @Nonnull + @Override + public RecordQueryPlan toEquivalentPlan(@Nonnull final PartialMatch partialMatch, + @Nonnull final PlanContext planContext, + @Nonnull final Memoizer memoizer, + @Nonnull final List comparisonRanges, + final boolean reverseScanOrder) { + final var matchInfo = partialMatch.getRegularMatchInfo(); + final var comparisons = extractRankComparison(comparisonRanges); + final var baseRecordType = + Type.Record.fromFieldDescriptorsMap(RecordMetaData.getFieldDescriptorMapFromTypes(queriedRecordTypes)); + return tryFetchCoveringIndexScan(partialMatch, planContext, memoizer, comparisonRanges, reverseScanOrder, baseRecordType) + .orElseGet(() -> + new RecordQueryIndexPlan(index.getName(), + primaryKey, + VectorIndexScanComparisons.byValue(toScanComparisons(comparisons.getLeft()), comparisons.getRight(), null), + planContext.getPlannerConfiguration().getIndexFetchMethod(), + RecordQueryFetchFromPartialRecordPlan.FetchIndexRecords.PRIMARY_KEY, + reverseScanOrder, + false, + partialMatch.getMatchCandidate(), + baseRecordType, + matchInfo.getConstraint())); + } + + private static NonnullPair, Comparisons.DistanceRankValueComparison> extractRankComparison(@Nonnull final List comparisonRanges) { + final var nonRankComparisonsBuilder = ImmutableList.builder(); + final var rankComparisonsBuilder = ImmutableList.builder(); + for (final var comparisonRange : comparisonRanges) { + if (comparisonRange.isEquality() && comparisonRange.getEqualityComparison() instanceof Comparisons.DistanceRankValueComparison) { + rankComparisonsBuilder.add((Comparisons.DistanceRankValueComparison)comparisonRange.getEqualityComparison()); + } else if (comparisonRange.isInequality() && comparisonRange.getInequalityComparisons().stream().allMatch(comp -> comp instanceof Comparisons.DistanceRankValueComparison)) { + rankComparisonsBuilder.addAll(comparisonRange.getInequalityComparisons().stream().map(comp -> (Comparisons.DistanceRankValueComparison)comp).collect(ImmutableList.toImmutableList())); + } else { + nonRankComparisonsBuilder.add(comparisonRange); + } + } + final var rankComparisons = rankComparisonsBuilder.build(); + Verify.verify(rankComparisons.size() == 1); + return NonnullPair.of(nonRankComparisonsBuilder.build(), rankComparisons.get(0)); + } + + @Nonnull + private Optional tryFetchCoveringIndexScan(@Nonnull final PartialMatch partialMatch, + @Nonnull final PlanContext planContext, + @Nonnull final Memoizer memoizer, + @Nonnull final List comparisonRanges, + final boolean isReverse, + @Nonnull Type.Record baseRecordType) { + final var indexEntryToLogicalRecordOptional = getIndexEntryToLogicalRecordMaybe(); + if (indexEntryToLogicalRecordOptional.isEmpty()) { + return Optional.empty(); + } + final var indexEntryToLogicalRecord = indexEntryToLogicalRecordOptional.get(); + final var comparisons = extractRankComparison(comparisonRanges); + final var indexPlan = + new RecordQueryIndexPlan(index.getName(), + primaryKey, + VectorIndexScanComparisons.byValue(toScanComparisons(comparisons.getLeft()), comparisons.getRight(), null), + planContext.getPlannerConfiguration().getIndexFetchMethod(), + RecordQueryFetchFromPartialRecordPlan.FetchIndexRecords.PRIMARY_KEY, + isReverse, + false, + partialMatch.getMatchCandidate(), + baseRecordType, + partialMatch.getRegularMatchInfo().getConstraint()); + + final var coveringIndexPlan = new RecordQueryCoveringIndexPlan(indexPlan, + indexEntryToLogicalRecord.getQueriedRecordType().getName(), + AvailableFields.NO_FIELDS, // not used except for old planner properties + indexEntryToLogicalRecord.getIndexKeyValueToPartialRecord()); + + return Optional.of(new RecordQueryFetchFromPartialRecordPlan(Quantifier.physical(memoizer.memoizePlan(coveringIndexPlan)), + coveringIndexPlan::pushValueThroughFetch, baseRecordType, RecordQueryFetchFromPartialRecordPlan.FetchIndexRecords.PRIMARY_KEY)); + } + + @Nonnull + @Override + public Optional pushValueThroughFetch(@Nonnull final Value toBePushedValue, + @Nonnull final CorrelationIdentifier sourceAlias, + @Nonnull final CorrelationIdentifier targetAlias) { + final var indexEntryToLogicalRecord = + getIndexEntryToLogicalRecordMaybe().orElseThrow(() -> new RecordCoreException("need index entry to logical record")); + + return ScanWithFetchMatchCandidate.pushValueThroughFetch(toBePushedValue, + baseAlias, + sourceAlias, + targetAlias, + Iterables.concat(indexEntryToLogicalRecord.getLogicalKeyValues(), + indexEntryToLogicalRecord.getLogicalValueValues())); + } + + @Nonnull + private static ScanComparisons toScanComparisons(@Nonnull final List comparisonRanges) { + ScanComparisons.Builder builder = new ScanComparisons.Builder(); + for (ComparisonRange comparisonRange : comparisonRanges) { + builder.addComparisonRange(comparisonRange); + } + return builder.build(); + } + + @Nonnull + public Placeholder getDistancePredicatePlaceholder() { + return distancePredicatePlaceholder; + } +} diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/expressions/SelectExpression.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/expressions/SelectExpression.java index 2dfcdbb3d5..ad50724bcb 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/expressions/SelectExpression.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/expressions/SelectExpression.java @@ -630,7 +630,7 @@ public PlannerGraph rewriteInternalPlannerGraph(@Nonnull final List 30 ? String.format("%02x", predicateString.hashCode()) : predicateString; return PlannerGraph.fromNodeAndChildGraphs( new PlannerGraph.LogicalOperatorNode(this, "SELECT " + resultValue, diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/predicates/RangeConstraints.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/predicates/RangeConstraints.java index d0fe0d7103..13a9a54427 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/predicates/RangeConstraints.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/predicates/RangeConstraints.java @@ -763,6 +763,9 @@ private boolean canBeUsedInScanPrefix(@Nonnull final Comparisons.Comparison comp case STARTS_WITH: case NOT_NULL: case IS_NULL: + case DISTANCE_RANK_EQUALS: + case DISTANCE_RANK_LESS_THAN: + case DISTANCE_RANK_LESS_THAN_OR_EQUAL: return true; case TEXT_CONTAINS_ALL: case TEXT_CONTAINS_ALL_WITHIN: diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/rules/AbstractDataAccessRule.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/rules/AbstractDataAccessRule.java index c98ed6dd73..6ee0e32b70 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/rules/AbstractDataAccessRule.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/rules/AbstractDataAccessRule.java @@ -45,12 +45,14 @@ import com.apple.foundationdb.record.query.plan.cascades.RequestedOrdering; import com.apple.foundationdb.record.query.plan.cascades.RequestedOrderingConstraint; import com.apple.foundationdb.record.query.plan.cascades.ValueIndexScanMatchCandidate; +import com.apple.foundationdb.record.query.plan.cascades.VectorIndexScanMatchCandidate; import com.apple.foundationdb.record.query.plan.cascades.WithPrimaryKeyMatchCandidate; import com.apple.foundationdb.record.query.plan.cascades.debug.Debugger; import com.apple.foundationdb.record.query.plan.cascades.expressions.LogicalDistinctExpression; import com.apple.foundationdb.record.query.plan.cascades.expressions.LogicalIntersectionExpression; import com.apple.foundationdb.record.query.plan.cascades.expressions.RelationalExpression; import com.apple.foundationdb.record.query.plan.cascades.matching.structure.BindingMatcher; +import com.apple.foundationdb.record.query.plan.cascades.predicates.Placeholder; import com.apple.foundationdb.record.query.plan.cascades.properties.CardinalitiesProperty.Cardinality; import com.apple.foundationdb.record.query.plan.cascades.values.Value; import com.apple.foundationdb.record.query.plan.plans.RecordQueryIntersectionPlan; @@ -490,6 +492,16 @@ private static List> maximumCoverageMatches(@Nonnu final var isContained = findContainingAccess(singleMatchedAccesses, outerAccess); + final var partialMatch = outerAccess.getPartialMatch(); + final var probeBoundPlaceholders = partialMatch.getBoundPlaceholders().stream().map(Placeholder::getParameterAlias).collect(ImmutableSet.toImmutableSet()); + if (partialMatch.getMatchCandidate() instanceof VectorIndexScanMatchCandidate) { + final var vectorMatchCandidate = (VectorIndexScanMatchCandidate)partialMatch.getMatchCandidate(); + final var predicatePlaceHolder = vectorMatchCandidate.getDistancePredicatePlaceholder(); + if (!probeBoundPlaceholders.contains(predicatePlaceHolder.getParameterAlias())) { + continue; + } + } + if (!isContained) { maximumCoverageMatchesBuilder.add( Vectored.of(outerAccess, index)); diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java index b9a2ac21bb..bf23706de7 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/typing/Type.java @@ -162,6 +162,10 @@ default boolean isRecord() { return getTypeCode().equals(TypeCode.RECORD); } + default boolean isVector() { + return getTypeCode().equals(TypeCode.VECTOR); + } + /** * Checks whether a {@link Type} is {@link Relation}. * diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/ArithmeticValue.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/ArithmeticValue.java index d7bb5d2ddd..c8bd1e4cb4 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/ArithmeticValue.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/ArithmeticValue.java @@ -546,21 +546,6 @@ public enum PhysicalOperator { BITMAP_BIT_POSITION_II(LogicalOperator.BITMAP_BIT_POSITION, TypeCode.INT, TypeCode.INT, TypeCode.INT, (l, r) -> Math.subtractExact((int)l, Math.multiplyExact(Math.floorDiv((int)l, (int)r), (int)r))), EUCLIDEAN_DISTANCE_VV(LogicalOperator.EUCLIDEAN_DISTANCE, TypeCode.VECTOR, TypeCode.VECTOR, TypeCode.DOUBLE, ((l, r) -> new Metric.EuclideanMetric().distance(((Vector)l).toDoubleVector().getData(), ((Vector)r).toDoubleVector().getData()))), -// EUCLIDEAN_DISTANCE_VA(LogicalOperator.EUCLIDEAN_DISTANCE, TypeCode.VECTOR, TypeCode.ARRAY, TypeCode.DOUBLE, ((l, r) -> { -// final List rhs = (List)r; -// final Vector lhs = (Vector)l; -// Verify.verify(rhs.size() == lhs.size()); -// if (rhs.isEmpty()) { -// return 0.0d; -// } -// final var doubleArray = new Double[rhs.size()]; -// for (int i = 0; i < rhs.size(); i++) { -// final var item = rhs.get(i); -// Verify.verify(item instanceof Number); -// doubleArray[i] = ((Number)item).doubleValue(); -// } -// return new Metric.CosineMetric().distance(((Vector)l).toDoubleVector().getData(), doubleArray); -// })), COSINE_DISTANCE_VV(LogicalOperator.COSINE_DISTANCE, TypeCode.VECTOR, TypeCode.VECTOR, TypeCode.DOUBLE, ((l, r) -> new Metric.CosineMetric().distance(((Vector)l).toDoubleVector().getData(), ((Vector)r).toDoubleVector().getData()))); @Nonnull diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/EuclideanDistanceRankValue.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/EuclideanDistanceRankValue.java new file mode 100644 index 0000000000..c186a05640 --- /dev/null +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/EuclideanDistanceRankValue.java @@ -0,0 +1,114 @@ +/* + * EuclideanDistanceRankValue.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2015-2022 Apple Inc. and the FoundationDB project authors + * + * Licensed 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 com.apple.foundationdb.record.query.plan.cascades.values; + +import com.apple.foundationdb.annotation.API; +import com.apple.foundationdb.record.ObjectPlanHash; +import com.apple.foundationdb.record.PlanDeserializer; +import com.apple.foundationdb.record.PlanSerializationContext; +import com.apple.foundationdb.record.planprotos.PEuclideanDistanceRank; +import com.apple.foundationdb.record.planprotos.PValue; +import com.apple.foundationdb.record.query.plan.cascades.typing.Type; +import com.google.auto.service.AutoService; + +import javax.annotation.Nonnull; +import java.util.Objects; + +/** + * A windowed value that computes the RANK of a Euclidean distance of a list of expressions which can optionally be + * partitioned by expressions defining a window. + */ +@API(API.Status.EXPERIMENTAL) +public class EuclideanDistanceRankValue extends WindowedValue implements Value.IndexOnlyValue { + private static final String NAME = "EuclideanDistanceRank"; + private static final ObjectPlanHash BASE_HASH = new ObjectPlanHash(NAME + "-Value"); + + public EuclideanDistanceRankValue(@Nonnull final PlanSerializationContext serializationContext, + @Nonnull final PEuclideanDistanceRank rankValueProto) { + super(serializationContext, Objects.requireNonNull(rankValueProto.getSuper())); + } + + public EuclideanDistanceRankValue(@Nonnull Iterable partitioningValues, + @Nonnull Iterable argumentValues) { + super(partitioningValues, argumentValues); + } + + @Nonnull + @Override + public String getName() { + return NAME; + } + + @Override + public int planHash(@Nonnull final PlanHashMode mode) { + return basePlanHash(mode, BASE_HASH); + } + + @Nonnull + @Override + public Type getResultType() { + return Type.primitiveType(Type.TypeCode.LONG); + } + + @Nonnull + @Override + public EuclideanDistanceRankValue withChildren(final Iterable newChildren) { + final var childrenPair = splitNewChildren(newChildren); + return new EuclideanDistanceRankValue(childrenPair.getKey(), childrenPair.getValue()); + } + + @Nonnull + @Override + public PEuclideanDistanceRank toProto(@Nonnull final PlanSerializationContext serializationContext) { + return PEuclideanDistanceRank.newBuilder().setSuper(toWindowedValueProto(serializationContext)).build(); + } + + @Nonnull + @Override + public PValue toValueProto(@Nonnull final PlanSerializationContext serializationContext) { + return PValue.newBuilder().setEuclideanDistanceRank(toProto(serializationContext)).build(); + } + + @Nonnull + public static EuclideanDistanceRankValue fromProto(@Nonnull final PlanSerializationContext serializationContext, + @Nonnull final PEuclideanDistanceRank rankValueProto) { + return new EuclideanDistanceRankValue(serializationContext, rankValueProto); + } + + /** + * Deserializer. + */ + @AutoService(PlanDeserializer.class) + public static class Deserializer implements PlanDeserializer { + @Nonnull + @Override + public Class getProtoMessageClass() { + return PEuclideanDistanceRank.class; + } + + @Nonnull + @Override + public EuclideanDistanceRankValue fromProto(@Nonnull final PlanSerializationContext serializationContext, + @Nonnull final PEuclideanDistanceRank euclideanDistanceRank) { + return EuclideanDistanceRankValue.fromProto(serializationContext, euclideanDistanceRank); + } + } +} diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/FieldValue.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/FieldValue.java index 9d8d5ba679..d48caa1d93 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/FieldValue.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/FieldValue.java @@ -22,6 +22,7 @@ import com.apple.foundationdb.annotation.API; import com.apple.foundationdb.annotation.SpotBugsSuppressWarnings; +import com.apple.foundationdb.async.hnsw.Vector; import com.apple.foundationdb.record.EvaluationContext; import com.apple.foundationdb.record.ObjectPlanHash; import com.apple.foundationdb.record.PlanDeserializer; @@ -188,6 +189,10 @@ private static Object unwrapPrimitive(@Nonnull Type type, @Nullable Object field } else if (type.isUuid()) { Verify.verify(fieldValue instanceof UUID); return fieldValue; + } else if (type.isVector()) { + Verify.verify(fieldValue instanceof ByteString); + final var byteString = (ByteString) fieldValue; + return Vector.fromBytes(byteString.toByteArray(), ((Type.Vector)type).getPrecision()); } else { // This also may need to turn ByteString's into byte[] for Type.TypeCode.BYTES return fieldValue; diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RankValue.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RankValue.java index 78cac98e1d..8468147be4 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RankValue.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RankValue.java @@ -26,7 +26,10 @@ import com.apple.foundationdb.record.PlanSerializationContext; import com.apple.foundationdb.record.planprotos.PRankValue; import com.apple.foundationdb.record.planprotos.PValue; +import com.apple.foundationdb.record.query.expressions.Comparisons; import com.apple.foundationdb.record.query.plan.cascades.BuiltInFunction; +import com.apple.foundationdb.record.query.plan.cascades.predicates.QueryPredicate; +import com.apple.foundationdb.record.query.plan.cascades.predicates.ValuePredicate; import com.apple.foundationdb.record.query.plan.cascades.typing.Type; import com.apple.foundationdb.record.query.plan.cascades.typing.Typed; import com.google.auto.service.AutoService; @@ -36,6 +39,7 @@ import javax.annotation.Nonnull; import java.util.List; import java.util.Objects; +import java.util.Optional; /** * A windowed value that computes the RANK of a list of expressions which can optionally be partitioned by expressions @@ -62,6 +66,61 @@ public String getName() { return NAME; } + @Nonnull + @Override + public Optional absorbUpperPredicate(@Nonnull final Comparisons.Type comparisonType, @Nonnull final Value comparand) { + // this enables the rank predicate to consume an upper predicate to produce a distance rank comparison with the + // correct comparison ranges. + // Currently, only the following tree structure below is matched and transformed: + // + // RelOpComparison (<=) (supported comparison operators: <=, <, and =). + // / \ + // / \ => ValuePredicate(EDR(Prtn, FV'), DRVC(<=, Cov')) + // Rank(Prtn, [ED(Fv', CoV')]) 42' + // + // ED=EuclideanDistance ArithmeticValue + // EDR=EuclideanDistanceRank + // DRVC=DistanceRangeValueComparison + if (!(comparand instanceof ConstantObjectValue)) { + return Optional.empty(); + } + if (getArgumentValues().size() > 1) { + return Optional.empty(); + } + final var argument = getArgumentValues().get(0); + if (!(argument instanceof ArithmeticValue)) { + return Optional.empty(); + } + final var arithmeticValue = (ArithmeticValue)argument; + if (arithmeticValue.getLogicalOperator() != ArithmeticValue.LogicalOperator.EUCLIDEAN_DISTANCE) { + return Optional.empty(); + } + final var euclideanDistanceArgs = ImmutableList.copyOf(arithmeticValue.getChildren()); + final var firstArg = euclideanDistanceArgs.get(0); + final var secondArg = euclideanDistanceArgs.get(1); + if (!(firstArg instanceof FieldValue && (secondArg instanceof ConstantObjectValue || secondArg instanceof LiteralValue))) { + return Optional.empty(); + } + + final Comparisons.Type distanceRankComparisonType; + switch (comparisonType) { + case EQUALS: + distanceRankComparisonType = Comparisons.Type.DISTANCE_RANK_EQUALS; + break; + case LESS_THAN: + distanceRankComparisonType = Comparisons.Type.DISTANCE_RANK_LESS_THAN; + break; + case LESS_THAN_OR_EQUALS: + distanceRankComparisonType = Comparisons.Type.DISTANCE_RANK_LESS_THAN_OR_EQUAL; + break; + default: + return Optional.empty(); + } + final var value = new EuclideanDistanceRankValue(getPartitioningValues(), ImmutableList.of(firstArg)); + final var comparison = new Comparisons.DistanceRankValueComparison(distanceRankComparisonType, secondArg, comparand); + return Optional.of(new ValuePredicate(value, comparison)); + } + @Override public int planHash(@Nonnull final PlanHashMode mode) { return basePlanHash(mode, BASE_HASH); @@ -128,6 +187,7 @@ public RankFn() { } @Nonnull + @SuppressWarnings("PMD.UnusedFormalParameter") private static RankValue encapsulateInternal(@Nonnull BuiltInFunction builtInFunction, @Nonnull final List arguments) { Verify.verify(arguments.size() == 2); // ordering expressions must be present, ok for now. diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RecordConstructorValue.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RecordConstructorValue.java index 7839dd254c..01ea2d2f1e 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RecordConstructorValue.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RecordConstructorValue.java @@ -22,6 +22,7 @@ import com.apple.foundationdb.annotation.API; import com.apple.foundationdb.annotation.SpotBugsSuppressWarnings; +import com.apple.foundationdb.async.hnsw.Vector; import com.apple.foundationdb.record.EvaluationContext; import com.apple.foundationdb.record.ObjectPlanHash; import com.apple.foundationdb.record.PlanDeserializer; @@ -224,6 +225,8 @@ private static Object protoObjectForPrimitive(@Nonnull Type type, @Nonnull Objec } } else if (type.getTypeCode() == Type.TypeCode.VERSION) { return ZeroCopyByteString.wrap(((FDBRecordVersion)field).toBytes(false)); + } else if (type.getTypeCode() == Type.TypeCode.VECTOR) { + return ZeroCopyByteString.wrap(((Vector)field).getRawData()); } return field; } diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RelOpValue.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RelOpValue.java index 0ab1f9c693..802183eeaf 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RelOpValue.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/RelOpValue.java @@ -155,6 +155,24 @@ public Optional toQueryPredicate(@Nullable final TypeRepository // can be correlated (or not) to anything except the innermostAlias final Value leftChild = it.next(); final Value rightChild = it.next(); + + // this gives either child a chance of consuming current predicate tree, effectively doing a rotation-like + // transformation: + // ParentPredicate + // / \ => LeftComparand.transform(ParentPredicate, RightComparand) + // LeftComparand RightComparand + var absorbedMaybe = leftChild.absorbUpperPredicate(comparisonType, rightChild); + if (absorbedMaybe.isPresent()) { + return absorbedMaybe; + } + final var invertedComparison = Comparisons.invertComparisonType(comparisonType); + if (invertedComparison != null) { + absorbedMaybe = rightChild.absorbUpperPredicate(invertedComparison, leftChild); + if (absorbedMaybe.isPresent()) { + return absorbedMaybe; + } + } + final Set leftChildCorrelatedTo = leftChild.getCorrelatedTo(); final Set rightChildCorrelatedTo = rightChild.getCorrelatedTo(); diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/Value.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/Value.java index 94f840e713..3f2e8a8250 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/Value.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/Value.java @@ -731,6 +731,12 @@ default Optional> match })); } + @Nonnull + default Optional absorbUpperPredicate(@Nonnull final Comparisons.Type comparisonType, + @Nonnull final Value comparand) { + return Optional.empty(); + } + @Nonnull PValue toValueProto(@Nonnull PlanSerializationContext serializationContext); diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/plans/RecordQueryIndexPlan.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/plans/RecordQueryIndexPlan.java index bce54f1450..ef7c2e848b 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/plans/RecordQueryIndexPlan.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/plans/RecordQueryIndexPlan.java @@ -59,6 +59,7 @@ import com.apple.foundationdb.record.provider.foundationdb.IndexScanRange; import com.apple.foundationdb.record.provider.foundationdb.MultidimensionalIndexScanComparisons; import com.apple.foundationdb.record.provider.foundationdb.UnsupportedRemoteFetchIndexException; +import com.apple.foundationdb.record.provider.foundationdb.VectorIndexScanComparisons; import com.apple.foundationdb.record.query.plan.AvailableFields; import com.apple.foundationdb.record.query.plan.QueryPlanConstraint; import com.apple.foundationdb.record.query.plan.ScanComparisons; @@ -576,12 +577,24 @@ public void logPlanStructure(StoreTimer timer) { @Override public boolean hasScanComparisons() { - return scanParameters instanceof IndexScanComparisons; + return (scanParameters instanceof IndexScanComparisons) || (scanParameters instanceof VectorIndexScanComparisons); } @Nonnull @Override public ScanComparisons getScanComparisons() { + if (scanParameters instanceof VectorIndexScanComparisons) { + final var vectorIndexComparisons = (VectorIndexScanComparisons)scanParameters; + final var builder = new ScanComparisons.Builder(); + builder.addAll(vectorIndexComparisons.getPrefixScanComparisons()); + if (vectorIndexComparisons.getDistanceRankValueComparison().getType().isEquality()) { + builder.addEqualityComparison(vectorIndexComparisons.getDistanceRankValueComparison()); + } else { + builder.addInequalityComparison(vectorIndexComparisons.getDistanceRankValueComparison()); + } + builder.addAll(vectorIndexComparisons.getSuffixScanComparisons()); + return builder.build(); + } if (scanParameters instanceof IndexScanComparisons) { return ((IndexScanComparisons)scanParameters).getComparisons(); } else { diff --git a/fdb-record-layer-core/src/main/proto/record_query_plan.proto b/fdb-record-layer-core/src/main/proto/record_query_plan.proto index c2fe657786..ffd66124f1 100644 --- a/fdb-record-layer-core/src/main/proto/record_query_plan.proto +++ b/fdb-record-layer-core/src/main/proto/record_query_plan.proto @@ -262,6 +262,7 @@ message PValue { PRangeValue range_value = 48; PFirstOrDefaultStreamingValue first_or_default_streaming_value = 49; PEvaluatesToValue evaluates_to_value = 50; + PEuclideanDistanceRank euclidean_distance_rank = 51; } } @@ -644,6 +645,10 @@ message PRankValue { optional PWindowedValue super = 1; } +message PEuclideanDistanceRank { + optional PWindowedValue super = 1; +} + message PRecordConstructorValue { message PColumn { optional PType.PRecordType.PField field = 1; diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/MessageTuple.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/MessageTuple.java index 18ad52c85b..b782779a63 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/MessageTuple.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/MessageTuple.java @@ -21,6 +21,8 @@ package com.apple.foundationdb.relational.recordlayer; import com.apple.foundationdb.annotation.API; +import com.apple.foundationdb.async.hnsw.Vector; +import com.apple.foundationdb.record.RecordMetaDataOptionsProto; import com.apple.foundationdb.record.TupleFieldsProto; import com.apple.foundationdb.record.metadata.expressions.TupleFieldsHelper; import com.apple.foundationdb.relational.api.exceptions.InvalidColumnReferenceException; @@ -50,13 +52,19 @@ public Object getObject(int position) throws InvalidColumnReferenceException { throw InvalidColumnReferenceException.getExceptionForInvalidPositionNumber(position); } Descriptors.FieldDescriptor fieldDescriptor = message.getDescriptorForType().getFields().get(position); + final var recordTypeOptions = fieldDescriptor.getOptions().getExtension(RecordMetaDataOptionsProto.field); + final var fieldValue = message.getField(message.getDescriptorForType().getFields().get(position)); + if (recordTypeOptions.hasVectorOptions()) { + final var precision = recordTypeOptions.getVectorOptions().getPrecision(); + final var byteStringFieldValue = (ByteString)fieldValue; + return Vector.fromBytes(byteStringFieldValue.toByteArray(), precision); + } if (fieldDescriptor.isRepeated()) { - final var list = (List) message.getField(message.getDescriptorForType().getFields().get(position)); + final var list = (List) fieldValue; return list.stream().map(MessageTuple::sanitizeField).collect(Collectors.toList()); } if (message.hasField(fieldDescriptor)) { - final var field = message.getField(message.getDescriptorForType().getFields().get(position)); - return sanitizeField(field); + return sanitizeField(fieldValue); } else { return null; } diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java index 3c5d69f47e..84aeb85ff7 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/BaseVisitor.java @@ -1399,7 +1399,7 @@ public Expression visitAggregateFunctionCall(@Nonnull RelationalParser.Aggregate @Nonnull @Override - public Expression visitNonAggregateFunctionCall(final RelationalParser.NonAggregateFunctionCallContext ctx) { + public Expression visitNonAggregateFunctionCall(@Nonnull RelationalParser.NonAggregateFunctionCallContext ctx) { return expressionVisitor.visitNonAggregateFunctionCall(ctx) ; } @@ -1537,13 +1537,13 @@ public Identifier visitWindowName(@Nonnull RelationalParser.WindowNameContext ct @Nonnull @Override - public Object visitWindowSpec(final RelationalParser.WindowSpecContext ctx) { + public Object visitWindowSpec(@Nonnull final RelationalParser.WindowSpecContext ctx) { return visitChildren(ctx); } @Nonnull @Override - public Expressions visitPartitionClause(final RelationalParser.PartitionClauseContext ctx) { + public Expressions visitPartitionClause(@Nonnull final RelationalParser.PartitionClauseContext ctx) { return expressionVisitor.visitPartitionClause(ctx); } diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java index c3a6474ee7..7844283deb 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DelegatingVisitor.java @@ -1239,7 +1239,7 @@ public Expression visitAggregateFunctionCall(@Nonnull RelationalParser.Aggregate @Nonnull @Override - public Expression visitNonAggregateFunctionCall(final RelationalParser.NonAggregateFunctionCallContext ctx) { + public Expression visitNonAggregateFunctionCall(@Nonnull final RelationalParser.NonAggregateFunctionCallContext ctx) { return getDelegate().visitNonAggregateFunctionCall(ctx); } @@ -1377,13 +1377,13 @@ public Identifier visitWindowName(@Nonnull RelationalParser.WindowNameContext ct @Nonnull @Override - public Object visitWindowSpec(final RelationalParser.WindowSpecContext ctx) { + public Object visitWindowSpec(@Nonnull final RelationalParser.WindowSpecContext ctx) { return getDelegate().visitWindowSpec(ctx); } @Nonnull @Override - public Expressions visitPartitionClause(final RelationalParser.PartitionClauseContext ctx) { + public Expressions visitPartitionClause(@Nonnull final RelationalParser.PartitionClauseContext ctx) { return getDelegate().visitPartitionClause(ctx); } diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/ExpressionVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/ExpressionVisitor.java index 39d330c00b..fcc956e5fb 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/ExpressionVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/ExpressionVisitor.java @@ -264,7 +264,7 @@ public Expression visitAggregateFunctionCall(@Nonnull RelationalParser.Aggregate @Nonnull @Override - public Expression visitNonAggregateFunctionCall(final RelationalParser.NonAggregateFunctionCallContext ctx) { + public Expression visitNonAggregateFunctionCall(@Nonnull final RelationalParser.NonAggregateFunctionCallContext ctx) { return visitNonAggregateWindowedFunction(ctx.nonAggregateWindowedFunction()); } @@ -284,8 +284,7 @@ public Expression visitNonAggregateWindowedFunction(@Nonnull final RelationalPar final var orderByArray = AbstractArrayConstructorValue.LightArrayConstructorValue.of(orderByValues, Type.any()); final var arguments = Expressions.of(ImmutableList.of(Expression.ofUnnamed(partitionArray), Expression.ofUnnamed(orderByArray))); - final var result = getDelegate().getSemanticAnalyzer().resolveScalarFunction("rank", arguments, false); - return result; + return getDelegate().getSemanticAnalyzer().resolveScalarFunction("rank", arguments, false); } @Nonnull @@ -380,7 +379,7 @@ public NonnullPair> visitOverClause(@Nonnul @Nonnull @Override - public Expressions visitPartitionClause(final RelationalParser.PartitionClauseContext ctx) { + public Expressions visitPartitionClause(@Nonnull final RelationalParser.PartitionClauseContext ctx) { return Expressions.of(ctx.expression().stream().map(expContext -> Assert.castUnchecked(visit(expContext), Expression.class)).collect(ImmutableList.toImmutableList())); } diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java index 2ab3136a34..ac8bc05503 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/TypedVisitor.java @@ -776,7 +776,7 @@ public interface TypedVisitor extends RelationalParserVisitor { @Nonnull @Override - Expression visitNonAggregateFunctionCall(final RelationalParser.NonAggregateFunctionCallContext ctx); + Expression visitNonAggregateFunctionCall(@Nonnull RelationalParser.NonAggregateFunctionCallContext ctx); @Nonnull @Override @@ -868,11 +868,11 @@ public interface TypedVisitor extends RelationalParserVisitor { @Nonnull @Override - Object visitWindowSpec(final RelationalParser.WindowSpecContext ctx); + Object visitWindowSpec(@Nonnull RelationalParser.WindowSpecContext ctx); @Nonnull @Override - Expressions visitPartitionClause(RelationalParser.PartitionClauseContext ctx); + Expressions visitPartitionClause(@Nonnull RelationalParser.PartitionClauseContext ctx); @Nonnull @Override diff --git a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java index 6bee4c648a..4322359fd4 100644 --- a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java +++ b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java @@ -20,12 +20,14 @@ package com.apple.foundationdb.relational.recordlayer.query; +import com.apple.foundationdb.async.hnsw.Vector; import com.apple.foundationdb.relational.api.StructResultSetMetaData; import com.apple.foundationdb.relational.api.metadata.DataType; import com.apple.foundationdb.relational.recordlayer.EmbeddedRelationalExtension; import com.apple.foundationdb.relational.recordlayer.Utils; import com.apple.foundationdb.relational.utils.Ddl; import org.assertj.core.api.Assertions; +import org.assertj.core.data.Offset; import org.junit.jupiter.api.Order; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; @@ -75,15 +77,24 @@ void vectorTest(@Nonnull final String ddlType, @Nonnull final DataType expectedT @Test void selectFromHnsw() throws Exception { final String schemaTemplate = "create table photos(zone string, recordId string, name string," + - "embedding vector(3), primary key (zone, recordId), organized by hnsw(embedding partition by zone + 3, name) " + + "embedding vector(3), primary key (zone, recordId), organized by hnsw(embedding partition by zone, name) " + "with (hnsw_m = 10, hnsw_ef_construction = 5))"; try (var ddl = Ddl.builder().database(URI.create("/TEST/QT")).relationalExtension(relationalExtension).schemaTemplate(schemaTemplate).build()) { try (var statement = ddl.setSchemaAndGetConnection().createStatement()) { - final var result = statement.execute("SELECT * FROM photos WHERE RANK() OVER (PARTITION BY zone ORDER BY euclidean_distance(embedding, vector(1.2, -0.5, 3.14)) DESC) < 10"); - final var metadata = statement.getResultSet().getMetaData(); - Assertions.assertThat(metadata).isInstanceOf(StructResultSetMetaData.class); - final var relationalMetadata = (StructResultSetMetaData)metadata; - final var type = relationalMetadata.getRelationalDataType().getFields().get(1).getType(); + statement.executeUpdate("insert into photos values ('1', '100', 'DarthVader', vector(1.2h, -0.3H, 3.14H))"); + statement.execute("SELECT * FROM photos WHERE zone = '1' and name = 'DarthVader' and RANK() OVER (PARTITION BY zone, " + + "name ORDER BY euclidean_distance(embedding, vector(1.2H, -0.5H, 3.14H)) DESC) < 10"); + final var resultSet = statement.getResultSet(); + resultSet.next(); + Assertions.assertThat(resultSet.getString(1)).isEqualTo("1"); + Assertions.assertThat(resultSet.getString(2)).isEqualTo("100"); + Assertions.assertThat(resultSet.getString(3)).isEqualTo("DarthVader"); + Assertions.assertThat(resultSet.getObject(4)).isInstanceOf(Vector.HalfVector.class); + final var halfVector = (Vector.HalfVector)resultSet.getObject(4); + Assertions.assertThat(halfVector.getData().length).isEqualTo(3); + Assertions.assertThat(halfVector.getData()[0].floatValue()).isCloseTo(1.2f, Offset.offset(0.01f)); + Assertions.assertThat(halfVector.getData()[1].floatValue()).isCloseTo(-0.3f, Offset.offset(0.01f)); + Assertions.assertThat(halfVector.getData()[2].floatValue()).isCloseTo(3.14f, Offset.offset(0.01f)); } } } From 97c3368da4e0bcc3e4f6af0a7f73e55f4c7c392b Mon Sep 17 00:00:00 2001 From: Youssef Hatem Date: Mon, 18 Aug 2025 17:27:19 +0100 Subject: [PATCH 32/34] fixes. - add test for prepared inserts. --- .../record/query/expressions/Comparisons.java | 3 +- .../indexes/VectorIndexTestBase.java | 4 +- .../recordlayer/query/VectorTypeTest.java | 40 ++++++++++++++++++- .../src/test/java/YamlIntegrationTests.java | 2 +- 4 files changed, 43 insertions(+), 6 deletions(-) diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java index 731f499f09..d5daa2b1f0 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/expressions/Comparisons.java @@ -1938,8 +1938,7 @@ public PComparison toComparisonProto(@Nonnull final PlanSerializationContext ser @Nullable public Vector getVector(@Nullable final FDBRecordStoreBase store, final @Nullable EvaluationContext context) { - final Object comparand = getComparand(store, context); - return comparand == null ? null : Vector.HalfVector.halfVectorFromBytes((byte[])comparand); + return (Vector)getComparand(store, context); } public int getLimit(@Nullable final FDBRecordStoreBase store, final @Nullable EvaluationContext context) { diff --git a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java index c03f3691f3..072a4eeb05 100644 --- a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java +++ b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexTestBase.java @@ -213,7 +213,7 @@ void basicWriteIndexReadTest() throws Exception { final DistanceRankValueComparison distanceRankComparison = new DistanceRankValueComparison(Comparisons.Type.DISTANCE_RANK_LESS_THAN_OR_EQUAL, new LiteralValue<>(Type.Vector.of(false, 16, 128), - randomVectorData(random, 128)), + Vector.HalfVector.fromBytes(randomVectorData(random, 128), 16)), new LiteralValue<>(10)); final VectorIndexScanComparisons vectorIndexScanComparisons = @@ -251,7 +251,7 @@ void basicWriteIndexReadGroupedTest() throws Exception { final DistanceRankValueComparison distanceRankComparison = new DistanceRankValueComparison(Comparisons.Type.DISTANCE_RANK_LESS_THAN_OR_EQUAL, new LiteralValue<>(Type.Vector.of(false, 16, 128), - randomVectorData(random, 128)), + Vector.HalfVector.fromBytes(randomVectorData(random, 128), 16)), new LiteralValue<>(10)); final VectorIndexScanComparisons vectorIndexScanComparisons = diff --git a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java index 4322359fd4..4a62c829bd 100644 --- a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java +++ b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java @@ -20,12 +20,15 @@ package com.apple.foundationdb.relational.recordlayer.query; +import com.apple.foundationdb.async.hnsw.HNSWHelpers; import com.apple.foundationdb.async.hnsw.Vector; import com.apple.foundationdb.relational.api.StructResultSetMetaData; import com.apple.foundationdb.relational.api.metadata.DataType; import com.apple.foundationdb.relational.recordlayer.EmbeddedRelationalExtension; import com.apple.foundationdb.relational.recordlayer.Utils; import com.apple.foundationdb.relational.utils.Ddl; +import com.apple.foundationdb.relational.utils.SchemaTemplateRule; +import com.christianheina.langx.half4j.Half; import org.assertj.core.api.Assertions; import org.assertj.core.data.Offset; import org.junit.jupiter.api.Order; @@ -79,9 +82,44 @@ void selectFromHnsw() throws Exception { final String schemaTemplate = "create table photos(zone string, recordId string, name string," + "embedding vector(3), primary key (zone, recordId), organized by hnsw(embedding partition by zone, name) " + "with (hnsw_m = 10, hnsw_ef_construction = 5))"; - try (var ddl = Ddl.builder().database(URI.create("/TEST/QT")).relationalExtension(relationalExtension).schemaTemplate(schemaTemplate).build()) { + try (var ddl = Ddl.builder().database(URI.create("/TEST/QT")).relationalExtension(relationalExtension).schemaTemplate(schemaTemplate).schemaTemplateOptions((new SchemaTemplateRule.SchemaTemplateOptions(true, true))).build()) { try (var statement = ddl.setSchemaAndGetConnection().createStatement()) { statement.executeUpdate("insert into photos values ('1', '100', 'DarthVader', vector(1.2h, -0.3H, 3.14H))"); + } + try (var statement = ddl.setSchemaAndGetConnection().createStatement()) { + statement.execute("SELECT * FROM photos WHERE zone = '1' and name = 'DarthVader' and RANK() OVER (PARTITION BY zone, " + + "name ORDER BY euclidean_distance(embedding, vector(1.2H, -0.5H, 3.14H)) DESC) < 10"); + final var resultSet = statement.getResultSet(); + resultSet.next(); + Assertions.assertThat(resultSet.getString(1)).isEqualTo("1"); + Assertions.assertThat(resultSet.getString(2)).isEqualTo("100"); + Assertions.assertThat(resultSet.getString(3)).isEqualTo("DarthVader"); + Assertions.assertThat(resultSet.getObject(4)).isInstanceOf(Vector.HalfVector.class); + final var halfVector = (Vector.HalfVector)resultSet.getObject(4); + Assertions.assertThat(halfVector.getData().length).isEqualTo(3); + Assertions.assertThat(halfVector.getData()[0].floatValue()).isCloseTo(1.2f, Offset.offset(0.01f)); + Assertions.assertThat(halfVector.getData()[1].floatValue()).isCloseTo(-0.3f, Offset.offset(0.01f)); + Assertions.assertThat(halfVector.getData()[2].floatValue()).isCloseTo(3.14f, Offset.offset(0.01f)); + } + } + } + + @Test + void insertPreparedVector() throws Exception { + final String schemaTemplate = "create table photos(zone string, recordId string, name string," + + "embedding vector(3), primary key (zone, recordId), organized by hnsw(embedding partition by zone, name) " + + "with (hnsw_m = 10, hnsw_ef_construction = 5))"; + try (var ddl = Ddl.builder().database(URI.create("/TEST/QT")).relationalExtension(relationalExtension).schemaTemplate(schemaTemplate).schemaTemplateOptions((new SchemaTemplateRule.SchemaTemplateOptions(true, true))).build()) { + try (var statement = ddl.setSchemaAndGetConnection().prepareStatement("insert into photos values (?, ?, ?, ?)")) { + statement.setString(1, "1"); + statement.setString(2, "100"); + statement.setString(3, "DarthVader"); + + final Half[] componentData = new Half[] {HNSWHelpers.halfValueOf(1.2f), HNSWHelpers.halfValueOf(-0.3f), HNSWHelpers.halfValueOf(3.14f)}; + statement.setObject(4, new Vector.HalfVector(componentData)); + statement.executeUpdate(); + } + try (var statement = ddl.setSchemaAndGetConnection().createStatement()) { statement.execute("SELECT * FROM photos WHERE zone = '1' and name = 'DarthVader' and RANK() OVER (PARTITION BY zone, " + "name ORDER BY euclidean_distance(embedding, vector(1.2H, -0.5H, 3.14H)) DESC) < 10"); final var resultSet = statement.getResultSet(); diff --git a/yaml-tests/src/test/java/YamlIntegrationTests.java b/yaml-tests/src/test/java/YamlIntegrationTests.java index e877d72973..b18575bdb8 100644 --- a/yaml-tests/src/test/java/YamlIntegrationTests.java +++ b/yaml-tests/src/test/java/YamlIntegrationTests.java @@ -38,7 +38,7 @@ public class YamlIntegrationTests { public void showcasingTests(YamlTest.Runner runner) throws Exception { runner.runYamsql("showcasing-tests.yamsql"); } - + @TestTemplate public void groupByTests(YamlTest.Runner runner) throws Exception { runner.runYamsql("groupby-tests.yamsql"); From 015979bfe6e89d6cf4d8ba1158072904c02a3c95 Mon Sep 17 00:00:00 2001 From: Youssef Hatem Date: Tue, 19 Aug 2025 14:31:19 +0100 Subject: [PATCH 33/34] - add HNSW_EF_SEARCH to query options. --- .../record/ExecuteProperties.java | 80 +++++++++++++++---- .../foundationdb/record/metadata/Index.java | 19 +++++ .../plan/plans/RecordQueryIndexPlan.java | 4 +- .../src/main/proto/record_query_plan.proto | 3 + .../foundationdb/relational/api/Options.java | 9 ++- .../src/main/antlr/RelationalLexer.g4 | 1 + .../src/main/antlr/RelationalParser.g4 | 1 + .../recordlayer/query/AstNormalizer.java | 3 + .../recordlayer/query/QueryPlan.java | 13 ++- .../recordlayer/query/VectorTypeTest.java | 33 ++++++++ 10 files changed, 145 insertions(+), 21 deletions(-) diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/ExecuteProperties.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/ExecuteProperties.java index 825fe56d3e..e8cdda25d9 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/ExecuteProperties.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/ExecuteProperties.java @@ -22,11 +22,15 @@ import com.apple.foundationdb.annotation.API; import com.apple.foundationdb.ReadTransaction; +import com.google.common.base.Verify; +import com.google.common.collect.ImmutableMap; import javax.annotation.Nonnull; import javax.annotation.Nullable; import java.util.ArrayList; import java.util.List; +import java.util.Map; +import java.util.Optional; /** * Limits on the execution of a query. @@ -37,6 +41,7 @@ *
  • limit on number of key-value pairs scanned
    • * */ +@SuppressWarnings("OptionalUsedAsFieldOrParameterType") @API(API.Status.UNSTABLE) public class ExecuteProperties { /** @@ -75,9 +80,13 @@ public class ExecuteProperties { private final CursorStreamingMode defaultCursorStreamingMode; + @Nonnull + private final Optional> overriddenIndexOptions; + @SuppressWarnings("java:S107") private ExecuteProperties(int skip, int rowLimit, @Nonnull IsolationLevel isolationLevel, long timeLimit, - @Nonnull ExecuteState state, boolean failOnScanLimitReached, @Nonnull CursorStreamingMode defaultCursorStreamingMode, boolean isDryRun) { + @Nonnull ExecuteState state, boolean failOnScanLimitReached, @Nonnull CursorStreamingMode defaultCursorStreamingMode, + boolean isDryRun, @Nonnull final Optional> overriddenIndexOptions) { this.skip = skip; this.rowLimit = rowLimit; this.isolationLevel = isolationLevel; @@ -86,6 +95,7 @@ private ExecuteProperties(int skip, int rowLimit, @Nonnull IsolationLevel isolat this.failOnScanLimitReached = failOnScanLimitReached; this.defaultCursorStreamingMode = defaultCursorStreamingMode; this.isDryRun = isDryRun; + this.overriddenIndexOptions = overriddenIndexOptions; } @Nonnull @@ -102,7 +112,7 @@ public ExecuteProperties setSkip(final int skip) { if (skip == this.skip) { return this; } - return copy(skip, rowLimit, timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun); + return copy(skip, rowLimit, timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun, overriddenIndexOptions); } public boolean isDryRun() { @@ -114,9 +124,31 @@ public ExecuteProperties setDryRun(final boolean isDryRun) { if (isDryRun == this.isDryRun) { return this; } - return copy(skip, rowLimit, timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun); + return copy(skip, rowLimit, timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun, overriddenIndexOptions); + } + + public boolean hasOverriddenIndexOptions() { + return overriddenIndexOptions.isPresent(); } + @Nonnull + public Map getOverriddenIndexOptions() { + Verify.verify(overriddenIndexOptions.isPresent()); + return overriddenIndexOptions.get(); + } + + @Nonnull + public Optional> getOverriddenIndexOptionsMaybe() { + return overriddenIndexOptions; + } + + public ExecuteProperties setOverriddenIndexOptions(@Nonnull final Map overriddenIndexOptions) { + Verify.verify(!overriddenIndexOptions.isEmpty()); + if (this.overriddenIndexOptions.isPresent() && this.overriddenIndexOptions.get().equals(overriddenIndexOptions)) { + return this; + } + return copy(skip, rowLimit, timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun, Optional.of(overriddenIndexOptions)); + } /** * Get the limit on the number of rows that will be returned as it would be passed to FDB. @@ -137,7 +169,7 @@ public ExecuteProperties setReturnedRowLimit(final int rowLimit) { if (newLimit == this.rowLimit) { return this; } - return copy(skip, newLimit, timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun); + return copy(skip, newLimit, timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun, overriddenIndexOptions); } /** @@ -184,7 +216,7 @@ public ExecuteState getState() { */ @Nonnull public ExecuteProperties setState(@Nonnull ExecuteState newState) { - return copy(skip, rowLimit, timeLimit, isolationLevel, newState, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun); + return copy(skip, rowLimit, timeLimit, isolationLevel, newState, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun, overriddenIndexOptions); } /** @@ -193,7 +225,7 @@ public ExecuteProperties setState(@Nonnull ExecuteState newState) { */ @Nonnull public ExecuteProperties clearState() { - return copy(skip, rowLimit, timeLimit, isolationLevel, new ExecuteState(), failOnScanLimitReached, defaultCursorStreamingMode, isDryRun); + return copy(skip, rowLimit, timeLimit, isolationLevel, new ExecuteState(), failOnScanLimitReached, defaultCursorStreamingMode, isDryRun, overriddenIndexOptions); } /** @@ -209,7 +241,7 @@ public ExecuteProperties setFailOnScanLimitReached(boolean failOnScanLimitReache if (failOnScanLimitReached == this.failOnScanLimitReached) { return this; } - return copy(skip, rowLimit, timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun); + return copy(skip, rowLimit, timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun, overriddenIndexOptions); } @Nonnull @@ -217,7 +249,7 @@ public ExecuteProperties clearReturnedRowLimit() { if (getReturnedRowLimit() == ReadTransaction.ROW_LIMIT_UNLIMITED) { return this; } - return copy(skip, ReadTransaction.ROW_LIMIT_UNLIMITED, timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun); + return copy(skip, ReadTransaction.ROW_LIMIT_UNLIMITED, timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun, overriddenIndexOptions); } /** @@ -229,7 +261,7 @@ public ExecuteProperties clearRowAndTimeLimits() { if (getTimeLimit() == UNLIMITED_TIME && getReturnedRowLimit() == ReadTransaction.ROW_LIMIT_UNLIMITED) { return this; } - return copy(skip, ReadTransaction.ROW_LIMIT_UNLIMITED, UNLIMITED_TIME, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun); + return copy(skip, ReadTransaction.ROW_LIMIT_UNLIMITED, UNLIMITED_TIME, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun, overriddenIndexOptions); } /** @@ -241,7 +273,7 @@ public ExecuteProperties clearSkipAndLimit() { if (skip == 0 && rowLimit == ReadTransaction.ROW_LIMIT_UNLIMITED) { return this; } - return copy(0, ReadTransaction.ROW_LIMIT_UNLIMITED, timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun); + return copy(0, ReadTransaction.ROW_LIMIT_UNLIMITED, timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun, overriddenIndexOptions); } /** @@ -254,7 +286,7 @@ public ExecuteProperties clearSkipAndAdjustLimit() { return this; } return copy(0, rowLimit == ReadTransaction.ROW_LIMIT_UNLIMITED ? ReadTransaction.ROW_LIMIT_UNLIMITED : rowLimit + skip, - timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun); + timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun, overriddenIndexOptions); } /** @@ -305,7 +337,7 @@ public ExecuteProperties setDefaultCursorStreamingMode(CursorStreamingMode defau if (defaultCursorStreamingMode == this.defaultCursorStreamingMode) { return this; } - return copy(skip, rowLimit, timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun); + return copy(skip, rowLimit, timeLimit, isolationLevel, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun, overriddenIndexOptions); } /** @@ -315,7 +347,7 @@ public ExecuteProperties setDefaultCursorStreamingMode(CursorStreamingMode defau */ @Nonnull public ExecuteProperties resetState() { - return copy(skip, rowLimit, timeLimit, isolationLevel, state.reset(), failOnScanLimitReached, defaultCursorStreamingMode, isDryRun); + return copy(skip, rowLimit, timeLimit, isolationLevel, state.reset(), failOnScanLimitReached, defaultCursorStreamingMode, isDryRun, overriddenIndexOptions); } /** @@ -328,13 +360,16 @@ public ExecuteProperties resetState() { * @param failOnScanLimitReached fail on scan limit reached * @param defaultCursorStreamingMode default streaming mode * @param isDryRun whether it is dry run + * @param overriddenIndexOptions overridden index options * @return a new properties with the given fields changed and other fields copied from this properties */ @SuppressWarnings("java:S107") @Nonnull protected ExecuteProperties copy(int skip, int rowLimit, long timeLimit, @Nonnull IsolationLevel isolationLevel, - @Nonnull ExecuteState state, boolean failOnScanLimitReached, CursorStreamingMode defaultCursorStreamingMode, boolean isDryRun) { - return new ExecuteProperties(skip, rowLimit, isolationLevel, timeLimit, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun); + @Nonnull ExecuteState state, boolean failOnScanLimitReached, CursorStreamingMode defaultCursorStreamingMode, + boolean isDryRun, @Nonnull final Optional> overriddenIndexOptions) { + return new ExecuteProperties(skip, rowLimit, isolationLevel, timeLimit, state, failOnScanLimitReached, defaultCursorStreamingMode, + isDryRun, overriddenIndexOptions); } @Nonnull @@ -408,6 +443,7 @@ public static class Builder { private ExecuteState executeState = null; private boolean failOnScanLimitReached = false; private boolean isDryRun = false; + private Optional> overriddenIndexOptions = Optional.empty(); private CursorStreamingMode defaultCursorStreamingMode = CursorStreamingMode.ITERATOR; private Builder() { @@ -455,6 +491,12 @@ public Builder setDryRun(boolean isDryRun) { return this; } + @Nonnull + public Builder setOverriddenIndexOptions(@Nonnull final Map overriddenIndexOptions) { + this.overriddenIndexOptions = Optional.of(ImmutableMap.copyOf(overriddenIndexOptions)); + return this; + } + @Nonnull public Builder setReturnedRowLimit(int rowLimit) { this.rowLimit = validateAndNormalizeRowLimit(rowLimit); @@ -485,6 +527,12 @@ public Builder clearSkipAndAdjustLimit() { return this; } + @Nonnull + public Builder clearOverriddenIndexOptions() { + this.overriddenIndexOptions = Optional.empty(); + return this; + } + @Nonnull public Builder setTimeLimit(long timeLimit) { this.timeLimit = validateAndNormalizeTimeLimit(timeLimit); @@ -607,7 +655,7 @@ public ExecuteProperties build() { } else { state = new ExecuteState(RecordScanLimiterFactory.enforce(scannedRecordsLimit), ByteScanLimiterFactory.enforce(scannedBytesLimit)); } - return new ExecuteProperties(skip, rowLimit, isolationLevel, timeLimit, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun); + return new ExecuteProperties(skip, rowLimit, isolationLevel, timeLimit, state, failOnScanLimitReached, defaultCursorStreamingMode, isDryRun, overriddenIndexOptions); } } } diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/Index.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/Index.java index 0a689d7c20..1850e7b1c6 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/Index.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/metadata/Index.java @@ -190,6 +190,25 @@ public Index(@Nonnull Index orig, @Nullable final IndexPredicate predicate) { this.lastModifiedVersion = orig.lastModifiedVersion; } + /** + * Copy constructor. This will create an index that is identical to the current Index with a given + * set of index options. + * @param orig original index to copy + * @param indexOptions the index options. + */ + public Index(@Nonnull Index orig, @Nonnull final Map indexOptions) { + this(orig.name, orig.rootExpression, orig.type, ImmutableMap.copyOf(indexOptions), orig.predicate); + if (orig.primaryKeyComponentPositions != null) { + this.primaryKeyComponentPositions = Arrays.copyOf(orig.primaryKeyComponentPositions, orig.primaryKeyComponentPositions.length); + } else { + this.primaryKeyComponentPositions = null; + } + this.subspaceKey = normalizeSubspaceKey(name, orig.subspaceKey); + this.useExplicitSubspaceKey = orig.useExplicitSubspaceKey; + this.addedVersion = orig.addedVersion; + this.lastModifiedVersion = orig.lastModifiedVersion; + } + @SuppressWarnings({"deprecation", "squid:CallToDeprecatedMethod", "java:S3776"}) // Old (deprecated) index type needs grouping compatibility @SpotBugsSuppressWarnings("NP_NONNULL_FIELD_NOT_INITIALIZED_IN_CONSTRUCTOR") public Index(@Nonnull RecordMetaDataProto.Index proto) throws KeyExpression.DeserializationException { diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/plans/RecordQueryIndexPlan.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/plans/RecordQueryIndexPlan.java index ef7c2e848b..c068fdbf55 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/plans/RecordQueryIndexPlan.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/plans/RecordQueryIndexPlan.java @@ -302,7 +302,9 @@ private RecordCursor executeUsingRemoteFetch(@N public RecordCursor executeEntries(@Nonnull FDBRecordStoreBase store, @Nonnull EvaluationContext context, @Nullable byte[] continuation, @Nonnull ExecuteProperties executeProperties) { final RecordMetaData metaData = store.getRecordMetaData(); - final Index index = metaData.getIndex(indexName); + final Index index = executeProperties.getOverriddenIndexOptionsMaybe() + .map(overriddenOptions -> new Index(metaData.getIndex(indexName), overriddenOptions)) + .orElse(metaData.getIndex(indexName)); final IndexScanBounds scanBounds = scanParameters.bind(store, index, context); if (!IndexScanType.BY_VALUE_OVER_SCAN.equals(getScanType())) { return store.scanIndex(index, scanBounds, continuation, executeProperties.asScanProperties(reverse)); diff --git a/fdb-record-layer-core/src/main/proto/record_query_plan.proto b/fdb-record-layer-core/src/main/proto/record_query_plan.proto index ffd66124f1..cf723df58b 100644 --- a/fdb-record-layer-core/src/main/proto/record_query_plan.proto +++ b/fdb-record-layer-core/src/main/proto/record_query_plan.proto @@ -409,6 +409,9 @@ message PArithmeticValue { BITMAP_BUCKET_NUMBER_II = 105; BITMAP_BIT_POSITION_LI = 106; BITMAP_BIT_POSITION_II = 107; + + EUCLIDEAN_DISTANCE_VV = 108; + COSINE_DISTANCE_VV = 109; } optional PPhysicalOperator operator = 1; diff --git a/fdb-relational-api/src/main/java/com/apple/foundationdb/relational/api/Options.java b/fdb-relational-api/src/main/java/com/apple/foundationdb/relational/api/Options.java index 210ce835b8..cafcaf9f42 100644 --- a/fdb-relational-api/src/main/java/com/apple/foundationdb/relational/api/Options.java +++ b/fdb-relational-api/src/main/java/com/apple/foundationdb/relational/api/Options.java @@ -221,7 +221,12 @@ public enum Name { * operations interacting with FDB. * Scope: Engine */ - ASYNC_OPERATIONS_TIMEOUT_MILLIS + ASYNC_OPERATIONS_TIMEOUT_MILLIS, + + /** + * Sets the EF_SEARCH parameter that controls the search space of queries leveraging the HNSW index. + */ + EF_SEARCH } public enum IndexFetchMethod { @@ -258,6 +263,7 @@ public enum IndexFetchMethod { builder.put(Name.CASE_SENSITIVE_IDENTIFIERS, false); builder.put(Name.CONTINUATIONS_CONTAIN_COMPILED_STATEMENTS, true); builder.put(Name.ASYNC_OPERATIONS_TIMEOUT_MILLIS, 10_000L); + builder.put(Name.EF_SEARCH, 100); OPTIONS_DEFAULT_VALUES = builder.build(); } @@ -425,6 +431,7 @@ private static Map> makeContracts() { data.put(Name.VALID_PLAN_HASH_MODES, List.of(TypeContract.stringType())); data.put(Name.CONTINUATIONS_CONTAIN_COMPILED_STATEMENTS, List.of(TypeContract.booleanType())); data.put(Name.ASYNC_OPERATIONS_TIMEOUT_MILLIS, List.of(TypeContract.longType(), RangeContract.of(0L, Long.MAX_VALUE))); + data.put(Name.EF_SEARCH, List.of(TypeContract.intType(), RangeContract.of(1, 100))); return Collections.unmodifiableMap(data); } diff --git a/fdb-relational-core/src/main/antlr/RelationalLexer.g4 b/fdb-relational-core/src/main/antlr/RelationalLexer.g4 index b85df08791..eb2da6b58f 100644 --- a/fdb-relational-core/src/main/antlr/RelationalLexer.g4 +++ b/fdb-relational-core/src/main/antlr/RelationalLexer.g4 @@ -513,6 +513,7 @@ HNSW_M: 'HNSW_M'; HNSW_MMAX: 'HNSW_MMAX'; HNSW_MMAX0: 'HNSW_MMAX0'; HNSW_EF_CONSTRUCTION: 'HNSW_EF_CONSTRUCTION'; +HNSW_EF_SEARCH: 'HNSW_EF_SEARCH'; IDENTIFIED: 'IDENTIFIED'; IGNORE_SERVER_IDS: 'IGNORE_SERVER_IDS'; IMPORT: 'IMPORT'; diff --git a/fdb-relational-core/src/main/antlr/RelationalParser.g4 b/fdb-relational-core/src/main/antlr/RelationalParser.g4 index fb8ba98963..27355a00e2 100644 --- a/fdb-relational-core/src/main/antlr/RelationalParser.g4 +++ b/fdb-relational-core/src/main/antlr/RelationalParser.g4 @@ -544,6 +544,7 @@ queryOption | LOG QUERY | DRY RUN | CONTINUATION CONTAINS COMPILED STATEMENT + | key=HNSW_EF_SEARCH '=' value=DECIMAL_LITERAL ; // Transaction's Statements diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/AstNormalizer.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/AstNormalizer.java index d434fb8b11..a43e682bac 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/AstNormalizer.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/AstNormalizer.java @@ -304,6 +304,9 @@ public Object visitQueryOption(@Nonnull RelationalParser.QueryOptionContext ctx) if (ctx.CONTINUATION() != null) { queryOptions.withOption(Options.Name.CONTINUATIONS_CONTAIN_COMPILED_STATEMENTS, true); } + if (ctx.HNSW_EF_SEARCH() != null) { + queryOptions.withOption(Options.Name.EF_SEARCH, ParseHelpers.parseDecimal(ctx.value.getText())); + } return null; } catch (SQLException e) { throw ExceptionUtil.toRelationalException(e).toUncheckedWrappedException(); diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/QueryPlan.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/QueryPlan.java index 47ea0ffa34..d56fbc6dee 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/QueryPlan.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/QueryPlan.java @@ -27,6 +27,7 @@ import com.apple.foundationdb.record.RecordCoreException; import com.apple.foundationdb.record.RecordCursor; import com.apple.foundationdb.record.RecordMetaData; +import com.apple.foundationdb.record.metadata.IndexOptions; import com.apple.foundationdb.record.provider.foundationdb.FDBRecordStoreBase; import com.apple.foundationdb.record.query.IndexQueryabilityFilter; import com.apple.foundationdb.record.query.plan.QueryPlanConstraint; @@ -79,6 +80,7 @@ import com.google.common.base.Suppliers; import com.google.common.base.Verify; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.protobuf.InvalidProtocolBufferException; import com.google.protobuf.Message; @@ -369,10 +371,15 @@ private RelationalResultSet executePhysicalPlan(@Nonnull final RecordLayerSchema validatePlanAgainstEnvironment(parsedContinuation, fdbRecordStore, executionContext, OptionsUtils.getValidPlanHashModes(options)); final RecordCursor cursor; - final var executeProperties = connection.getExecuteProperties().toBuilder() + final var executePropertiesBuilder = connection.getExecuteProperties().toBuilder() .setReturnedRowLimit(options.getOption(Options.Name.MAX_ROWS)) - .setDryRun(options.getOption(Options.Name.DRY_RUN)) - .build(); + .setDryRun(options.getOption(Options.Name.DRY_RUN)); + + final int efSearch = options.getOption(Options.Name.EF_SEARCH); + if (!options.getOption(Options.Name.EF_SEARCH).equals(Options.defaultOptions().get(Options.Name.EF_SEARCH))) { + executePropertiesBuilder.setOverriddenIndexOptions(ImmutableMap.of(IndexOptions.HNSW_EF_SEARCH, String.valueOf(efSearch))); + } + final var executeProperties = executePropertiesBuilder.build(); cursor = executionContext.metricCollector.clock( RelationalMetric.RelationalEvent.EXECUTE_RECORD_QUERY_PLAN, () -> recordQueryPlan.executePlan(fdbRecordStore, evaluationContext, parsedContinuation.getExecutionState(), diff --git a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java index 4a62c829bd..7d5a20cabf 100644 --- a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java +++ b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java @@ -136,4 +136,37 @@ void insertPreparedVector() throws Exception { } } } + + @Test + void insertPreparedVectorWithIndexHint() throws Exception { + final String schemaTemplate = "create table photos(zone string, recordId string, name string," + + "embedding vector(3), primary key (zone, recordId), organized by hnsw(embedding partition by zone, name) " + + "with (hnsw_m = 10, hnsw_ef_construction = 5))"; + try (var ddl = Ddl.builder().database(URI.create("/TEST/QT")).relationalExtension(relationalExtension).schemaTemplate(schemaTemplate).schemaTemplateOptions((new SchemaTemplateRule.SchemaTemplateOptions(true, true))).build()) { + try (var statement = ddl.setSchemaAndGetConnection().prepareStatement("insert into photos values (?, ?, ?, ?)")) { + statement.setString(1, "1"); + statement.setString(2, "100"); + statement.setString(3, "DarthVader"); + + final Half[] componentData = new Half[] {HNSWHelpers.halfValueOf(1.2f), HNSWHelpers.halfValueOf(-0.3f), HNSWHelpers.halfValueOf(3.14f)}; + statement.setObject(4, new Vector.HalfVector(componentData)); + statement.executeUpdate(); + } + try (var statement = ddl.setSchemaAndGetConnection().createStatement()) { + statement.execute("SELECT * FROM photos WHERE zone = '1' and name = 'DarthVader' and RANK() OVER (PARTITION BY zone, " + + "name ORDER BY euclidean_distance(embedding, vector(1.2H, -0.5H, 3.14H)) DESC) < 10 options (HNsw_ef_search = 50)"); + final var resultSet = statement.getResultSet(); + resultSet.next(); + Assertions.assertThat(resultSet.getString(1)).isEqualTo("1"); + Assertions.assertThat(resultSet.getString(2)).isEqualTo("100"); + Assertions.assertThat(resultSet.getString(3)).isEqualTo("DarthVader"); + Assertions.assertThat(resultSet.getObject(4)).isInstanceOf(Vector.HalfVector.class); + final var halfVector = (Vector.HalfVector)resultSet.getObject(4); + Assertions.assertThat(halfVector.getData().length).isEqualTo(3); + Assertions.assertThat(halfVector.getData()[0].floatValue()).isCloseTo(1.2f, Offset.offset(0.01f)); + Assertions.assertThat(halfVector.getData()[1].floatValue()).isCloseTo(-0.3f, Offset.offset(0.01f)); + Assertions.assertThat(halfVector.getData()[2].floatValue()).isCloseTo(3.14f, Offset.offset(0.01f)); + } + } + } } From 812a3203c8612a9b31b27d8e0b2e75eb588f86f9 Mon Sep 17 00:00:00 2001 From: Youssef Hatem Date: Tue, 19 Aug 2025 15:46:34 +0100 Subject: [PATCH 34/34] - minor fixes for ef_search. --- .../indexes/VectorIndexMaintainer.java | 2 +- .../recordlayer/query/AstNormalizer.java | 1 + .../recordlayer/query/VectorTypeTest.java | 18 +++++++++++++++++- 3 files changed, 19 insertions(+), 2 deletions(-) diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java index b48abdedd5..2c55ee91f1 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/indexes/VectorIndexMaintainer.java @@ -176,7 +176,7 @@ private CompletableFuture> kNearestNeighborSearch(@Null @Nonnull final HNSW hnsw, @Nonnull final ReadTransaction transaction, @Nonnull final VectorIndexScanBounds vectorIndexScanBounds) { - return hnsw.kNearestNeighborsSearch(transaction, vectorIndexScanBounds.getAdjustedLimit(), 100, + return hnsw.kNearestNeighborsSearch(transaction, vectorIndexScanBounds.getAdjustedLimit(), config.getEfSearch(), Objects.requireNonNull(vectorIndexScanBounds.getQueryVector()).toHalfVector()) .thenApply(nearestNeighbors -> { final ImmutableList.Builder nearestNeighborEntriesBuilder = ImmutableList.builder(); diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/AstNormalizer.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/AstNormalizer.java index a43e682bac..9e59a6d0da 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/AstNormalizer.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/AstNormalizer.java @@ -306,6 +306,7 @@ public Object visitQueryOption(@Nonnull RelationalParser.QueryOptionContext ctx) } if (ctx.HNSW_EF_SEARCH() != null) { queryOptions.withOption(Options.Name.EF_SEARCH, ParseHelpers.parseDecimal(ctx.value.getText())); + queryCachingFlags.add(NormalizationResult.QueryCachingFlags.WITH_NO_CACHE_OPTION); } return null; } catch (SQLException e) { diff --git a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java index 7d5a20cabf..8ffc16e733 100644 --- a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java +++ b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/VectorTypeTest.java @@ -154,7 +154,23 @@ void insertPreparedVectorWithIndexHint() throws Exception { } try (var statement = ddl.setSchemaAndGetConnection().createStatement()) { statement.execute("SELECT * FROM photos WHERE zone = '1' and name = 'DarthVader' and RANK() OVER (PARTITION BY zone, " + - "name ORDER BY euclidean_distance(embedding, vector(1.2H, -0.5H, 3.14H)) DESC) < 10 options (HNsw_ef_search = 50)"); + "name ORDER BY euclidean_distance(embedding, vector(1.2H, -0.5H, 3.14H)) DESC) < 10 options (HNsw_ef_search = 94)"); + final var resultSet = statement.getResultSet(); + resultSet.next(); + Assertions.assertThat(resultSet.getString(1)).isEqualTo("1"); + Assertions.assertThat(resultSet.getString(2)).isEqualTo("100"); + Assertions.assertThat(resultSet.getString(3)).isEqualTo("DarthVader"); + Assertions.assertThat(resultSet.getObject(4)).isInstanceOf(Vector.HalfVector.class); + final var halfVector = (Vector.HalfVector)resultSet.getObject(4); + Assertions.assertThat(halfVector.getData().length).isEqualTo(3); + Assertions.assertThat(halfVector.getData()[0].floatValue()).isCloseTo(1.2f, Offset.offset(0.01f)); + Assertions.assertThat(halfVector.getData()[1].floatValue()).isCloseTo(-0.3f, Offset.offset(0.01f)); + Assertions.assertThat(halfVector.getData()[2].floatValue()).isCloseTo(3.14f, Offset.offset(0.01f)); + } + + try (var statement = ddl.setSchemaAndGetConnection().createStatement()) { + statement.execute("SELECT * FROM photos WHERE zone = '1' and name = 'DarthVader' and RANK() OVER (PARTITION BY zone, " + + "name ORDER BY euclidean_distance(embedding, vector(1.2H, -0.5H, 3.14H)) DESC) < 10 options (HNsw_ef_search = 56)"); final var resultSet = statement.getResultSet(); resultSet.next(); Assertions.assertThat(resultSet.getString(1)).isEqualTo("1");