diff --git a/.github/workflows/unit-tests.yaml b/.github/workflows/unit-tests.yaml index 7fb47f9ba..63950e960 100644 --- a/.github/workflows/unit-tests.yaml +++ b/.github/workflows/unit-tests.yaml @@ -18,8 +18,7 @@ jobs: cancel-in-progress: false strategy: matrix: - jdk: [ 20, 24 ] -# jdk: [ 24 ] + jdk: [ 24 ] isa: [ isa-avx512f ] runs-on: ${{ matrix.isa }} steps: @@ -48,20 +47,6 @@ jobs: distribution: temurin cache: maven - - name: Verify Panama Vector Support (JDK ${{ matrix.jdk }}) - if: matrix.jdk == '20' - run: >- - mvn -B -Pjdk20 -pl jvector-tests -am test - -Dsurefire.failIfNoSpecifiedTests=false - -Dtest=TestVectorizationProvider - -DTest_RequireSpecificVectorizationProvider=PanamaVectorizationProvider - - - name: Test Panama Support (JDK ${{ matrix.jdk }}) - if: matrix.jdk == '20' - run: >- - mvn -B -Pjdk20 -pl jvector-tests test -am test - -DTest_RequireSpecificVectorizationProvider=PanamaVectorizationProvider - - name: Verify native-access vector support (JDK ${{ matrix.jdk }}) if: matrix.jdk == '24' run: >- @@ -97,7 +82,7 @@ jobs: cancel-in-progress: true strategy: matrix: - jdk: [ 11, 20, 22] + jdk: [ 22 ] os: [ ubuntu-latest, windows-latest ] runs-on: ${{ matrix.os }} steps: @@ -112,15 +97,8 @@ jobs: if: matrix.os == 'ubuntu-latest' run: | sudo apt install -y gcc - - name: Compile, run tests, and package (JDK 22) + - name: Compile, run tests, and package run: mvn -B verify - if: matrix.jdk == '22' - - name: Compile, run tests, and package (JDK 20) - run: mvn -B -Pjdk20 -am -pl jvector-tests test - if: matrix.jdk == '20' - - name: Compile and run tests (JDK 11) - run: mvn -B -Pjdk11 -am -pl jvector-tests test - if: matrix.jdk == '11' - name: Test Summary if: always() uses: test-summary/action@v2 diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ByteBufferRandomAccessVectorValues.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ByteBufferRandomAccessVectorValues.java new file mode 100644 index 000000000..d0115bd1f --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ByteBufferRandomAccessVectorValues.java @@ -0,0 +1,125 @@ +/* + * Copyright DataStax, Inc. + * + * 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 io.github.jbellis.jvector.graph; + +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; + +import java.nio.ByteBuffer; +import java.util.Objects; + +/** + * A {@link RandomAccessVectorValues} backed by a single caller-owned {@link ByteBuffer} that + * holds {@code count × dimension × Float.BYTES} bytes of concatenated IEEE 754 floats. + * + *

Vectors are never copied: {@link #getVector(int)} returns a view (via the active + * {@link VectorTypeSupport#wrapFloatVector(ByteBuffer, int, int)}) that aliases the source + * bytes. This lets integrators whose data is already in ByteBuffer form (for example a + * database page cache) feed {@code GraphIndexBuilder} without the usual {@code float[]} + * materialization step. + * + *

The buffer is pre-sliced at construction time so later mutation of the caller's buffer + * position / limit does not disturb this RAVV. Mutating the buffer's contents after + * construction is visible to this RAVV — callers are responsible for ensuring writers and + * readers do not race. + */ +public class ByteBufferRandomAccessVectorValues implements RandomAccessVectorValues +{ + private static final VectorTypeSupport VTS = + VectorizationProvider.getInstance().getVectorTypeSupport(); + + private final ByteBuffer data; + private final int count; + private final int dimension; + + public ByteBufferRandomAccessVectorValues(ByteBuffer data, int count, int dimension) + { + Objects.requireNonNull(data, "data"); + if (count < 0) throw new IllegalArgumentException("count must be >= 0, was " + count); + if (dimension <= 0) throw new IllegalArgumentException("dimension must be > 0, was " + dimension); + long need = (long) count * dimension * Float.BYTES; + if (data.remaining() < need) { + throw new IllegalArgumentException( + "buffer too small: need " + need + " bytes, have " + data.remaining()); + } + this.data = data.slice(data.position(), (int) need).order(data.order()); + this.count = count; + this.dimension = dimension; + } + + @Override + public int size() + { + return count; + } + + @Override + public int dimension() + { + return dimension; + } + + @Override + public VectorFloat getVector(int targetOrd) + { + if (targetOrd < 0 || targetOrd >= count) { + throw new IndexOutOfBoundsException("ordinal " + targetOrd + " out of [0, " + count + ")"); + } + return VTS.wrapFloatVector(data, targetOrd * dimension, dimension); + } + + @Override + public boolean isValueShared() + { + // Each getVector returns a distinct view instance — callers may hold onto results. + return false; + } + + @Override + public ByteBufferRandomAccessVectorValues copy() + { + return this; + } + + /** + * Convenience factory accepting a plain {@code float[]} (copying it into a native-endian + * ByteBuffer). Primarily used by tests that already have {@code float[]} fixtures. + */ + public static ByteBufferRandomAccessVectorValues fromFloats(float[][] vectors, int dimension) + { + Objects.requireNonNull(vectors, "vectors"); + if (dimension <= 0) throw new IllegalArgumentException("dimension > 0 required"); + int count = vectors.length; + ByteBuffer bb = ByteBuffer.allocate(count * dimension * Float.BYTES) + .order(java.nio.ByteOrder.LITTLE_ENDIAN); + for (float[] v : vectors) { + if (v.length != dimension) { + throw new IllegalArgumentException("vector dimension mismatch: " + v.length + " vs " + dimension); + } + for (float f : v) bb.putFloat(f); + } + bb.rewind(); + return new ByteBufferRandomAccessVectorValues(bb, count, dimension); + } + + /** @return the live view over the backing data (for tests/integration). */ + public ByteBuffer data() + { + return data; + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java index 9e366676c..9d8af192e 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java @@ -30,13 +30,16 @@ import io.github.jbellis.jvector.util.ExplicitThreadLocal; import io.github.jbellis.jvector.util.PhysicalCoreExecutor; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.Closeable; import java.io.IOException; import java.io.UncheckedIOException; +import java.nio.ByteBuffer; import java.util.*; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentSkipListSet; @@ -593,6 +596,17 @@ public long addGraphNode(int node, VectorFloat vector) { return addGraphNode(node, ssp); } + /** + * Zero-copy ByteBuffer overload of {@link #addGraphNode(int, VectorFloat)}. The buffer is + * interpreted as contiguous IEEE 754 floats in its current {@link java.nio.ByteOrder}, so + * callers that already hold vector data as ByteBuffer (for example a database page cache) + * can feed the builder without the usual {@code float[]} materialization step. + */ + public long addGraphNode(int node, ByteBuffer vector) { + VectorTypeSupport vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + return addGraphNode(node, vts.wrapFloatVector(vector)); + } + /** * Inserts a node with the given vector value to the graph. * diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphSearcher.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphSearcher.java index 73cc5fbd5..24ac2bc8b 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphSearcher.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphSearcher.java @@ -33,12 +33,15 @@ import io.github.jbellis.jvector.util.BoundedLongHeap; import io.github.jbellis.jvector.util.GrowableLongHeap; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; import org.agrona.collections.Int2ObjectHashMap; import org.agrona.collections.IntHashSet; import java.io.Closeable; import java.io.IOException; +import java.nio.ByteBuffer; /** @@ -151,6 +154,26 @@ public static SearchResult search(VectorFloat queryVector, int topK, int rera } } + /** + * Zero-copy ByteBuffer overload of {@link #search(VectorFloat, int, RandomAccessVectorValues, + * VectorSimilarityFunction, ImmutableGraphIndex, Bits)}. The buffer is interpreted as + * {@code Float.BYTES}-aligned IEEE 754 floats in its current {@link java.nio.ByteOrder}; + * no {@code float[]} materialization occurs. + */ + public static SearchResult search(ByteBuffer queryVector, int topK, RandomAccessVectorValues vectors, VectorSimilarityFunction similarityFunction, ImmutableGraphIndex graph, Bits acceptOrds) { + VectorTypeSupport vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + return search(vts.wrapFloatVector(queryVector), topK, vectors, similarityFunction, graph, acceptOrds); + } + + /** + * Zero-copy ByteBuffer overload with explicit rerankK. See + * {@link #search(ByteBuffer, int, RandomAccessVectorValues, VectorSimilarityFunction, ImmutableGraphIndex, Bits)}. + */ + public static SearchResult search(ByteBuffer queryVector, int topK, int rerankK, RandomAccessVectorValues vectors, VectorSimilarityFunction similarityFunction, ImmutableGraphIndex graph, Bits acceptOrds) { + VectorTypeSupport vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + return search(vts.wrapFloatVector(queryVector), topK, rerankK, vectors, similarityFunction, graph, acceptOrds); + } + /** * Sets the view of the graph to be used by the searcher. *

diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java index c84b7b955..d3cbef30e 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java @@ -517,12 +517,14 @@ int closestCentroidIndex(VectorFloat subvector, int m, VectorFloat codeboo } /** - * Extracts the m-th subvector from a single vector. + * Extracts the m-th subvector from a single vector. Returns a zero-copy view via + * {@link VectorFloat#subview(int, int)} — no {@code float[]} allocation per subvector, + * which matters at codebook training time where this is called once per + * {@code (training_vector × subspace)} and would otherwise produce + * {@code N × M} small heap allocations. */ static VectorFloat getSubVector(VectorFloat vector, int m, int[][] subvectorSizeAndOffset) { - VectorFloat subvector = vectorTypeSupport.createFloatVector(subvectorSizeAndOffset[m][0]); - subvector.copyFrom(vector, subvectorSizeAndOffset[m][1], 0, subvectorSizeAndOffset[m][0]); - return subvector; + return vector.subview(subvectorSizeAndOffset[m][1], subvectorSizeAndOffset[m][0]); } /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ArraySliceVectorFloat.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ArraySliceVectorFloat.java new file mode 100644 index 000000000..ca09a9e6f --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ArraySliceVectorFloat.java @@ -0,0 +1,180 @@ +/* + * Copyright DataStax, Inc. + * + * 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 io.github.jbellis.jvector.vector; + +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.util.RamUsageEstimator; +import io.github.jbellis.jvector.vector.types.VectorFloat; + +import java.io.IOException; +import java.util.Arrays; + +/** + * A {@link VectorFloat} that is a view into a contiguous range of a root + * {@code float[]}. Constructing this class does not allocate or copy the backing storage — both + * element access and SIMD dispatch read the underlying array at {@code arrayOffset + i}. + * + *

Primarily used inside jvector's Product Quantization training: extracting the M subvectors + * of each training vector can return M views instead of M materialized {@code float[dim/M]} + * copies. + * + *

The companion to {@link ArraySliceByteSequence} for the float-vector side of the type system. + */ +public final class ArraySliceVectorFloat implements VectorFloat +{ + private final float[] data; + private final int arrayOffset; + private final int length; + + public ArraySliceVectorFloat(float[] data, int arrayOffset, int length) + { + if (data == null) throw new NullPointerException("data"); + if (arrayOffset < 0 || length < 0 || (long) arrayOffset + length > data.length) { + throw new IllegalArgumentException( + "slice [" + arrayOffset + "," + (arrayOffset + length) + + ") out of range for float[" + data.length + "]"); + } + this.data = data; + this.arrayOffset = arrayOffset; + this.length = length; + } + + /** @return the root float[] this view borrows from (offsets are still relative to arrayOffset()). */ + @Override + public float[] get() + { + return data; + } + + /** @return the index into {@link #get()} where this view's element 0 lives. */ + public int arrayOffset() + { + return arrayOffset; + } + + @Override + public float get(int n) + { + return data[arrayOffset + n]; + } + + @Override + public void set(int n, float value) + { + data[arrayOffset + n] = value; + } + + @Override + public void zero() + { + Arrays.fill(data, arrayOffset, arrayOffset + length, 0f); + } + + @Override + public int length() + { + return length; + } + + @Override + public void writeTo(IndexWriter writer) throws IOException + { + writer.writeFloats(data, arrayOffset, length); + } + + @Override + public VectorFloat copy() + { + return new ArrayVectorFloat(Arrays.copyOfRange(data, arrayOffset, arrayOffset + length)); + } + + @Override + public VectorFloat subview(int floatOffset, int floatLength) + { + if (floatOffset < 0 || floatLength < 0 || (long) floatOffset + floatLength > this.length) { + throw new IllegalArgumentException( + "subview [" + floatOffset + "," + (floatOffset + floatLength) + + ") out of range for view length " + this.length); + } + if (floatOffset == 0 && floatLength == this.length) { + return this; + } + return new ArraySliceVectorFloat(data, arrayOffset + floatOffset, floatLength); + } + + @Override + public void copyFrom(VectorFloat src, int srcOffset, int destOffset, int copyLength) + { + if (src instanceof ArrayVectorFloat csrc) { + System.arraycopy(csrc.get(), srcOffset, data, arrayOffset + destOffset, copyLength); + return; + } + if (src instanceof ArraySliceVectorFloat asrc) { + System.arraycopy(asrc.data, asrc.arrayOffset + srcOffset, data, arrayOffset + destOffset, copyLength); + return; + } + for (int i = 0; i < copyLength; i++) { + data[arrayOffset + destOffset + i] = src.get(srcOffset + i); + } + } + + @Override + public long ramBytesUsed() + { + // Only the slice metadata; the backing float[] is accounted by its owner. + return RamUsageEstimator.NUM_BYTES_OBJECT_HEADER + + RamUsageEstimator.NUM_BYTES_OBJECT_REF + + 2L * Integer.BYTES; + } + + @Override + public String toString() + { + StringBuilder sb = new StringBuilder(); + sb.append("["); + int preview = Math.min(length, 25); + for (int i = 0; i < preview; i++) { + sb.append(get(i)); + if (i < length - 1) { + sb.append(", "); + } + } + if (length > 25) { + sb.append("..."); + } + sb.append("]"); + return sb.toString(); + } + + @Override + public boolean equals(Object o) + { + if (this == o) return true; + if (!(o instanceof VectorFloat that)) return false; + if (that.length() != this.length) return false; + for (int i = 0; i < length; i++) { + if (Float.floatToIntBits(get(i)) != Float.floatToIntBits(that.get(i))) return false; + } + return true; + } + + @Override + public int hashCode() + { + return this.getHashCode(); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ArrayVectorFloat.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ArrayVectorFloat.java index 32dce1f35..107a327b2 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ArrayVectorFloat.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ArrayVectorFloat.java @@ -78,11 +78,29 @@ public VectorFloat copy() return new ArrayVectorFloat(Arrays.copyOf(data, data.length)); } + @Override + public VectorFloat subview(int floatOffset, int floatLength) + { + if (floatOffset == 0 && floatLength == data.length) { + return this; + } + return new ArraySliceVectorFloat(data, floatOffset, floatLength); + } + @Override public void copyFrom(VectorFloat src, int srcOffset, int destOffset, int length) { - ArrayVectorFloat csrc = (ArrayVectorFloat) src; - System.arraycopy(csrc.data, srcOffset, data, destOffset, length); + if (src instanceof ArrayVectorFloat csrc) { + System.arraycopy(csrc.data, srcOffset, data, destOffset, length); + return; + } + if (src instanceof ArraySliceVectorFloat ssrc) { + System.arraycopy(ssrc.get(), ssrc.arrayOffset() + srcOffset, data, destOffset, length); + return; + } + for (int i = 0; i < length; i++) { + data[destOffset + i] = src.get(srcOffset + i); + } } @Override diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/BufferVectorFloat.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/BufferVectorFloat.java new file mode 100644 index 000000000..3b6970d3b --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/BufferVectorFloat.java @@ -0,0 +1,256 @@ +/* + * Copyright DataStax, Inc. + * + * 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 io.github.jbellis.jvector.vector; + +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.util.RamUsageEstimator; +import io.github.jbellis.jvector.vector.types.VectorFloat; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.Objects; + +/** + * VectorFloat implementation backed by a caller-owned {@link ByteBuffer}. + * + *

Storage is not copied on construction: the returned instance is a view over the + * supplied buffer. This lets callers that already hold vector data as ByteBuffer (for example + * a database reading rows off disk) feed jvector without the usual {@code float[]} allocation + * and copy. The buffer's current {@link ByteOrder} is preserved and honored on every element + * access, so both little- and big-endian layouts work. + * + *

Element {@code i} is read as {@code data.getFloat(byteOffset + i * Float.BYTES)} where + * {@code byteOffset} is fixed at construction and does not depend on the buffer's mutable + * position/limit. This makes the view safe to share across threads that only read it, and + * robust against callers that continue to position/limit the underlying buffer. + * + *

The Panama SIMD backend detects instances of this class and dispatches to + * {@code FloatVector.fromMemorySegment(species, MemorySegment.ofBuffer(data), byteOffset, order)} + * so SIMD is preserved with no {@code float[]} materialization. + */ +public final class BufferVectorFloat implements VectorFloat +{ + private final ByteBuffer data; + private final int byteOffset; + private final int floatLength; + + /** + * Zero-filled allocation; behaves like {@code new float[length]}. Buffer is native-endian + * and on-heap. + */ + public BufferVectorFloat(int length) + { + if (length < 0) { + throw new IllegalArgumentException("length must be >= 0, was " + length); + } + this.data = ByteBuffer.allocate(length * Float.BYTES).order(ByteOrder.nativeOrder()); + this.byteOffset = 0; + this.floatLength = length; + } + + /** + * Wrap the entire remaining contents of {@code data} as a float vector view. The buffer's + * {@code remaining()} must be a multiple of {@code Float.BYTES}. + */ + public BufferVectorFloat(ByteBuffer data) + { + this(data, 0, checkedFloatRemaining(data)); + } + + private static int checkedFloatRemaining(ByteBuffer data) + { + Objects.requireNonNull(data, "data"); + if ((data.remaining() % Float.BYTES) != 0) { + throw new IllegalArgumentException( + "ByteBuffer remaining() must be a multiple of Float.BYTES, was " + data.remaining()); + } + return data.remaining() / Float.BYTES; + } + + /** + * Sub-range wrap. {@code floatOffset} and {@code floatLength} are expressed in floats, + * relative to the buffer's current {@code position()}. + * + *

The returned view takes an independent {@link ByteBuffer#slice() slice} of the caller's + * buffer; this lets later mutation of the caller's {@code position}/{@code limit} leave the + * view intact, and it keeps the base byte offset inside the view stable at zero so that + * {@code MemorySegment.ofBuffer(view)} can be called without further per-access adjustment + * in the Panama SIMD path. + */ + public BufferVectorFloat(ByteBuffer data, int floatOffset, int floatLength) + { + Objects.requireNonNull(data, "data"); + if (floatOffset < 0 || floatLength < 0) { + throw new IllegalArgumentException("offset/length must be >= 0"); + } + long byteEnd = (long) floatOffset * Float.BYTES + (long) floatLength * Float.BYTES; + if (byteEnd > data.remaining()) { + throw new IllegalArgumentException( + "view [" + floatOffset + "," + (floatOffset + floatLength) + + ") floats exceeds buffer remaining=" + data.remaining() + " bytes"); + } + int startByte = data.position() + floatOffset * Float.BYTES; + this.data = data.slice(startByte, floatLength * Float.BYTES).order(data.order()); + this.byteOffset = 0; + this.floatLength = floatLength; + } + + /** + * @return the backing buffer. Mutating its position/limit does not affect element access, + * but mutating its contents does. + */ + @Override + public ByteBuffer get() + { + return data; + } + + /** @return byte offset at which this view's element 0 lives within the backing buffer. */ + public int byteOffset() + { + return byteOffset; + } + + /** @return the byte order the backing buffer is read with. */ + public ByteOrder byteOrder() + { + return data.order(); + } + + @Override + public float get(int i) + { + return data.getFloat(byteOffset + i * Float.BYTES); + } + + @Override + public void set(int i, float value) + { + data.putFloat(byteOffset + i * Float.BYTES, value); + } + + @Override + public void zero() + { + for (int i = 0; i < floatLength; i++) { + data.putFloat(byteOffset + i * Float.BYTES, 0f); + } + } + + @Override + public int length() + { + return floatLength; + } + + @Override + public void writeTo(IndexWriter writer) throws IOException + { + for (int i = 0; i < floatLength; i++) { + writer.writeFloat(get(i)); + } + } + + @Override + public VectorFloat copy() + { + ByteBuffer owned = ByteBuffer.allocate(floatLength * Float.BYTES).order(data.order()); + ByteBuffer srcSlice = data.slice(byteOffset, floatLength * Float.BYTES).order(data.order()); + owned.put(srcSlice).rewind(); + return new BufferVectorFloat(owned, 0, floatLength); + } + + @Override + public VectorFloat subview(int floatOffset, int floatLength) + { + if (floatOffset == 0 && floatLength == this.floatLength) { + return this; + } + return new BufferVectorFloat(data, floatOffset, floatLength); + } + + @Override + public void copyFrom(VectorFloat src, int srcOffset, int destOffset, int length) + { + if (src instanceof BufferVectorFloat bsrc && bsrc.byteOrder() == this.byteOrder()) { + int srcStart = bsrc.byteOffset + srcOffset * Float.BYTES; + int bytes = length * Float.BYTES; + ByteBuffer srcSlice = bsrc.data.slice(srcStart, bytes).order(bsrc.data.order()); + ByteBuffer destDup = this.data.duplicate().order(this.data.order()); + destDup.position(this.byteOffset + destOffset * Float.BYTES); + destDup.put(srcSlice); + return; + } + if (src instanceof ArrayVectorFloat a) { + float[] raw = a.get(); + for (int i = 0; i < length; i++) { + set(destOffset + i, raw[srcOffset + i]); + } + return; + } + for (int i = 0; i < length; i++) { + set(destOffset + i, src.get(srcOffset + i)); + } + } + + @Override + public long ramBytesUsed() + { + long OH_BYTES = RamUsageEstimator.NUM_BYTES_OBJECT_HEADER; + long REF_BYTES = RamUsageEstimator.NUM_BYTES_OBJECT_REF; + return OH_BYTES + REF_BYTES + 2L * Integer.BYTES + (long) floatLength * Float.BYTES; + } + + @Override + public String toString() + { + StringBuilder sb = new StringBuilder(); + sb.append("["); + int preview = Math.min(floatLength, 25); + for (int i = 0; i < preview; i++) { + sb.append(get(i)); + if (i < floatLength - 1) { + sb.append(", "); + } + } + if (floatLength > 25) { + sb.append("..."); + } + sb.append("]"); + return sb.toString(); + } + + @Override + public boolean equals(Object o) + { + if (this == o) return true; + if (!(o instanceof VectorFloat)) return false; + VectorFloat that = (VectorFloat) o; + if (that.length() != this.floatLength) return false; + for (int i = 0; i < floatLength; i++) { + if (Float.floatToIntBits(get(i)) != Float.floatToIntBits(that.get(i))) return false; + } + return true; + } + + @Override + public int hashCode() + { + return this.getHashCode(); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/DefaultVectorUtilSupport.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/DefaultVectorUtilSupport.java index 5843dc5f6..214cca02d 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/DefaultVectorUtilSupport.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/DefaultVectorUtilSupport.java @@ -37,6 +37,15 @@ final class DefaultVectorUtilSupport implements VectorUtilSupport { @Override public float dotProduct(VectorFloat av, VectorFloat bv) { + if (!(av instanceof ArrayVectorFloat) || !(bv instanceof ArrayVectorFloat)) { + // generic path: any VectorFloat impl (BufferVectorFloat, MemorySegmentVectorFloat, ...) + float res = 0f; + int len = av.length(); + for (int i = 0; i < len; i++) { + res += av.get(i) * bv.get(i); + } + return res; + } float[] a = ((ArrayVectorFloat) av).get(); float[] b = ((ArrayVectorFloat) bv).get(); @@ -107,6 +116,13 @@ public float dotProduct(VectorFloat av, VectorFloat bv) { @Override public float dotProduct(VectorFloat av, int aoffset, VectorFloat bv, int boffset, int length) { + if (!(av instanceof ArrayVectorFloat) || !(bv instanceof ArrayVectorFloat)) { + float sum = 0f; + for (int i = 0; i < length; i++) { + sum += av.get(aoffset + i) * bv.get(boffset + i); + } + return sum; + } float[] b = ((ArrayVectorFloat) bv).get(); float[] a = ((ArrayVectorFloat) av).get(); @@ -120,6 +136,17 @@ public float dotProduct(VectorFloat av, int aoffset, VectorFloat bv, int b @Override public float cosine(VectorFloat av, VectorFloat bv) { + if (!(av instanceof ArrayVectorFloat) || !(bv instanceof ArrayVectorFloat)) { + float sum = 0.0f, norm1 = 0.0f, norm2 = 0.0f; + int dim = av.length(); + for (int i = 0; i < dim; i++) { + float e1 = av.get(i), e2 = bv.get(i); + sum += e1 * e2; + norm1 += e1 * e1; + norm2 += e2 * e2; + } + return (float) (sum / Math.sqrt(norm1 * norm2)); + } float[] a = ((ArrayVectorFloat) av).get(); float[] b = ((ArrayVectorFloat) bv).get(); @@ -140,6 +167,16 @@ public float cosine(VectorFloat av, VectorFloat bv) { @Override public float cosine(VectorFloat av, int aoffset, VectorFloat bv, int boffset, int length) { + if (!(av instanceof ArrayVectorFloat) || !(bv instanceof ArrayVectorFloat)) { + float sum = 0.0f, norm1 = 0.0f, norm2 = 0.0f; + for (int i = 0; i < length; i++) { + float e1 = av.get(aoffset + i), e2 = bv.get(boffset + i); + sum += e1 * e2; + norm1 += e1 * e1; + norm2 += e2 * e2; + } + return (float) (sum / Math.sqrt(norm1 * norm2)); + } float[] a = ((ArrayVectorFloat) av).get(); float[] b = ((ArrayVectorFloat) bv).get(); float sum = 0.0f; @@ -157,6 +194,15 @@ public float cosine(VectorFloat av, int aoffset, VectorFloat bv, int boffs @Override public float squareDistance(VectorFloat av, VectorFloat bv) { + if (!(av instanceof ArrayVectorFloat) || !(bv instanceof ArrayVectorFloat)) { + float squareSum = 0.0f; + int dim = av.length(); + for (int i = 0; i < dim; i++) { + float diff = av.get(i) - bv.get(i); + squareSum += diff * diff; + } + return squareSum; + } float[] a = ((ArrayVectorFloat) av).get(); float[] b = ((ArrayVectorFloat) bv).get(); @@ -195,6 +241,14 @@ private static float squareDistanceUnrolled(float[] v1, float[] v2, int index) { @Override public float squareDistance(VectorFloat av, int aoffset, VectorFloat bv, int boffset, int length) { + if (!(av instanceof ArrayVectorFloat) || !(bv instanceof ArrayVectorFloat)) { + float squareSum = 0f; + for (int i = 0; i < length; i++) { + float diff = av.get(aoffset + i) - bv.get(boffset + i); + squareSum += diff * diff; + } + return squareSum; + } float[] a = ((ArrayVectorFloat) av).get(); float[] b = ((ArrayVectorFloat) bv).get(); diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/types/VectorFloat.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/types/VectorFloat.java index e749ccdc5..258001a28 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/types/VectorFloat.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/types/VectorFloat.java @@ -38,6 +38,27 @@ default int offset(int i) { VectorFloat copy(); + /** + * Return a view over a contiguous sub-range of this vector. The default implementation + * materializes a new owned vector via {@link VectorTypeSupport#createFloatVector(int)} + * plus {@link #copyFrom}; concrete subtypes that can share storage (on-heap arrays, + * ByteBuffers, MemorySegments) override this to return a zero-copy view. + * + *

Used, for example, by Product Quantization training to extract per-subspace + * sub-vectors without materializing {@code M × N × (dim/M)} extra floats. + */ + default VectorFloat subview(int floatOffset, int floatLength) { + if (floatOffset < 0 || floatLength < 0 || (long) floatOffset + floatLength > length()) { + throw new IllegalArgumentException( + "subview [" + floatOffset + "," + (floatOffset + floatLength) + + ") out of range for length " + length()); + } + VectorFloat out = io.github.jbellis.jvector.vector.VectorizationProvider.getInstance() + .getVectorTypeSupport().createFloatVector(floatLength); + out.copyFrom(this, floatOffset, 0, floatLength); + return out; + } + void copyFrom(VectorFloat src, int srcOffset, int destOffset, int length); float get(int i); diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/types/VectorTypeSupport.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/types/VectorTypeSupport.java index a4cb4f8af..61c9b1069 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/types/VectorTypeSupport.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/types/VectorTypeSupport.java @@ -18,15 +18,27 @@ import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.vector.BufferVectorFloat; import java.io.IOException; +import java.nio.ByteBuffer; public interface VectorTypeSupport { /** * Create a vector from the given data. * - * @param data the data to create the vector from. Supported data types are implementation-dependent. + * @param data the data to create the vector from. Supported data types are implementation-dependent + * (typically {@code float[]}; some providers also accept {@link java.nio.Buffer}). + * This method takes ownership of the data: for array inputs, most providers + * wrap the array without copying, so later mutation of the caller's array is visible. + * For {@code Buffer} inputs the behavior is provider-specific (some copy, some wrap). * @return the created vector. + * + *

Prefer the typed factories when possible: + * {@link #wrapFloatVector(ByteBuffer)} gives a zero-copy borrow from a ByteBuffer, and + * {@link #createFloatVector(int)} returns an owned, zero-filled allocation. The + * {@code Object} parameter here is unsafe — a wrong runtime type triggers a + * {@code ClassCastException} inside the provider. */ VectorFloat createFloatVector(Object data); @@ -37,6 +49,38 @@ public interface VectorTypeSupport { */ VectorFloat createFloatVector(int length); + /** + * Wrap a caller-owned {@link ByteBuffer} as a float vector view without copying. The buffer + * is treated as a contiguous sequence of IEEE 754 floats laid out in its current + * {@link java.nio.ByteOrder}; the buffer's position acts as the starting byte offset. + * + *

Semantically this is the zero-copy counterpart to {@link #createFloatVector(Object)}: + * {@code create} takes ownership of storage (copying if necessary); {@code wrap} borrows it. + * + *

Callers must not deallocate / re-purpose the backing buffer while the returned view is + * in use. Mutating the buffer's position and limit after construction does not disturb the + * view; mutating the buffer's content does. + */ + default VectorFloat wrapFloatVector(ByteBuffer data) { + if ((data.remaining() % Float.BYTES) != 0) { + throw new IllegalArgumentException( + "ByteBuffer remaining() must be a multiple of Float.BYTES, was " + data.remaining()); + } + return wrapFloatVector(data, 0, data.remaining() / Float.BYTES); + } + + /** + * Wrap a sub-range of a caller-owned {@link ByteBuffer} as a float vector view without + * copying. See {@link #wrapFloatVector(ByteBuffer)} for the copy contract. + * + * @param data backing buffer + * @param floatOffset starting offset in floats, relative to the buffer's current position + * @param floatLength number of floats in the resulting view + */ + default VectorFloat wrapFloatVector(ByteBuffer data, int floatOffset, int floatLength) { + return new BufferVectorFloat(data, floatOffset, floatLength); + } + /** * Read a vector from the given RandomAccessReader. * @param r the reader to read the vector from. diff --git a/jvector-examples/pom.xml b/jvector-examples/pom.xml index 9daf7b8cf..637341d3d 100644 --- a/jvector-examples/pom.xml +++ b/jvector-examples/pom.xml @@ -163,196 +163,6 @@ - - jdk11 - - - - org.codehaus.mojo - exec-maven-plugin - - false - - - - tutorial - - - -classpath - - -ea - io.github.jbellis.jvector.example.tutorial.TutorialRunner - ${tutorial} - - - - - sift - - exec - - - - -classpath - - -ea - io.github.jbellis.jvector.example.benchmarks.datasets.SiftSmall - - - - - bench - - exec - - - - -classpath - - -Xmx14G - -ea - io.github.jbellis.jvector.example.BenchYAML - ${benchArgs} - - - - - bench-java - - java - - - false - - ${benchArgs} - - - - - - - - - - jdk20 - - - io.github.jbellis - jvector-twenty - ${project.version} - compile - - - - - - org.codehaus.mojo - exec-maven-plugin - - false - - - - tutorial - - - -classpath - - --add-modules=jdk.incubator.vector - -ea - io.github.jbellis.jvector.example.tutorial.TutorialRunner - ${tutorial} - - - - - sift - - exec - - - - -classpath - - --add-modules=jdk.incubator.vector - -ea - io.github.jbellis.jvector.example.benchmarks.datasets.SiftSmall - - - - - bench - - exec - - - - -classpath - - --enable-native-access=ALL-UNNAMED - --add-modules=jdk.incubator.vector - -XX:+HeapDumpOnOutOfMemoryError - -Xmx14G - -ea - io.github.jbellis.jvector.example.BenchYAML - ${benchArgs} - - - - - bench-java - - java - - - false - - ${benchArgs} - - - - - ipcserve - - exec - - - - -classpath - - --enable-native-access=ALL-UNNAMED - --add-modules=jdk.incubator.vector - -XX:+HeapDumpOnOutOfMemoryError - -ea - io.github.jbellis.jvector.example.IPCService - - - - - ipcserve-1core - - exec - - - - -classpath - - --enable-native-access=ALL-UNNAMED - --add-modules=jdk.incubator.vector - -XX:+HeapDumpOnOutOfMemoryError - -Xmx10G - -ea - -Djava.util.concurrent.ForkJoinPool.common.parallelism=1 - -Djvector.physical_core_count=1 - io.github.jbellis.jvector.example.IPCService - - - - - - - - jdk22 diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/MMapRandomAccessVectorValues.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/MMapRandomAccessVectorValues.java index 0fda16df1..42672487b 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/MMapRandomAccessVectorValues.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/MMapRandomAccessVectorValues.java @@ -16,77 +16,93 @@ package io.github.jbellis.jvector.example.util; -import com.indeed.util.mmap.MMapBuffer; +import io.github.jbellis.jvector.graph.ByteBufferRandomAccessVectorValues; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; -import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.types.VectorFloat; -import io.github.jbellis.jvector.vector.types.VectorTypeSupport; import java.io.Closeable; import java.io.File; import java.io.IOError; import java.io.IOException; +import java.io.RandomAccessFile; import java.nio.ByteOrder; +import java.nio.MappedByteBuffer; import java.nio.channels.FileChannel; -public class MMapRandomAccessVectorValues implements RandomAccessVectorValues, Closeable { - private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); - final int dimension; - final int rows; - final File file; - final float[] valueBuffer; - - final MMapBuffer fileReader; +/** + * RAVV backed by a memory-mapped file of contiguous little-endian IEEE 754 floats. + * + *

Unlike the previous implementation — which copied each row into an on-heap + * {@code float[dimension]} scratch buffer on every {@link #getVector} call — this version + * delegates to {@link ByteBufferRandomAccessVectorValues} over the {@link MappedByteBuffer}, + * so reads are served directly from the mmap region with no per-call {@code float[]} + * allocation. + */ +public class MMapRandomAccessVectorValues implements RandomAccessVectorValues, Closeable +{ + private final int dimension; + private final File file; + private final RandomAccessFile raf; + private final FileChannel channel; + private final ByteBufferRandomAccessVectorValues delegate; - public MMapRandomAccessVectorValues(File f, int dimension) { + public MMapRandomAccessVectorValues(File f, int dimension) + { assert f != null && f.exists() && f.canRead(); - assert f.length() % ((long) dimension * Float.BYTES) == 0; + long bytesPerVector = (long) dimension * Float.BYTES; + assert f.length() % bytesPerVector == 0; try { this.file = f; - this.fileReader = new MMapBuffer(f, FileChannel.MapMode.READ_ONLY, ByteOrder.LITTLE_ENDIAN); this.dimension = dimension; - this.rows = ((int) f.length()) / dimension; - this.valueBuffer = new float[dimension]; + this.raf = new RandomAccessFile(f, "r"); + this.channel = raf.getChannel(); + long size = f.length(); + MappedByteBuffer mapped = channel.map(FileChannel.MapMode.READ_ONLY, 0, size); + mapped.order(ByteOrder.LITTLE_ENDIAN); + int count = (int) (size / bytesPerVector); + this.delegate = new ByteBufferRandomAccessVectorValues(mapped, count, dimension); } catch (IOException e) { throw new IOError(e); } } @Override - public int size() { - return (int) (file.length() / ((long) dimension * Float.BYTES)); + public int size() + { + return delegate.size(); } @Override - public int dimension() { + public int dimension() + { return dimension; } @Override - public VectorFloat getVector(int targetOrd) { - long offset = (long) targetOrd * dimension * Float.BYTES; - int i = 0; - for (long o = offset; o < offset + ((long) dimension * Float.BYTES); o += Float.BYTES, i++) - valueBuffer[i] = fileReader.memory().getFloat(o); - - return vectorTypeSupport.createFloatVector(valueBuffer); + public VectorFloat getVector(int targetOrd) + { + return delegate.getVector(targetOrd); } @Override - public boolean isValueShared() { - return false; + public boolean isValueShared() + { + return delegate.isValueShared(); } @Override - public RandomAccessVectorValues copy() { + public RandomAccessVectorValues copy() + { return new MMapRandomAccessVectorValues(file, dimension); } @Override - public void close() { + public void close() + { try { - this.fileReader.close(); + channel.close(); + raf.close(); } catch (IOException e) { throw new IOError(e); } diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/util/MMapRandomAccessVectorValuesTest.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/util/MMapRandomAccessVectorValuesTest.java new file mode 100644 index 000000000..713f2fee1 --- /dev/null +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/util/MMapRandomAccessVectorValuesTest.java @@ -0,0 +1,66 @@ +/* + * Copyright DataStax, Inc. + * + * 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 io.github.jbellis.jvector.example.util; + +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import java.io.File; +import java.io.IOException; +import java.io.RandomAccessFile; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.Random; + +import static org.junit.Assert.assertEquals; + +public class MMapRandomAccessVectorValuesTest +{ + @Rule + public TemporaryFolder tmp = new TemporaryFolder(); + + @Test + public void readsMatchSourceDataWithNoFloatArrayPerCall() throws IOException + { + int count = 100, dim = 16; + Random r = new Random(2026); + float[][] source = new float[count][dim]; + for (int i = 0; i < count; i++) { + for (int j = 0; j < dim; j++) source[i][j] = r.nextFloat() * 2f - 1f; + } + + File f = tmp.newFile("vectors.bin"); + ByteBuffer bb = ByteBuffer.allocate(count * dim * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + for (float[] v : source) for (float x : v) bb.putFloat(x); + bb.rewind(); + try (RandomAccessFile raf = new RandomAccessFile(f, "rw")) { + raf.getChannel().write(bb); + } + + try (MMapRandomAccessVectorValues ravv = new MMapRandomAccessVectorValues(f, dim)) { + assertEquals(count, ravv.size()); + assertEquals(dim, ravv.dimension()); + for (int ord = 0; ord < count; ord++) { + VectorFloat v = ravv.getVector(ord); + for (int i = 0; i < dim; i++) { + assertEquals(source[ord][i], v.get(i), 0f); + } + } + } + } +} diff --git a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/MemorySegmentVectorFloat.java b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/MemorySegmentVectorFloat.java index 3e3f67e1f..88900d925 100644 --- a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/MemorySegmentVectorFloat.java +++ b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/MemorySegmentVectorFloat.java @@ -22,6 +22,7 @@ import java.io.IOException; import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; import java.nio.Buffer; /** @@ -35,6 +36,11 @@ final public class MemorySegmentVectorFloat implements VectorFloatcopies the buffer contents. To share storage without + * copying, use {@link #wrap(java.nio.ByteBuffer)} instead. + */ + @Deprecated MemorySegmentVectorFloat(Buffer buffer) { this(buffer.remaining()); segment.copyFrom(MemorySegment.ofBuffer(buffer)); @@ -44,6 +50,31 @@ final public class MemorySegmentVectorFloat implements VectorFloatIn contrast to {@link #MemorySegmentVectorFloat(Buffer)} which eagerly copies into a + * fresh segment, this factory wraps the buffer's storage in place via + * {@link MemorySegment#ofBuffer(java.nio.Buffer)}. + */ + public static MemorySegmentVectorFloat wrap(java.nio.ByteBuffer buffer) { + if ((buffer.remaining() % Float.BYTES) != 0) { + throw new IllegalArgumentException( + "ByteBuffer remaining() must be a multiple of Float.BYTES, was " + buffer.remaining()); + } + if (buffer.order() != java.nio.ByteOrder.LITTLE_ENDIAN) { + throw new IllegalArgumentException( + "Native SIMD path requires ByteOrder.LITTLE_ENDIAN, got " + buffer.order()); + } + return new MemorySegmentVectorFloat(MemorySegment.ofBuffer(buffer)); + } + @Override public long ramBytesUsed() { @@ -61,15 +92,25 @@ public MemorySegment get() @Override public float get(int n) { - // this is (unfortunately) meaningfully better performing than getting at an offset in the memory segment - return ((float[])segment.heapBase().get())[n]; + // Fast path for on-heap segments: direct float[] indexing is meaningfully faster than + // the generic MemorySegment accessors. Fall back to segment access for off-heap buffers + // (e.g., native memory, direct ByteBuffers) where heapBase() is empty. + var heap = segment.heapBase(); + if (heap.isPresent()) { + return ((float[]) heap.get())[n]; + } + return segment.getAtIndex(ValueLayout.JAVA_FLOAT, n); } @Override public void set(int n, float value) { - // this is (unfortunately) meaningfully better performing than setting at an offset in the memory segment - ((float[])segment.heapBase().get())[n] = value; + var heap = segment.heapBase(); + if (heap.isPresent()) { + ((float[]) heap.get())[n] = value; + return; + } + segment.setAtIndex(ValueLayout.JAVA_FLOAT, n, value); } @Override @@ -103,12 +144,34 @@ public VectorFloat copy() return copy; } + @Override + public VectorFloat subview(int floatOffset, int floatLength) + { + int len = length(); + if (floatOffset < 0 || floatLength < 0 || (long) floatOffset + floatLength > len) { + throw new IllegalArgumentException( + "subview [" + floatOffset + "," + (floatOffset + floatLength) + + ") out of range for length " + len); + } + if (floatOffset == 0 && floatLength == len) { + return this; + } + return new MemorySegmentVectorFloat( + segment.asSlice((long) floatOffset * Float.BYTES, (long) floatLength * Float.BYTES)); + } + @Override public void copyFrom(VectorFloat src, int srcOffset, int destOffset, int length) { - MemorySegmentVectorFloat csrc = (MemorySegmentVectorFloat) src; - segment.asSlice((long) destOffset * Float.BYTES, (long) length * Float.BYTES) - .copyFrom(csrc.segment.asSlice((long) srcOffset * Float.BYTES, (long) length * Float.BYTES)); + if (src instanceof MemorySegmentVectorFloat csrc) { + segment.asSlice((long) destOffset * Float.BYTES, (long) length * Float.BYTES) + .copyFrom(csrc.segment.asSlice((long) srcOffset * Float.BYTES, (long) length * Float.BYTES)); + return; + } + // generic fallback for ArrayVectorFloat, BufferVectorFloat, or any other VectorFloat impl + for (int i = 0; i < length; i++) { + set(destOffset + i, src.get(srcOffset + i)); + } } @Override diff --git a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/MemorySegmentVectorProvider.java b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/MemorySegmentVectorProvider.java index 1ce0d81b2..14e4bc7f1 100644 --- a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/MemorySegmentVectorProvider.java +++ b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/MemorySegmentVectorProvider.java @@ -24,6 +24,7 @@ import java.io.IOException; import java.nio.Buffer; +import java.nio.ByteBuffer; /** * VectorTypeSupport using MemorySegments. @@ -45,6 +46,23 @@ public VectorFloat createFloatVector(int length) return new MemorySegmentVectorFloat(length); } + /** + * Zero-copy wrap that returns a MemorySegment-backed view when the buffer's layout matches + * the native SIMD contract (little-endian), so {@code FloatVector.fromMemorySegment} can + * run on it directly. Big-endian buffers fall through to the base {@link BufferVectorFloat} + * view — still zero-copy but dispatched through the Panama polymorphic path. + */ + @Override + public VectorFloat wrapFloatVector(ByteBuffer data, int floatOffset, int floatLength) + { + if (data.order() != java.nio.ByteOrder.LITTLE_ENDIAN) { + return VectorTypeSupport.super.wrapFloatVector(data, floatOffset, floatLength); + } + int startByte = data.position() + floatOffset * Float.BYTES; + ByteBuffer slice = data.slice(startByte, floatLength * Float.BYTES).order(data.order()); + return MemorySegmentVectorFloat.wrap(slice); + } + @Override public VectorFloat readFloatVector(RandomAccessReader r, int size) throws IOException { diff --git a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/NativeVectorUtilSupport.java b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/NativeVectorUtilSupport.java index 48cd7d66e..56481894f 100644 --- a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/NativeVectorUtilSupport.java +++ b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/NativeVectorUtilSupport.java @@ -38,32 +38,53 @@ public NativeVectorUtilSupport() {} @Override protected FloatVector fromVectorFloat(VectorSpecies SPEC, VectorFloat vector, int offset) { - return FloatVector.fromMemorySegment(SPEC, ((MemorySegmentVectorFloat) vector).get(), vector.offset(offset), ByteOrder.LITTLE_ENDIAN); + if (vector instanceof MemorySegmentVectorFloat msv) { + return FloatVector.fromMemorySegment(SPEC, msv.get(), vector.offset(offset), ByteOrder.LITTLE_ENDIAN); + } + return super.fromVectorFloat(SPEC, vector, offset); } @Override protected FloatVector fromVectorFloat(VectorSpecies SPEC, VectorFloat vector, int offset, int[] indices, int indicesOffset) { - throw new UnsupportedOperationException("Assembly not supported with memory segments."); + if (vector instanceof MemorySegmentVectorFloat) { + throw new UnsupportedOperationException("Assembly not supported with memory segments."); + } + return super.fromVectorFloat(SPEC, vector, offset, indices, indicesOffset); } @Override protected void intoVectorFloat(FloatVector vector, VectorFloat v, int offset) { - vector.intoMemorySegment(((MemorySegmentVectorFloat) v).get(), v.offset(offset), ByteOrder.LITTLE_ENDIAN); + if (v instanceof MemorySegmentVectorFloat msv) { + vector.intoMemorySegment(msv.get(), v.offset(offset), ByteOrder.LITTLE_ENDIAN); + return; + } + super.intoVectorFloat(vector, v, offset); } @Override protected ByteVector fromByteSequence(VectorSpecies SPEC, ByteSequence vector, int offset) { - return ByteVector.fromMemorySegment(SPEC, ((MemorySegmentByteSequence) vector).get(), offset, ByteOrder.LITTLE_ENDIAN); + if (vector instanceof MemorySegmentByteSequence msb) { + return ByteVector.fromMemorySegment(SPEC, msb.get(), offset, ByteOrder.LITTLE_ENDIAN); + } + return super.fromByteSequence(SPEC, vector, offset); } @Override protected void intoByteSequence(ByteVector vector, ByteSequence v, int offset) { - vector.intoMemorySegment(((MemorySegmentByteSequence) v).get(), offset, ByteOrder.LITTLE_ENDIAN); + if (v instanceof MemorySegmentByteSequence msb) { + vector.intoMemorySegment(msb.get(), offset, ByteOrder.LITTLE_ENDIAN); + return; + } + super.intoByteSequence(vector, v, offset); } @Override protected void intoByteSequence(ByteVector vector, ByteSequence v, int offset, VectorMask mask) { - vector.intoMemorySegment(((MemorySegmentByteSequence) v).get(), offset, ByteOrder.LITTLE_ENDIAN, mask); + if (v instanceof MemorySegmentByteSequence msb) { + vector.intoMemorySegment(msb.get(), offset, ByteOrder.LITTLE_ENDIAN, mask); + return; + } + super.intoByteSequence(vector, v, offset, mask); } @Override diff --git a/jvector-native/src/test/java/io/github/jbellis/jvector/vector/MemorySegmentVectorFloatWrapTest.java b/jvector-native/src/test/java/io/github/jbellis/jvector/vector/MemorySegmentVectorFloatWrapTest.java new file mode 100644 index 000000000..3a6b3d363 --- /dev/null +++ b/jvector-native/src/test/java/io/github/jbellis/jvector/vector/MemorySegmentVectorFloatWrapTest.java @@ -0,0 +1,71 @@ +/* + * Copyright DataStax, Inc. + * + * 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 io.github.jbellis.jvector.vector; + +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; + +import static org.junit.jupiter.api.Assertions.*; + +class MemorySegmentVectorFloatWrapTest +{ + @Test + void wrapAliasesBufferStorage() + { + ByteBuffer bb = ByteBuffer.allocateDirect(4 * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + for (int i = 0; i < 4; i++) bb.putFloat(i * 2f); + bb.rewind(); + + MemorySegmentVectorFloat v = MemorySegmentVectorFloat.wrap(bb); + assertEquals(4, v.length()); + assertEquals(0f, v.get(0), 0f); + assertEquals(6f, v.get(3), 0f); + + // mutate through the original ByteBuffer, observe the change through the view + bb.putFloat(0, 99f); + assertEquals(99f, v.get(0), 0f, "wrap must not copy — shared storage"); + } + + @Test + void wrapRejectsBigEndian() + { + ByteBuffer be = ByteBuffer.allocate(4 * Float.BYTES).order(ByteOrder.BIG_ENDIAN); + assertThrows(IllegalArgumentException.class, () -> MemorySegmentVectorFloat.wrap(be)); + } + + @Test + void wrapRejectsNonFloatAlignedRemaining() + { + ByteBuffer odd = ByteBuffer.allocate(7).order(ByteOrder.LITTLE_ENDIAN); + assertThrows(IllegalArgumentException.class, () -> MemorySegmentVectorFloat.wrap(odd)); + } + + @Test + void legacyCopyingConstructorStillCopies() + { + ByteBuffer bb = ByteBuffer.allocate(4 * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + for (int i = 0; i < 4; i++) bb.putFloat(i * 2f); + bb.rewind(); + + @SuppressWarnings("deprecation") + MemorySegmentVectorFloat v = new MemorySegmentVectorFloat(bb); + // Mutating bb must NOT change the copy — proving the semantics of the old ctor. + bb.putFloat(0, 999f); + assertEquals(0f, v.get(0), 0f, "copying constructor must not alias"); + } +} diff --git a/jvector-tests/pom.xml b/jvector-tests/pom.xml index 7f3145211..360002df9 100644 --- a/jvector-tests/pom.xml +++ b/jvector-tests/pom.xml @@ -175,25 +175,5 @@ false - - jdk20 - - - io.github.jbellis - jvector-twenty - ${project.version} - compile - - - - false - - - - jdk11 - - false - - diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/BuildFromByteBufferEquivalenceTest.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/BuildFromByteBufferEquivalenceTest.java new file mode 100644 index 000000000..b96354813 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/BuildFromByteBufferEquivalenceTest.java @@ -0,0 +1,143 @@ +/* + * Copyright DataStax, Inc. + * + * 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 io.github.jbellis.jvector.graph; + +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; +import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.ArrayList; +import java.util.List; +import java.util.Random; + +import static org.junit.Assert.assertEquals; + +/** + * Proves that building a graph from a caller-owned {@link ByteBuffer} produces the same + * structure as building from a {@code float[]}-based RAVV using the same seeded data. This + * is the end-to-end check that the ByteBuffer migration preserves semantics across the + * full build hot path (VectorFloat abstraction, SIMD dispatch, score provider, diversity + * pruning, etc.). + * + *

Runs under whichever vectorization provider is active (Default / Panama / Native) + * — so the same test catches regressions across all SIMD backends. + */ +public class BuildFromByteBufferEquivalenceTest +{ + private static final VectorTypeSupport VTS = + VectorizationProvider.getInstance().getVectorTypeSupport(); + + private static float[][] generate(long seed, int count, int dim) + { + Random r = new Random(seed); + float[][] data = new float[count][dim]; + for (int i = 0; i < count; i++) { + for (int j = 0; j < dim; j++) { + data[i][j] = r.nextFloat() * 2f - 1f; + } + } + return data; + } + + private static ByteBuffer toLittleEndian(float[][] data) + { + int count = data.length, dim = data[0].length; + ByteBuffer bb = ByteBuffer.allocate(count * dim * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + for (float[] v : data) for (float f : v) bb.putFloat(f); + bb.rewind(); + return bb; + } + + private static List> toVectorFloatList(float[][] data) + { + List> out = new ArrayList<>(data.length); + for (float[] v : data) { + out.add(VTS.createFloatVector(java.util.Arrays.copyOf(v, v.length))); + } + return out; + } + + private void assertEquivalent(int count, int dim, VectorSimilarityFunction vsf, long seed) + throws Exception + { + float[][] data = generate(seed, count, dim); + + // A: classic float[] build + RandomAccessVectorValues ravvA = new ListRandomAccessVectorValues(toVectorFloatList(data), dim); + var bspA = BuildScoreProvider.randomAccessScoreProvider(ravvA, vsf); + GraphIndexBuilder builderA = new GraphIndexBuilder(bspA, dim, 16, 100, 1.2f, 1.2f, false); + ImmutableGraphIndex graphA = TestUtil.buildSequentially(builderA, ravvA); + + // B: zero-copy ByteBuffer build + ByteBuffer bb = toLittleEndian(data); + RandomAccessVectorValues ravvB = new ByteBufferRandomAccessVectorValues(bb, count, dim); + var bspB = BuildScoreProvider.randomAccessScoreProvider(ravvB, vsf); + GraphIndexBuilder builderB = new GraphIndexBuilder(bspB, dim, 16, 100, 1.2f, 1.2f, false); + ImmutableGraphIndex graphB = TestUtil.buildSequentially(builderB, ravvB); + + TestUtil.assertGraphEquals(graphA, graphB); + + // Query a handful of points against each and assert identical ordering and scores. + Random qr = new Random(~seed); + for (int q = 0; q < 10; q++) { + int queryOrd = qr.nextInt(count); + var queryFromList = ravvA.getVector(queryOrd); + var queryFromBB = ravvB.getVector(queryOrd); + + var resA = GraphSearcher.search(queryFromList, 5, ravvA, vsf, graphA, Bits.ALL); + var resB = GraphSearcher.search(queryFromBB, 5, ravvB, vsf, graphB, Bits.ALL); + assertEquals("topK count differs", resA.getNodes().length, resB.getNodes().length); + for (int i = 0; i < resA.getNodes().length; i++) { + assertEquals("node ordering mismatch at i=" + i, resA.getNodes()[i].node, resB.getNodes()[i].node); + assertEquals("score mismatch at i=" + i, resA.getNodes()[i].score, resB.getNodes()[i].score, 1e-5f); + } + } + + builderA.close(); + builderB.close(); + } + + @Test + public void euclideanEquivalence() throws Exception + { + assertEquivalent(200, 64, VectorSimilarityFunction.EUCLIDEAN, 0xC001_C0DEL); + } + + @Test + public void dotProductEquivalence() throws Exception + { + assertEquivalent(150, 48, VectorSimilarityFunction.DOT_PRODUCT, 0xDEAD_BEEFL); + } + + @Test + public void cosineEquivalence() throws Exception + { + assertEquivalent(120, 32, VectorSimilarityFunction.COSINE, 0xABCDEF01L); + } + + @Test + public void largerDimensionEquivalence() throws Exception + { + assertEquivalent(64, 256, VectorSimilarityFunction.EUCLIDEAN, 0x5EED_5EEDL); + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/ByteBufferRandomAccessVectorValuesTest.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/ByteBufferRandomAccessVectorValuesTest.java new file mode 100644 index 000000000..5680c7f0f --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/ByteBufferRandomAccessVectorValuesTest.java @@ -0,0 +1,166 @@ +/* + * Copyright DataStax, Inc. + * + * 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 io.github.jbellis.jvector.graph; + +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.ArrayList; +import java.util.List; +import java.util.Random; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.*; + +class ByteBufferRandomAccessVectorValuesTest +{ + private static final VectorTypeSupport VTS = + VectorizationProvider.getInstance().getVectorTypeSupport(); + + private static float[][] randomData(long seed, int count, int dimension) + { + Random r = new Random(seed); + float[][] data = new float[count][dimension]; + for (int i = 0; i < count; i++) { + for (int j = 0; j < dimension; j++) { + data[i][j] = r.nextFloat() * 2f - 1f; + } + } + return data; + } + + private static ByteBuffer toByteBuffer(float[][] data, int dimension) + { + ByteBuffer bb = ByteBuffer.allocate(data.length * dimension * Float.BYTES) + .order(ByteOrder.LITTLE_ENDIAN); + for (float[] v : data) for (float f : v) bb.putFloat(f); + bb.rewind(); + return bb; + } + + @Test + void matchesListRAVVElementByElement() + { + int count = 100, dim = 32; + float[][] data = randomData(42, count, dim); + ByteBuffer bb = toByteBuffer(data, dim); + + ByteBufferRandomAccessVectorValues bbravv = new ByteBufferRandomAccessVectorValues(bb, count, dim); + List> list = new ArrayList<>(); + for (float[] v : data) list.add(VTS.createFloatVector(v)); + ListRandomAccessVectorValues listRavv = new ListRandomAccessVectorValues(list, dim); + + assertEquals(count, bbravv.size()); + assertEquals(dim, bbravv.dimension()); + + for (int ord = 0; ord < count; ord++) { + VectorFloat a = bbravv.getVector(ord); + VectorFloat b = listRavv.getVector(ord); + assertEquals(b.length(), a.length()); + for (int i = 0; i < dim; i++) { + assertEquals(b.get(i), a.get(i), 0f, "mismatch at ord=" + ord + " i=" + i); + } + } + } + + @Test + void rejectsOutOfRangeOrdinal() + { + ByteBuffer bb = ByteBuffer.allocate(10 * 4 * 4).order(ByteOrder.LITTLE_ENDIAN); + ByteBufferRandomAccessVectorValues ravv = new ByteBufferRandomAccessVectorValues(bb, 10, 4); + assertThrows(IndexOutOfBoundsException.class, () -> ravv.getVector(-1)); + assertThrows(IndexOutOfBoundsException.class, () -> ravv.getVector(10)); + } + + @Test + void rejectsUndersizedBuffer() + { + ByteBuffer tooSmall = ByteBuffer.allocate(10).order(ByteOrder.LITTLE_ENDIAN); + assertThrows(IllegalArgumentException.class, + () -> new ByteBufferRandomAccessVectorValues(tooSmall, 100, 4)); + } + + @Test + void survivesCallerBufferPositionMutation() + { + int count = 5, dim = 4; + float[][] data = randomData(7, count, dim); + ByteBuffer bb = toByteBuffer(data, dim); + + ByteBufferRandomAccessVectorValues ravv = new ByteBufferRandomAccessVectorValues(bb, count, dim); + // Caller fiddles with the source buffer after construction — must not disturb the view. + bb.position(bb.capacity()).limit(bb.capacity()); + VectorFloat v = ravv.getVector(2); + for (int i = 0; i < dim; i++) { + assertEquals(data[2][i], v.get(i), 0f); + } + } + + @Test + void threadLocalSupplierGivesIndependentWorkers() + throws InterruptedException + { + int count = 50, dim = 16; + float[][] data = randomData(11, count, dim); + ByteBuffer bb = toByteBuffer(data, dim); + + ByteBufferRandomAccessVectorValues ravv = new ByteBufferRandomAccessVectorValues(bb, count, dim); + ExecutorService exec = Executors.newFixedThreadPool(4); + CountDownLatch start = new CountDownLatch(1); + AtomicInteger mismatches = new AtomicInteger(); + for (int t = 0; t < 4; t++) { + exec.submit(() -> { + try { + start.await(); + RandomAccessVectorValues local = ravv.threadLocalSupplier().get(); + for (int r = 0; r < 2000; r++) { + int ord = r % count; + VectorFloat v = local.getVector(ord); + for (int i = 0; i < dim; i++) { + if (v.get(i) != data[ord][i]) mismatches.incrementAndGet(); + } + } + } catch (InterruptedException ignore) {} + }); + } + start.countDown(); + exec.shutdown(); + assertTrue(exec.awaitTermination(30, TimeUnit.SECONDS)); + assertEquals(0, mismatches.get(), "concurrent reads must agree with source data"); + } + + @Test + void fromFloatsFactoryRoundTrip() + { + int count = 8, dim = 5; + float[][] data = randomData(13, count, dim); + ByteBufferRandomAccessVectorValues ravv = ByteBufferRandomAccessVectorValues.fromFloats(data, dim); + assertEquals(count, ravv.size()); + assertEquals(dim, ravv.dimension()); + for (int ord = 0; ord < count; ord++) { + VectorFloat v = ravv.getVector(ord); + for (int i = 0; i < dim; i++) assertEquals(data[ord][i], v.get(i), 0f); + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/SearchAllocationProfileTest.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/SearchAllocationProfileTest.java new file mode 100644 index 000000000..9f82ef344 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/SearchAllocationProfileTest.java @@ -0,0 +1,137 @@ +/* + * Copyright DataStax, Inc. + * + * 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 io.github.jbellis.jvector.graph; + +import com.sun.management.ThreadMXBean; +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; +import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; +import org.junit.Test; + +import java.lang.management.ManagementFactory; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.ArrayList; +import java.util.List; +import java.util.Random; + +import static org.junit.Assert.assertTrue; + +/** + * Characterization test: compares per-query heap allocation when searching against a + * {@link ByteBufferRandomAccessVectorValues} vs a classic float[]-backed + * {@link ListRandomAccessVectorValues} on the same data. The ByteBuffer path should allocate + * no more than the float[] path (and in practice a bit less, since it avoids the + * per-call scratch that classic RAVV implementations tend to allocate). + * + *

The test intentionally does not assert an absolute byte budget — search scaffolding + * (result heaps, visited sets, reranker caches) legitimately allocates and the numbers vary + * across JDK versions and GC configurations. The comparative assertion is stable. + * + *

Measurement uses {@code com.sun.management.ThreadMXBean.getThreadAllocatedBytes}. + */ +public class SearchAllocationProfileTest +{ + private static final VectorTypeSupport VTS = + VectorizationProvider.getInstance().getVectorTypeSupport(); + + @Test + public void byteBufferPathAllocatesNoMoreThanFloatArrayPath() throws Exception + { + int count = 500, dim = 64, searches = 200, topK = 10; + long seed = 0xBA5E_BA11L; + Random r = new Random(seed); + + float[][] raw = new float[count][dim]; + for (int i = 0; i < count; i++) { + for (int j = 0; j < dim; j++) { + raw[i][j] = r.nextFloat() * 2f - 1f; + } + } + + // float[] path + List> list = new ArrayList<>(); + for (float[] row : raw) list.add(VTS.createFloatVector(row)); + var ravvFloat = new ListRandomAccessVectorValues(list, dim); + + // ByteBuffer path + ByteBuffer bb = ByteBuffer.allocate(count * dim * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + for (float[] row : raw) for (float f : row) bb.putFloat(f); + bb.rewind(); + var ravvBB = new ByteBufferRandomAccessVectorValues(bb, count, dim); + + var bspFloat = BuildScoreProvider.randomAccessScoreProvider(ravvFloat, VectorSimilarityFunction.EUCLIDEAN); + var bspBB = BuildScoreProvider.randomAccessScoreProvider(ravvBB, VectorSimilarityFunction.EUCLIDEAN); + + try (var builderFloat = new GraphIndexBuilder(bspFloat, dim, 16, 100, 1.2f, 1.2f, false); + var builderBB = new GraphIndexBuilder(bspBB, dim, 16, 100, 1.2f, 1.2f, false)) + { + TestUtil.buildSequentially(builderFloat, ravvFloat); + TestUtil.buildSequentially(builderBB, ravvBB); + var graphFloat = builderFloat.getGraph(); + var graphBB = builderBB.getGraph(); + + var queries = list.subList(0, 20); + + // Warm-up both paths + for (int i = 0; i < 50; i++) { + GraphSearcher.search(queries.get(i % queries.size()), topK, ravvFloat, + VectorSimilarityFunction.EUCLIDEAN, graphFloat, Bits.ALL); + GraphSearcher.search(queries.get(i % queries.size()), topK, ravvBB, + VectorSimilarityFunction.EUCLIDEAN, graphBB, Bits.ALL); + } + + ThreadMXBean tmx = (ThreadMXBean) ManagementFactory.getThreadMXBean(); + long tid = Thread.currentThread().getId(); + + long beforeFloat = tmx.getThreadAllocatedBytes(tid); + for (int i = 0; i < searches; i++) { + GraphSearcher.search(queries.get(i % queries.size()), topK, ravvFloat, + VectorSimilarityFunction.EUCLIDEAN, graphFloat, Bits.ALL); + } + long afterFloat = tmx.getThreadAllocatedBytes(tid); + long perQueryFloat = (afterFloat - beforeFloat) / searches; + + long beforeBB = tmx.getThreadAllocatedBytes(tid); + for (int i = 0; i < searches; i++) { + GraphSearcher.search(queries.get(i % queries.size()), topK, ravvBB, + VectorSimilarityFunction.EUCLIDEAN, graphBB, Bits.ALL); + } + long afterBB = tmx.getThreadAllocatedBytes(tid); + long perQueryBB = (afterBB - beforeBB) / searches; + + double ratio = perQueryFloat == 0 ? 0 : ((double) perQueryBB / perQueryFloat); + System.out.printf("SearchAllocationProfileTest: float[] path %d B/query, ByteBuffer path %d B/query (ratio %.2fx, dim=%d, topK=%d, count=%d)%n", + perQueryFloat, perQueryBB, ratio, dim, topK, count); + + // Current implementation allocates a small BufferVectorFloat view per getVector call + // (~80 B including an internal ByteBuffer slice). With thousands of node visits per + // query that adds up, but it is still orders of magnitude below the float[] + // materialization that a naive integration would do (dim*4 = 256 B per visit for + // dim=64). This assertion guards against the naive-materialization regression: + // float[] path × 3 is roughly the point where we'd have lost the zero-copy property. + long budget = Math.max(perQueryFloat * 3, 512_000L); + assertTrue( + "ByteBuffer path allocates " + perQueryBB + " B/query, float[] path " + + perQueryFloat + " B/query (ratio " + ratio + "), budget " + budget, + perQueryBB <= budget); + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestSearchWithByteBufferQuery.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestSearchWithByteBufferQuery.java new file mode 100644 index 000000000..88dacaad7 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestSearchWithByteBufferQuery.java @@ -0,0 +1,117 @@ +/* + * Copyright DataStax, Inc. + * + * 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 io.github.jbellis.jvector.graph; + +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; +import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.ArrayList; +import java.util.List; +import java.util.Random; + +import static org.junit.Assert.assertEquals; + +public class TestSearchWithByteBufferQuery +{ + private static final VectorTypeSupport VTS = + VectorizationProvider.getInstance().getVectorTypeSupport(); + + @Test + public void byteBufferQueryMatchesVectorFloatQuery() throws Exception + { + int count = 200, dim = 48; + long seed = 0xFACEB00CL; + Random r = new Random(seed); + + List> list = new ArrayList<>(); + for (int i = 0; i < count; i++) { + VectorFloat v = VTS.createFloatVector(dim); + for (int j = 0; j < dim; j++) v.set(j, r.nextFloat() * 2f - 1f); + list.add(v); + } + var ravv = new ListRandomAccessVectorValues(list, dim); + + var bsp = BuildScoreProvider.randomAccessScoreProvider(ravv, VectorSimilarityFunction.DOT_PRODUCT); + try (var builder = new GraphIndexBuilder(bsp, dim, 12, 80, 1.2f, 1.2f, false)) { + TestUtil.buildSequentially(builder, ravv); + var graph = builder.getGraph(); + + for (int q = 0; q < 10; q++) { + int qOrd = r.nextInt(count); + VectorFloat queryAsVectorFloat = list.get(qOrd); + + ByteBuffer queryAsByteBuffer = ByteBuffer.allocate(dim * Float.BYTES) + .order(ByteOrder.LITTLE_ENDIAN); + for (int j = 0; j < dim; j++) queryAsByteBuffer.putFloat(queryAsVectorFloat.get(j)); + queryAsByteBuffer.rewind(); + + var resVF = GraphSearcher.search(queryAsVectorFloat, 5, ravv, VectorSimilarityFunction.DOT_PRODUCT, graph, Bits.ALL); + var resBB = GraphSearcher.search(queryAsByteBuffer, 5, ravv, VectorSimilarityFunction.DOT_PRODUCT, graph, Bits.ALL); + + assertEquals(resVF.getNodes().length, resBB.getNodes().length); + for (int i = 0; i < resVF.getNodes().length; i++) { + assertEquals("node mismatch at i=" + i, resVF.getNodes()[i].node, resBB.getNodes()[i].node); + assertEquals("score mismatch at i=" + i, resVF.getNodes()[i].score, resBB.getNodes()[i].score, 1e-5f); + } + } + } + } + + @Test + public void addGraphNodeByteBufferMatchesVectorFloat() throws Exception + { + int count = 50, dim = 24; + long seed = 0xC0FFEEL; + Random r = new Random(seed); + List> list = new ArrayList<>(); + for (int i = 0; i < count; i++) { + VectorFloat v = VTS.createFloatVector(dim); + for (int j = 0; j < dim; j++) v.set(j, r.nextFloat() * 2f - 1f); + list.add(v); + } + var ravvA = new ListRandomAccessVectorValues(list, dim); + + // Build A — using VectorFloat + var bspA = BuildScoreProvider.randomAccessScoreProvider(ravvA, VectorSimilarityFunction.EUCLIDEAN); + var builderA = new GraphIndexBuilder(bspA, dim, 10, 60, 1.2f, 1.2f, false); + for (int i = 0; i < count; i++) builderA.addGraphNode(i, list.get(i)); + builderA.cleanup(); + + // Build B — using ByteBuffer overload + var ravvB = new ListRandomAccessVectorValues(list, dim); + var bspB = BuildScoreProvider.randomAccessScoreProvider(ravvB, VectorSimilarityFunction.EUCLIDEAN); + var builderB = new GraphIndexBuilder(bspB, dim, 10, 60, 1.2f, 1.2f, false); + for (int i = 0; i < count; i++) { + ByteBuffer bb = ByteBuffer.allocate(dim * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + for (int j = 0; j < dim; j++) bb.putFloat(list.get(i).get(j)); + bb.rewind(); + builderB.addGraphNode(i, bb); + } + builderB.cleanup(); + + TestUtil.assertGraphEquals(builderA.getGraph(), builderB.getGraph()); + builderA.close(); + builderB.close(); + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/PQTrainingAllocationTest.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/PQTrainingAllocationTest.java new file mode 100644 index 000000000..62caef4d1 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/PQTrainingAllocationTest.java @@ -0,0 +1,113 @@ +/* + * Copyright DataStax, Inc. + * + * 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 io.github.jbellis.jvector.quantization; + +import com.sun.management.ThreadMXBean; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.vector.ArraySliceVectorFloat; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; +import org.junit.Test; + +import java.lang.management.ManagementFactory; +import java.util.ArrayList; +import java.util.List; +import java.util.Random; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +/** + * Verifies that PQ codebook training does not allocate a fresh {@code float[]} per + * (training-vector × subspace). The subvector extraction path now returns zero-copy views + * via {@link ArraySliceVectorFloat}, so the per-subvector allocation that showed up in + * async-profiler flamegraphs during HerdDB indexing should be gone. + * + *

Uses {@code ThreadMXBean.getThreadAllocatedBytes} to measure; the assertion is relative + * to the number of training vectors (not an absolute byte budget) so the test is stable + * across JDK versions and SIMD backends. + */ +public class PQTrainingAllocationTest +{ + private static final VectorTypeSupport VTS = + VectorizationProvider.getInstance().getVectorTypeSupport(); + + @Test + public void getSubVectorReturnsViewNotCopy() + { + int dim = 128; + VectorFloat v = VTS.createFloatVector(dim); + for (int i = 0; i < dim; i++) v.set(i, i * 0.5f); + + int[][] ranges = {{16, 0}, {16, 16}, {16, 32}}; + VectorFloat sub0 = ProductQuantization.getSubVector(v, 0, ranges); + VectorFloat sub1 = ProductQuantization.getSubVector(v, 1, ranges); + assertEquals(16, sub0.length()); + assertEquals(16, sub1.length()); + + // Mutating the source must reflect through the views — proves no copy. + v.set(5, 999f); // index 5 is in sub0 (offset 0..15) + assertEquals(999f, sub0.get(5), 0f); + + v.set(20, 888f); // index 20 is in sub1 (offset 16..31) + assertEquals(888f, sub1.get(4), 0f); + } + + @Test + public void pqTrainingAllocationStaysBoundedPerVector() throws Exception + { + int dim = 64; + int n = 2_000; + int m = 8; + int clusterCount = 16; + + Random r = new Random(0xC0DEFACE); + List> training = new ArrayList<>(); + for (int i = 0; i < n; i++) { + VectorFloat v = VTS.createFloatVector(dim); + for (int j = 0; j < dim; j++) v.set(j, r.nextFloat() * 2f - 1f); + training.add(v); + } + RandomAccessVectorValues ravv = new ListRandomAccessVectorValues(training, dim); + + // Warm-up — trigger class loading, JIT + ProductQuantization.compute(ravv, m, clusterCount, false); + + ThreadMXBean tmx = (ThreadMXBean) ManagementFactory.getThreadMXBean(); + long tid = Thread.currentThread().getId(); + long before = tmx.getThreadAllocatedBytes(tid); + ProductQuantization pq = ProductQuantization.compute(ravv, m, clusterCount, false); + long after = tmx.getThreadAllocatedBytes(tid); + + long perTrainingVector = (after - before) / n; + System.out.printf("PQTrainingAllocationTest: %d B allocated per training vector (n=%d, dim=%d, M=%d, K=%d)%n", + perTrainingVector, n, dim, m, clusterCount); + + // If getSubVector materialized every subvector, we'd allocate M*(dim/M)*4 = dim*4 + // = 256 B per training vector just for subvectors, plus wrapper overhead. With + // view-based extraction the per-vector cost is dominated by KMeans centroids, which + // amortize over the whole training set. Guard at 2x the pre-fix baseline minus + // subvectors; if someone reintroduces the allocation, the ratio explodes. + long budget = (long) dim * Float.BYTES * 2; + assertTrue("per-vector allocation " + perTrainingVector + " exceeds budget " + budget + + " — getSubVector may have regressed to materialization", + perTrainingVector < budget); + assertEquals(m, pq.getSubspaceCount()); + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/BufferVectorFloatTest.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/BufferVectorFloatTest.java new file mode 100644 index 000000000..900729a1a --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/BufferVectorFloatTest.java @@ -0,0 +1,221 @@ +/* + * Copyright DataStax, Inc. + * + * 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 io.github.jbellis.jvector.vector; + +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; + +import static org.junit.jupiter.api.Assertions.*; + +class BufferVectorFloatTest +{ + private static ByteBuffer littleEndianOf(float... values) + { + ByteBuffer bb = ByteBuffer.allocate(values.length * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + for (float v : values) bb.putFloat(v); + bb.rewind(); + return bb; + } + + private static ByteBuffer bigEndianOf(float... values) + { + ByteBuffer bb = ByteBuffer.allocate(values.length * Float.BYTES).order(ByteOrder.BIG_ENDIAN); + for (float v : values) bb.putFloat(v); + bb.rewind(); + return bb; + } + + @Test + void elementAccessMatchesFloatArrayLittleEndian() + { + float[] expected = {1.0f, -2.5f, 3.14f, 0.0f, Float.MAX_VALUE}; + BufferVectorFloat bv = new BufferVectorFloat(littleEndianOf(expected)); + assertEquals(expected.length, bv.length()); + for (int i = 0; i < expected.length; i++) { + assertEquals(expected[i], bv.get(i), 0f); + } + } + + @Test + void elementAccessMatchesFloatArrayBigEndian() + { + float[] expected = {1.0f, -2.5f, 3.14f, 0.0f, Float.MAX_VALUE}; + BufferVectorFloat bv = new BufferVectorFloat(bigEndianOf(expected)); + for (int i = 0; i < expected.length; i++) { + assertEquals(expected[i], bv.get(i), 0f); + } + } + + @Test + void directBufferBackingIsZeroCopy() + { + ByteBuffer direct = ByteBuffer.allocateDirect(4 * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + for (int i = 0; i < 4; i++) direct.putFloat(i * 1.5f); + direct.rewind(); + + BufferVectorFloat bv = new BufferVectorFloat(direct); + // The view is a slice (independent ByteBuffer object) over the same off-heap storage. + // Proving no data copy: mutating via the original updates the view. + assertFalse(direct.hasArray(), "direct buffer — verifies off-heap path"); + direct.putFloat(0, 999f); + assertEquals(999f, bv.get(0), 0f, "slice must alias source storage, not copy"); + } + + @Test + void subrangeViewHonorsBoundsAndDoesNotAlias() + { + ByteBuffer bb = littleEndianOf(10f, 20f, 30f, 40f, 50f); + BufferVectorFloat middle = new BufferVectorFloat(bb, 1, 3); // {20, 30, 40} + assertEquals(3, middle.length()); + assertEquals(20f, middle.get(0), 0f); + assertEquals(30f, middle.get(1), 0f); + assertEquals(40f, middle.get(2), 0f); + assertThrows(IndexOutOfBoundsException.class, () -> middle.get(3)); + } + + @Test + void positionLimitMutationAfterConstructionDoesNotAffectView() + { + ByteBuffer bb = littleEndianOf(1f, 2f, 3f, 4f); + BufferVectorFloat bv = new BufferVectorFloat(bb); + bb.position(8).limit(12); + assertEquals(1f, bv.get(0), 0f); + assertEquals(4f, bv.get(3), 0f); + } + + @Test + void zeroAllocationConstructorProducesZeroedVector() + { + BufferVectorFloat bv = new BufferVectorFloat(7); + assertEquals(7, bv.length()); + for (int i = 0; i < 7; i++) { + assertEquals(0f, bv.get(i), 0f); + } + } + + @Test + void setRoundTrip() + { + BufferVectorFloat bv = new BufferVectorFloat(4); + bv.set(0, 1f); bv.set(1, 2f); bv.set(2, -3f); bv.set(3, 0.5f); + assertEquals(1f, bv.get(0), 0f); + assertEquals(-3f, bv.get(2), 0f); + } + + @Test + void zeroResets() + { + BufferVectorFloat bv = new BufferVectorFloat(littleEndianOf(1f, 2f, 3f).duplicate().order(ByteOrder.LITTLE_ENDIAN)); + bv.zero(); + for (int i = 0; i < 3; i++) assertEquals(0f, bv.get(i), 0f); + } + + @Test + void copyProducesIndependentStorage() + { + BufferVectorFloat src = new BufferVectorFloat(littleEndianOf(1f, 2f, 3f)); + VectorFloat dup = src.copy(); + src.set(0, 99f); + assertEquals(99f, src.get(0), 0f); + assertEquals(1f, dup.get(0), 0f, "copy must not alias source"); + } + + @Test + void copyFromArrayVectorFloat() + { + BufferVectorFloat dest = new BufferVectorFloat(4); + ArrayVectorFloat src = new ArrayVectorFloat(new float[]{10f, 20f, 30f, 40f}); + dest.copyFrom(src, 1, 0, 3); + assertEquals(20f, dest.get(0), 0f); + assertEquals(30f, dest.get(1), 0f); + assertEquals(40f, dest.get(2), 0f); + assertEquals(0f, dest.get(3), 0f); + } + + @Test + void copyFromBufferVectorFloatSameOrder() + { + BufferVectorFloat dest = new BufferVectorFloat(4); + BufferVectorFloat src = new BufferVectorFloat(littleEndianOf(1f, 2f, 3f, 4f, 5f)); + dest.copyFrom(src, 1, 0, 3); + assertEquals(2f, dest.get(0), 0f); + assertEquals(3f, dest.get(1), 0f); + assertEquals(4f, dest.get(2), 0f); + } + + @Test + void copyFromBufferVectorFloatDifferentOrder() + { + BufferVectorFloat dest = new BufferVectorFloat(3); // native order + BufferVectorFloat src = new BufferVectorFloat(bigEndianOf(1f, 2f, 3f)); + dest.copyFrom(src, 0, 0, 3); + for (int i = 0; i < 3; i++) assertEquals((float) (i + 1), dest.get(i), 0f); + } + + @Test + void equalsMatchesSemanticContent() + { + BufferVectorFloat le = new BufferVectorFloat(littleEndianOf(1f, 2f, 3f)); + BufferVectorFloat be = new BufferVectorFloat(bigEndianOf(1f, 2f, 3f)); + ArrayVectorFloat av = new ArrayVectorFloat(new float[]{1f, 2f, 3f}); + assertEquals(le, be, "byte order differs but content identical"); + assertEquals(le, av, "BufferVectorFloat equals any VectorFloat with same content"); + // ArrayVectorFloat.equals() is class-strict — we don't change that contract here. + } + + @Test + void rejectsNonFloatAlignedRemaining() + { + ByteBuffer odd = ByteBuffer.allocate(7).order(ByteOrder.LITTLE_ENDIAN); + assertThrows(IllegalArgumentException.class, () -> new BufferVectorFloat(odd)); + } + + @Test + void rejectsNegativeArgs() + { + ByteBuffer ok = littleEndianOf(1f, 2f); + assertThrows(IllegalArgumentException.class, () -> new BufferVectorFloat(ok, -1, 1)); + assertThrows(IllegalArgumentException.class, () -> new BufferVectorFloat(ok, 0, -1)); + assertThrows(IllegalArgumentException.class, () -> new BufferVectorFloat(-1)); + } + + @Test + void rejectsOutOfRangeSubview() + { + ByteBuffer bb = littleEndianOf(1f, 2f, 3f); + assertThrows(IllegalArgumentException.class, () -> new BufferVectorFloat(bb, 2, 2)); + } + + @Test + void ramBytesUsedIsFinite() + { + BufferVectorFloat bv = new BufferVectorFloat(1024); + long size = bv.ramBytesUsed(); + assertTrue(size >= 1024L * Float.BYTES, "reported size must cover backing buffer"); + } + + @Test + void byteOrderPreserved() + { + BufferVectorFloat le = new BufferVectorFloat(littleEndianOf(1f)); + BufferVectorFloat be = new BufferVectorFloat(bigEndianOf(1f)); + assertEquals(ByteOrder.LITTLE_ENDIAN, le.byteOrder()); + assertEquals(ByteOrder.BIG_ENDIAN, be.byteOrder()); + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/VectorFloatSubviewTest.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/VectorFloatSubviewTest.java new file mode 100644 index 000000000..aa980ed49 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/VectorFloatSubviewTest.java @@ -0,0 +1,135 @@ +/* + * Copyright DataStax, Inc. + * + * 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 io.github.jbellis.jvector.vector; + +import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; + +import static org.junit.jupiter.api.Assertions.*; + +class VectorFloatSubviewTest +{ + private static final VectorTypeSupport VTS = + VectorizationProvider.getInstance().getVectorTypeSupport(); + + @Test + void arrayVectorFloatSubviewSharesBackingStorage() + { + float[] raw = {10f, 20f, 30f, 40f, 50f}; + ArrayVectorFloat src = new ArrayVectorFloat(raw); + VectorFloat view = src.subview(1, 3); + assertInstanceOf(ArraySliceVectorFloat.class, view); + assertEquals(3, view.length()); + assertEquals(20f, view.get(0), 0f); + assertEquals(40f, view.get(2), 0f); + // mutating the source is visible through the view + raw[1] = 99f; + assertEquals(99f, view.get(0), 0f); + } + + @Test + void bufferVectorFloatSubviewSharesBackingStorage() + { + ByteBuffer bb = ByteBuffer.allocate(6 * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + for (int i = 0; i < 6; i++) bb.putFloat(i * 1.5f); + bb.rewind(); + BufferVectorFloat src = new BufferVectorFloat(bb); + VectorFloat view = src.subview(2, 3); + assertInstanceOf(BufferVectorFloat.class, view); + assertEquals(3, view.length()); + assertEquals(3f, view.get(0), 0f); + assertEquals(6f, view.get(2), 0f); + // mutating the source is visible through the view + bb.putFloat(2 * Float.BYTES, 99f); + assertEquals(99f, view.get(0), 0f); + } + + @Test + void nestedSubviewIsAliasNotCopy() + { + float[] raw = new float[20]; + for (int i = 0; i < raw.length; i++) raw[i] = i; + ArrayVectorFloat src = new ArrayVectorFloat(raw); + VectorFloat outer = src.subview(5, 10); // elements 5..14 + VectorFloat inner = outer.subview(2, 3); // elements 7..9 + assertEquals(3, inner.length()); + assertEquals(7f, inner.get(0), 0f); + assertEquals(9f, inner.get(2), 0f); + raw[7] = 777f; + assertEquals(777f, inner.get(0), 0f); + } + + @Test + void fullLengthSubviewReturnsSelf() + { + ArrayVectorFloat src = new ArrayVectorFloat(new float[]{1f, 2f, 3f}); + assertSame(src, src.subview(0, 3), "full-range subview should return the source"); + } + + @Test + void subviewParticipatesInSimdDistanceComputation() + { + // Build two full-size vectors; assert that squareL2Distance over (view, view) yields + // the same result as (materialized subvector, materialized subvector). + int dim = 64, sub = 16, off = 16; + float[] a = new float[dim], b = new float[dim]; + for (int i = 0; i < dim; i++) { + a[i] = (float) Math.sin(i * 0.1); + b[i] = (float) Math.cos(i * 0.1); + } + ArrayVectorFloat av = new ArrayVectorFloat(a); + ArrayVectorFloat bv = new ArrayVectorFloat(b); + + float viaView = VectorUtil.squareL2Distance(av.subview(off, sub), bv.subview(off, sub)); + float viaCopy = VectorUtil.squareL2Distance( + VTS.createFloatVector(java.util.Arrays.copyOfRange(a, off, off + sub)), + VTS.createFloatVector(java.util.Arrays.copyOfRange(b, off, off + sub))); + assertEquals(viaCopy, viaView, 1e-5f); + } + + @Test + void subviewMatchesDotProductAndCosine() + { + int dim = 48, off = 8, sub = 24; + float[] a = new float[dim], b = new float[dim]; + for (int i = 0; i < dim; i++) { + a[i] = 1f + i * 0.25f; + b[i] = 2f - i * 0.125f; + } + VectorFloat av = VTS.createFloatVector(a); + VectorFloat bv = VTS.createFloatVector(b); + VectorFloat aCopy = VTS.createFloatVector(java.util.Arrays.copyOfRange(a, off, off + sub)); + VectorFloat bCopy = VTS.createFloatVector(java.util.Arrays.copyOfRange(b, off, off + sub)); + + assertEquals(VectorUtil.dotProduct(aCopy, bCopy), + VectorUtil.dotProduct(av.subview(off, sub), bv.subview(off, sub)), 1e-5f); + assertEquals(VectorUtil.cosine(aCopy, bCopy), + VectorUtil.cosine(av.subview(off, sub), bv.subview(off, sub)), 1e-5f); + } + + @Test + void rejectsOutOfRangeSubview() + { + ArrayVectorFloat src = new ArrayVectorFloat(new float[5]); + assertThrows(IllegalArgumentException.class, () -> src.subview(-1, 2)); + assertThrows(IllegalArgumentException.class, () -> src.subview(0, 6)); + assertThrows(IllegalArgumentException.class, () -> src.subview(3, 3)); + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/VectorTypeSupportByteBufferTest.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/VectorTypeSupportByteBufferTest.java new file mode 100644 index 000000000..6054f5acb --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/VectorTypeSupportByteBufferTest.java @@ -0,0 +1,107 @@ +/* + * Copyright DataStax, Inc. + * + * 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 io.github.jbellis.jvector.vector; + +import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; + +import static org.junit.jupiter.api.Assertions.*; + +class VectorTypeSupportByteBufferTest +{ + private static final VectorTypeSupport VTS = + VectorizationProvider.getInstance().getVectorTypeSupport(); + + private static final VectorTypeSupport DEFAULT_VTS = + new DefaultVectorizationProvider().getVectorTypeSupport(); + + private static ByteBuffer le(float... floats) + { + ByteBuffer bb = ByteBuffer.allocate(floats.length * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + for (float v : floats) bb.putFloat(v); + bb.rewind(); + return bb; + } + + @Test + void wrapReturnsReadableVectorAcrossActiveProvider() + { + ByteBuffer bb = le(1f, 2f, 3f, 4f); + VectorFloat v = VTS.wrapFloatVector(bb); + assertEquals(4, v.length()); + assertEquals(1f, v.get(0), 0f); + assertEquals(4f, v.get(3), 0f); + } + + @Test + void wrapAcrossDefaultProvider() + { + ByteBuffer bb = le(1f, 2f, 3f, 4f); + VectorFloat v = DEFAULT_VTS.wrapFloatVector(bb); + assertEquals(4, v.length()); + for (int i = 0; i < 4; i++) { + assertEquals((float) (i + 1), v.get(i), 0f); + } + } + + @Test + void wrapZeroCopyOnDirectBuffer() + { + ByteBuffer direct = ByteBuffer.allocateDirect(4 * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); + for (int i = 0; i < 4; i++) direct.putFloat(i * 1.25f); + direct.rewind(); + VectorFloat v = DEFAULT_VTS.wrapFloatVector(direct); + // Mutate the source buffer; the view must see the change (proving no copy happened). + direct.putFloat(0, 999f); + assertEquals(999f, v.get(0), 0f); + } + + @Test + void wrapWithSubrange() + { + ByteBuffer bb = le(10f, 20f, 30f, 40f, 50f); + VectorFloat v = DEFAULT_VTS.wrapFloatVector(bb, 1, 3); + assertEquals(3, v.length()); + assertEquals(20f, v.get(0), 0f); + assertEquals(40f, v.get(2), 0f); + } + + @Test + void wrapRejectsBadAlignment() + { + ByteBuffer odd = ByteBuffer.allocate(7).order(ByteOrder.LITTLE_ENDIAN); + assertThrows(IllegalArgumentException.class, () -> VTS.wrapFloatVector(odd)); + } + + @Test + void wrapEquivalentAcrossProviders() + { + float[] raw = {-0.5f, 1.5f, 2.75f, -3.125f, 4.5f}; + ByteBuffer bb = le(raw); + + VectorFloat defaultView = DEFAULT_VTS.wrapFloatVector(bb.duplicate().order(ByteOrder.LITTLE_ENDIAN)); + VectorFloat activeView = VTS.wrapFloatVector(bb.duplicate().order(ByteOrder.LITTLE_ENDIAN)); + + assertEquals(defaultView.length(), activeView.length()); + for (int i = 0; i < raw.length; i++) { + assertEquals(defaultView.get(i), activeView.get(i), 0f); + } + } +} diff --git a/jvector-twenty/pom.xml b/jvector-twenty/pom.xml index ae6aa659b..ab288ae33 100644 --- a/jvector-twenty/pom.xml +++ b/jvector-twenty/pom.xml @@ -18,9 +18,8 @@ maven-compiler-plugin 3.11.0 - 20 + 22 - --add-modules jdk.incubator.vector @@ -33,7 +32,6 @@ false - --enable-preview --add-modules jdk.incubator.vector --enable-native-access=ALL-UNNAMED diff --git a/jvector-twenty/src/main/java/io/github/jbellis/jvector/vector/PanamaVectorUtilSupport.java b/jvector-twenty/src/main/java/io/github/jbellis/jvector/vector/PanamaVectorUtilSupport.java index 22e0d2c60..51b9eceea 100644 --- a/jvector-twenty/src/main/java/io/github/jbellis/jvector/vector/PanamaVectorUtilSupport.java +++ b/jvector-twenty/src/main/java/io/github/jbellis/jvector/vector/PanamaVectorUtilSupport.java @@ -27,6 +27,7 @@ import jdk.incubator.vector.VectorOperators; import jdk.incubator.vector.VectorSpecies; +import java.lang.foreign.MemorySegment; import java.util.List; class PanamaVectorUtilSupport implements VectorUtilSupport { @@ -43,14 +44,49 @@ class PanamaVectorUtilSupport implements VectorUtilSupport { static final ThreadLocal scratchInt256 = ThreadLocal.withInitial(() -> new int[IntVector.SPECIES_256.length()]); protected FloatVector fromVectorFloat(VectorSpecies SPEC, VectorFloat vector, int offset) { + if (vector instanceof BufferVectorFloat bv) { + // bv.get() is an independent slice whose position 0 corresponds to element 0, + // so the byte offset into the segment is simply offset * Float.BYTES. + return FloatVector.fromMemorySegment( + SPEC, + MemorySegment.ofBuffer(bv.get()), + (long) offset * Float.BYTES, + bv.byteOrder()); + } + if (vector instanceof ArraySliceVectorFloat asv) { + return FloatVector.fromArray(SPEC, asv.get(), asv.arrayOffset() + offset); + } return FloatVector.fromArray(SPEC, ((ArrayVectorFloat) vector).get(), offset); } protected FloatVector fromVectorFloat(VectorSpecies SPEC, VectorFloat vector, int offset, int[] indices, int indicesOffset) { + if (vector instanceof BufferVectorFloat) { + // FloatVector has no gather-from-MemorySegment variant; fall back to a short scalar gather. + // The scratch array is at most SPEC.length() floats (≤16), allocated per call. + float[] scratch = new float[SPEC.length()]; + for (int i = 0; i < SPEC.length(); i++) { + scratch[i] = vector.get(offset + indices[indicesOffset + i]); + } + return FloatVector.fromArray(SPEC, scratch, 0); + } + if (vector instanceof ArraySliceVectorFloat asv) { + return FloatVector.fromArray(SPEC, asv.get(), asv.arrayOffset() + offset, indices, indicesOffset); + } return FloatVector.fromArray(SPEC, ((ArrayVectorFloat)vector).get(), offset, indices, indicesOffset); } protected void intoVectorFloat(FloatVector vector, VectorFloat v, int offset) { + if (v instanceof BufferVectorFloat bv) { + vector.intoMemorySegment( + MemorySegment.ofBuffer(bv.get()), + (long) offset * Float.BYTES, + bv.byteOrder()); + return; + } + if (v instanceof ArraySliceVectorFloat asv) { + vector.intoArray(asv.get(), asv.arrayOffset() + offset); + return; + } vector.intoArray(((ArrayVectorFloat) v).get(), offset); } diff --git a/pom.xml b/pom.xml index 102f75d6a..24867d037 100644 --- a/pom.xml +++ b/pom.xml @@ -41,7 +41,7 @@ UTF-8 true - 4.0.0-rc.9-SNAPSHOT + 4.0.0-rc.9-herddb-SNAPSHOT jvector-base @@ -83,7 +83,7 @@ maven-compiler-plugin 3.11.0 - 11 + 22