You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by be...@apache.org on 2025/05/27 19:38:04 UTC
(lucene) branch branch_10x updated: Reduce NeighborArray heap memory (#14527)
This is an automated email from the ASF dual-hosted git repository.
benwtrent pushed a commit to branch branch_10x
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/branch_10x by this push:
new be23f36590a Reduce NeighborArray heap memory (#14527)
be23f36590a is described below
commit be23f36590a581b91abc8edbc4e80f02bbcbbe7f
Author: weizijun <we...@gmail.com>
AuthorDate: Wed May 28 03:36:57 2025 +0800
Reduce NeighborArray heap memory (#14527)
* use FloatArrayList\IntArrayList to replace float[]\int[]
* use getScores(int i) to replace scores()
* add more tests
* add change log and change the init value
* update OnHeapHnswGraph ramBytesUsed method
* improve
* add MaxSizedIntArrayList
* add MaxSizedFloatArrayList
* add MaxSizedFloatArrayList
* fixed tests
* revert
---
lucene/CHANGES.txt | 2 +
.../internal/hppc/MaxSizedFloatArrayList.java | 102 ++++++++++
.../lucene/internal/hppc/MaxSizedIntArrayList.java | 102 ++++++++++
.../src/java/org/apache/lucene/util/ArrayUtil.java | 24 +++
.../apache/lucene/util/hnsw/HnswGraphBuilder.java | 9 +-
.../org/apache/lucene/util/hnsw/NeighborArray.java | 103 ++++++----
.../apache/lucene/util/hnsw/OnHeapHnswGraph.java | 47 +++--
.../internal/hppc/TestMaxSizedFloatArrayList.java | 178 +++++++++++++++++
.../internal/hppc/TestMaxSizedIntArrayList.java | 176 +++++++++++++++++
.../test/org/apache/lucene/util/TestArrayUtil.java | 30 +++
.../apache/lucene/util/hnsw/TestNeighborArray.java | 210 ++++++++++++++++++++-
11 files changed, 929 insertions(+), 54 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 51e83b7d730..186084731a8 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -47,6 +47,8 @@ Optimizations
* GITHUB#14447: Compute the doc range more efficiently when flushing doc block. (Pan Guixin)
+* GITHUB#14527: Reduce NeighborArray heap memory. (weizijun)
+
* GITHUB#14529, GITHUB#14555, GITHUB#14618: Impl intoBitset for IndexedDISI and Docvalues. (Guo Feng)
* GITHUB#14552: Speed up flush of softdelete by intoBitset. (Guo Feng)
diff --git a/lucene/core/src/java/org/apache/lucene/internal/hppc/MaxSizedFloatArrayList.java b/lucene/core/src/java/org/apache/lucene/internal/hppc/MaxSizedFloatArrayList.java
new file mode 100644
index 00000000000..31799af4ca7
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/internal/hppc/MaxSizedFloatArrayList.java
@@ -0,0 +1,102 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.internal.hppc;
+
+import static org.apache.lucene.internal.hppc.HashContainers.DEFAULT_EXPECTED_ELEMENTS;
+
+import org.apache.lucene.util.ArrayUtil;
+import org.apache.lucene.util.RamUsageEstimator;
+
+/**
+ * An array-backed list of {@code float} with a maximum size limit.
+ *
+ * @lucene.internal
+ */
+public class MaxSizedFloatArrayList extends FloatArrayList {
+ private static final long BASE_RAM_BYTES_USED =
+ RamUsageEstimator.shallowSizeOfInstance(MaxSizedFloatArrayList.class);
+
+ final int maxSize;
+
+ /** New instance with sane defaults. */
+ public MaxSizedFloatArrayList(int maxSize) {
+ this(maxSize, DEFAULT_EXPECTED_ELEMENTS);
+ }
+
+ /**
+ * New instance with sane defaults.
+ *
+ * @param maxSize The maximum size this list can grow to
+ * @param expectedElements The expected number of elements guaranteed not to cause buffer
+ * expansion (inclusive).
+ */
+ public MaxSizedFloatArrayList(int maxSize, int expectedElements) {
+ super(expectedElements);
+ assert expectedElements <= maxSize
+ : "expectedElements (" + expectedElements + ") must be <= maxSize (" + maxSize + ")";
+ this.maxSize = maxSize;
+ }
+
+ /** Creates a new list from the elements of another list in its iteration order. */
+ public MaxSizedFloatArrayList(MaxSizedFloatArrayList list) {
+ super(list.size());
+ this.maxSize = list.maxSize;
+ addAll(list);
+ }
+
+ @Override
+ protected void ensureBufferSpace(int expectedAdditions) {
+ if (elementsCount + expectedAdditions > maxSize) {
+ throw new IllegalStateException("Cannot grow beyond maxSize: " + maxSize);
+ }
+ if (elementsCount + expectedAdditions > buffer.length) {
+ this.buffer = ArrayUtil.growInRange(buffer, elementsCount + expectedAdditions, maxSize);
+ }
+ }
+
+ @Override
+ public int hashCode() {
+ int h = 1, max = elementsCount;
+ h = 31 * h + maxSize;
+ for (int i = 0; i < max; i++) {
+ h = 31 * h + BitMixer.mix(this.buffer[i]);
+ }
+ return h;
+ }
+
+ /**
+ * Returns <code>true</code> only if the other object is an instance of the same class and with
+ * the same elements and maxSize.
+ */
+ @Override
+ public boolean equals(Object obj) {
+ if (this == obj) {
+ return true;
+ }
+ if (obj == null || getClass() != obj.getClass()) {
+ return false;
+ }
+ MaxSizedFloatArrayList other = (MaxSizedFloatArrayList) obj;
+ return maxSize == other.maxSize && super.equals(obj);
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return BASE_RAM_BYTES_USED + RamUsageEstimator.sizeOf(buffer);
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/internal/hppc/MaxSizedIntArrayList.java b/lucene/core/src/java/org/apache/lucene/internal/hppc/MaxSizedIntArrayList.java
new file mode 100644
index 00000000000..8640b3fb40a
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/internal/hppc/MaxSizedIntArrayList.java
@@ -0,0 +1,102 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.internal.hppc;
+
+import static org.apache.lucene.internal.hppc.HashContainers.DEFAULT_EXPECTED_ELEMENTS;
+
+import org.apache.lucene.util.ArrayUtil;
+import org.apache.lucene.util.RamUsageEstimator;
+
+/**
+ * An array-backed list of {@code int} with a maximum size limit.
+ *
+ * @lucene.internal
+ */
+public class MaxSizedIntArrayList extends IntArrayList {
+ private static final long BASE_RAM_BYTES_USED =
+ RamUsageEstimator.shallowSizeOfInstance(MaxSizedIntArrayList.class);
+
+ final int maxSize;
+
+ /** New instance with sane defaults. */
+ public MaxSizedIntArrayList(int maxSize) {
+ this(maxSize, DEFAULT_EXPECTED_ELEMENTS);
+ }
+
+ /**
+ * New instance with sane defaults.
+ *
+ * @param maxSize The maximum size this list can grow to
+ * @param expectedElements The expected number of elements guaranteed not to cause buffer
+ * expansion (inclusive).
+ */
+ public MaxSizedIntArrayList(int maxSize, int expectedElements) {
+ super(expectedElements);
+ assert expectedElements <= maxSize
+ : "expectedElements (" + expectedElements + ") must be <= maxSize (" + maxSize + ")";
+ this.maxSize = maxSize;
+ }
+
+ /** Creates a new list from the elements of another list in its iteration order. */
+ public MaxSizedIntArrayList(MaxSizedIntArrayList list) {
+ super(list.size());
+ this.maxSize = list.maxSize;
+ addAll(list);
+ }
+
+ @Override
+ protected void ensureBufferSpace(int expectedAdditions) {
+ if (elementsCount + expectedAdditions > maxSize) {
+ throw new IllegalStateException("Cannot grow beyond maxSize: " + maxSize);
+ }
+ if (elementsCount + expectedAdditions > buffer.length) {
+ this.buffer = ArrayUtil.growInRange(buffer, elementsCount + expectedAdditions, maxSize);
+ }
+ }
+
+ @Override
+ public int hashCode() {
+ int h = 1, max = elementsCount;
+ h = 31 * h + maxSize;
+ for (int i = 0; i < max; i++) {
+ h = 31 * h + BitMixer.mix(this.buffer[i]);
+ }
+ return h;
+ }
+
+ /**
+ * Returns <code>true</code> only if the other object is an instance of the same class and with
+ * the same elements and maxSize.
+ */
+ @Override
+ public boolean equals(Object obj) {
+ if (this == obj) {
+ return true;
+ }
+ if (obj == null || getClass() != obj.getClass()) {
+ return false;
+ }
+ MaxSizedIntArrayList other = (MaxSizedIntArrayList) obj;
+ return maxSize == other.maxSize && super.equals(obj);
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return BASE_RAM_BYTES_USED + RamUsageEstimator.sizeOf(buffer);
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/util/ArrayUtil.java b/lucene/core/src/java/org/apache/lucene/util/ArrayUtil.java
index 2cd3cb63cfe..2fee4e52b8b 100644
--- a/lucene/core/src/java/org/apache/lucene/util/ArrayUtil.java
+++ b/lucene/core/src/java/org/apache/lucene/util/ArrayUtil.java
@@ -346,6 +346,30 @@ public final class ArrayUtil {
return growExact(array, Math.min(maxLength, potentialLength));
}
+ /**
+ * Returns an array whose size is at least {@code minLength} but not over {@code maxLength},
+ * growing exponentially if it needs to grow.
+ */
+ public static float[] growInRange(float[] array, int minLength, int maxLength) {
+ assert minLength >= 0
+ : "minLength must be positive (got " + minLength + "): likely integer overflow?";
+
+ if (minLength > maxLength) {
+ throw new IllegalArgumentException(
+ "requested minimum array length "
+ + minLength
+ + " is larger than requested maximum array length "
+ + maxLength);
+ }
+
+ if (array.length >= minLength) {
+ return array;
+ }
+
+ int potentialLength = oversize(minLength, Float.BYTES);
+ return growExact(array, Math.min(maxLength, potentialLength));
+ }
+
/**
* Returns an array whose size is at least {@code minSize}, generally over-allocating
* exponentially
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java
index e6569e3b6b2..2cc4cb061be 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java
@@ -370,13 +370,13 @@ public class HnswGraphBuilder implements HnswBuilder {
Lock lock = hnswLock.write(level, nbr);
try {
NeighborArray nbrsOfNbr = getGraph().getNeighbors(level, nbr);
- nbrsOfNbr.addAndEnsureDiversity(node, candidates.scores()[i], nbr, scorer);
+ nbrsOfNbr.addAndEnsureDiversity(node, candidates.getScores(i), nbr, scorer);
} finally {
lock.unlock();
}
} else {
NeighborArray nbrsOfNbr = hnsw.getNeighbors(level, nbr);
- nbrsOfNbr.addAndEnsureDiversity(node, candidates.scores()[i], nbr, scorer);
+ nbrsOfNbr.addAndEnsureDiversity(node, candidates.getScores(i), nbr, scorer);
}
}
}
@@ -397,7 +397,7 @@ public class HnswGraphBuilder implements HnswBuilder {
// compare each neighbor (in distance order) against the closer neighbors selected so far,
// only adding it if it is closer to the target than to any of the other selected neighbors
int cNode = candidates.nodes()[i];
- float cScore = candidates.scores()[i];
+ float cScore = candidates.getScores(i);
assert cNode <= hnsw.maxNodeId();
scorer.setScoringOrdinal(cNode);
if (diversityCheck(cScore, neighbors, scorer)) {
@@ -452,6 +452,7 @@ public class HnswGraphBuilder implements HnswBuilder {
// see: https://github.com/apache/lucene/issues/14214
// connectComponents();
frozen = true;
+ hnsw.finishBuild();
}
@SuppressWarnings("unused")
@@ -546,7 +547,7 @@ public class HnswGraphBuilder implements HnswBuilder {
// must subtract 1 here since the nodes array is one larger than the configured
// max neighbors (M / 2M).
// We should have taken care of this check by searching for not-full nodes
- int maxConn = nbr0.nodes().length - 1;
+ int maxConn = nbr0.maxSize() - 1;
assert notFullyConnected.get(n0);
assert nbr0.size() < maxConn : "node " + n0 + " is full, has " + nbr0.size() + " friends";
nbr0.addOutOfOrder(n1, score);
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/NeighborArray.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/NeighborArray.java
index c7745120a5f..07d96d353db 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/NeighborArray.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/NeighborArray.java
@@ -19,7 +19,10 @@ package org.apache.lucene.util.hnsw;
import java.io.IOException;
import java.util.Arrays;
-import org.apache.lucene.util.ArrayUtil;
+import org.apache.lucene.internal.hppc.MaxSizedFloatArrayList;
+import org.apache.lucene.internal.hppc.MaxSizedIntArrayList;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.RamUsageEstimator;
/**
* NeighborArray encodes the neighbors of a node and their mutual scores in the HNSW graph as a pair
@@ -29,16 +32,21 @@ import org.apache.lucene.util.ArrayUtil;
*
* @lucene.internal
*/
-public class NeighborArray {
+public class NeighborArray implements Accountable {
+ private static final long BASE_RAM_BYTES_USED =
+ RamUsageEstimator.shallowSizeOfInstance(NeighborArray.class);
+
private final boolean scoresDescOrder;
private int size;
- private final float[] scores;
- private final int[] nodes;
+ private final int maxSize;
+ private final MaxSizedFloatArrayList scores;
+ private final MaxSizedIntArrayList nodes;
private int sortedNodeSize;
public NeighborArray(int maxSize, boolean descOrder) {
- nodes = new int[maxSize];
- scores = new float[maxSize];
+ this.maxSize = maxSize;
+ nodes = new MaxSizedIntArrayList(maxSize, maxSize / 8);
+ scores = new MaxSizedFloatArrayList(maxSize, maxSize / 8);
this.scoresDescOrder = descOrder;
}
@@ -48,32 +56,32 @@ public class NeighborArray {
*/
public void addInOrder(int newNode, float newScore) {
assert size == sortedNodeSize : "cannot call addInOrder after addOutOfOrder";
- if (size == nodes.length) {
+ if (size == maxSize) {
throw new IllegalStateException("No growth is allowed");
}
if (size > 0) {
- float previousScore = scores[size - 1];
+ float previousScore = scores.get(size - 1);
assert ((scoresDescOrder && (previousScore >= newScore))
|| (scoresDescOrder == false && (previousScore <= newScore)))
: "Nodes are added in the incorrect order! Comparing "
+ newScore
+ " to "
- + Arrays.toString(ArrayUtil.copyOfSubArray(scores, 0, size));
+ + Arrays.toString(scores.toArray());
}
- nodes[size] = newNode;
- scores[size] = newScore;
+ nodes.add(newNode);
+ scores.add(newScore);
++size;
++sortedNodeSize;
}
/** Add node and newScore but do not insert as sorted */
public void addOutOfOrder(int newNode, float newScore) {
- if (size == nodes.length) {
+ if (size == maxSize) {
throw new IllegalStateException("No growth is allowed");
}
- scores[size] = newScore;
- nodes[size] = newNode;
+ nodes.add(newNode);
+ scores.add(newScore);
size++;
}
@@ -90,13 +98,13 @@ public class NeighborArray {
int newNode, float newScore, int nodeId, UpdateableRandomVectorScorer scorer)
throws IOException {
addOutOfOrder(newNode, newScore);
- if (size < nodes.length) {
+ if (size < maxSize) {
return;
}
// we're oversize, need to do diversity check and pop out the least diverse neighbour
scorer.setScoringOrdinal(nodeId);
removeIndex(findWorstNonDiverse(scorer));
- assert size == nodes.length - 1;
+ assert size == maxSize - 1;
}
/**
@@ -134,8 +142,8 @@ public class NeighborArray {
/** insert the first unsorted node into its sorted position */
private int insertSortedInternal(RandomVectorScorer scorer) throws IOException {
assert sortedNodeSize < size : "Call this method only when there's unsorted node";
- int tmpNode = nodes[sortedNodeSize];
- float tmpScore = scores[sortedNodeSize];
+ int tmpNode = nodes.get(sortedNodeSize);
+ float tmpScore = scores.get(sortedNodeSize);
if (Float.isNaN(tmpScore)) {
tmpScore = scorer.score(tmpNode);
@@ -146,11 +154,19 @@ public class NeighborArray {
? descSortFindRightMostInsertionPoint(tmpScore, sortedNodeSize)
: ascSortFindRightMostInsertionPoint(tmpScore, sortedNodeSize);
System.arraycopy(
- nodes, insertionPoint, nodes, insertionPoint + 1, sortedNodeSize - insertionPoint);
+ nodes.buffer,
+ insertionPoint,
+ nodes.buffer,
+ insertionPoint + 1,
+ sortedNodeSize - insertionPoint);
System.arraycopy(
- scores, insertionPoint, scores, insertionPoint + 1, sortedNodeSize - insertionPoint);
- nodes[insertionPoint] = tmpNode;
- scores[insertionPoint] = tmpScore;
+ scores.buffer,
+ insertionPoint,
+ scores.buffer,
+ insertionPoint + 1,
+ sortedNodeSize - insertionPoint);
+ nodes.buffer[insertionPoint] = tmpNode;
+ scores.buffer[insertionPoint] = tmpScore;
++sortedNodeSize;
return insertionPoint;
}
@@ -171,19 +187,29 @@ public class NeighborArray {
* @lucene.internal
*/
public int[] nodes() {
- return nodes;
+ return nodes.buffer;
}
- public float[] scores() {
- return scores;
+ /**
+ * Get the score at the given index
+ *
+ * @param i index of the score to get
+ * @return the score at the given index
+ */
+ public float getScores(int i) {
+ return scores.get(i);
}
public void clear() {
size = 0;
sortedNodeSize = 0;
+ nodes.clear();
+ scores.clear();
}
void removeLast() {
+ nodes.removeLast();
+ scores.removeLast();
size--;
sortedNodeSize = Math.min(sortedNodeSize, size);
}
@@ -193,8 +219,8 @@ public class NeighborArray {
removeLast();
return;
}
- System.arraycopy(nodes, idx + 1, nodes, idx, size - idx - 1);
- System.arraycopy(scores, idx + 1, scores, idx, size - idx - 1);
+ nodes.removeAt(idx);
+ scores.removeAt(idx);
if (idx < sortedNodeSize) {
sortedNodeSize--;
}
@@ -207,11 +233,11 @@ public class NeighborArray {
}
private int ascSortFindRightMostInsertionPoint(float newScore, int bound) {
- int insertionPoint = Arrays.binarySearch(scores, 0, bound, newScore);
+ int insertionPoint = Arrays.binarySearch(scores.buffer, 0, bound, newScore);
if (insertionPoint >= 0) {
// find the right most position with the same score
while ((insertionPoint < bound - 1)
- && (scores[insertionPoint + 1] == scores[insertionPoint])) {
+ && (scores.get(insertionPoint + 1) == scores.get(insertionPoint))) {
insertionPoint++;
}
insertionPoint++;
@@ -226,7 +252,7 @@ public class NeighborArray {
int end = bound - 1;
while (start <= end) {
int mid = (start + end) / 2;
- if (scores[mid] < newScore) end = mid - 1;
+ if (scores.get(mid) < newScore) end = mid - 1;
else start = mid + 1;
}
return start;
@@ -245,7 +271,7 @@ public class NeighborArray {
// no unchecked node left
break;
}
- scorer.setScoringOrdinal(nodes[i]);
+ scorer.setScoringOrdinal(nodes.get(i));
if (isWorstNonDiverse(i, uncheckedIndexes, uncheckedCursor, scorer)) {
return i;
}
@@ -259,11 +285,11 @@ public class NeighborArray {
private boolean isWorstNonDiverse(
int candidateIndex, int[] uncheckedIndexes, int uncheckedCursor, RandomVectorScorer scorer)
throws IOException {
- float minAcceptedSimilarity = scores[candidateIndex];
+ float minAcceptedSimilarity = scores.get(candidateIndex);
if (candidateIndex == uncheckedIndexes[uncheckedCursor]) {
// the candidate itself is unchecked
for (int i = candidateIndex - 1; i >= 0; i--) {
- float neighborSimilarity = scorer.score(nodes[i]);
+ float neighborSimilarity = scorer.score(nodes.get(i));
// candidate node is too similar to node i given its score relative to the base node
if (neighborSimilarity >= minAcceptedSimilarity) {
return true;
@@ -274,7 +300,7 @@ public class NeighborArray {
// inserted) unchecked nodes
assert candidateIndex > uncheckedIndexes[uncheckedCursor];
for (int i = uncheckedCursor; i >= 0; i--) {
- float neighborSimilarity = scorer.score(nodes[uncheckedIndexes[i]]);
+ float neighborSimilarity = scorer.score(nodes.get(uncheckedIndexes[i]));
// candidate node is too similar to node i given its score relative to the base node
if (neighborSimilarity >= minAcceptedSimilarity) {
return true;
@@ -283,4 +309,13 @@ public class NeighborArray {
}
return false;
}
+
+ public int maxSize() {
+ return maxSize;
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ return BASE_RAM_BYTES_USED + nodes.ramBytesUsed() + scores.ramBytesUsed();
+ }
}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/OnHeapHnswGraph.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/OnHeapHnswGraph.java
index 532bb8d08b1..f577dae0106 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/OnHeapHnswGraph.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/OnHeapHnswGraph.java
@@ -63,6 +63,8 @@ public final class OnHeapHnswGraph extends HnswGraph implements Accountable {
private int upto;
private NeighborArray cur;
+ private volatile long graphRamBytesUsed;
+
/**
* ctor
*
@@ -155,6 +157,15 @@ public final class OnHeapHnswGraph extends HnswGraph implements Accountable {
nonZeroLevelSize.incrementAndGet();
}
maxNodeId.accumulateAndGet(node, Math::max);
+ // update graphRamBytesUsed every 1000 nodes
+ if (level == 0 && node % 1000 == 0) {
+ updateGraphRamBytesUsed();
+ }
+ }
+
+ /** Finish building the graph. */
+ public void finishBuild() {
+ updateGraphRamBytesUsed();
}
@Override
@@ -285,23 +296,29 @@ public final class OnHeapHnswGraph extends HnswGraph implements Accountable {
lastFreezeSize = size();
}
+ /** Update the estimated ram bytes used for the neighbor array. */
+ public void updateGraphRamBytesUsed() {
+ long currentRamBytesUsedEstimate = RamUsageEstimator.NUM_BYTES_ARRAY_HEADER;
+ for (int node = 0; node < graph.length; node++) {
+ if (graph[node] == null) {
+ continue;
+ }
+
+ for (int i = 0; i < graph[node].length; i++) {
+ if (graph[node][i] == null) {
+ continue;
+ }
+ currentRamBytesUsedEstimate += graph[node][i].ramBytesUsed();
+ }
+
+ currentRamBytesUsedEstimate += RamUsageEstimator.NUM_BYTES_OBJECT_HEADER;
+ }
+ graphRamBytesUsed = currentRamBytesUsedEstimate;
+ }
+
@Override
public long ramBytesUsed() {
- long neighborArrayBytes0 =
- (long) nsize0 * (Integer.BYTES + Float.BYTES)
- + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER * 2L
- + RamUsageEstimator.NUM_BYTES_OBJECT_REF * 2L
- + Integer.BYTES * 3;
- long neighborArrayBytes =
- (long) nsize * (Integer.BYTES + Float.BYTES)
- + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER * 2L
- + RamUsageEstimator.NUM_BYTES_OBJECT_REF * 2L
- + Integer.BYTES * 3;
- long total = 0;
- total +=
- size() * (neighborArrayBytes0 + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER)
- + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER; // for graph and level 0;
- total += nonZeroLevelSize.get() * neighborArrayBytes; // for non-zero level
+ long total = graphRamBytesUsed; // all NeighborArray
total += 4 * Integer.BYTES; // all int fields
total += 1; // field: noGrowth
total +=
diff --git a/lucene/core/src/test/org/apache/lucene/internal/hppc/TestMaxSizedFloatArrayList.java b/lucene/core/src/test/org/apache/lucene/internal/hppc/TestMaxSizedFloatArrayList.java
new file mode 100644
index 00000000000..484a602b0bc
--- /dev/null
+++ b/lucene/core/src/test/org/apache/lucene/internal/hppc/TestMaxSizedFloatArrayList.java
@@ -0,0 +1,178 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.internal.hppc;
+
+import org.junit.Test;
+
+/** Tests for {@link MaxSizedFloatArrayList}. */
+public class TestMaxSizedFloatArrayList extends TestFloatArrayList {
+ private static final int MAX_SIZE = 5;
+
+ @Test
+ public void testMaxSizeConstructor() {
+ // Test valid maxSize
+ MaxSizedFloatArrayList list = new MaxSizedFloatArrayList(MAX_SIZE, 3);
+ assertEquals(3, list.buffer.length);
+ assertEquals(MAX_SIZE, list.maxSize);
+
+ // Test invalid maxSize (expectedElements > maxSize)
+ expectThrows(
+ AssertionError.class,
+ () -> {
+ new MaxSizedFloatArrayList(2, 3);
+ });
+ }
+
+ @Test
+ public void testMaxSizeLimit() {
+ MaxSizedFloatArrayList list = new MaxSizedFloatArrayList(MAX_SIZE);
+
+ // Fill up to maxSize
+ for (int i = 0; i < MAX_SIZE; i++) {
+ list.add(i + 0.5f);
+ }
+ assertEquals(MAX_SIZE, list.size());
+
+ // Try to add beyond maxSize
+ expectThrows(
+ IllegalStateException.class,
+ () -> {
+ list.add(MAX_SIZE + 0.5f);
+ });
+
+ // Try to add array beyond maxSize
+ expectThrows(
+ IllegalStateException.class,
+ () -> {
+ list.add(new float[] {1.5f, 2.5f, 3.5f});
+ });
+
+ // Try to addAll beyond maxSize
+ FloatArrayList otherList = new FloatArrayList();
+ otherList.add(1.5f);
+ expectThrows(
+ IllegalStateException.class,
+ () -> {
+ list.addAll(otherList);
+ });
+ }
+
+ @Test
+ public void testMaxSizeWithInsert() {
+ MaxSizedFloatArrayList list = new MaxSizedFloatArrayList(MAX_SIZE);
+
+ // Fill up to maxSize - 1
+ for (int i = 0; i < MAX_SIZE - 1; i++) {
+ list.add(i + 0.5f);
+ }
+ assertEquals(MAX_SIZE - 1, list.size());
+
+ // Insert at the end (should succeed)
+ list.insert(MAX_SIZE - 1, (MAX_SIZE - 1) + 0.5f);
+ assertEquals(MAX_SIZE, list.size());
+
+ // Try to insert beyond maxSize
+ expectThrows(
+ IllegalStateException.class,
+ () -> {
+ list.insert(0, MAX_SIZE + 0.5f);
+ });
+ }
+
+ @Test
+ public void testMaxSizeWithResize() {
+ MaxSizedFloatArrayList list = new MaxSizedFloatArrayList(MAX_SIZE);
+
+ // Try to resize beyond maxSize
+ expectThrows(
+ IllegalStateException.class,
+ () -> {
+ list.resize(MAX_SIZE + 1);
+ });
+
+ // Resize to maxSize (should succeed)
+ list.resize(MAX_SIZE);
+ assertEquals(MAX_SIZE, list.size());
+ }
+
+ @Test
+ public void testMaxSizeWithEnsureCapacity() {
+ MaxSizedFloatArrayList list = new MaxSizedFloatArrayList(MAX_SIZE);
+
+ // Try to ensure capacity beyond maxSize
+ expectThrows(
+ IllegalStateException.class,
+ () -> {
+ list.ensureCapacity(MAX_SIZE + 1);
+ });
+
+ // Ensure capacity to maxSize (should succeed)
+ list.ensureCapacity(MAX_SIZE);
+ // Buffer can be larger than maxSize due to growth strategy, but not smaller if capacity is
+ // ensured to maxSize
+ assertTrue(list.buffer.length >= MAX_SIZE || list.buffer.length == 0 && MAX_SIZE == 0);
+ }
+
+ @Test
+ public void testMaxSizeWithCopyConstructor() {
+ MaxSizedFloatArrayList original = new MaxSizedFloatArrayList(MAX_SIZE);
+ original.add(1.5f, 2.5f, 3.5f);
+
+ MaxSizedFloatArrayList copy = new MaxSizedFloatArrayList(original);
+ assertEquals(original.maxSize, copy.maxSize);
+ assertEquals(original.size(), copy.size());
+ assertArrayEquals(original.toArray(), copy.toArray(), 0.001f);
+ }
+
+ @Test
+ public void testHashCodeWithMaxSize() {
+ MaxSizedFloatArrayList list1 = new MaxSizedFloatArrayList(5);
+ MaxSizedFloatArrayList list2 = new MaxSizedFloatArrayList(10);
+
+ // Same elements but different maxSize should have different hash codes
+ list1.add(1.5f, 2.5f, 3.5f);
+ list2.add(1.5f, 2.5f, 3.5f);
+ assertNotEquals(list1.hashCode(), list2.hashCode());
+
+ // Same maxSize and same elements should have same hash code
+ MaxSizedFloatArrayList list3 = new MaxSizedFloatArrayList(5);
+ list3.add(1.5f, 2.5f, 3.5f);
+ assertEquals(list1.hashCode(), list3.hashCode());
+ }
+
+ @Test
+ public void testEqualsWithMaxSize() {
+ MaxSizedFloatArrayList list1 = new MaxSizedFloatArrayList(5);
+ MaxSizedFloatArrayList list2 = new MaxSizedFloatArrayList(10);
+ MaxSizedFloatArrayList list3 = new MaxSizedFloatArrayList(5);
+
+ // Same elements but different maxSize should not be equal
+ list1.add(1.5f, 2.5f, 3.5f);
+ list2.add(1.5f, 2.5f, 3.5f);
+ assertNotEquals(list1, list2);
+
+ // Same maxSize and same elements should be equal
+ list3.add(1.5f, 2.5f, 3.5f);
+ assertEquals(list1, list3);
+
+ // Different elements but same maxSize should not be equal
+ list3.clear();
+ list3.add(1.5f, 2.5f, 4.5f);
+ assertNotEquals(list1, list3);
+ }
+}
diff --git a/lucene/core/src/test/org/apache/lucene/internal/hppc/TestMaxSizedIntArrayList.java b/lucene/core/src/test/org/apache/lucene/internal/hppc/TestMaxSizedIntArrayList.java
new file mode 100644
index 00000000000..dc567fd2557
--- /dev/null
+++ b/lucene/core/src/test/org/apache/lucene/internal/hppc/TestMaxSizedIntArrayList.java
@@ -0,0 +1,176 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.internal.hppc;
+
+import org.junit.Test;
+
+/** Tests for {@link MaxSizedIntArrayList}. */
+public class TestMaxSizedIntArrayList extends TestIntArrayList {
+ private static final int MAX_SIZE = 5;
+
+ @Test
+ public void testMaxSizeConstructor() {
+ // Test valid maxSize
+ MaxSizedIntArrayList list = new MaxSizedIntArrayList(MAX_SIZE, 3);
+ assertEquals(3, list.buffer.length);
+ assertEquals(MAX_SIZE, list.maxSize);
+
+ // Test invalid maxSize (expectedElements > maxSize)
+ expectThrows(
+ AssertionError.class,
+ () -> {
+ new MaxSizedIntArrayList(2, 3);
+ });
+ }
+
+ @Test
+ public void testMaxSizeLimit() {
+ MaxSizedIntArrayList list = new MaxSizedIntArrayList(MAX_SIZE);
+
+ // Fill up to maxSize
+ for (int i = 0; i < MAX_SIZE; i++) {
+ list.add(i);
+ }
+ assertEquals(MAX_SIZE, list.size());
+
+ // Try to add beyond maxSize
+ expectThrows(
+ IllegalStateException.class,
+ () -> {
+ list.add(MAX_SIZE);
+ });
+
+ // Try to add array beyond maxSize
+ expectThrows(
+ IllegalStateException.class,
+ () -> {
+ list.add(new int[] {1, 2, 3});
+ });
+
+ // Try to addAll beyond maxSize
+ IntArrayList otherList = new IntArrayList();
+ otherList.add(1);
+ expectThrows(
+ IllegalStateException.class,
+ () -> {
+ list.addAll(otherList);
+ });
+ }
+
+ @Test
+ public void testMaxSizeWithInsert() {
+ MaxSizedIntArrayList list = new MaxSizedIntArrayList(MAX_SIZE);
+
+ // Fill up to maxSize - 1
+ for (int i = 0; i < MAX_SIZE - 1; i++) {
+ list.add(i);
+ }
+ assertEquals(MAX_SIZE - 1, list.size());
+
+ // Insert at the end (should succeed)
+ list.insert(MAX_SIZE - 1, MAX_SIZE - 1);
+ assertEquals(MAX_SIZE, list.size());
+
+ // Try to insert beyond maxSize
+ expectThrows(
+ IllegalStateException.class,
+ () -> {
+ list.insert(0, MAX_SIZE);
+ });
+ }
+
+ @Test
+ public void testMaxSizeWithResize() {
+ MaxSizedIntArrayList list = new MaxSizedIntArrayList(MAX_SIZE);
+
+ // Try to resize beyond maxSize
+ expectThrows(
+ IllegalStateException.class,
+ () -> {
+ list.resize(MAX_SIZE + 1);
+ });
+
+ // Resize to maxSize (should succeed)
+ list.resize(MAX_SIZE);
+ assertEquals(MAX_SIZE, list.size());
+ }
+
+ @Test
+ public void testMaxSizeWithEnsureCapacity() {
+ MaxSizedIntArrayList list = new MaxSizedIntArrayList(MAX_SIZE);
+
+ // Try to ensure capacity beyond maxSize
+ expectThrows(
+ IllegalStateException.class,
+ () -> {
+ list.ensureCapacity(MAX_SIZE + 1);
+ });
+
+ // Ensure capacity to maxSize (should succeed)
+ list.ensureCapacity(MAX_SIZE);
+ assertEquals(MAX_SIZE, list.buffer.length);
+ }
+
+ @Test
+ public void testMaxSizeWithCopyConstructor() {
+ MaxSizedIntArrayList original = new MaxSizedIntArrayList(MAX_SIZE);
+ original.add(1, 2, 3);
+
+ MaxSizedIntArrayList copy = new MaxSizedIntArrayList(original);
+ assertEquals(original.maxSize, copy.maxSize);
+ assertEquals(original.size(), copy.size());
+ assertArrayEquals(original.toArray(), copy.toArray());
+ }
+
+ @Test
+ public void testHashCodeWithMaxSize() {
+ MaxSizedIntArrayList list1 = new MaxSizedIntArrayList(5);
+ MaxSizedIntArrayList list2 = new MaxSizedIntArrayList(10);
+
+ // Same elements but different maxSize should have different hash codes
+ list1.add(1, 2, 3);
+ list2.add(1, 2, 3);
+ assertNotEquals(list1.hashCode(), list2.hashCode());
+
+ // Same maxSize and same elements should have same hash code
+ MaxSizedIntArrayList list3 = new MaxSizedIntArrayList(5);
+ list3.add(1, 2, 3);
+ assertEquals(list1.hashCode(), list3.hashCode());
+ }
+
+ @Test
+ public void testEqualsWithMaxSize() {
+ MaxSizedIntArrayList list1 = new MaxSizedIntArrayList(5);
+ MaxSizedIntArrayList list2 = new MaxSizedIntArrayList(10);
+ MaxSizedIntArrayList list3 = new MaxSizedIntArrayList(5);
+
+ // Same elements but different maxSize should not be equal
+ list1.add(1, 2, 3);
+ list2.add(1, 2, 3);
+ assertNotEquals(list1, list2);
+
+ // Same maxSize and same elements should be equal
+ list3.add(1, 2, 3);
+ assertEquals(list1, list3);
+
+ // Different elements but same maxSize should not be equal
+ list3.clear();
+ list3.add(1, 2, 4);
+ assertNotEquals(list1, list3);
+ }
+}
diff --git a/lucene/core/src/test/org/apache/lucene/util/TestArrayUtil.java b/lucene/core/src/test/org/apache/lucene/util/TestArrayUtil.java
index 90e5bafc9fa..d09b2b7aa36 100644
--- a/lucene/core/src/test/org/apache/lucene/util/TestArrayUtil.java
+++ b/lucene/core/src/test/org/apache/lucene/util/TestArrayUtil.java
@@ -398,6 +398,36 @@ public class TestArrayUtil extends LuceneTestCase {
assertEquals(minLength, growInRange(new int[] {1, 2, 3}, minLength, minLength).length);
}
+ public void testGrowInRangeFloat() {
+ float[] array = new float[] {1f, 2f, 3f};
+
+ // If minLength is negative, maxLength does not matter
+ expectThrows(AssertionError.class, () -> growInRange(array, -1, 4));
+ expectThrows(AssertionError.class, () -> growInRange(array, -1, 0));
+ expectThrows(AssertionError.class, () -> growInRange(array, -1, -1));
+
+ // If minLength > maxLength, we throw an exception
+ expectThrows(IllegalArgumentException.class, () -> growInRange(array, 1, 0));
+ expectThrows(IllegalArgumentException.class, () -> growInRange(array, 4, 3));
+ expectThrows(IllegalArgumentException.class, () -> growInRange(array, 5, 4));
+
+ // If minLength is sufficient, we return the array
+ assertSame(array, growInRange(array, 1, 4));
+ assertSame(array, growInRange(array, 1, 2));
+ assertSame(array, growInRange(array, 1, 1));
+
+ int minLength = 4;
+ int maxLength = Integer.MAX_VALUE;
+
+ // The array grows normally if maxLength permits
+ assertEquals(
+ oversize(minLength, Float.BYTES),
+ growInRange(new float[] {1f, 2f, 3f}, minLength, maxLength).length);
+
+ // The array grows to maxLength if maxLength is limiting
+ assertEquals(minLength, growInRange(new float[] {1f, 2f, 3f}, minLength, minLength).length);
+ }
+
public void testCopyOfSubArray() {
short[] shortArray = {1, 2, 3};
assertArrayEquals(new short[] {1}, copyOfSubArray(shortArray, 0, 1));
diff --git a/lucene/core/src/test/org/apache/lucene/util/hnsw/TestNeighborArray.java b/lucene/core/src/test/org/apache/lucene/util/hnsw/TestNeighborArray.java
index d23f7dd61fe..4d281b9011e 100644
--- a/lucene/core/src/test/org/apache/lucene/util/hnsw/TestNeighborArray.java
+++ b/lucene/core/src/test/org/apache/lucene/util/hnsw/TestNeighborArray.java
@@ -213,9 +213,217 @@ public class TestNeighborArray extends LuceneTestCase {
assertScoresEqual(new float[] {7, 6, 5, 4, 3, 2, 1}, neighbors);
}
+ public void testMaxSizeAddInOrder() throws IOException {
+ NeighborArray neighbors = new NeighborArray(3, true);
+ neighbors.addInOrder(0, 1.0f);
+ neighbors.addInOrder(1, 0.8f);
+ neighbors.addInOrder(2, 0.6f);
+
+ // Should throw when trying to add beyond max size
+ expectThrows(IllegalStateException.class, () -> neighbors.addInOrder(3, 0.4f));
+
+ assertScoresEqual(new float[] {1.0f, 0.8f, 0.6f}, neighbors);
+ assertNodesEqual(new int[] {0, 1, 2}, neighbors);
+ }
+
+ public void testMaxSizeAddOutOfOrder() throws IOException {
+ NeighborArray neighbors = new NeighborArray(3, true);
+ neighbors.addOutOfOrder(0, 0.8f);
+ neighbors.addOutOfOrder(1, 1.0f);
+ neighbors.addOutOfOrder(2, 0.6f);
+
+ // Should throw when trying to add beyond max size
+ expectThrows(IllegalStateException.class, () -> neighbors.addOutOfOrder(3, 0.4f));
+
+ // Sort to verify contents
+ neighbors.sort(null);
+ assertScoresEqual(new float[] {1.0f, 0.8f, 0.6f}, neighbors);
+ assertNodesEqual(new int[] {1, 0, 2}, neighbors);
+ }
+
+ public void testMaxSizeInsertSorted() throws IOException {
+ NeighborArray neighbors = new NeighborArray(3, true);
+ neighbors.insertSorted(0, 1.0f);
+ neighbors.insertSorted(1, 0.8f);
+ neighbors.insertSorted(2, 0.6f);
+
+ // Should throw when trying to insert beyond max size
+ expectThrows(IllegalStateException.class, () -> neighbors.insertSorted(3, 0.4f));
+
+ assertScoresEqual(new float[] {1.0f, 0.8f, 0.6f}, neighbors);
+ assertNodesEqual(new int[] {0, 1, 2}, neighbors);
+ }
+
+ public void testMaxSizeMixedOperations() throws IOException {
+ NeighborArray neighbors = new NeighborArray(3, true);
+
+ // Add first node
+ neighbors.addInOrder(0, 1.0f);
+
+ // Add second node out of order
+ neighbors.addOutOfOrder(1, 0.8f);
+
+ // Try to add third node in order - should fail because we used addOutOfOrder
+ expectThrows(AssertionError.class, () -> neighbors.addInOrder(2, 0.6f));
+
+ // Add third node out of order
+ neighbors.addOutOfOrder(2, 0.6f);
+
+ // Try to add fourth node - should fail due to max size
+ expectThrows(IllegalStateException.class, () -> neighbors.addOutOfOrder(3, 0.4f));
+
+ // Sort to verify contents
+ neighbors.sort(null);
+ assertScoresEqual(new float[] {1.0f, 0.8f, 0.6f}, neighbors);
+ assertNodesEqual(new int[] {0, 1, 2}, neighbors);
+ }
+
+ public void testBoundaryValues() throws IOException {
+ // Test with empty array
+ NeighborArray neighbors = new NeighborArray(3, true);
+ assertEquals(0, neighbors.size());
+
+ // Test with single element
+ neighbors.addInOrder(0, Float.MAX_VALUE);
+ assertScoresEqual(new float[] {Float.MAX_VALUE}, neighbors);
+ assertNodesEqual(new int[] {0}, neighbors);
+
+ // Test with NaN values
+ neighbors.clear();
+ neighbors.addOutOfOrder(0, Float.NaN);
+ neighbors.addOutOfOrder(1, Float.NaN);
+ assertScoresEqual(new float[] {Float.NaN, Float.NaN}, neighbors);
+ assertNodesEqual(new int[] {0, 1}, neighbors);
+ }
+
+ public void testSortOrder() throws IOException {
+ // Test ascending order
+ NeighborArray ascNeighbors = new NeighborArray(5, false);
+ ascNeighbors.addInOrder(0, 0.1f);
+ ascNeighbors.addInOrder(1, 0.2f);
+ ascNeighbors.addInOrder(2, 0.3f);
+ assertScoresEqual(new float[] {0.1f, 0.2f, 0.3f}, ascNeighbors);
+
+ // Test descending order
+ NeighborArray descNeighbors = new NeighborArray(5, true);
+ descNeighbors.addInOrder(0, 0.3f);
+ descNeighbors.addInOrder(1, 0.2f);
+ descNeighbors.addInOrder(2, 0.1f);
+ assertScoresEqual(new float[] {0.3f, 0.2f, 0.1f}, descNeighbors);
+
+ // Test equal scores
+ NeighborArray equalNeighbors = new NeighborArray(5, true);
+ equalNeighbors.addInOrder(0, 0.5f);
+ equalNeighbors.addInOrder(1, 0.5f);
+ equalNeighbors.addInOrder(2, 0.5f);
+ assertScoresEqual(new float[] {0.5f, 0.5f, 0.5f}, equalNeighbors);
+ }
+
+ public void testRemoveOperations() throws IOException {
+ NeighborArray neighbors = new NeighborArray(5, true);
+ neighbors.addInOrder(0, 1.0f);
+ neighbors.addInOrder(1, 0.8f);
+ neighbors.addInOrder(2, 0.6f);
+ neighbors.addInOrder(3, 0.4f);
+
+ // Test remove last
+ neighbors.removeLast();
+ assertScoresEqual(new float[] {1.0f, 0.8f, 0.6f}, neighbors);
+ assertNodesEqual(new int[] {0, 1, 2}, neighbors);
+
+ // Test remove middle
+ neighbors.removeIndex(1);
+ assertScoresEqual(new float[] {1.0f, 0.6f}, neighbors);
+ assertNodesEqual(new int[] {0, 2}, neighbors);
+
+ // Test remove first
+ neighbors.removeIndex(0);
+ assertScoresEqual(new float[] {0.6f}, neighbors);
+ assertNodesEqual(new int[] {2}, neighbors);
+
+ // Test remove last remaining
+ neighbors.removeLast();
+ assertEquals(0, neighbors.size());
+ }
+
+ public void testClearOperation() throws IOException {
+ NeighborArray neighbors = new NeighborArray(5, true);
+ neighbors.addInOrder(0, 1.0f);
+ neighbors.addInOrder(1, 0.8f);
+ neighbors.addInOrder(2, 0.6f);
+
+ // Test clear
+ neighbors.clear();
+ assertEquals(0, neighbors.size());
+
+ // Test adding after clear
+ neighbors.addInOrder(3, 1.0f);
+ assertScoresEqual(new float[] {1.0f}, neighbors);
+ assertNodesEqual(new int[] {3}, neighbors);
+ }
+
+ public void testComplexOperations() throws IOException {
+ NeighborArray neighbors = new NeighborArray(5, true);
+
+ // Add some nodes in order
+ neighbors.addInOrder(0, 1.0f);
+ neighbors.addInOrder(1, 0.8f);
+
+ // Add some nodes out of order
+ neighbors.addOutOfOrder(2, 0.9f);
+ neighbors.addOutOfOrder(3, 0.7f);
+
+ // Sort and verify
+ neighbors.sort(null);
+ assertScoresEqual(new float[] {1.0f, 0.9f, 0.8f, 0.7f}, neighbors);
+ assertNodesEqual(new int[] {0, 2, 1, 3}, neighbors);
+
+ // Remove some nodes
+ neighbors.removeIndex(1);
+ neighbors.removeLast();
+
+ // Add more nodes
+ neighbors.insertSorted(4, 0.85f);
+ neighbors.insertSorted(5, 0.75f);
+
+ // Final verification
+ assertScoresEqual(new float[] {1.0f, 0.85f, 0.8f, 0.75f}, neighbors);
+ assertNodesEqual(new int[] {0, 4, 1, 5}, neighbors);
+ }
+
+ public void testScorerOperations() throws IOException {
+ NeighborArray neighbors = new NeighborArray(5, true);
+ TestRandomVectorScorer scorer =
+ nodeId -> {
+ switch (nodeId) {
+ case 0:
+ return 1.0f;
+ case 1:
+ return 0.8f;
+ case 2:
+ return 0.9f;
+ case 3:
+ return 0.7f;
+ default:
+ return 0.0f;
+ }
+ };
+
+ // Add nodes with NaN scores
+ neighbors.addOutOfOrder(0, Float.NaN);
+ neighbors.addOutOfOrder(1, Float.NaN);
+ neighbors.addOutOfOrder(2, Float.NaN);
+ neighbors.addOutOfOrder(3, Float.NaN);
+
+ // Sort using scorer
+ neighbors.sort(scorer);
+ assertScoresEqual(new float[] {1.0f, 0.9f, 0.8f, 0.7f}, neighbors);
+ assertNodesEqual(new int[] {0, 2, 1, 3}, neighbors);
+ }
+
private void assertScoresEqual(float[] scores, NeighborArray neighbors) {
for (int i = 0; i < scores.length; i++) {
- assertEquals(scores[i], neighbors.scores()[i], 0.01f);
+ assertEquals(scores[i], neighbors.getScores(i), 0.01f);
}
}