save point -- in the middle of just mess

Author: normen662

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<N extends Neighbor> implements Node<N> {
     @Nonnull
     private final Tuple primaryKey;
 
-    @Nonnull
-    private final Vector<Half> vector;
-
     @Nonnull
     private final List<N> neighbors;
 
-    protected AbstractNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector<Half> vector,
+    protected AbstractNode(@Nonnull final Tuple primaryKey,
                            @Nonnull final List<N> neighbors) {
         this.primaryKey = primaryKey;
-        this.vector = vector;
         this.neighbors = ImmutableList.copyOf(neighbors);
     }
 
@@ -54,11 +49,6 @@ public Tuple getPrimaryKey() {
         return primaryKey;
     }
 
-    @Nonnull
-    public Vector<Half> getVector() {
-        return vector;
-    }
-
     @Nonnull
     @Override
     public List<N> getNeighbors() {

fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/ByNodeStorageAdapter.java

@@ -75,21 +75,22 @@ public ByNodeStorageAdapter(@Nonnull final HNSW.Config config, @Nonnull final Su
     }
 
     @Override
-    public CompletableFuture<NodeKeyWithLayer> fetchEntryNodeKey(@Nonnull final ReadTransaction readTransaction) {
+    public CompletableFuture<EntryPointAndLayer> fetchEntryNodeKey(@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 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<Node> 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 <N extends Neighbor> Node<N> nodeFromTuple(@Nonnull final Node.NodeCreator<N> 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 <N extends Neighbor> Node<N> dataNodeFromTuples(@Nonnull final Node.NodeCreator<N> creator,
+                                                            @Nonnull final Tuple primaryKey,
+                                                            @Nonnull final Tuple vectorTuple,
+                                                            @Nonnull final Tuple neighborsTuple) {
+        final Vector<Half> vector = vectorFromTuple(vectorTuple);
+
+        List<Neighbor> 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<NeighborWithVector> neighborsWithVectors = null;
+
+        return creator.create(NodeKind.DATA, primaryKey, vector, neighbors);
+    }
+
+    @Nonnull
+    private <N extends Neighbor> Node<N> intermediateNodeFromTuples(@Nonnull final Node.NodeCreator<N> creator,
+                                                                    @Nonnull final Tuple primaryKey,
+                                                                    @Nonnull final Tuple neighborsTuple) {
+        List<NeighborWithVector> neighborsWithVectors = Lists.newArrayListWithExpectedSize(neighborsTuple.size());
         Half[] neighborVectorHalfs = null;
-        List<Neighbor> 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<Half> 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 <S extends NodeSlot, N extends AbstractNode<S, N>> AbstractChangeSet<S, N>

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<Neighbor> {
+    @Nonnull
+    private final Vector<Half> vector;
+
     public DataNode(@Nonnull final Tuple primaryKey, @Nonnull final Vector<Half> vector,
                     @Nonnull final List<Neighbor> 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<Half> getVector() {
+        return vector;
+    }
+
     @Nonnull
     @Override
     public DataNode asDataNode() {
@@ -71,9 +80,9 @@ public NodeCreator<Neighbor> sameCreator() {
     @SuppressWarnings("unchecked")
     public static Node<Neighbor> creator(@Nonnull final NodeKind nodeKind,
                                          @Nonnull final Tuple primaryKey,
-                                         @Nonnull final Vector<Half> vector,
+                                         @Nullable final Vector<Half> vector,
                                          @Nonnull final List<? extends Neighbor> neighbors) {
         Verify.verify(nodeKind == NodeKind.INTERMEDIATE);
-        return new DataNode(primaryKey, vector, (List<Neighbor>)neighbors);
+        return new DataNode(primaryKey, Objects.requireNonNull(vector), (List<Neighbor>)neighbors);
     }
 }

fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayerAndDistance.java renamed to 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<Tuple> primaryKeys(@Nonnull List<Element> elements) {
+        return () -> elements.stream()
+                .map(Element::getPrimaryKey)
+                .iterator();
+    }
 }

fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/NodeKeyWithLayer.java renamed to 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<Half> vector;
 
-    public NodeKeyWithLayer(final int layer, @Nonnull final Tuple primaryKey) {
+    public EntryPointAndLayer(final int layer, @Nonnull final Tuple primaryKey, @Nonnull final Vector<Half> 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<Half> getVector() {
+        return vector;
+    }
 }

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());
+    }
+}