diff --git a/UPGRADING.md b/UPGRADING.md index c99de606f..999e0c448 100644 --- a/UPGRADING.md +++ b/UPGRADING.md @@ -1,3 +1,16 @@ +# Upgrading from 4.0.x to 4.1.x + +## On-disk format version 7 +- Version 7 of the on-disk graph format adds one feature, inline int8 vectors (`FeatureId.INLINE_BYTE_VECTORS`): + each node's byte vector stored inline, one byte per dimension. Apart from that feature, a version 7 file has + the version 6 layout. +- JVector 4.0.x reads formats up to 6 and rejects version 7 files. `OnDiskGraphIndex.CURRENT_VERSION`, + the default for `OnDiskGraphIndexWriter.Builder` and for `CompressedVectors.write(out)`, is now 7. +- To keep a rollback to 4.0.x possible, pass the version explicitly until every reader runs 4.1: + `withVersion(6)` on graph writer builders and `write(out, 6)` for compressors and compressed vectors. + A version 6 writer rejects `INLINE_BYTE_VECTORS`. +- `OnDiskGraphIndexCompactor` always writes `CURRENT_VERSION`; it has no option to write an older version. + # Upgrading from 3.0.x to 4.0.x ## New features diff --git a/docs/release notes/4.1.0/709.feature.md b/docs/release notes/4.1.0/709.feature.md new file mode 100644 index 000000000..dcd6a71c0 --- /dev/null +++ b/docs/release notes/4.1.0/709.feature.md @@ -0,0 +1,53 @@ +### Native Int8 (Byte) Vector Support + +**Description** +JVector can now build, store and search signed int8 vectors natively, without converting them to float32. +This is useful when an embedding model already emits int8 vectors, or when vectors have been scaled to +int8 up front: they take 4x less space on disk and in memory than float32, and all scoring is done +byte-by-byte. + +* `RandomAccessByteVectorValues` (and `ListRandomAccessByteVectorValues`) provide random access to + `ByteSequence` vectors. It and `RandomAccessVectorValues` now share a common `VectorValues` + interface, which also provides `getVectorInto(node, destination, offset)` for reading a + vector into a caller-owned buffer. +* `ByteVectorSimilarityFunction` scores two int8 vectors: `EUCLIDEAN` is `1 / (1 + squaredL2)`, + `DOT_PRODUCT` is `dot < 0 ? 1 / (1 - dot) : dot + 1`, and `COSINE` is `(1 + cosine) / 2`. The scores are + not normalised by dimension, so distinct integer distances stay distinct at high dimensions. + `DOT_PRODUCT` scores are positive but not bounded above by 1. +* `GraphIndexBuilder.builder(RandomAccessByteVectorValues, ByteVectorSimilarityFunction, M)` builds a graph + from int8 vectors, and `GraphIndexBuilder.build(VectorValues)` accepts either vector type. +* The `InlineByteVectors` on-disk feature (`FeatureId.INLINE_BYTE_VECTORS`) stores each vector as `dimension` + raw bytes. `OnDiskGraphIndex.View.byteVectorRerankerFor(...)` scores candidates directly from disk, and + `getByteVectorInto(...)` reads a vector into a reusable buffer. +* Int8 distance kernels (squared L2, dot product, cosine) are available in the default, Panama and native + SIMD implementations, and accumulate in integers so all execution paths return identical results. +* `BenchYAML` and `AutoBenchYAML` support int8 datasets: mark a dataset with `is_int8: true` in + `dataset-metadata.yml`. A runnable example is in `Int8Example` (`-Dtutorial=int8`). + +**How to Enable** +No configuration is required; the float32 APIs are unchanged. To use int8 vectors, build with the byte-vector +builder and write the index with the `InlineByteVectors` feature: + +```java +RandomAccessByteVectorValues vectors = new ListRandomAccessByteVectorValues(byteVectors, dimension); +try (var builder = GraphIndexBuilder.builder(vectors, ByteVectorSimilarityFunction.EUCLIDEAN, 32).build()) { + var graph = builder.build(vectors); + try (var writer = GraphIndexWriter.getBuilderFor(GraphIndexWriterTypes.RANDOM_ACCESS_PARALLEL, graph, path) + .with(new InlineByteVectors(dimension)).build()) { + writer.write(Map.of(FeatureId.INLINE_BYTE_VECTORS, n -> new InlineByteVectors.State(vectors.getVector(n)))); + } +} +``` + +**Notes** +* The new interfaces and classes are marked `@Experimental` and the API may change in later releases. +* Only signed int8 (-128..127) is supported. Unsigned uint8 data, such as conventional SIFT1B/BIGANN `.bvecs` + files, must be converted first; it is not detected, and reading it as signed silently changes distances. +* Quantization and compression (PQ, BQ, NVQ, fused graphs) are not supported for int8 vectors. Traversal + already scores exactly, so a separate reranking step is not needed. `rerankK` still controls how much of the + graph is searched. +* Zero vectors are not supported with `COSINE` (the result is undefined); do not index them. +* Existing float32 callers keep working: `RandomAccessVectorValues.getVector` keeps its `VectorFloat` + signature and `GraphIndexBuilder.build(RandomAccessVectorValues)` is retained as a forwarding overload. +* Int8 indexes require on-disk format version 7 or later. JVector 4.0.x reads formats up to 6 and rejects + version 7 files; a version 6 writer rejects `INLINE_BYTE_VECTORS`. 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 ec7ad410a..7b2d30ceb 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 @@ -32,7 +32,9 @@ import io.github.jbellis.jvector.quantization.PQVectors; import io.github.jbellis.jvector.quantization.ProductQuantization; import io.github.jbellis.jvector.util.*; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.types.ByteSequence; import io.github.jbellis.jvector.vector.types.VectorFloat; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -114,6 +116,11 @@ private static BuildScoreProvider getBuildScoreProvider(RandomAccessVectorValues } } + private static BuildScoreProvider getBuildScoreProvider(RandomAccessByteVectorValues vectorValues, ByteVectorSimilarityFunction similarityFunction) { + // Byte vectors do not support PQ/BQ build-time compression; always use the direct provider. + return BuildScoreProvider.byteVectorScoreProvider(vectorValues, similarityFunction); + } + /** * Reads all the vectors from vector values, builds a graph connecting them by their dense * ordinals, using the given hyperparameter settings, and returns the resulting graph. @@ -407,6 +414,31 @@ public static Builder builder(RandomAccessVectorValues vectorValues, VectorSimil return builder(vectorValues, similarityFunction, List.of(M)); } + /** + * Entry point for the fluent builder for int8 byte vectors, resolving the {@link BuildScoreProvider} + * from raw byte vectors via {@link BuildScoreProvider#byteVectorScoreProvider}. + * + * @param vectorValues the int8 byte vectors whose relations are represented by the graph + * @param similarityFunction the byte-vector similarity function to score vectors with + * @param maxDegrees the maximum number of connections a node can have in each layer; if fewer entries + * are specified than the number of layers, the last entry is used for all remaining layers. + */ + public static Builder builder(RandomAccessByteVectorValues vectorValues, ByteVectorSimilarityFunction similarityFunction, List maxDegrees) { + return new Builder(getBuildScoreProvider(vectorValues, similarityFunction), vectorValues.dimension(), maxDegrees); + } + + /** + * Entry point for the fluent builder for int8 byte vectors, for the common case of a single + * (non-hierarchical) max degree. + * + * @param vectorValues the int8 byte vectors whose relations are represented by the graph + * @param similarityFunction the byte-vector similarity function to score vectors with + * @param M the maximum number of connections a node can have + */ + public static Builder builder(RandomAccessByteVectorValues vectorValues, ByteVectorSimilarityFunction similarityFunction, int M) { + return builder(vectorValues, similarityFunction, List.of(M)); + } + /** * Entry point for the fluent builder, building from an existing {@link MutableGraphIndex} (e.g. one just * loaded from disk) rather than constructing a fresh {@link OnHeapGraphIndex}. {@code addHierarchy} is not @@ -695,15 +727,42 @@ public static GraphIndexBuilder rescore(GraphIndexBuilder other, BuildScoreProvi return newBuilder; } + /** + * Retained for binary compatibility with callers compiled against the float-only signature. + * Equivalent to {@link #build(VectorValues)}. + */ public ImmutableGraphIndex build(RandomAccessVectorValues ravv) { - var vv = ravv.threadLocalSupplier(); + return build((VectorValues) ravv); + } + + + /** + * Builds the graph from any {@link VectorValues} source — works for both + * {@link RandomAccessVectorValues} (float32) and {@link RandomAccessByteVectorValues} (int8). + * The score provider supplied at construction time determines how vectors are compared. + */ + public ImmutableGraphIndex build(VectorValues ravv) { int size = ravv.size(); - simdExecutor.submit(() -> { - IntStream.range(0, size).parallel().forEach(node -> { - addGraphNode(node, vv.get().getVector(node)); - }); - }).join(); + if (ravv instanceof RandomAccessVectorValues) { + // Float sources may share (borrowed) vectors; read each through a per-thread view. + var vv = ((RandomAccessVectorValues) ravv).threadLocalSupplier(); + simdExecutor.submit(() -> { + IntStream.range(0, size).parallel().forEach(node -> { + addGraphNode(node, vv.get().getVector(node)); + }); + }).join(); + } else if (ravv instanceof RandomAccessByteVectorValues) { + // Same for int8 sources: read the supplied vectors rather than assuming the score provider's match. + var vv = ((RandomAccessByteVectorValues) ravv).threadLocalSupplier(); + simdExecutor.submit(() -> { + IntStream.range(0, size).parallel().forEach(node -> { + addGraphNode(node, vv.get().getVector(node)); + }); + }).join(); + } else { + throw new IllegalArgumentException("Unsupported VectorValues type: " + ravv.getClass().getName()); + } cleanup(); return graph; @@ -852,6 +911,19 @@ public long addGraphNode(int node, VectorFloat vector) { return addGraphNode(node, ssp); } + /** + * Inserts a node with the given int8 byte vector into the graph. + * + * @param node the node ID to add + * @param vector the byte vector to add + * @return an estimate of the number of extra bytes used by the graph after adding the given node + * @throws UnsupportedOperationException if this builder was not constructed with a byte-vector score provider + */ + public long addGraphNode(int node, ByteSequence vector) { + var ssp = scoreProvider.searchProviderFor(vector); + return addGraphNode(node, ssp); + } + /** * Inserts a node with the given vector value to the graph. * @@ -1343,4 +1415,4 @@ public static ImmutableGraphIndex buildAndMergeNewNodes(RandomAccessReader in, return builder.getGraph(); } } -} \ No newline at end of file +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ListRandomAccessByteVectorValues.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ListRandomAccessByteVectorValues.java new file mode 100644 index 000000000..861a5e71e --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ListRandomAccessByteVectorValues.java @@ -0,0 +1,72 @@ +/* + * 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.annotations.Experimental; +import io.github.jbellis.jvector.vector.types.ByteSequence; + +import java.util.List; + +/** + * A List-backed implementation of the {@link RandomAccessByteVectorValues} interface. + *

+ * It is acceptable to provide this class to a GraphBuilder, and then continue + * to add vectors to the backing List as you add to the graph. + *

+ * This will be as threadsafe as the provided List. + */ +@Experimental +public class ListRandomAccessByteVectorValues implements RandomAccessByteVectorValues { + private final List> vectors; + private final int dimension; + + /** + * Construct a new instance of {@link ListRandomAccessByteVectorValues}. + * + * @param vectors a (potentially mutable) list of byte vectors. + * @param dimension the dimension of the vectors. + */ + public ListRandomAccessByteVectorValues(List> vectors, int dimension) { + this.vectors = vectors; + this.dimension = dimension; + } + + @Override + public int size() { + return vectors.size(); + } + + @Override + public int dimension() { + return dimension; + } + + @Override + public ByteSequence getVector(int nodeId) { + return vectors.get(nodeId); + } + + @Override + public boolean isValueShared() { + return false; + } + + @Override + public ListRandomAccessByteVectorValues copy() { + return this; + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessByteVectorValues.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessByteVectorValues.java new file mode 100644 index 000000000..21a12399c --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessByteVectorValues.java @@ -0,0 +1,68 @@ +/* + * 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.annotations.Experimental; +import io.github.jbellis.jvector.util.ExplicitThreadLocal; +import io.github.jbellis.jvector.vector.types.ByteSequence; + +import java.util.function.Supplier; + +/** + * Provides random access to byte (int8) vectors by dense ordinal. + *

+ * This is the byte-vector parallel to {@link RandomAccessVectorValues}. + * Both extend the common super-interface {@link VectorValues}. + * It is used by graph-based index builders and searchers that operate natively + * on int8 vectors without a float32 round-trip. + */ +@Experimental +public interface RandomAccessByteVectorValues extends VectorValues> { + + /** + * Retrieve the vector associated with a given node, and store it in the destination vector at the given offset. + * The default implementation copies from {@link #getVector}; implementations that can read directly into the + * destination (e.g. from disk) should override it. + */ + @Override + default void getVectorInto(int node, ByteSequence destinationVector, int offset) { + destinationVector.copyFrom(getVector(node), 0, offset, dimension()); + } + + /** + * Creates a new copy of this {@link RandomAccessByteVectorValues}. + * Un-shared implementations may simply return {@code this}. + */ + @Override + RandomAccessByteVectorValues copy(); + + /** + * Returns a supplier of thread-local copies of the RAVV. + */ + @Override + default Supplier threadLocalSupplier() { + if (!isValueShared()) { + return () -> this; + } + + if (this instanceof AutoCloseable) { + VectorValues.LOG.warning("VectorValues is shared and implements AutoCloseable; threadLocalSupplier() may lead to leaks"); + } + var tl = ExplicitThreadLocal.withInitial(this::copy); + return tl::get; + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessVectorValues.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessVectorValues.java index eb8f6df24..aa3660d58 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessVectorValues.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessVectorValues.java @@ -30,14 +30,15 @@ import io.github.jbellis.jvector.vector.types.VectorFloat; import java.util.function.Supplier; -import java.util.logging.Logger; /** - * Provides random access to vectors by dense ordinal. This interface is used by graph-based + * Provides random access to float32 vectors by dense ordinal. This interface is used by graph-based * implementations of KNN search. + *

+ * For int8 vectors see {@link RandomAccessByteVectorValues}. + * Both extend the common super-interface {@link VectorValues}. */ -public interface RandomAccessVectorValues { - Logger LOG = Logger.getLogger(RandomAccessVectorValues.class.getName()); +public interface RandomAccessVectorValues extends VectorValues> { /** * Return the number of vector values. @@ -46,21 +47,22 @@ public interface RandomAccessVectorValues { * (1) implementing a threadsafe, un-shared RAVV, where `copy` returns `this`, or * (2) implementing a fixed-size RAVV. */ + @Override int size(); - /** Return the dimension of the returned vector values */ - int dimension(); - /** * Return the vector value indexed at the given ordinal. - * - *

For performance, implementations are free to re-use the same object across invocations. - * That is, you will get back the same VectorFloat<?> - * reference (for instance) for every requested ordinal. If you want to use those values across - * calls, you should make a copy. + *

+ * Redeclared here (rather than only inherited from {@link VectorValues}) so that the method keeps its + * {@code VectorFloat getVector(int)} binary signature; the generic inherited declaration erases to + * {@code Object getVector(int)}, which would break previously compiled callers and implementations. + *

+ * For performance, implementations are free to re-use the same object across invocations. + * If you need to retain the value across calls, make a copy. * * @param nodeId a valid ordinal, ≥ 0 and < {@link #size()}. */ + @Override VectorFloat getVector(int nodeId); @Deprecated @@ -74,16 +76,11 @@ default VectorFloat vectorValue(int targetOrd) { * @param destinationVector the vector to store the result in * @param offset the offset in the destination vector to store the result */ + @Override default void getVectorInto(int node, VectorFloat destinationVector, int offset) { destinationVector.copyFrom(getVector(node), 0, offset, dimension()); } - /** - * @return true iff the vector returned by `getVector` is shared. A shared vector will - * only be valid until the next call to getVector overwrites it. - */ - boolean isValueShared(); - /** * Creates a new copy of this {@link RandomAccessVectorValues}. This is helpful when you need to * access different values at once, to avoid overwriting the underlying float vector returned by @@ -91,18 +88,20 @@ default void getVectorInto(int node, VectorFloat destinationVector, int offse *

* Un-shared implementations may simply return `this`. */ + @Override RandomAccessVectorValues copy(); /** * Returns a supplier of thread-local copies of the RAVV. */ + @Override default Supplier threadLocalSupplier() { if (!isValueShared()) { return () -> this; } if (this instanceof AutoCloseable) { - LOG.warning("RAVV is shared and implements AutoCloseable; threadLocalSupplier() may lead to leaks"); + VectorValues.LOG.warning("VectorValues is shared and implements AutoCloseable; threadLocalSupplier() may lead to leaks"); } var tl = ExplicitThreadLocal.withInitial(this::copy); return tl::get; diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/VectorValues.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/VectorValues.java new file mode 100644 index 000000000..d42c6d260 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/VectorValues.java @@ -0,0 +1,90 @@ +/* + * 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.annotations.Experimental; +import io.github.jbellis.jvector.util.ExplicitThreadLocal; + +import java.util.function.Supplier; +import java.util.logging.Logger; + +/** + * Common super-interface for random access to vectors by dense ordinal. + *

+ * {@code V} is the vector element type — {@code VectorFloat} for float32 vectors + * (see {@link RandomAccessVectorValues}) and {@code ByteSequence} for int8 vectors + * (see {@link RandomAccessByteVectorValues}). + */ +@Experimental +public interface VectorValues { + Logger LOG = Logger.getLogger(VectorValues.class.getName()); + + /** Return the number of vector values. */ + int size(); + + /** Return the dimension of the returned vector values. */ + int dimension(); + + /** + * Return the vector value indexed at the given ordinal. + *

+ * For performance, implementations are free to re-use the same object across invocations. + * If you need to retain the value across calls, make a copy. + * + * @param nodeId a valid ordinal, ≥ 0 and < {@link #size()}. + */ + V getVector(int nodeId); + + /** + * Retrieve the vector associated with a given node, and store it in the destination vector at the given offset. + * Unlike {@link #getVector}, the result is owned by the caller, so this is the way to take a private copy + * (and, for sources that read from disk, to avoid an intermediate allocation). + * + * @param node the node to retrieve + * @param destinationVector the vector to store the result in + * @param offset the offset in the destination vector to store the result + */ + void getVectorInto(int node, V destinationVector, int offset); + + /** + * @return true iff the vector returned by {@link #getVector} is shared across calls. + * A shared vector is only valid until the next call to {@link #getVector} overwrites it. + */ + boolean isValueShared(); + + /** + * Creates a new copy of this instance. + * Un-shared implementations may simply return {@code this}. + */ + VectorValues copy(); + + /** + * Returns a supplier of thread-local copies of this instance. + */ + @SuppressWarnings("unchecked") + default Supplier> threadLocalSupplier() { + if (!isValueShared()) { + return () -> this; + } + + if (this instanceof AutoCloseable) { + LOG.warning("VectorValues is shared and implements AutoCloseable; threadLocalSupplier() may lead to leaks"); + } + var tl = ExplicitThreadLocal.withInitial(this::copy); + return tl::get; + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java index d43497944..12bb1d28a 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java @@ -381,7 +381,7 @@ public int commonHeaderSize() { public void writeHeaderFeatures(IndexWriter out, Map features) throws IOException { // we restrict pre-version-3 writers to INLINE_VECTORS features, so we don't need additional version-handling here for (Feature writer : features.values()) { - writer.writeHeader(out); + writer.writeHeader(out, getVersion()); } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java index 305b0850f..a68ab9bc0 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java @@ -20,6 +20,8 @@ import io.github.jbellis.jvector.graph.ImmutableGraphIndex; import io.github.jbellis.jvector.graph.disk.feature.Feature; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.FusedFeature; +import io.github.jbellis.jvector.graph.disk.feature.InlineByteVectors; import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; import io.github.jbellis.jvector.graph.disk.feature.NVQ; import io.github.jbellis.jvector.graph.disk.feature.SeparatedFeature; @@ -259,6 +261,8 @@ public K build() throws IOException { int dimension; if (features.containsKey(FeatureId.INLINE_VECTORS)) { dimension = ((InlineVectors) features.get(FeatureId.INLINE_VECTORS)).dimension(); + } else if (features.containsKey(FeatureId.INLINE_BYTE_VECTORS)) { + dimension = ((InlineByteVectors) features.get(FeatureId.INLINE_BYTE_VECTORS)).dimension(); } else if (features.containsKey(FeatureId.NVQ_VECTORS)) { dimension = ((NVQ) features.get(FeatureId.NVQ_VECTORS)).dimension(); } else if (features.containsKey(FeatureId.SEPARATED_VECTORS)) { diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatFactory.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatFactory.java index 12010bc1e..b7c0014a1 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatFactory.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatFactory.java @@ -35,7 +35,8 @@ public class GraphIndexFormatFactory { 3, new GraphIndexFormatV3(), 4, new GraphIndexFormatV4(), 5, new GraphIndexFormatV5(), - 6, new GraphIndexFormatV6() + 6, new GraphIndexFormatV6(), + 7, new GraphIndexFormatV7() ); /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java index c9d965aa4..a9a46ed5c 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java @@ -26,6 +26,7 @@ import java.util.EnumMap; import java.util.LinkedHashMap; import java.util.Map; +import java.util.Set; import java.util.function.IntFunction; import java.util.stream.Collectors; @@ -44,7 +45,12 @@ class GraphIndexFormatV6 extends GraphIndexFormatV5 { /** Creates the singleton format for version 6. */ GraphIndexFormatV6() { - super(6, allFeatures()); + this(6, allFeatures()); + } + + /** Allows later feature extensions to reuse the V6 record and hierarchy layout. */ + protected GraphIndexFormatV6(int version, Set supportedFeatures) { + super(version, supportedFeatures); } /** Places fused features last so non-fused inline features occupy a contiguous prefix. */ @@ -112,7 +118,7 @@ public void writeHeaderFeatures(IndexWriter out, Map supportedFeatures() { + var features = EnumSet.copyOf(allFeatures()); + features.add(FeatureId.INLINE_BYTE_VECTORS); + return features; + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java index f43357c27..c67461162 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java @@ -27,6 +27,7 @@ import io.github.jbellis.jvector.graph.disk.feature.FeatureSource; import io.github.jbellis.jvector.graph.disk.feature.FusedPQ; import io.github.jbellis.jvector.graph.disk.feature.FusedFeature; +import io.github.jbellis.jvector.graph.disk.feature.InlineByteVectors; import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; import io.github.jbellis.jvector.graph.disk.feature.NVQ; import io.github.jbellis.jvector.graph.disk.feature.SeparatedFeature; @@ -36,8 +37,10 @@ import java.util.ArrayList; import io.github.jbellis.jvector.util.Bits; import io.github.jbellis.jvector.util.RamUsageEstimator; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.ByteSequence; import io.github.jbellis.jvector.vector.types.VectorFloat; import io.github.jbellis.jvector.vector.types.VectorTypeSupport; @@ -68,7 +71,7 @@ public class OnDiskGraphIndex implements ImmutableGraphIndex, AutoCloseable, Accountable { private static final Logger logger = LoggerFactory.getLogger(OnDiskGraphIndex.class); - public static final int CURRENT_VERSION = 6; + public static final int CURRENT_VERSION = 7; static final int MAGIC = 0xFFFF0D61; // FFFF to distinguish from old graphs, which should never start with a negative size "ODGI" static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); final ReaderSupplier readerSupplier; @@ -578,6 +581,64 @@ public void getVectorInto(int node, VectorFloat vector, int offset) { } } + /** + * Returns the signed int8 vector stored for {@code node} via the + * {@link FeatureId#INLINE_BYTE_VECTORS} feature. + * + * @throws UnsupportedOperationException if the graph was not written with + * {@link InlineByteVectors} + */ + public ByteSequence getByteVector(int node) { + ByteSequence vec = vectorTypeSupport.createByteSequence(dimension); + getByteVectorInto(node, vec); + return vec; + } + + /** + * Reads the signed int8 vector stored for {@code node} into {@code destination}, which must have + * exactly {@link #dimension()} elements. Use this instead of {@link #getByteVector(int)} on hot paths + * to avoid allocating a new vector per read. + * + * @throws UnsupportedOperationException if the graph was not written with + * {@link InlineByteVectors} + */ + public void getByteVectorInto(int node, ByteSequence destination) { + if (!features.containsKey(FeatureId.INLINE_BYTE_VECTORS)) { + throw new UnsupportedOperationException("No inline byte vectors in this graph"); + } + try { + long diskOffset = offsetFor(node, FeatureId.INLINE_BYTE_VECTORS); + reader.seek(diskOffset); + vectorTypeSupport.readByteSequence(reader, destination); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + } + + /** + * Returns a {@link ScoreFunction.ExactScoreFunction} that scores candidates by reading + * their int8 vectors from disk and comparing them byte×byte against {@code queryBytes}. + * The resulting function reuses a single scratch vector and is NOT thread-safe. + * + * @throws UnsupportedOperationException if the graph was not written with + * {@link InlineByteVectors} + */ + public ScoreFunction.ExactScoreFunction byteVectorRerankerFor(ByteSequence queryBytes, + ByteVectorSimilarityFunction bvsf) { + if (!features.containsKey(FeatureId.INLINE_BYTE_VECTORS)) { + throw new UnsupportedOperationException("No inline byte vectors in this graph"); + } + return new ScoreFunction.ExactScoreFunction() { + private final ByteSequence scratch = vectorTypeSupport.createByteSequence(dimension); + + @Override + public float similarityTo(int node) { + getByteVectorInto(node, scratch); + return bvsf.compare(queryBytes, scratch); + } + }; + } + public NodesIterator getNeighborsIterator(int level, int node) { try { int[] stored; diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/Feature.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/Feature.java index 03dd2ec98..1602aa2fa 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/Feature.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/Feature.java @@ -36,8 +36,27 @@ default boolean isFused() { int featureSize(); + /** + * Writes this feature's header at {@link io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex#CURRENT_VERSION}. + * + * @deprecated a header embedded in a graph must be written at the graph's version; use + * {@link #writeHeader(IndexWriter, int)} + */ + @Deprecated void writeHeader(IndexWriter out) throws IOException; + /** + * Writes this feature's header for a graph written at {@code version}. Features that serialize a + * compressor write it at this version, so a graph written at an older version is laid out as that + * version wrote it. The default ignores the version, for headers that don't depend on it. + * + * @param out the output to write to + * @param version the on-disk format version of the graph being written + */ + default void writeHeader(IndexWriter out, int version) throws IOException { + writeHeader(out); + } + default void writeInline(IndexWriter out, State state) throws IOException { // default no-op } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FeatureId.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FeatureId.java index 131c4c8ee..37b62e43a 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FeatureId.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FeatureId.java @@ -33,7 +33,8 @@ public enum FeatureId { FUSED_PQ(FusedPQ::load), NVQ_VECTORS(NVQ::load), SEPARATED_VECTORS(SeparatedVectors::load), - SEPARATED_NVQ(SeparatedNVQ::load); + SEPARATED_NVQ(SeparatedNVQ::load), + INLINE_BYTE_VECTORS(InlineByteVectors::load); private final BiFunction loader; diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java index 3dc404481..a8d39aa1b 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java @@ -135,9 +135,15 @@ public QuantizationCompactionStrategy createCompactionStrategy(CompactionContext return new FusedCompactionStrategy(ctx, this, retrainer); } + @Deprecated @Override public void writeHeader(IndexWriter out) throws IOException { - pq.write(out, OnDiskGraphIndex.CURRENT_VERSION); + writeHeader(out, OnDiskGraphIndex.CURRENT_VERSION); + } + + @Override + public void writeHeader(IndexWriter out, int version) throws IOException { + pq.write(out, version); } // this is an awkward fit for the Feature.State design since we need to diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/InlineByteVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/InlineByteVectors.java new file mode 100644 index 000000000..981999b95 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/InlineByteVectors.java @@ -0,0 +1,99 @@ +/* + * 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.disk.feature; + +import io.github.jbellis.jvector.annotations.Experimental; +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.disk.CommonHeader; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.ByteSequence; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; + +import java.io.IOException; + +/** + * Stores signed int8 (byte) vectors inline in an {@link io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex}. + *

+ * Each vector occupies exactly {@code dimension} bytes on disk — 4× smaller than + * the float32 {@link InlineVectors} representation. The on-disk layout is otherwise + * identical: one contiguous byte block per node record, written and read via + * {@link VectorTypeSupport#writeByteSequence} and {@link VectorTypeSupport#readByteSequence}. + *

+ * Use {@link io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex.View#getByteVector(int)} + * to retrieve a stored vector at search time. + */ +@Experimental +public class InlineByteVectors extends AbstractFeature { + private static final VectorTypeSupport vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + + private final int dimension; + + public InlineByteVectors(int dimension) { + this.dimension = dimension; + } + + @Override + public FeatureId id() { + return FeatureId.INLINE_BYTE_VECTORS; + } + + /** No extra header bytes — dimension is already in the {@link CommonHeader}. */ + @Override + public int headerSize() { + return 0; + } + + /** One byte per component. */ + @Override + public int featureSize() { + return dimension; + } + + public int dimension() { + return dimension; + } + + static InlineByteVectors load(CommonHeader header, RandomAccessReader reader) { + return new InlineByteVectors(header.dimension); + } + + @Deprecated + @Override + public void writeHeader(IndexWriter out) { + writeHeader(out, OnDiskGraphIndex.CURRENT_VERSION); + } + + @Override + public void writeHeader(IndexWriter out, int version) { + // common header carries dimension; nothing extra needed + } + + @Override + public void writeInline(IndexWriter out, Feature.State state) throws IOException { + vts.writeByteSequence(out, ((State) state).vector); + } + + public static class State implements Feature.State { + public final ByteSequence vector; + + public State(ByteSequence vector) { + this.vector = vector; + } + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/InlineVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/InlineVectors.java index 005103ddb..cac33f6f2 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/InlineVectors.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/InlineVectors.java @@ -19,6 +19,7 @@ import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.disk.RandomAccessReader; import io.github.jbellis.jvector.graph.disk.CommonHeader; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.types.VectorFloat; import io.github.jbellis.jvector.vector.types.VectorTypeSupport; @@ -58,8 +59,14 @@ static InlineVectors load(CommonHeader header, RandomAccessReader reader) { return new InlineVectors(header.dimension); } + @Deprecated @Override public void writeHeader(IndexWriter out) { + writeHeader(out, OnDiskGraphIndex.CURRENT_VERSION); + } + + @Override + public void writeHeader(IndexWriter out, int version) { // common header contains dimension, which is sufficient } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/NVQ.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/NVQ.java index 492bb0dfe..863780bf1 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/NVQ.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/NVQ.java @@ -69,9 +69,15 @@ static NVQ load(CommonHeader header, RandomAccessReader reader) { } } + @Deprecated @Override public void writeHeader(IndexWriter out) throws IOException { - nvq.write(out, OnDiskGraphIndex.CURRENT_VERSION); + writeHeader(out, OnDiskGraphIndex.CURRENT_VERSION); + } + + @Override + public void writeHeader(IndexWriter out, int version) throws IOException { + nvq.write(out, version); } @Override diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/SeparatedNVQ.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/SeparatedNVQ.java index b65370632..a0e3ddd21 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/SeparatedNVQ.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/SeparatedNVQ.java @@ -67,9 +67,15 @@ public int featureSize() { return nvq.compressedVectorSize(); } + @Deprecated @Override public void writeHeader(IndexWriter out) throws IOException { - nvq.write(out, OnDiskGraphIndex.CURRENT_VERSION); + writeHeader(out, OnDiskGraphIndex.CURRENT_VERSION); + } + + @Override + public void writeHeader(IndexWriter out, int version) throws IOException { + nvq.write(out, version); out.writeLong(offset); } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/SeparatedVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/SeparatedVectors.java index 017aef53f..b399574af 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/SeparatedVectors.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/SeparatedVectors.java @@ -19,6 +19,7 @@ import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.disk.RandomAccessReader; import io.github.jbellis.jvector.graph.disk.CommonHeader; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.types.VectorTypeSupport; @@ -60,8 +61,14 @@ public int featureSize() { return dimension * Float.BYTES; } + @Deprecated @Override public void writeHeader(IndexWriter out) throws IOException { + writeHeader(out, OnDiskGraphIndex.CURRENT_VERSION); + } + + @Override + public void writeHeader(IndexWriter out, int version) throws IOException { out.writeLong(offset); } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProvider.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProvider.java index 8e1f91cf8..5ae9fc533 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProvider.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProvider.java @@ -16,15 +16,20 @@ package io.github.jbellis.jvector.graph.similarity; +import io.github.jbellis.jvector.graph.RandomAccessByteVectorValues; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.RemappedRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.VectorValues; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; import io.github.jbellis.jvector.quantization.BQVectors; import io.github.jbellis.jvector.quantization.PQVectors; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import io.github.jbellis.jvector.vector.VectorUtil; import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.ByteSequence; import io.github.jbellis.jvector.vector.types.VectorFloat; import io.github.jbellis.jvector.vector.types.VectorTypeSupport; +import java.util.function.Supplier; /** * Encapsulates comparing node distances for GraphIndexBuilder. @@ -59,6 +64,20 @@ public interface BuildScoreProvider { */ SearchScoreProvider searchProviderFor(VectorFloat vector); + /** + * Create a search score provider to use *internally* during construction, for a byte (int8) query vector. + *

+ * The default implementation throws {@link UnsupportedOperationException}; only + * {@link #byteVectorScoreProvider} overrides it. + * + * @param vector the int8 query vector to provide similarity scores against + */ + default SearchScoreProvider searchProviderFor(ByteSequence vector) { + throw new UnsupportedOperationException( + "This BuildScoreProvider does not support byte-vector queries; " + + "use byteVectorScoreProvider() to construct a byte-vector builder"); + } + /** * Create a search score provider to use *internally* during construction. *

@@ -106,8 +125,8 @@ static BuildScoreProvider randomAccessScoreProvider(RandomAccessVectorValues rav static BuildScoreProvider randomAccessScoreProvider(RandomAccessVectorValues ravv, VectorSimilarityFunction similarityFunction) { // We need two sources of vectors in order to perform diversity check comparisons without // colliding. ThreadLocalSupplier makes this a no-op if the RAVV is actually un-shared. - var vectors = ravv.threadLocalSupplier(); - var vectorsCopy = ravv.threadLocalSupplier(); + Supplier vectors = ravv.threadLocalSupplier(); + Supplier vectorsCopy = ravv.threadLocalSupplier(); return new BuildScoreProvider() { /// Exact scores against `node1`. @@ -123,7 +142,10 @@ private ScoreFunction.ExactScoreFunction scoreFunctionFor(int node1) { @Override public float similarityTo(int node2) { var a = anchor; - if (a == null) anchor = a = vectors.get().getVector(node1).copy(); + if (a == null) { + anchor = a = vts.createFloatVector(ravv.dimension()); + vectors.get().getVectorInto(node1, a, 0); + } return similarityFunction.compare(a, vectorsCopy.get().getVector(node2)); } }; @@ -225,6 +247,93 @@ public VectorFloat approximateCentroid() { }; } + + /** + * Returns a BSP that performs exact score comparisons using the given + * {@link RandomAccessByteVectorValues} and {@link ByteVectorSimilarityFunction}. + * All scoring is byte×byte with no float32 round-trip. + */ + static BuildScoreProvider byteVectorScoreProvider(RandomAccessByteVectorValues ravv, ByteVectorSimilarityFunction bvsf) { + var vectors = ravv.threadLocalSupplier(); + var vectorsCopy = ravv.threadLocalSupplier(); + + return new BuildScoreProvider() { + /// Exact scores against `node1`. + /// + /// For a value-shared source the anchor is read on the first comparison and kept as + /// a private copy, so creating a score function reads nothing, later reads through the + /// same views can't overwrite the anchor, and each comparison reads only `node2`. + private ScoreFunction.ExactScoreFunction scoreFunctionFor(int node1) { + if (ravv.isValueShared()) { + return new ScoreFunction.ExactScoreFunction() { + private ByteSequence anchor; + + @Override + public float similarityTo(int node2) { + var a = anchor; + if (a == null) { + anchor = a = vts.createByteSequence(ravv.dimension()); + vectors.get().getVectorInto(node1, a, 0); + } + return bvsf.compare(a, vectorsCopy.get().getVector(node2)); + } + }; + } + var vector = vectors.get().getVector(node1); + var vc = vectorsCopy.get(); + return node2 -> bvsf.compare(vector, vc.getVector(node2)); + } + + @Override + public boolean isExact() { + return true; + } + + @Override + public VectorFloat approximateCentroid() { + var vv = vectors.get(); + var centroid = vts.createFloatVector(vv.dimension()); + for (int i = 0; i < vv.size(); i++) { + var v = vv.getVector(i); + if (v == null) continue; + for (int d = 0; d < vv.dimension(); d++) { + centroid.set(d, centroid.get(d) + v.get(d)); + } + } + VectorUtil.scale(centroid, 1.0f / vv.size()); + return centroid; + } + + @Override + public SearchScoreProvider searchProviderFor(VectorFloat vector) { + throw new UnsupportedOperationException( + "byteVectorScoreProvider does not support float query vectors; use searchProviderFor(int node)"); + } + + @Override + public SearchScoreProvider searchProviderFor(ByteSequence vector) { + var vc = vectorsCopy.get(); + var sf = (ScoreFunction.ExactScoreFunction) node2 -> bvsf.compare(vector, vc.getVector(node2)); + return new DefaultSearchScoreProvider(sf); + } + + @Override + public SearchScoreProvider searchProviderFor(int node1) { + return new DefaultSearchScoreProvider(scoreFunctionFor(node1)); + } + + @Override + public SearchScoreProvider diversityProviderFor(int node1) { + return searchProviderFor(node1); + } + + @Override + public ScoreFunction diversityScoreFunctionFor(int node1) { + return scoreFunctionFor(node1); + } + }; + } + static BuildScoreProvider bqBuildScoreProvider(BQVectors bqv) { return new BuildScoreProvider() { @Override diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ByteVectorSimilarityFunction.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ByteVectorSimilarityFunction.java new file mode 100644 index 000000000..913fe186b --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ByteVectorSimilarityFunction.java @@ -0,0 +1,84 @@ +/* + * 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.annotations.Experimental; +import io.github.jbellis.jvector.vector.types.ByteSequence; + +/** + * Vector similarity function for signed int8 (byte) vectors; parallel to + * {@link VectorSimilarityFunction} but operating on {@link ByteSequence}. + *

+ * Bytes are treated as signed int8 values (Java's {@code byte} is already signed, range −128..127). + * Return-value conventions match {@link VectorSimilarityFunction}: higher is more similar. + */ +@Experimental +public enum ByteVectorSimilarityFunction { + + /** + * Euclidean similarity in {@code (0, 1]}: {@code 1 / (1 + squaredL2)}, the same mapping as + * {@link VectorSimilarityFunction#EUCLIDEAN}. + */ + EUCLIDEAN { + @Override + public float compare(ByteSequence v1, ByteSequence v2) { + // Do not divide the squared distance by a dimension-dependent maximum (e.g. n * 255^2) here. + // That squeezes every score into a band just below 1.0 and, in float32, collapses nearby + // distances to the same value (at 1024 dimensions squared distances 0, 1, 2 and 3 all score 1.0). + // Graph construction prunes with `similarity(other) > score * alpha`, which relies on scores + // keeping their relative spacing, so collapsed scores erase ranking distinctions and break + // pruning of longer links. + return 1.0f / (1.0f + VectorUtil.squareL2Distance(v1, v2)); + } + }, + + /** + * Dot product mapped to a positive, monotonically increasing score: + * {@code dot < 0 ? 1 / (1 - dot) : dot + 1}, the same mapping Lucene uses for unnormalised dot products. + *

+ * The raw integer dot product is deliberately not divided by a dimension-dependent maximum (e.g. + * {@code n * 128^2}): at 1,024 dimensions that makes adjacent dot products differ by less than float32 + * resolution, collapsing distinct scores. This mapping keeps distinct integer dot products distinct. + * Scores are in {@code (0, inf)}, with 1 meaning orthogonal. For unit-norm int8 vectors prefer + * {@link #COSINE}. + */ + DOT_PRODUCT { + @Override + public float compare(ByteSequence v1, ByteSequence v2) { + float dot = VectorUtil.dotProduct(v1, v2); + return dot < 0 ? 1 / (1 - dot) : dot + 1; + } + }, + + /** Cosine similarity: {@code (1 + cosine(v1, v2)) / 2} */ + COSINE { + @Override + public float compare(ByteSequence v1, ByteSequence v2) { + return (1.0f + VectorUtil.cosine(v1, v2)) / 2.0f; + } + }; + + /** + * Calculates a similarity score between the two int8 vectors. + * Higher values correspond to closer vectors. + * + * @param v1 a byte vector + * @param v2 another byte vector, of the same dimension + * @return the similarity score + */ + public abstract float compare(ByteSequence v1, ByteSequence v2); +} 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..5502e4508 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 @@ -338,6 +338,41 @@ public float assembleAndSumPQ( return res; } + // The int8 methods below accumulate in integers and convert to float only at the end, matching the + // Panama and native implementations, so that all execution paths produce identical results. + // A float accumulator stops being exact above 2^24, which int8 products reach at ~1,000 dimensions. + + @Override + public float dotProduct(ByteSequence a, ByteSequence b) { + int sum = 0; + for (int i = 0; i < a.length(); i++) { + sum += a.get(i) * b.get(i); + } + return sum; + } + + @Override + public float squareDistance(ByteSequence a, ByteSequence b) { + int sum = 0; + for (int i = 0; i < a.length(); i++) { + int diff = a.get(i) - b.get(i); + sum += diff * diff; + } + return sum; + } + + @Override + public float cosine(ByteSequence a, ByteSequence b) { + long dot = 0, normA = 0, normB = 0; + for (int i = 0; i < a.length(); i++) { + int ai = a.get(i), bi = b.get(i); + dot += ai * bi; + normA += ai * ai; + normB += bi * bi; + } + return (float) (dot / Math.sqrt((double) normA * normB)); + } + @Override public int hammingDistance(long[] v1, long[] v2) { int hd = 0; diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtil.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtil.java index 4be2004c7..5f719ae28 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtil.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtil.java @@ -235,6 +235,47 @@ public static float assembleAndSumPQ(VectorFloat data, int subspaceCount, Byt return impl.assembleAndSumPQ(data, subspaceCount, dataOffsets1, dataOffsetsOffset1, dataOffsets2, dataOffsetsOffset2, clusterCount); } + /** + * Returns the dot product of two signed int8 byte vectors. + * + * @throws IllegalArgumentException if the vectors' dimensions differ. + */ + public static float dotProduct(ByteSequence a, ByteSequence b) { + if (a.length() != b.length()) { + throw new IllegalArgumentException("vector dimensions differ: " + a.length() + "!=" + b.length()); + } + return impl.dotProduct(a, b); + } + + /** + * Returns the sum of squared differences of two signed int8 byte vectors. + * + * @throws IllegalArgumentException if the vectors' dimensions differ. + */ + public static float squareL2Distance(ByteSequence a, ByteSequence b) { + if (a.length() != b.length()) { + throw new IllegalArgumentException("vector dimensions differ: " + a.length() + "!=" + b.length()); + } + return impl.squareDistance(a, b); + } + + /** + * Returns the cosine similarity of two signed int8 byte vectors. + *

+ * As with the float overload, zero vectors (all components zero) are not supported: their norm is zero, + * so the result is NaN. Callers must not pass zero vectors; this is checked with an assertion. + * + * @throws IllegalArgumentException if the vectors' dimensions differ. + */ + public static float cosine(ByteSequence a, ByteSequence b) { + if (a.length() != b.length()) { + throw new IllegalArgumentException("vector dimensions differ: " + a.length() + "!=" + b.length()); + } + float r = impl.cosine(a, b); + assert Float.isFinite(r) : String.format("cosine(%s, %s) = %s", a, b, r); + return r; + } + public static int hammingDistance(long[] v1, long[] v2) { return impl.hammingDistance(v1, v2); } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtilSupport.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtilSupport.java index 9d663d10c..ff041274c 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtilSupport.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtilSupport.java @@ -173,6 +173,46 @@ default boolean tryUpdateKMeansDistances(VectorFloat[] points, VectorFloat */ float assembleAndSumPQ(VectorFloat codebookPartialSums, int subspaceCount, ByteSequence vector1Ordinals, int vector1OrdinalOffset, ByteSequence node2Ordinals, int node2OrdinalOffset, int clusterCount); + /** + * Calculates the dot product of two signed int8 byte vectors. + * The default is a scalar fallback; implementations should override it with a faster version. + */ + default float dotProduct(ByteSequence a, ByteSequence b) { + int sum = 0; + for (int i = 0; i < a.length(); i++) { + sum += a.get(i) * b.get(i); + } + return sum; + } + + /** + * Returns the sum of squared differences of two signed int8 byte vectors. + * The default is a scalar fallback; implementations should override it with a faster version. + */ + default float squareDistance(ByteSequence a, ByteSequence b) { + int sum = 0; + for (int i = 0; i < a.length(); i++) { + int diff = a.get(i) - b.get(i); + sum += diff * diff; + } + return sum; + } + + /** + * Returns the cosine similarity of two signed int8 byte vectors. + * The default is a scalar fallback; implementations should override it with a faster version. + */ + default float cosine(ByteSequence a, ByteSequence b) { + long dot = 0, normA = 0, normB = 0; + for (int i = 0; i < a.length(); i++) { + int ai = a.get(i), bi = b.get(i); + dot += ai * bi; + normA += ai * ai; + normB += bi * bi; + } + return (float) (dot / Math.sqrt((double) normA * normB)); + } + int hammingDistance(long[] v1, long[] v2); void calculatePartialSums(VectorFloat codebook, int codebookIndex, int size, int clusterCount, VectorFloat query, int offset, VectorSimilarityFunction vsf, VectorFloat partialSums); diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java index 2d48791bc..d3e46d3c3 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java @@ -21,6 +21,8 @@ import io.github.jbellis.jvector.example.util.BenchmarkSummarizer.SummaryStats; import io.github.jbellis.jvector.example.util.CheckpointManager; import io.github.jbellis.jvector.example.benchmarks.datasets.BaseVectorLoading; +import io.github.jbellis.jvector.example.benchmarks.datasets.ByteDataSet; +import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets; import io.github.jbellis.jvector.example.yaml.DatasetCollection; import io.github.jbellis.jvector.example.yaml.MultiConfig; @@ -147,9 +149,12 @@ public static void main(String[] args) throws IOException { try (var ds = DataSets.loadDataSet(datasetName, loading).orElseThrow( () -> new RuntimeException("Dataset " + datasetName + " not found") ).getDataSet()) { - logger.info("Dataset loaded: {} with {} vectors", datasetName, ds.getBaseVectors().size()); + logger.info("Dataset loaded: {} with {} vectors", datasetName, ds.size()); + if (ds instanceof ByteDataSet) { + config.validateForInt8Dataset(); + } - List datasetResults = Grid.runAndCollectResults(ds, GridSpec.fromConfig(config)); + List datasetResults = Grid.runAndCollectResults(ds, GridSpec.fromConfig(config, ds)); results.addAll(datasetResults); logger.info("Benchmark completed for dataset: {}", datasetName); @@ -158,24 +163,28 @@ public static void main(String[] args) throws IOException { // Compaction regression is opt-in; failures remain non-fatal. if (config.isCompactionEnabled()) { - try { - logger.info("Running compaction benchmark for dataset: {}", datasetName); - List datasetCompactionResults; - if (loading == BaseVectorLoading.BUFFERED) { - // This legacy validation queries raw originals; keep its resident source. - try (var resident = DataSets.loadDataSet(datasetName, BaseVectorLoading.PRELOAD) - .orElseThrow(() -> new RuntimeException("Dataset " + datasetName + " not found")) - .getDataSet()) { - datasetCompactionResults = CompactionBench.run(resident); + if (ds instanceof ByteDataSet) { + logger.info("Skipping compaction benchmark for int8 dataset: {}", datasetName); + } else { + try { + logger.info("Running compaction benchmark for dataset: {}", datasetName); + List datasetCompactionResults; + if (loading == BaseVectorLoading.BUFFERED) { + // This legacy validation queries raw originals; keep its resident source. + try (var resident = DataSets.loadDataSet(datasetName, BaseVectorLoading.PRELOAD) + .orElseThrow(() -> new RuntimeException("Dataset " + datasetName + " not found")) + .getDataSet()) { + datasetCompactionResults = CompactionBench.run(resident); + } + } else { + datasetCompactionResults = CompactionBench.run((DataSet) ds); } - } else { - datasetCompactionResults = CompactionBench.run(ds); + compactionResults.addAll(datasetCompactionResults); + logger.info("Compaction benchmark completed for dataset: {} ({} configs)", + datasetName, datasetCompactionResults.size()); + } catch (Exception e) { + logger.error("Compaction benchmark failed for dataset {}", datasetName, e); } - compactionResults.addAll(datasetCompactionResults); - logger.info("Compaction benchmark completed for dataset: {} ({} configs)", - datasetName, datasetCompactionResults.size()); - } catch (Exception e) { - logger.error("Compaction benchmark failed for dataset {}", datasetName, e); } } else { logger.info("Skipping compaction benchmark for dataset: {} (withCompaction not enabled)", datasetName); diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java index 8eb6ee09f..2c6e4022d 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java @@ -17,6 +17,7 @@ package io.github.jbellis.jvector.example; import io.github.jbellis.jvector.example.benchmarks.datasets.BaseVectorLoading; +import io.github.jbellis.jvector.example.benchmarks.datasets.ByteDataSet; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets; import io.github.jbellis.jvector.example.reporting.RunArtifacts; import io.github.jbellis.jvector.example.reporting.SearchReportingCatalog; @@ -127,7 +128,11 @@ public static void main(String[] args) throws IOException { // Register dataset info the first time we actually load the dataset for benchmarking artifacts.registerDataset(datasetName, ds); - Grid.run(ds, GridSpec.fromConfig(config), artifacts); + if (ds instanceof ByteDataSet) { + config.validateForInt8Dataset(); + } + + Grid.run(ds, GridSpec.fromConfig(config, ds), artifacts); } } } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java index 4c494a916..cf52f847a 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java @@ -121,7 +121,7 @@ public static List run(DataSet ds) throws Exception { private static BenchResult runConfig(DataSet ds, PartitionConfig cfg) throws Exception { String datasetName = ds.getName(); logger.info("Compaction bench [{}] config {}: {} vectors", - datasetName, cfg.dirName(), ds.getBaseVectors().size()); + datasetName, cfg.dirName(), ds.size()); // 1. Fetch pre-built partitions from S3 (cached locally). List partitionPaths = CompactionPartitionSource.ensurePartitions( diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/FeatureProviders.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/FeatureProviders.java index 5a138a4bf..7d4b46c89 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/FeatureProviders.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/FeatureProviders.java @@ -17,6 +17,7 @@ package io.github.jbellis.jvector.example; import io.github.jbellis.jvector.graph.ImmutableGraphIndex; +import io.github.jbellis.jvector.graph.RandomAccessByteVectorValues; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; import io.github.jbellis.jvector.quantization.PQVectors; @@ -34,6 +35,7 @@ final class FeatureProviders { static { register(new InlineVectorsProvider()); + register(new InlineByteVectorsProvider()); register(new NVQProvider()); register(new FusedPQProvider()); } @@ -96,4 +98,17 @@ static FeatureWriteContext configureWriter(Set features, } return ctx; } + + /** As {@link #configureWriter}, for an int8 dataset. */ + static FeatureWriteContext configureWriter(Set features, + ImmutableGraphIndex graph, + Path outPath, + RandomAccessByteVectorValues byteVectors) throws FileNotFoundException { + var writer = GraphBuilderFactory.writerBuilder(graph, outPath, byteVectors.size()); + var ctx = new FeatureWriteContext(writer, graph, byteVectors); + for (var id : features) { + require(id).configure(ctx); + } + return ctx; + } } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/FeatureWriteContext.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/FeatureWriteContext.java index 583318656..3aff536f1 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/FeatureWriteContext.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/FeatureWriteContext.java @@ -17,6 +17,7 @@ package io.github.jbellis.jvector.example; import io.github.jbellis.jvector.graph.ImmutableGraphIndex; +import io.github.jbellis.jvector.graph.RandomAccessByteVectorValues; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.disk.RandomAccessOnDiskGraphIndexWriter; import io.github.jbellis.jvector.graph.disk.feature.Feature; @@ -35,7 +36,9 @@ final class FeatureWriteContext { private final RandomAccessOnDiskGraphIndexWriter.Builder writer; private final ImmutableGraphIndex graph; + // Exactly one of vectors and byteVectors is set, according to the dataset's element type. private final RandomAccessVectorValues vectors; + private final RandomAccessByteVectorValues byteVectors; // File-backed sources reuse scratch vectors; ordinal suppliers need a view per worker. private final Supplier writerVectors; private final PQVectors buildPQ; @@ -58,16 +61,51 @@ final class FeatureWriteContext { this.writer = writer; this.graph = graph; this.vectors = vectors; + this.byteVectors = null; this.writerVectors = vectors.threadLocalSupplier(); this.buildPQ = buildPQ; this.metrics = metrics; } + /** A context for an int8 dataset: no float vectors, no PQ, no metrics. */ + FeatureWriteContext(RandomAccessOnDiskGraphIndexWriter.Builder writer, + ImmutableGraphIndex graph, + RandomAccessByteVectorValues byteVectors) { + this.writer = writer; + this.graph = graph; + this.vectors = null; + this.byteVectors = byteVectors; + this.writerVectors = null; + this.buildPQ = null; + this.metrics = null; + } + /** The graph being written; complete only for in-memory builds. */ ImmutableGraphIndex graph() { return graph; } - /** The full-precision base vectors. */ - RandomAccessVectorValues vectors() { return vectors; } + /** + * The full-precision base vectors. + * + * @throws IllegalStateException for an int8 dataset, whose features cannot be built from float vectors + */ + RandomAccessVectorValues vectors() { + if (vectors == null) { + throw new IllegalStateException("This is an int8 dataset; the feature needs float vectors"); + } + return vectors; + } + + /** + * The int8 base vectors. + * + * @throws IllegalStateException for a float dataset + */ + RandomAccessByteVectorValues byteVectors() { + if (byteVectors == null) { + throw new IllegalStateException("This is a float dataset; the feature needs int8 vectors"); + } + return byteVectors; + } /** PQ-encoded build vectors, or null for full-precision and non-PQ builds. */ PQVectors buildPQ() { return buildPQ; } @@ -85,6 +123,11 @@ void supplyFromVector(FeatureId id, Function, Feature.State> enco ordinalSuppliers.put(id, ordinal -> encode.apply(writerVectors.get().getVector(ordinal))); } + /** Registers a feature whose per-node state is computed from the node's ordinal alone. */ + void supplyByOrdinal(FeatureId id, IntFunction supplier) { + ordinalSuppliers.put(id, supplier); + } + /** Registers a feature whose state needs the completed graph, such as fused neighbor codes. */ void supplyAfterBuild(FeatureId id, Function> supplier) { afterBuildSuppliers.put(id, supplier); diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/GraphBuilderFactory.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/GraphBuilderFactory.java index 491893847..14e6a4ef2 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/GraphBuilderFactory.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/GraphBuilderFactory.java @@ -19,6 +19,7 @@ import io.github.jbellis.jvector.example.util.FilteredForkJoinPool; import io.github.jbellis.jvector.graph.GraphIndexBuilder; import io.github.jbellis.jvector.graph.ImmutableGraphIndex; +import io.github.jbellis.jvector.graph.RandomAccessByteVectorValues; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.disk.OrdinalMapper; import io.github.jbellis.jvector.graph.disk.RandomAccessOnDiskGraphIndexWriter; @@ -27,6 +28,7 @@ import io.github.jbellis.jvector.quantization.CompressedVectors; import io.github.jbellis.jvector.quantization.PQVectors; import io.github.jbellis.jvector.util.PhysicalCoreExecutor; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import java.io.FileNotFoundException; @@ -56,6 +58,23 @@ static GraphIndexBuilder forFullPrecision(RandomAccessVectorValues vectors, FilteredForkJoinPool.createFilteredPool()); } + /** A builder scoring with exact similarity over int8 vectors. */ + static GraphIndexBuilder forByteVectors(RandomAccessByteVectorValues vectors, + ByteVectorSimilarityFunction similarity, + BuildConfig build) { + var bsp = BuildScoreProvider.byteVectorScoreProvider(vectors, similarity); + return new GraphIndexBuilder(bsp, + vectors.dimension(), + build.M, + build.efConstruction, + build.neighborOverflow, + build.alpha, + build.addHierarchy, + build.refineFinalGraph, + PhysicalCoreExecutor.pool(), + FilteredForkJoinPool.createFilteredPool()); + } + /** A builder scoring with PQ- or BQ-encoded vectors. */ static GraphIndexBuilder forCompressed(CompressedVectors encoded, VectorSimilarityFunction similarity, diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java index 444d64947..78148559c 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java @@ -26,6 +26,8 @@ import io.github.jbellis.jvector.example.benchmarks.ThroughputBenchmark; import io.github.jbellis.jvector.example.benchmarks.diagnostics.BenchmarkDiagnostics; import io.github.jbellis.jvector.example.benchmarks.diagnostics.DiagnosticLevel; +import io.github.jbellis.jvector.example.benchmarks.datasets.BaseDataSet; +import io.github.jbellis.jvector.example.benchmarks.datasets.ByteDataSet; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; import io.github.jbellis.jvector.example.reporting.RunArtifacts; import io.github.jbellis.jvector.example.util.CompressorParameters; @@ -42,6 +44,8 @@ import io.github.jbellis.jvector.quantization.CompressedVectors; import io.github.jbellis.jvector.quantization.VectorCompressor; import io.github.jbellis.jvector.util.ExplicitThreadLocal; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; +import io.github.jbellis.jvector.vector.types.ByteSequence; import io.github.jbellis.jvector.vector.types.VectorFloat; import java.io.IOException; @@ -85,7 +89,7 @@ public static Double getIndexBuildTimeSeconds(String datasetName) { * Runs {@code spec} against {@code ds}, printing console tables and writing experiments.csv rows as * selected by {@code artifacts}. */ - public static void run(DataSet ds, GridSpec spec, RunArtifacts artifacts) throws IOException { + public static void run(BaseDataSet ds, GridSpec spec, RunArtifacts artifacts) throws IOException { runGrid(ds, spec, () -> setupBenchmarks(artifacts.benchmarksToCompute()), reportingListeners(artifacts)); } @@ -93,7 +97,7 @@ public static void run(DataSet ds, GridSpec spec, RunArtifacts artifacts) throws * Runs {@code spec} against {@code ds} with the fixed regression benchmarks, printing console tables * and returning one result per measurement. */ - public static List runAndCollectResults(DataSet ds, GridSpec spec) throws IOException { + public static List runAndCollectResults(BaseDataSet ds, GridSpec spec) throws IOException { var collector = new BenchResultCollector(); runGrid(ds, spec, Grid::regressionBenchmarks, List.of(new ConsoleReportListener(RunArtifacts.disabled()), collector)); @@ -111,7 +115,7 @@ private static List reportingListeners(RunArtifacts artifacts) { * @param benchmarks creates the query benchmarks for each index and search compressor * @param listeners receive every measurement, in order */ - private static void runGrid(DataSet ds, + private static void runGrid(BaseDataSet ds, GridSpec spec, Supplier> benchmarks, List listeners) throws IOException @@ -191,7 +195,7 @@ static void runOneGraph(OnDiskGraphIndexCache cache, Map> topKGrid, List usePruningGrid, RunArtifacts artifacts, - DataSet ds, + BaseDataSet ds, Path workDirectory) throws IOException { var build = BuildConfig.builder() @@ -215,7 +219,7 @@ private static void runOneGraph(OnDiskGraphIndexCache cache, List> featureSets, BuildConfig build, SearchConfig search, - DataSet ds, + BaseDataSet ds, Path workDirectory, Supplier> benchmarks, List listeners) throws IOException @@ -228,26 +232,37 @@ private static void runOneGraph(OnDiskGraphIndexCache cache, diagnostics.startMonitoring("testDirectory", workDirectory); diagnostics.startMonitoring("indexCache", Paths.get(indexCacheDir)); diagnostics.capturePrePhaseSnapshot("Graph Build"); - System.out.printf("%s: Dataset similarity function is %s%n", ds.getName(), ds.getSimilarityFunction()); + String sfLabel = (ds instanceof ByteDataSet) + ? ((ByteDataSet) ds).getSimilarityFunction().toString() + : ((DataSet) ds).getSimilarityFunction().toString(); + System.out.printf("%s: Dataset similarity function is %s%n", ds.getName(), sfLabel); // Resolve build compressor (and label quant type) so we can record compute time VectorCompressor buildCompressorObj = null; String buildQuantType = null; - if (build.buildCompressor != null) { - var buildParams = build.buildCompressor.apply(ds); + // Int8 datasets only support compression type "None" (enforced by MultiConfig.validateForInt8Dataset), + // so there is no build compressor for them. + if (build.buildCompressor != null && !(ds instanceof ByteDataSet)) { + var floatDs = (DataSet) ds; + var buildParams = build.buildCompressor.apply(floatDs); buildQuantType = quantTypeOf(buildParams); // "PQ", "BQ", or null - buildCompressorObj = getCompressor(build.buildCompressor, ds, constructionMetrics, Phase.INDEX, buildQuantType); + buildCompressorObj = getCompressor(build.buildCompressor, floatDs, constructionMetrics, Phase.INDEX, buildQuantType); } // Used in logging String buildCompressorString = (buildCompressorObj == null) ? "None" : String.valueOf(buildCompressorObj); - IndexBuildStrategy strategy = (buildCompressorObj == null) - ? new InMemoryBuild(workDirectory) - : new CompressedOnDiskBuild(cache, workDirectory, buildCompressorObj, constructionMetrics); - var built = strategy.build(featureSets, build, ds); + IndexBuildStrategy.BuiltIndexes built; + if (ds instanceof ByteDataSet) { + built = new InMemoryBuild(workDirectory).buildInt8(featureSets, build, (ByteDataSet) ds); + } else { + IndexBuildStrategy strategy = (buildCompressorObj == null) + ? new InMemoryBuild(workDirectory) + : new CompressedOnDiskBuild(cache, workDirectory, buildCompressorObj, constructionMetrics); + built = strategy.build(featureSets, build, (DataSet) ds); + } // Null when every index came from the index cache. Read it from here, not the per-dataset map, // which still holds the previous build config's time in that case. Double indexBuildTimeS = built.buildTimeS; @@ -283,10 +298,14 @@ private static void runOneGraph(OnDiskGraphIndexCache cache, if (FeatureProviders.anyMatch(featureSetForIndex, FeatureProvider::providesApproximateScoring)) { cv = null; System.out.format("%s: configured to use FUSED PQ, skipping vector compression%n", ds.getName()); + } else if (ds instanceof ByteDataSet) { + cv = null; + System.out.format("%s: No search compressor configured, FULL PRECISION vectors will be used for search%n", ds.getName()); } else { - var searchParams = cpSupplier.apply(ds); + var floatDs = (DataSet) ds; + var searchParams = cpSupplier.apply(floatDs); String searchQuantType = quantTypeOf(searchParams); // "PQ", "BQ", or null - var compressor = getCompressor(cpSupplier, ds, constructionMetrics, Phase.SEARCH, searchQuantType); + var compressor = getCompressor(cpSupplier, floatDs, constructionMetrics, Phase.SEARCH, searchQuantType); if (compressor == null) { cv = null; @@ -294,14 +313,14 @@ private static void runOneGraph(OnDiskGraphIndexCache cache, } else { long start = System.nanoTime(); cv = constructionMetrics.search(searchQuantType) - .timeEncode(() -> compressor.encodeAll(ds.getBaseRavv())); + .timeEncode(() -> compressor.encodeAll(floatDs.getBaseRavv())); double encodingTimeS = (System.nanoTime() - start) / 1_000_000_000.0; if (cv == null) { throw new IllegalStateException(String.format( "Compressor '%s' was provided but failed to encode vectors for dataset '%s'. " + "Aborting to prevent false recall results.", compressor, ds.getName())); } - System.out.format("%s: %s encoded %d vectors [%.2f MB] in %.2fs%n", ds.getName(), compressor, ds.getBaseVectors().size(), (cv.ramBytesUsed() / 1024f / 1024f), encodingTimeS); + System.out.format("%s: %s encoded %d vectors [%.2f MB] in %.2fs%n", ds.getName(), compressor, ds.size(), (cv.ramBytesUsed() / 1024f / 1024f), encodingTimeS); } } @@ -638,7 +657,7 @@ private static void saveToCache(String fname, VectorCompressor compressor) { } public static class ConfiguredSystem implements AutoCloseable { - DataSet ds; + BaseDataSet ds; ImmutableGraphIndex index; CompressedVectors cv; Set features; @@ -650,38 +669,50 @@ public static class ConfiguredSystem implements AutoCloseable { return new GraphSearcher(index); }); - ConfiguredSystem(DataSet ds, ImmutableGraphIndex index, CompressedVectors cv, Set features) { + ConfiguredSystem(BaseDataSet ds, ImmutableGraphIndex index, CompressedVectors cv, Set features) { this.ds = ds; this.index = index; this.cv = cv; this.features = features; this.scoresFromIndex = FeatureProviders.anyMatch(features, FeatureProvider::providesApproximateScoring); // File-backed sources borrow scratch storage; concurrent scorers need separate views. - this.exactVectors = ds.getBaseRavv().threadLocalSupplier(); + // Int8 datasets do not have float base vectors; leave exactVectors null. + this.exactVectors = ds instanceof DataSet ? ((DataSet) ds).getBaseRavv().threadLocalSupplier() : null; } + /** Score provider for float query vectors (float datasets). */ public SearchScoreProvider scoreProviderFor(VectorFloat queryVector, ImmutableGraphIndex.View view) { + var floatDs = (DataSet) ds; var scoringView = (ImmutableGraphIndex.ScoringView) view; ScoreFunction.ApproximateScoreFunction asf; if (scoresFromIndex) { - asf = scoringView.approximateScoreFunctionFor(queryVector, ds.getSimilarityFunction()); + asf = scoringView.approximateScoreFunctionFor(queryVector, floatDs.getSimilarityFunction()); } else { // if we're not compressing then just use the exact score function if (cv == null) { - return DefaultSearchScoreProvider.exact(queryVector, ds.getSimilarityFunction(), exactVectors.get()); + return DefaultSearchScoreProvider.exact(queryVector, floatDs.getSimilarityFunction(), exactVectors.get()); } - asf = cv.precomputedScoreFunctionFor(queryVector, ds.getSimilarityFunction()); + asf = cv.precomputedScoreFunctionFor(queryVector, floatDs.getSimilarityFunction()); } - var rr = scoringView.rerankerFor(queryVector, ds.getSimilarityFunction()); + var rr = scoringView.rerankerFor(queryVector, floatDs.getSimilarityFunction()); return new DefaultSearchScoreProvider(asf, rr); } + /** Score provider for int8 byte query vectors (int8 datasets). */ + public SearchScoreProvider scoreProviderFor(ByteSequence queryBytes, ImmutableGraphIndex.View view) { + var byteDs = (ByteDataSet) ds; + ByteVectorSimilarityFunction bvsf = byteDs.getSimilarityFunction(); + var diskView = (OnDiskGraphIndex.View) view; + ScoreFunction.ExactScoreFunction sf = diskView.byteVectorRerankerFor(queryBytes, bvsf); + return new DefaultSearchScoreProvider(sf); + } + public GraphSearcher getSearcher() { return searchers.get(); } - public DataSet getDataSet() { + public BaseDataSet getDataSet() { return ds; } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/GridSpec.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/GridSpec.java index 82bf3eacd..65c17e4fb 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/GridSpec.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/GridSpec.java @@ -16,6 +16,8 @@ package io.github.jbellis.jvector.example; +import io.github.jbellis.jvector.example.benchmarks.datasets.BaseDataSet; +import io.github.jbellis.jvector.example.benchmarks.datasets.ByteDataSet; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; import io.github.jbellis.jvector.example.util.CompressorParameters; import io.github.jbellis.jvector.example.yaml.MultiConfig; @@ -51,6 +53,14 @@ private GridSpec(Builder b) { /** The sweep a dataset's YAML config describes. */ public static GridSpec fromConfig(MultiConfig config) { + return fromConfig(config, null); + } + + /** + * Like {@link #fromConfig(MultiConfig)}, but for an int8 dataset the feature sets are the fixed + * int8 set, since reranking and fused-graph options do not apply to byte vectors. + */ + public static GridSpec fromConfig(MultiConfig config, BaseDataSet ds) { var construction = config.construction; var search = config.search; return fromGrids(construction.addHierarchy, @@ -59,7 +69,7 @@ public static GridSpec fromConfig(MultiConfig config) { construction.neighborOverflow, construction.efConstruction, construction.getCompressorParameters(), - construction.getFeatureSets(), + ds instanceof ByteDataSet ? construction.getInt8FeatureSets() : construction.getFeatureSets(), SearchConfig.builder() .searchCompressors(search.getCompressorParameters()) .topKGrid(search.topKOverquery) diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/InMemoryBuild.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/InMemoryBuild.java index 346818c2d..e847c52fb 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/InMemoryBuild.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/InMemoryBuild.java @@ -17,6 +17,7 @@ package io.github.jbellis.jvector.example; import io.github.jbellis.jvector.disk.ReaderSupplierFactory; +import io.github.jbellis.jvector.example.benchmarks.datasets.ByteDataSet; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; import io.github.jbellis.jvector.graph.ImmutableGraphIndex; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; @@ -85,6 +86,38 @@ public BuiltIndexes build(List> featureSets, BuildConfi return new BuiltIndexes(indexes, new LinkedHashMap<>(), buildTimeS); } + /** Builds an int8 index in memory and writes each requested feature set; int8 supports INLINE_BYTE_VECTORS. */ + BuiltIndexes buildInt8(List> featureSets, BuildConfig build, ByteDataSet ds) throws IOException { + var byteVectors = ds.getBaseRavv(); + Map, ImmutableGraphIndex> indexes = new LinkedHashMap<>(); + var builder = GraphBuilderFactory.forByteVectors(byteVectors, ds.getSimilarityFunction(), build); + long start = System.nanoTime(); + var onHeapGraph = builder.build(byteVectors); + double buildTimeS = (System.nanoTime() - start) / 1_000_000_000.0; + System.out.format("Build (int8) M=%d overflow=%.2f ef=%d in %.2fs%n", + build.M, build.neighborOverflow, build.efConstruction, buildTimeS); + for (int i = 0; i <= onHeapGraph.getMaxLevel(); i++) { + System.out.format(" L%d: %d nodes, %.2f avg degree%n", + i, onHeapGraph.size(i), onHeapGraph.getAverageDegree(i)); + } + int n = 0; + for (var features : featureSets) { + if (FeatureProviders.anyMatch(features, p -> !p.supportsInMemoryBuild())) { + System.out.println("Skipping " + features + " when building in memory: it needs a compressed build"); + continue; + } + var graphPath = workDirectory.resolve("graph" + n++); + var ctx = FeatureProviders.configureWriter(features, onHeapGraph, graphPath, byteVectors); + try (var writer = ctx.buildWriter()) { + start = System.nanoTime(); + writer.write(ctx.ordinalSuppliers()); + System.out.format("Wrote %s in %.2fs%n", features, (System.nanoTime() - start) / 1_000_000_000.0); + } + indexes.put(features, OnDiskGraphIndex.load(ReaderSupplierFactory.open(graphPath))); + } + return new BuiltIndexes(indexes, new LinkedHashMap<>(), buildTimeS); + } + /** @param view a view of the completed graph, or null when no feature needs one */ private static Map> writeSuppliers(FeatureWriteContext ctx, ImmutableGraphIndex.View view) { diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/InlineByteVectorsProvider.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/InlineByteVectorsProvider.java new file mode 100644 index 000000000..7b4d877af --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/InlineByteVectorsProvider.java @@ -0,0 +1,41 @@ +/* + * 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; + +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.InlineByteVectors; + +/** Signed int8 vectors stored alongside each node, used for exact reranking of int8 datasets. */ +final class InlineByteVectorsProvider implements FeatureProvider { + @Override + public FeatureId id() { + return FeatureId.INLINE_BYTE_VECTORS; + } + + @Override + public void configure(FeatureWriteContext ctx) { + var byteVectors = ctx.byteVectors(); + ctx.add(new InlineByteVectors(byteVectors.dimension())); + ctx.supplyByOrdinal(FeatureId.INLINE_BYTE_VECTORS, + ordinal -> new InlineByteVectors.State(byteVectors.getVector(ordinal))); + } + + @Override + public boolean providesRerankVectors() { + return true; + } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/SearchContext.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/SearchContext.java index 9371abc5a..18015c293 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/SearchContext.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/SearchContext.java @@ -16,6 +16,7 @@ package io.github.jbellis.jvector.example; +import io.github.jbellis.jvector.example.benchmarks.datasets.BaseDataSet; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; import io.github.jbellis.jvector.example.benchmarks.diagnostics.SystemMonitor; import io.github.jbellis.jvector.example.util.CompressorParameters; @@ -30,7 +31,7 @@ *

To give listeners more data, add a field here and set it where {@link Grid} creates the context.

*/ final class SearchContext { - final DataSet ds; + final BaseDataSet ds; final BuildConfig build; final SearchConfig search; /** The feature set the index was requested with. */ @@ -49,7 +50,7 @@ final class SearchContext { /** Memory state captured right after the build, or null if none was captured. */ final SystemMonitor.SystemSnapshot buildSnapshot; - SearchContext(DataSet ds, + SearchContext(BaseDataSet ds, BuildConfig build, SearchConfig search, Set requestedFeatures, diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/AccuracyBenchmark.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/AccuracyBenchmark.java index 897eec550..ae707948b 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/AccuracyBenchmark.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/AccuracyBenchmark.java @@ -88,7 +88,7 @@ public List runBenchmark( throw new RuntimeException("At least one metric must be displayed"); } - int totalQueries = cs.getDataSet().getQueryVectors().size(); + int totalQueries = cs.getDataSet().getQueryCount(); // execute all queries in parallel and collect results List results = IntStream.range(0, totalQueries) diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/CountBenchmark.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/CountBenchmark.java index 1e10c102d..3d0adafb4 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/CountBenchmark.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/CountBenchmark.java @@ -105,7 +105,7 @@ public List runBenchmark( LongAdder nodesVisited = new LongAdder(); LongAdder nodesExpanded = new LongAdder(); LongAdder nodesExpandedBaseLayer = new LongAdder(); - int totalQueries = cs.getDataSet().getQueryVectors().size(); + int totalQueries = cs.getDataSet().getQueryCount(); for (int run = 0; run < queryRuns; run++) { IntStream.range(0, totalQueries) diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/ExecutionTimeBenchmark.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/ExecutionTimeBenchmark.java index 1899b99cd..29a1c1d5a 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/ExecutionTimeBenchmark.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/ExecutionTimeBenchmark.java @@ -56,7 +56,7 @@ public List runBenchmark( boolean usePruning, int queryRuns) { - int totalQueries = cs.getDataSet().getQueryVectors().size(); + int totalQueries = cs.getDataSet().getQueryCount(); double totalRuntime = 0; for (int run = 0; run < queryRuns; run++) { diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/LatencyBenchmark.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/LatencyBenchmark.java index d04a9c1bc..969ddbec5 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/LatencyBenchmark.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/LatencyBenchmark.java @@ -104,7 +104,7 @@ public List runBenchmark( throw new IllegalArgumentException("At least one parameter must be set to true"); } - int totalQueries = cs.getDataSet().getQueryVectors().size(); + int totalQueries = cs.getDataSet().getQueryCount(); double mean = 0.0; double m2 = 0.0; int count = 0; diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/QueryExecutor.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/QueryExecutor.java index 45094bbd6..f20df1a1c 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/QueryExecutor.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/QueryExecutor.java @@ -17,13 +17,17 @@ package io.github.jbellis.jvector.example.benchmarks; import io.github.jbellis.jvector.example.Grid.ConfiguredSystem; +import io.github.jbellis.jvector.example.benchmarks.datasets.ByteDataSet; +import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; import io.github.jbellis.jvector.graph.SearchResult; import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.vector.types.ByteSequence; import io.github.jbellis.jvector.vector.types.VectorFloat; public class QueryExecutor { /** * Executes the query at index i using the given parameters. + * Dispatches to the int8 or float path depending on the dataset type. * * @param cs Configured system that contains the query vectors. * @param topK Number of top results. @@ -33,17 +37,32 @@ public class QueryExecutor { * @return the SearchResult for query i. */ public static SearchResult executeQuery(ConfiguredSystem cs, int topK, int rerankK, boolean usePruning, int i) { - var queryVector = cs.getDataSet().getQueryVectors().get(i); + if (cs.getDataSet() instanceof ByteDataSet) { + ByteDataSet byteDs = (ByteDataSet) cs.getDataSet(); + return executeByteQuery(cs, topK, rerankK, usePruning, byteDs.getQueryVectors().get(i)); + } + var queryVector = ((DataSet) cs.getDataSet()).getQueryVectors().get(i); return executeQuery(cs, topK, rerankK, usePruning, queryVector); } // Overload to allow single query injection (e.g., for warm-up with random vectors) - public static SearchResult executeQuery(ConfiguredSystem cs, int topK, int rerankK, boolean usePruning, VectorFloat queryVector - ) { + public static SearchResult executeQuery(ConfiguredSystem cs, int topK, int rerankK, boolean usePruning, VectorFloat queryVector) { var searcher = cs.getSearcher(); searcher.usePruning(usePruning); var sf = cs.scoreProviderFor(queryVector, searcher.getView()); return searcher.search(sf, topK, rerankK, 0.0f, 0.0f, Bits.ALL); } + + // Overload to allow single int8 query injection (e.g., for warm-up with random vectors) + public static SearchResult executeQuery(ConfiguredSystem cs, int topK, int rerankK, boolean usePruning, ByteSequence queryBytes) { + return executeByteQuery(cs, topK, rerankK, usePruning, queryBytes); + } + + private static SearchResult executeByteQuery(ConfiguredSystem cs, int topK, int rerankK, boolean usePruning, ByteSequence queryBytes) { + var searcher = cs.getSearcher(); + searcher.usePruning(usePruning); + var sf = cs.scoreProviderFor(queryBytes, searcher.getView()); + return searcher.search(sf, topK, rerankK, 0.0f, 0.0f, Bits.ALL); + } } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/ThroughputBenchmark.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/ThroughputBenchmark.java index eaca57a17..99f8f0aab 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/ThroughputBenchmark.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/ThroughputBenchmark.java @@ -17,7 +17,9 @@ package io.github.jbellis.jvector.example.benchmarks; import io.github.jbellis.jvector.example.Grid.ConfiguredSystem; +import io.github.jbellis.jvector.example.benchmarks.datasets.ByteDataSet; import io.github.jbellis.jvector.graph.SearchResult; +import io.github.jbellis.jvector.vector.types.ByteSequence; import io.github.jbellis.jvector.vector.types.VectorFloat; import io.github.jbellis.jvector.vector.types.VectorTypeSupport; import io.github.jbellis.jvector.vector.VectorUtil; @@ -140,7 +142,7 @@ public List runBenchmark( } try (BenchmarkDiagnostics diagnostics = new BenchmarkDiagnostics(diagnosticLevel)) { - int totalQueries = cs.getDataSet().getQueryVectors().size(); + int totalQueries = cs.getDataSet().getQueryCount(); int dim = cs.getDataSet().getDimension(); // Warmup Phase with diagnostics @@ -154,14 +156,24 @@ public List runBenchmark( .forEach(k -> { long queryStart = System.nanoTime(); - // Generate a random vector - VectorFloat randQ = vts.createFloatVector(dim); - for (int j = 0; j < dim; j++) { - randQ.set(j, ThreadLocalRandom.current().nextFloat()); + SearchResult sr; + if (cs.getDataSet() instanceof ByteDataSet) { + // Int8 dataset: random signed byte query + ByteSequence randB = vts.createByteSequence(dim); + for (int j = 0; j < dim; j++) { + randB.set(j, (byte) ThreadLocalRandom.current().nextInt(-128, 128)); + } + sr = QueryExecutor.executeQuery(cs, topK, rerankK, usePruning, randB); + } else { + // Generate a random vector + VectorFloat randQ = vts.createFloatVector(dim); + for (int j = 0; j < dim; j++) { + randQ.set(j, ThreadLocalRandom.current().nextFloat()); + } + VectorUtil.l2normalize(randQ); + sr = QueryExecutor.executeQuery( + cs, topK, rerankK, usePruning, randQ); } - VectorUtil.l2normalize(randQ); - SearchResult sr = QueryExecutor.executeQuery( - cs, topK, rerankK, usePruning, randQ); SINK += sr.getVisitedCount(); long queryEnd = System.nanoTime(); diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/BaseDataSet.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/BaseDataSet.java new file mode 100644 index 000000000..b069cc0bb --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/BaseDataSet.java @@ -0,0 +1,67 @@ +/* + * 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.benchmarks.datasets; + +import java.io.IOException; +import java.util.List; + +/** + * Common base interface for all vector datasets, regardless of element type (float or byte). + */ +public interface BaseDataSet extends AutoCloseable { + + /** Release file-backed input resources after all benchmark consumers stop. */ + @Override + default void close() throws IOException {} + + /** + * Get dimensions of the vectors in this dataset. + * @return the dimensionality + */ + int getDimension(); + + /** True only for the harness's explicit bounded-input mode. */ + default boolean useDatasetPrefetching() { return false; } + + /** + * The symbolic name of this dataset, used for dataset selection and result labeling. + * @return the dataset name + */ + String getName(); + + /** + * The ground truth as a list. + * Each major index corresponds to the query index. + * Each minor index within represents the corresponding base vector ordinal. + * @return the ground truth list + */ + List> getGroundTruth(); + + /** + * The number of base vectors in this dataset. + * @return the base vector count + */ + int size(); + + /** + * The number of query vectors available. Works for all dataset types. + * @return the query vector count + */ + default int getQueryCount() { + return getGroundTruth().size(); + } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/ByteDataSet.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/ByteDataSet.java new file mode 100644 index 000000000..e8a9330fa --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/ByteDataSet.java @@ -0,0 +1,63 @@ +/* + * 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.benchmarks.datasets; + +import io.github.jbellis.jvector.graph.RandomAccessByteVectorValues; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; +import io.github.jbellis.jvector.vector.types.ByteSequence; + +import java.util.List; + +/** + * A dataset backed by native signed int8 (byte) vectors. + */ +public interface ByteDataSet extends BaseDataSet { + + /** + * Get a random-access view of base byte vectors. + * @return base byte vectors + */ + RandomAccessByteVectorValues getBaseRavv(); + + /** + * The similarity function used for byte vector comparison. + * @return the byte similarity function + */ + ByteVectorSimilarityFunction getSimilarityFunction(); + + /** + * The base byte vectors as a list. + * @return a list of base byte vectors + */ + List> getBaseVectors(); + + /** + * The query byte vectors as a list. + * @return a list of query byte vectors + */ + List> getQueryVectors(); + + @Override + default int size() { + return getBaseVectors().size(); + } + + @Override + default int getQueryCount() { + return getQueryVectors().size(); + } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/ByteSimpleDataSet.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/ByteSimpleDataSet.java new file mode 100644 index 000000000..5bc7e1729 --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/ByteSimpleDataSet.java @@ -0,0 +1,109 @@ +/* + * 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.benchmarks.datasets; + +import io.github.jbellis.jvector.graph.ListRandomAccessByteVectorValues; +import io.github.jbellis.jvector.graph.RandomAccessByteVectorValues; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; +import io.github.jbellis.jvector.vector.types.ByteSequence; + +import java.util.List; + +/** + * A {@link ByteDataSet} backed by native signed int8 (byte) vectors. + */ +public class ByteSimpleDataSet implements ByteDataSet { + + private final String name; + private final ByteVectorSimilarityFunction similarityFunction; + private final List> baseVectors; + private final List> queryVectors; + private final List> groundTruth; + private final int dimension; + private RandomAccessByteVectorValues baseRavv; + + public ByteSimpleDataSet(String name, + ByteVectorSimilarityFunction similarityFunction, + List> baseVectors, + List> queryVectors, + List> groundTruth) + { + if (baseVectors.isEmpty()) { + throw new IllegalArgumentException("Base vectors must not be empty"); + } + if (queryVectors.isEmpty()) { + throw new IllegalArgumentException("Query vectors must not be empty"); + } + if (groundTruth.isEmpty()) { + throw new IllegalArgumentException("Ground truth must not be empty"); + } + this.dimension = baseVectors.get(0).length(); + if (queryVectors.get(0).length() != dimension) { + throw new IllegalArgumentException("Base and query vectors must have the same dimensionality"); + } + if (queryVectors.size() != groundTruth.size()) { + throw new IllegalArgumentException("Query and ground truth lists must be the same size"); + } + + this.name = name; + this.similarityFunction = similarityFunction; + this.baseVectors = baseVectors; + this.queryVectors = queryVectors; + this.groundTruth = groundTruth; + + System.out.format("%n%s: %d base and %d query int8 vectors created, dimensions %d%n", + name, baseVectors.size(), queryVectors.size(), dimension); + } + + @Override + public int getDimension() { + return dimension; + } + + @Override + public String getName() { + return name; + } + + @Override + public ByteVectorSimilarityFunction getSimilarityFunction() { + return similarityFunction; + } + + @Override + public RandomAccessByteVectorValues getBaseRavv() { + if (baseRavv == null) { + baseRavv = new ListRandomAccessByteVectorValues(baseVectors, dimension); + } + return baseRavv; + } + + @Override + public List> getBaseVectors() { + return baseVectors; + } + + @Override + public List> getQueryVectors() { + return queryVectors; + } + + @Override + public List> getGroundTruth() { + return groundTruth; + } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSet.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSet.java index daa91d1d0..f528eaa27 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSet.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSet.java @@ -20,26 +20,12 @@ import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import io.github.jbellis.jvector.vector.types.VectorFloat; -import java.io.IOException; -import java.util.*; +import java.util.List; /** - * This provides a uniform way to access vector test data, regardless of where it comes from or how it is implemented. + * A dataset backed by single-precision floating point vectors (default dataset type in JVector). */ -public interface DataSet extends AutoCloseable { - - /** Release file-backed input resources after all benchmark consumers stop. */ - @Override - default void close() throws IOException {} - - /** - * Get dimensions of the vectors in this dataset. - * @return the dimensionality - */ - int getDimension(); - - /** True only for the harness's explicit bounded-input mode. */ - default boolean useDatasetPrefetching() { return false; } +public interface DataSet extends BaseDataSet { /** * Get a random-access view of base vectors. @@ -47,12 +33,6 @@ default void close() throws IOException {} */ RandomAccessVectorValues getBaseRavv(); - /** - * The symbolic name of this dataset, used for dataset selection and result labeling. - * @return the dataset name - */ - String getName(); - /** * The similarity function originally used to build this dataset, and the one that should be used for testing * during indexing and traversal. @@ -74,11 +54,13 @@ default void close() throws IOException {} */ List> getQueryVectors(); - /** - * The ground truth as a list. - * Each major index corresponds to the self-same index from {@link #getQueryVectors()}. - * Each minor index within represents the corresponding ordinal from {@link #getBaseVectors()} and {@link #getBaseRavv()}. - * @return a list of query vectors. - */ - List> getGroundTruth(); + @Override + default int size() { + return getBaseVectors().size(); + } + + @Override + default int getQueryCount() { + return getQueryVectors().size(); + } } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetInfo.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetInfo.java index 94eb9f011..bb517b053 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetInfo.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetInfo.java @@ -49,9 +49,9 @@ /// @see DataSetLoader /// @see DataSets public class DataSetInfo implements DataSetProperties { - private final Supplier loader; + private final Supplier loader; private final DataSetProperties baseProperties; - private volatile DataSet cached; + private volatile BaseDataSet cached; /// Creates a new dataset info handle. /// @@ -61,7 +61,7 @@ public class DataSetInfo implements DataSetProperties { /// /// @param baseProperties the dataset properties (name, similarity function, etc.) /// @param loader a supplier that performs the deferred load; invoked at most once - public DataSetInfo(DataSetProperties baseProperties, Supplier loader) { + public DataSetInfo(DataSetProperties baseProperties, Supplier loader) { this.baseProperties = baseProperties; this.loader = loader; } @@ -124,7 +124,8 @@ public boolean isDuplicateVectorFree() { /// completes, after which all callers share the same cached instance. /// /// @return the ready-to-use {@link DataSet} - public DataSet getDataSet() { + @SuppressWarnings("unchecked") + public T getDataSet() { if (cached == null) { synchronized (this) { if (cached == null) { @@ -132,6 +133,6 @@ public DataSet getDataSet() { } } } - return cached; + return (T) cached; } } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFD.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFD.java index d32918635..e64783789 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFD.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFD.java @@ -16,6 +16,7 @@ package io.github.jbellis.jvector.example.benchmarks.datasets; import io.github.jbellis.jvector.example.util.SiftLoader; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.yaml.snakeyaml.Yaml; @@ -452,8 +453,18 @@ public Optional loadDataSet(String dataSetName, BaseVectorLoading l return Optional.of(new DataSetInfo(props, () -> { if (loading == BaseVectorLoading.BUFFERED && props.loadBehavior() != DataSetProperties.LoadBehavior.NO_SCRUB) throw new IllegalArgumentException("Buffered loading requires NO_SCRUB metadata for " + dataSetName); - var queryVectors = SiftLoader.readFvecs(effectiveCacheDir.resolve(queryFile).toString()); var gtVectors = SiftLoader.readIvecs(effectiveCacheDir.resolve(gtFile).toString()); + if (props.isInt8()) { + if (loading == BaseVectorLoading.BUFFERED) + throw new IllegalArgumentException("Buffered loading is not supported for int8 datasets: " + dataSetName); + var bvsf = ByteVectorSimilarityFunction.valueOf( + props.similarityFunction().orElseThrow().name()); + // is_int8 means signed int8 payloads; unsigned uint8 .bvecs files (e.g. SIFT1B/BIGANN) are not supported + var baseVectors = SiftLoader.readInt8Bvecs(effectiveCacheDir.resolve(baseFile).toString()); + var queryVectors = SiftLoader.readInt8Bvecs(effectiveCacheDir.resolve(queryFile).toString()); + return new ByteSimpleDataSet(dataSetName, bvsf, baseVectors, queryVectors, gtVectors); + } + var queryVectors = SiftLoader.readFvecs(effectiveCacheDir.resolve(queryFile).toString()); if (loading == BaseVectorLoading.BUFFERED) { try { return new BufferedDataSet(dataSetName, props.similarityFunction().orElseThrow(), diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetProperties.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetProperties.java index 6d017d899..4c3ff7e02 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetProperties.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetProperties.java @@ -54,6 +54,10 @@ public interface DataSetProperties { /// Canonical key for how benchmark loaders should treat the dataset at load time. String KEY_LOAD_BEHAVIOR = "load_behavior"; + /// Canonical key for whether the dataset uses native **signed int8** (-128..127) vectors ({@link Boolean}). + /// Unsigned uint8 datasets (e.g. conventional SIFT1B/BIGANN `.bvecs` files) are not supported. + String KEY_IS_INT8 = "is_int8"; + /** * Controls benchmark-loader behavior for this dataset. * @@ -124,6 +128,16 @@ default LoadBehavior loadBehavior() { return LoadBehavior.LEGACY_SCRUB; } + /** + * Returns true if this dataset stores native signed int8 (-128..127) vectors. + * Unsigned uint8 data (e.g. conventional SIFT1B/BIGANN {@code .bvecs} files, 0..255) is not supported, + * because reading it as signed would silently change distances. + * When true, loaders must use {@code readInt8Bvecs} and return a {@code ByteSimpleDataSet}. + */ + default boolean isInt8() { + return false; + } + /** * A convenience method to capture the notion of a valid dataset. * As any additional qualifiers are added to this data carrier, this method should be updated accordingly. @@ -261,5 +275,10 @@ public LoadBehavior loadBehavior() { } return LoadBehavior.LEGACY_SCRUB; } + + @Override + public boolean isInt8() { + return Boolean.TRUE.equals(properties.get(KEY_IS_INT8)); + } } } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/DatasetInfoWriter.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/DatasetInfoWriter.java index fd3b94ce3..fe91ef5e7 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/DatasetInfoWriter.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/DatasetInfoWriter.java @@ -16,6 +16,8 @@ package io.github.jbellis.jvector.example.reporting; +import io.github.jbellis.jvector.example.benchmarks.datasets.BaseDataSet; +import io.github.jbellis.jvector.example.benchmarks.datasets.ByteDataSet; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; import java.io.IOException; @@ -114,17 +116,20 @@ public static Row fromDataSet(String datasetName, String basePath, String queryPath, String groundTruthPath, - DataSet ds) { + BaseDataSet ds) { + String sf = (ds instanceof ByteDataSet) + ? ((ByteDataSet) ds).getSimilarityFunction().toString() + : ((DataSet) ds).getSimilarityFunction().toString(); return new Row( datasetName, basePath, queryPath, groundTruthPath, - ds.getBaseVectors().size(), - ds.getQueryVectors().size(), + ds.size(), + ds.getQueryCount(), ds.getGroundTruth().size(), ds.getDimension(), - ds.getSimilarityFunction().toString() + sf ); } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/RunArtifacts.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/RunArtifacts.java index da8459645..fc6b37261 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/RunArtifacts.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/RunArtifacts.java @@ -16,7 +16,7 @@ package io.github.jbellis.jvector.example.reporting; -import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; +import io.github.jbellis.jvector.example.benchmarks.datasets.BaseDataSet; import io.github.jbellis.jvector.example.benchmarks.Metric; import io.github.jbellis.jvector.example.yaml.MultiConfig; import io.github.jbellis.jvector.example.yaml.MetricSelection; @@ -261,7 +261,7 @@ public void logRow(String datasetName, public Map> benchmarksToLog() { return benchmarksToLog; } public MetricSelection metricsToLog() { return metricsToLog; } - public void registerDataset(String datasetName, DataSet ds) throws IOException { + public void registerDataset(String datasetName, BaseDataSet ds) throws IOException { if (datasetInfoWriter == null) { return; // disabled } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/Int8Example.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/Int8Example.java new file mode 100644 index 000000000..71c4e6f69 --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/Int8Example.java @@ -0,0 +1,203 @@ +/* + * 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.tutorial; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Random; + +import io.github.jbellis.jvector.disk.ReaderSupplier; +import io.github.jbellis.jvector.disk.ReaderSupplierFactory; +import io.github.jbellis.jvector.example.util.SiftLoader; +import io.github.jbellis.jvector.graph.GraphIndexBuilder; +import io.github.jbellis.jvector.graph.GraphSearcher; +import io.github.jbellis.jvector.graph.ImmutableGraphIndex; +import io.github.jbellis.jvector.graph.ListRandomAccessByteVectorValues; +import io.github.jbellis.jvector.graph.RandomAccessByteVectorValues; +import io.github.jbellis.jvector.graph.SearchResult; +import io.github.jbellis.jvector.graph.disk.GraphIndexWriter; +import io.github.jbellis.jvector.graph.disk.GraphIndexWriterTypes; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.InlineByteVectors; +import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider; +import io.github.jbellis.jvector.graph.similarity.SearchScoreProvider; +import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.ByteSequence; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; + +/** + * Demonstrates end-to-end INT8 (signed byte) vector support in JVector: + * + *

How to run: + *

+ *   # The siftsmall dataset is checked into the repository root, so run from the
+ *   # repository root (siftsmall/siftsmall_base.fvecs must be present relative to
+ *   # the working directory).
+ *
+ *   # Build and run — jdk22 is the default profile (active when no -P flag is given):
+ *   ./mvnw compile -pl jvector-examples -am
+ *   ./mvnw exec:exec@tutorial -pl jvector-examples -Dtutorial=int8
+ *
+ *   # Use -Pjdk20 or -Pjdk11 if you are on an older JDK:
+ *   ./mvnw exec:exec@tutorial -pl jvector-examples -Pjdk20 -Dtutorial=int8
+ * 
+ * + *
    + *
  1. Read the siftsmall dataset from .fvecs files (float32).
  2. + *
  3. Convert each float32 vector to a signed int8 {@link ByteSequence}.
  4. + *
  5. Build a graph index using only byte×byte distance calculations.
  6. + *
  7. Save the graph to disk with {@link InlineByteVectors} — 1 byte per component on disk.
  8. + *
  9. Load the index back from disk.
  10. + *
  11. Generate random int8 query vectors and search, scoring directly from disk byte vectors.
  12. + *
+ * + *

The siftsmall dataset ships with the repository at {@code siftsmall/siftsmall_base.fvecs}; + * the working directory must be the repository root (or pass the dataset directory as the first argument). + */ +public class Int8Example { + + private static final VectorTypeSupport vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + + public static void main(String[] args) throws IOException { + String siftDir = args.length > 0 ? args[0] : "siftsmall"; + + // ── Step 1: Read siftsmall base vectors (.fvecs) ───────────────────────────── + System.out.println("Loading siftsmall base vectors..."); + List> floatVectors = SiftLoader.readFvecs(siftDir + "/siftsmall_base.fvecs"); + int dimension = floatVectors.get(0).length(); + System.out.printf("Loaded %d vectors of dimension %d%n", floatVectors.size(), dimension); + + // ── Step 2: Convert float32 vectors to signed int8 (ByteSequence) ──────────── + // SIFT base vectors store gradient histogram bins as unsigned bytes in [0, 255]. + // We subtract 128 to shift them into the signed range [-128, 127] that + // ByteVectorSimilarityFunction and the underlying SIMD routines expect. + // For other float32 datasets you would typically scale by a dataset-specific + // factor and then clamp before casting. + System.out.println("Converting float32 vectors to int8..."); + List> byteVectors = new ArrayList<>(floatVectors.size()); + for (VectorFloat fv : floatVectors) { + ByteSequence bv = vts.createByteSequence(dimension); + for (int i = 0; i < dimension; i++) { + // SIFT component in [0,255] → shift to signed [-128, 127] + bv.set(i, (byte) ((int) fv.get(i) - 128)); + } + byteVectors.add(bv); + } + + // Wrap the list in a RandomAccessByteVectorValues (RABVV) — + // the byte-vector analogue of RandomAccessVectorValues. + RandomAccessByteVectorValues rabvv = new ListRandomAccessByteVectorValues(byteVectors, dimension); + + // ── Step 3: Build the graph index using int8 vectors ───────────────────────── + // Use the fluent builder API. addHierarchy and refineFinalGraph are controlled + // via GraphIndexBuilderConfig (JMX/system property); no need to pass them here. + // The BuildScoreProvider wires up byte×byte scoring — no float32 round-trip + // occurs during construction. + int M = 32; + int efConstruction = 100; + float neighborOverflow = 1.2f; + float alpha = 1.2f; + + System.out.println("Building graph index from int8 vectors..."); + ImmutableGraphIndex heapGraph; + try (GraphIndexBuilder builder = GraphIndexBuilder.builder(rabvv, ByteVectorSimilarityFunction.EUCLIDEAN, M) + .withBeamWidth(efConstruction) + .withNeighborOverflow(neighborOverflow) + .withAlpha(alpha) + .build()) + { + heapGraph = builder.build(rabvv); + } + System.out.printf("Graph built with %d nodes%n", heapGraph.size(0)); + + // ── Step 4: Save the graph to disk with native int8 storage ────────────────── + // InlineByteVectors stores each vector as `dimension` raw bytes on disk — + // 4× more compact than the float32 InlineVectors alternative. + Path graphPath = Files.createTempFile("jvector-int8-example", null); + System.out.printf("Writing graph to disk (%d bytes/vector): %s%n", dimension, graphPath); + try (GraphIndexWriter writer = GraphIndexWriter + .getBuilderFor(GraphIndexWriterTypes.RANDOM_ACCESS_PARALLEL, heapGraph, graphPath) + .with(new InlineByteVectors(dimension)) + .build()) + { + writer.write(Map.of( + FeatureId.INLINE_BYTE_VECTORS, + nodeId -> new InlineByteVectors.State(rabvv.getVector(nodeId)) + )); + } + + // ── Step 5: Load the index from disk ───────────────────────────────────────── + System.out.println("Loading graph from disk..."); + ReaderSupplier readerSupplier = ReaderSupplierFactory.open(graphPath); + OnDiskGraphIndex diskGraph = OnDiskGraphIndex.load(readerSupplier); + + // ── Step 6: Search with random int8 query vectors ──────────────────────────── + // The score function reads each candidate's byte vector directly from disk — + // true int8 end-to-end, no float conversion anywhere in the search path. + int numQueries = 10; + int topK = 5; + Random rng = new Random(42); + + System.out.printf("%nSearching with %d random int8 query vectors (top-%d):%n", numQueries, topK); + + try (GraphSearcher searcher = new GraphSearcher(diskGraph)) { + OnDiskGraphIndex.View view = (OnDiskGraphIndex.View) searcher.getView(); + + for (int q = 0; q < numQueries; q++) { + ByteSequence queryBytes = randomInt8Vector(dimension, rng); + + // byteVectorRerankerFor reads each candidate's int8 vector from disk + // and scores it byte×byte — no float32 anywhere in the hot path. + SearchScoreProvider ssp = new DefaultSearchScoreProvider( + view.byteVectorRerankerFor(queryBytes, ByteVectorSimilarityFunction.EUCLIDEAN)); + + SearchResult result = searcher.search(ssp, topK, Bits.ALL); + + System.out.printf("Query %2d → top neighbors: ", q); + for (SearchResult.NodeScore ns : result.getNodes()) { + System.out.printf("(id=%d, score=%.4f) ", ns.node, ns.score); + } + System.out.println(); + } + } + + // cleanup + readerSupplier.close(); + Files.deleteIfExists(graphPath); + } + + /** + * Returns a random signed-byte vector as a {@link ByteSequence}. + * Each component is independently and uniformly drawn from [-128, 127]. + */ + private static ByteSequence randomInt8Vector(int dimension, Random rng) { + ByteSequence v = vts.createByteSequence(dimension); + for (int i = 0; i < dimension; i++) { + // nextInt(256) gives [0, 255]; subtract 128 → [-128, 127] + v.set(i, (byte) (rng.nextInt(256) - 128)); + } + return v; + } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/NvqExample.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/NvqExample.java index 661d29cb2..8ccc64f1b 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/NvqExample.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/NvqExample.java @@ -24,6 +24,7 @@ import java.util.stream.IntStream; import io.github.jbellis.jvector.disk.ReaderSupplierFactory; +import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets; import io.github.jbellis.jvector.example.util.AccuracyMetrics; import io.github.jbellis.jvector.graph.GraphIndexBuilder; @@ -51,7 +52,7 @@ public class NvqExample { public static void main(String[] args) throws IOException { // Load a preconfigured dataset String datasetName = args.length == 0 ? "ada002-100k" : args[0]; - var ds = DataSets.loadDataSet(datasetName).orElseThrow(() -> + DataSet ds = DataSets.loadDataSet(datasetName).orElseThrow(() -> new RuntimeException("dataset not found")) .getDataSet(); var dim = ds.getDimension(); diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/TutorialRunner.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/TutorialRunner.java index c675f90c8..0e914e65d 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/TutorialRunner.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/TutorialRunner.java @@ -41,6 +41,9 @@ public static void main(String[] args) throws IOException { case "nvq": NvqExample.main(forwardArgs); break; + case "int8": + Int8Example.main(forwardArgs); + break; default: throw new IllegalArgumentException("Unknown example" + args[0]); } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/SiftLoader.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/SiftLoader.java index a491d0c9e..2a7442eb5 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/SiftLoader.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/SiftLoader.java @@ -17,6 +17,7 @@ package io.github.jbellis.jvector.example.util; import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.ByteSequence; import io.github.jbellis.jvector.vector.types.VectorFloat; import io.github.jbellis.jvector.vector.types.VectorTypeSupport; @@ -60,6 +61,38 @@ public static List> readFvecs(String filePath) { return vectors; } + /// Reads a `.bvecs`-layout file whose components are **signed int8** values (-128..127). + /// + /// The layout is the same as the TexMex `.bvecs` format (a little-endian int dimension followed by + /// `dimension` bytes, per vector), but the payload is interpreted as signed. + /// **Unsigned uint8 payloads, which is what conventional `.bvecs` files such as SIFT1B/BIGANN contain + /// (0..255), are not supported.** + /// Bytes are copied as-is, so an unsigned file would be silently misread (200 becomes -56) + /// and distances would change. Convert unsigned data to signed int8 before loading it. + /// + /// @param filePath path to a signed-int8 `.bvecs`-layout file + /// @return the vectors, as signed int8 byte sequences + public static List> readInt8Bvecs(String filePath) { + var vectors = new ArrayList>(); + try (var dis = new DataInputStream(new BufferedInputStream(new FileInputStream(filePath)))) { + while (dis.available() > 0) { + var dimension = Integer.reverseBytes(dis.readInt()); + if (dimension <= 0) { + throw new IOException("Corrupt bvecs file: negative or zero dimension " + dimension + " (possible file corruption or wrong format)"); + } + if (dimension > 100_000) { + throw new IOException("Unreasonable dimension " + dimension + " in bvecs file (possible file corruption or wrong format)"); + } + var data = new byte[dimension]; + dis.readFully(data); + vectors.add(vectorTypeSupport.createByteSequence(data)); + } + } catch (IOException ex) { + throw new UncheckedIOException(ex); + } + return vectors; + } + public static List> readIvecs(String filename) { var groundTruthTopK = new ArrayList>(); diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/Compression.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/Compression.java index a8277508a..41f3ad967 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/Compression.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/Compression.java @@ -44,7 +44,7 @@ public Function getCompressorParameters() { if (strCenterData == null) { centerData = ds.getSimilarityFunction() == VectorSimilarityFunction.EUCLIDEAN; } else { - centerData = strCenterData.equals("Yes");; + centerData = strCenterData.equals("Yes"); } if (parameters.containsKey("m")) { diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/ConstructionParameters.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/ConstructionParameters.java index 5177fdf4a..9494dadd4 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/ConstructionParameters.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/ConstructionParameters.java @@ -35,6 +35,12 @@ public class ConstructionParameters extends CommonParameters { public List> getFeatureSets() { List> featureSets = null; + if (reranking == null || reranking.isEmpty()) { + throw new IllegalArgumentException("construction.reranking is required (e.g. [FP] or [NVQ])"); + } + if (fusedGraph == null || fusedGraph.isEmpty()) { + throw new IllegalArgumentException("construction.fusedGraph is required (e.g. [No] or [Yes])"); + } for (var fusedItem : fusedGraph) { var newFeatures = reranking.stream().map(item -> { EnumSet features; @@ -69,4 +75,12 @@ public List> getFeatureSets() { return featureSets; } -} \ No newline at end of file + + /** + * Feature sets for int8 datasets, which are stored natively as bytes and have no reranking or fused-graph + * options, so there is nothing to configure. + */ + public List> getInt8FeatureSets() { + return List.of(EnumSet.of(FeatureId.INLINE_BYTE_VECTORS)); + } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/MultiConfig.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/MultiConfig.java index 851d2b2ab..7aeb9addf 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/MultiConfig.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/MultiConfig.java @@ -185,6 +185,38 @@ public boolean isCompactionEnabled() { return withCompaction != null && withCompaction.contains(Boolean.TRUE); } + /** + * Int8 datasets are stored and scored natively as bytes; reranking and compression are not supported. + * Fails with an explicit error unless reranking is absent and every compression type is "None" and fusedGraph is not enabled. + */ + public void validateForInt8Dataset() { + String name = dataset; + if (construction != null) { + if (construction.reranking != null && !construction.reranking.isEmpty()) { + throw new IllegalArgumentException("Int8 dataset '" + name + "' does not support construction.reranking " + + construction.reranking + "; remove 'reranking' from the YAML config."); + } + if (construction.fusedGraph != null && construction.fusedGraph.contains(true)) { + throw new IllegalArgumentException("Int8 dataset '" + name + "' does not support construction.fusedGraph: " + + "fused graphs require compression; use [No] or remove 'fusedGraph' from the YAML config."); + } + requireNoCompression(name, "construction", construction.compression); + } + if (search != null) { + requireNoCompression(name, "search", search.compression); + } + } + + private static void requireNoCompression(String dataset, String section, List compression) { + if (compression == null) return; + for (var c : compression) { + if (!"None".equals(c.type)) { + throw new IllegalArgumentException("Int8 dataset '" + dataset + "' does not support " + section + + ".compression type '" + c.type + "'; only 'None' is allowed."); + } + } + } + public int getYamlSchemaVersion() { return yamlSchemaVersion; } diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFDTest.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFDTest.java index bcc9032ac..a1198c3fd 100644 --- a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFDTest.java +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFDTest.java @@ -114,8 +114,8 @@ public void loadsDatasetFromLocalCatalogAndFiles() throws IOException { assertEquals("test-ds", info.get().getName()); var ds = info.get().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); - assertEquals(2, ds.getQueryVectors().size()); + assertEquals(5, ds.size()); + assertEquals(2, ds.getQueryCount()); assertEquals(2, ds.getGroundTruth().size()); assertEquals(4, ds.getDimension()); } @@ -126,7 +126,7 @@ public void bufferedModePreservesValuesAndClosesItsSource() throws Exception { writeTestDataFiles(cacheDir); var loader = new DataSetLoaderSimpleMFD(null, cacheDir.toString(), false, testMetadata); var info = DataSets.loadDataSet("test-ds", java.util.List.of(loader), BaseVectorLoading.BUFFERED).orElseThrow(); - var ds = info.getDataSet(); + DataSet ds = info.getDataSet(); assertSame(ds, info.getDataSet()); var values = ds.getBaseRavv(); var file = assertInstanceOf(io.github.jbellis.jvector.disk.FvecFileVectorValues.class, values); @@ -135,7 +135,7 @@ public void bufferedModePreservesValuesAndClosesItsSource() throws Exception { ds.getBaseVectors().get(1); assertEquals(1f, retained.get(0)); assertEquals(0f, retained.get(1)); - try (ds; var preload = loader.loadDataSet("test-ds").orElseThrow().getDataSet()) { + try (ds; var preload = (DataSet) loader.loadDataSet("test-ds").orElseThrow().getDataSet()) { assertEquals(preload.getGroundTruth(), ds.getGroundTruth()); assertEquals(preload.getSimilarityFunction(), ds.getSimilarityFunction()); var consumers = java.util.concurrent.Executors.newFixedThreadPool(4); @@ -145,7 +145,7 @@ public void bufferedModePreservesValuesAndClosesItsSource() throws Exception { tasks.add(() -> { try (var view = file.copy()) { for (int pass = 0; pass < 20; pass++) { - for (int ordinal = 0; ordinal < ds.getBaseVectors().size(); ordinal++) { + for (int ordinal = 0; ordinal < ds.size(); ordinal++) { var expected = preload.getBaseVectors().get(ordinal); var actual = view.getVector(ordinal); for (int col = 0; col < expected.length(); col++) @@ -300,7 +300,7 @@ public void loadsWithLocalPathAsYamlFile() throws IOException { ); var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -343,7 +343,7 @@ public void nullCatalogUrlWorksWithLocalCatalog() throws IOException { ); var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -356,7 +356,7 @@ public void emptyCatalogUrlWorksWithLocalCatalog() throws IOException { ); var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -410,8 +410,8 @@ public void catalogUrlFetchesRemoteCatalogWhenNoLocalCatalogExists() throws IOEx assertTrue(Files.exists(cacheDir.resolve("catalog_entries.yaml"))); var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); - assertEquals(2, ds.getQueryVectors().size()); + assertEquals(5, ds.size()); + assertEquals(2, ds.getQueryCount()); assertEquals(2, ds.getGroundTruth().size()); assertEquals(4, ds.getDimension()); @@ -537,7 +537,7 @@ public void subdirectoryDataFilesResolveRelativeToTheirCatalog() throws IOExcept // data files should resolve relative to subDir, not cacheDir var ds = loader.loadDataSet("sub-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -563,7 +563,7 @@ public void duplicateEntryAcrossCatalogsDoesNotFail() throws IOException { // should load without error — whichever catalog wins, the dataset is valid var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); assertNotNull(ds); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -669,7 +669,7 @@ public void base_urlOverrideIsUsedForDownload() throws IOException { ); var ds = loader.loadDataSet("private-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -711,7 +711,7 @@ public void base_urlWithoutTrailingSlashIsNormalized() throws IOException { // should load fine — base_url is normalized with trailing slash var ds = loader.loadDataSet("private-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -737,7 +737,7 @@ public void subdirectoryPathsInFileValuesResolveCorrectly() throws IOException { ); var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); assertEquals(4, ds.getDimension()); } @@ -783,7 +783,7 @@ public void defaultsAreFoldedIntoEntries() throws IOException { // files exist locally so base_url isn't hit, but the entry should load fine var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -804,7 +804,7 @@ public void entryOverridesDefaults() throws IOException { ); var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -852,7 +852,7 @@ public void cacheDirOverridesLocalDir() throws IOException { ); var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -873,7 +873,7 @@ public void cacheDirFromDefaults() throws IOException { ); var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -897,7 +897,7 @@ public void cacheDirEntryOverridesDefault() throws IOException { ); var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -969,7 +969,7 @@ public void nonExistentCacheDirWithLocalFilesPrePopulated() throws IOException { ); var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } // ======================================================================== @@ -994,7 +994,7 @@ public void envVarExpandedInBaseurl() throws IOException { // files exist locally so the expanded base_url isn't hit, but parsing should succeed var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -1015,7 +1015,7 @@ public void envVarExpandedInCacheDir() throws IOException { ); var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -1036,7 +1036,7 @@ public void envVarExpandedInDefaults() throws IOException { ); var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -1070,7 +1070,7 @@ public void envVarWithDefaultUsesDefault() throws IOException { ); var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -1147,7 +1147,7 @@ public void includeWithUnreachableRemoteWarnsButDoesNotFail() throws IOException // local entry should still work var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -1181,7 +1181,7 @@ public void includeWithMissingUrlFieldIsIgnored() throws IOException { ); var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -1204,7 +1204,7 @@ public void localEntryOverridesIncludedEntry() throws IOException { // local entry should work — the failed include shouldn't prevent it var ds = loader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, ds.getBaseVectors().size()); + assertEquals(5, ds.size()); } @Test @@ -1398,7 +1398,7 @@ public void includeOnlyCatalogLoadsOfflineFromCachedRemoteCatalog() throws IOExc ); var onlineDs = onlineLoader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, onlineDs.getBaseVectors().size()); + assertEquals(5, onlineDs.size()); assertTrue(Files.exists(cachedDataDir.resolve("test_base.fvecs"))); assertTrue(Files.exists(cachedDataDir.resolve("test_query.fvecs"))); assertTrue(Files.exists(cachedDataDir.resolve("test_gt.ivecs"))); @@ -1413,8 +1413,8 @@ public void includeOnlyCatalogLoadsOfflineFromCachedRemoteCatalog() throws IOExc ); var offlineDs = offlineLoader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(5, offlineDs.getBaseVectors().size()); - assertEquals(2, offlineDs.getQueryVectors().size()); + assertEquals(5, offlineDs.size()); + assertEquals(2, offlineDs.getQueryCount()); assertEquals(2, offlineDs.getGroundTruth().size()); assertEquals(4, offlineDs.getDimension()); } @@ -1456,7 +1456,7 @@ public void localCatalogOverridesCachedIncludedRemoteCatalogOffline() throws IOE ); var onlineDs = onlineLoader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(1, onlineDs.getBaseVectors().size()); + assertEquals(1, onlineDs.size()); } finally { server.stop(0); } @@ -1467,8 +1467,8 @@ public void localCatalogOverridesCachedIncludedRemoteCatalogOffline() throws IOE ); var offlineDs = offlineLoader.loadDataSet("test-ds").orElseThrow().getDataSet(); - assertEquals(1, offlineDs.getBaseVectors().size()); - assertEquals(1, offlineDs.getQueryVectors().size()); + assertEquals(1, offlineDs.size()); + assertEquals(1, offlineDs.getQueryCount()); assertEquals(1, offlineDs.getGroundTruth().size()); assertEquals(4, offlineDs.getDimension()); } diff --git a/jvector-examples/yaml-configs/dataset-metadata.yml b/jvector-examples/yaml-configs/dataset-metadata.yml index a0c840a8c..f670c520c 100644 --- a/jvector-examples/yaml-configs/dataset-metadata.yml +++ b/jvector-examples/yaml-configs/dataset-metadata.yml @@ -108,4 +108,18 @@ openai-3072-1m: load_behavior: NO_SCRUB openai-1536-1m: similarity_function: DOT_PRODUCT - load_behavior: NO_SCRUB \ No newline at end of file + load_behavior: NO_SCRUB +# is_int8: true means the vectors are SIGNED int8 (-128..127). Unsigned uint8 .bvecs data +# (e.g. conventional SIFT1B/BIGANN files, 0..255) is not supported and would be silently misread. +sift1m-128-euclidean-int8: + similarity_function: EUCLIDEAN + load_behavior: NO_SCRUB + is_int8: true +wikipedia-2023-11-en-mxbai-int8-1m: + similarity_function: DOT_PRODUCT + load_behavior: NO_SCRUB + is_int8: true +wiki-multilingual-v3-int8-1m: + similarity_function: DOT_PRODUCT + load_behavior: NO_SCRUB + is_int8: true diff --git a/jvector-examples/yaml-configs/datasets.yml b/jvector-examples/yaml-configs/datasets.yml index 60aa893f5..01ce18725 100644 --- a/jvector-examples/yaml-configs/datasets.yml +++ b/jvector-examples/yaml-configs/datasets.yml @@ -1,3 +1,6 @@ +int8: + - wikipedia-2023-11-en-mxbai-int8-1m + - wiki-multilingual-v3-int8-1m jvector-100k: - cohere-english-v3-100k - ada002-100k diff --git a/jvector-examples/yaml-configs/index-parameters/cap-10M.yml b/jvector-examples/yaml-configs/index-parameters/cap-10M.yml index c9c6c09d8..93b5b4eb9 100644 --- a/jvector-examples/yaml-configs/index-parameters/cap-10M.yml +++ b/jvector-examples/yaml-configs/index-parameters/cap-10M.yml @@ -1,5 +1,5 @@ yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 dataset: cap-10M useDatasetPrefetching: false diff --git a/jvector-examples/yaml-configs/index-parameters/cap-1M.yml b/jvector-examples/yaml-configs/index-parameters/cap-1M.yml index 6aa8454b5..08a8e38a4 100644 --- a/jvector-examples/yaml-configs/index-parameters/cap-1M.yml +++ b/jvector-examples/yaml-configs/index-parameters/cap-1M.yml @@ -1,5 +1,5 @@ yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 dataset: cap-1M useDatasetPrefetching: false diff --git a/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-10M.yml b/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-10M.yml index 17b3dd54a..97fb04c36 100644 --- a/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-10M.yml +++ b/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-10M.yml @@ -1,5 +1,5 @@ yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 dataset: cohere-english-v3-10M useDatasetPrefetching: false diff --git a/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-1M.yml b/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-1M.yml index 56fb2849f..dc9dcf7f2 100644 --- a/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-1M.yml +++ b/jvector-examples/yaml-configs/index-parameters/cohere-english-v3-1M.yml @@ -1,5 +1,5 @@ yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 dataset: cohere-english-v3-1M useDatasetPrefetching: false diff --git a/jvector-examples/yaml-configs/index-parameters/colbert-1M.yml b/jvector-examples/yaml-configs/index-parameters/colbert-1M.yml index 66b47a056..3bcf9668a 100644 --- a/jvector-examples/yaml-configs/index-parameters/colbert-1M.yml +++ b/jvector-examples/yaml-configs/index-parameters/colbert-1M.yml @@ -1,5 +1,5 @@ yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 dataset: colbert-1M useDatasetPrefetching: false diff --git a/jvector-examples/yaml-configs/index-parameters/default.yml b/jvector-examples/yaml-configs/index-parameters/default.yml index e0a4ed3a5..fdb3da125 100644 --- a/jvector-examples/yaml-configs/index-parameters/default.yml +++ b/jvector-examples/yaml-configs/index-parameters/default.yml @@ -1,5 +1,5 @@ yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 dataset: default # Prefetch bounded batches instead of preloading the entire dataset into memory. diff --git a/jvector-examples/yaml-configs/index-parameters/dpr-gemma-10M.yml b/jvector-examples/yaml-configs/index-parameters/dpr-gemma-10M.yml index 90eac34e8..16acbe6cf 100644 --- a/jvector-examples/yaml-configs/index-parameters/dpr-gemma-10M.yml +++ b/jvector-examples/yaml-configs/index-parameters/dpr-gemma-10M.yml @@ -1,5 +1,5 @@ yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 dataset: dpr-gemma-10M useDatasetPrefetching: false diff --git a/jvector-examples/yaml-configs/index-parameters/dpr-gemma-1M.yml b/jvector-examples/yaml-configs/index-parameters/dpr-gemma-1M.yml index 354c06436..a194ace14 100644 --- a/jvector-examples/yaml-configs/index-parameters/dpr-gemma-1M.yml +++ b/jvector-examples/yaml-configs/index-parameters/dpr-gemma-1M.yml @@ -1,5 +1,5 @@ yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 dataset: dpr-gemma-1M useDatasetPrefetching: false diff --git a/jvector-examples/yaml-configs/index-parameters/glove-100-angular.yml b/jvector-examples/yaml-configs/index-parameters/glove-100-angular.yml index 1b03b1edd..920fee6a8 100644 --- a/jvector-examples/yaml-configs/index-parameters/glove-100-angular.yml +++ b/jvector-examples/yaml-configs/index-parameters/glove-100-angular.yml @@ -1,6 +1,6 @@ # This is here for loader testing purposes only, use other datasets for realistic test data yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 dataset: glove-100-angular useDatasetPrefetching: false diff --git a/jvector-examples/yaml-configs/index-parameters/glove-200-angular.yml b/jvector-examples/yaml-configs/index-parameters/glove-200-angular.yml index 1630baaba..f39457919 100644 --- a/jvector-examples/yaml-configs/index-parameters/glove-200-angular.yml +++ b/jvector-examples/yaml-configs/index-parameters/glove-200-angular.yml @@ -1,6 +1,6 @@ # This is here for loader testing purposes only, use other datasets for realistic test data yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 dataset: glove-200-angular useDatasetPrefetching: false diff --git a/jvector-examples/yaml-configs/index-parameters/glove-25-angular.yml b/jvector-examples/yaml-configs/index-parameters/glove-25-angular.yml index bc2293a9d..e25426f0f 100644 --- a/jvector-examples/yaml-configs/index-parameters/glove-25-angular.yml +++ b/jvector-examples/yaml-configs/index-parameters/glove-25-angular.yml @@ -1,6 +1,6 @@ # This is here for loader testing purposes only, use other datasets for realistic test data yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 dataset: glove-25-angular useDatasetPrefetching: false diff --git a/jvector-examples/yaml-configs/index-parameters/glove-50-angular.yml b/jvector-examples/yaml-configs/index-parameters/glove-50-angular.yml index 309ef9769..cde4630f9 100644 --- a/jvector-examples/yaml-configs/index-parameters/glove-50-angular.yml +++ b/jvector-examples/yaml-configs/index-parameters/glove-50-angular.yml @@ -1,6 +1,6 @@ # This is here for loader testing purposes only, use other datasets for realistic test data yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 dataset: glove-50-angular useDatasetPrefetching: false diff --git a/jvector-examples/yaml-configs/index-parameters/lastfm-64-dot.yml b/jvector-examples/yaml-configs/index-parameters/lastfm-64-dot.yml index c2e2fb98a..5ac19c195 100644 --- a/jvector-examples/yaml-configs/index-parameters/lastfm-64-dot.yml +++ b/jvector-examples/yaml-configs/index-parameters/lastfm-64-dot.yml @@ -1,6 +1,6 @@ # This is here for loader testing purposes only, use other datasets for realistic test data yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 dataset: lastfm-64-dot useDatasetPrefetching: false diff --git a/jvector-examples/yaml-configs/index-parameters/nytimes-256-angular.yml b/jvector-examples/yaml-configs/index-parameters/nytimes-256-angular.yml index 4e2db21a7..50e1ba072 100644 --- a/jvector-examples/yaml-configs/index-parameters/nytimes-256-angular.yml +++ b/jvector-examples/yaml-configs/index-parameters/nytimes-256-angular.yml @@ -1,6 +1,6 @@ # This is here for loader testing purposes only, use other datasets for realistic test data yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 dataset: nytimes-256-angular useDatasetPrefetching: false diff --git a/jvector-examples/yaml-configs/index-parameters/sift-128-euclidean.yml b/jvector-examples/yaml-configs/index-parameters/sift-128-euclidean.yml index 94362ec70..93569a009 100644 --- a/jvector-examples/yaml-configs/index-parameters/sift-128-euclidean.yml +++ b/jvector-examples/yaml-configs/index-parameters/sift-128-euclidean.yml @@ -1,6 +1,6 @@ # This is here for loader testing purposes only, use other datasets for realistic test data yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 dataset: sift-128-euclidean useDatasetPrefetching: false diff --git a/jvector-examples/yaml-configs/index-parameters/wiki-multilingual-v3-int8-1m.yml b/jvector-examples/yaml-configs/index-parameters/wiki-multilingual-v3-int8-1m.yml new file mode 100644 index 000000000..e7d98b7dc --- /dev/null +++ b/jvector-examples/yaml-configs/index-parameters/wiki-multilingual-v3-int8-1m.yml @@ -0,0 +1,23 @@ +yamlSchemaVersion: 1 +onDiskIndexVersion: 7 + +dataset: wiki-multilingual-v3-int8-1m + +construction: + outDegree: [32] + efConstruction: [100] + neighborOverflow: [1.2f] + addHierarchy: [No] + refineFinalGraph: [No] + fusedGraph: [No] + compression: + - type: None + useSavedIndexIfExists: No + +search: + topKOverquery: + 10: [1.0] + 100: [1.0] + useSearchPruning: [Yes] + compression: + - type: None diff --git a/jvector-examples/yaml-configs/index-parameters/wikipedia-2023-11-en-mxbai-int8-1m.yml b/jvector-examples/yaml-configs/index-parameters/wikipedia-2023-11-en-mxbai-int8-1m.yml new file mode 100644 index 000000000..23b470734 --- /dev/null +++ b/jvector-examples/yaml-configs/index-parameters/wikipedia-2023-11-en-mxbai-int8-1m.yml @@ -0,0 +1,23 @@ +yamlSchemaVersion: 1 +onDiskIndexVersion: 7 + +dataset: wikipedia-2023-11-en-mxbai-int8-1m + +construction: + outDegree: [32] + efConstruction: [100] + neighborOverflow: [1.2f] + addHierarchy: [No] + refineFinalGraph: [No] + fusedGraph: [No] + compression: + - type: None + useSavedIndexIfExists: No + +search: + topKOverquery: + 10: [1.0] + 100: [1.0] + useSearchPruning: [Yes] + compression: + - type: None diff --git a/jvector-examples/yaml-configs/run-config.yml b/jvector-examples/yaml-configs/run-config.yml index ba5ee1077..04dd1e160 100644 --- a/jvector-examples/yaml-configs/run-config.yml +++ b/jvector-examples/yaml-configs/run-config.yml @@ -1,6 +1,6 @@ # Run-level controls for BenchYAML (applies to the whole invocation) yamlSchemaVersion: 1 -onDiskIndexVersion: 6 +onDiskIndexVersion: 7 # Benchmark stats to compute benchmarks: 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 442ae812f..f26fb7bb9 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 @@ -91,6 +91,30 @@ public VectorFloat nvqSub(VectorFloat a, int aOffset, VectorFloat b, in return NativeNVQArithmetic.sub(a, aOffset, b, bOffset, length); } + @Override + public float dotProduct(ByteSequence a, ByteSequence b) { + return NativeSimdOps.dot_product_i8( + ((MemorySegmentByteSequence) a).get(), (long) a.offset(), + ((MemorySegmentByteSequence) b).get(), (long) b.offset(), + (long) a.length()); + } + + @Override + public float squareDistance(ByteSequence a, ByteSequence b) { + return NativeSimdOps.euclidean_i8( + ((MemorySegmentByteSequence) a).get(), (long) a.offset(), + ((MemorySegmentByteSequence) b).get(), (long) b.offset(), + (long) a.length()); + } + + @Override + public float cosine(ByteSequence a, ByteSequence b) { + return NativeSimdOps.cosine_i8( + ((MemorySegmentByteSequence) a).get(), (long) a.offset(), + ((MemorySegmentByteSequence) b).get(), (long) b.offset(), + (long) a.length()); + } + @Override protected FloatVector fromVectorFloat(VectorSpecies SPEC, VectorFloat vector, int offset) { return FloatVector.fromMemorySegment(SPEC, ((MemorySegmentVectorFloat) vector).get(), vector.offset(offset), ByteOrder.LITTLE_ENDIAN); diff --git a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java index bbe8e0cda..68115ac53 100644 --- a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java +++ b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java @@ -3030,6 +3030,192 @@ public static void nvq_shuffle_query_in_place_8bit(MemorySegment vector, long le } } + private static class dot_product_i8 { + public static final FunctionDescriptor DESC = FunctionDescriptor.of( + NativeSimdOps.C_FLOAT, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG, + NativeSimdOps.C_LONG + ); + + public static final MemorySegment ADDR = NativeSimdOps.findOrThrow("dot_product_i8"); + + public static final MethodHandle HANDLE = Linker.nativeLinker().downcallHandle(ADDR, DESC, Linker.Option.critical(true)); + } + + /** + * Function descriptor for: + * {@snippet lang=c : + * float dot_product_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static FunctionDescriptor dot_product_i8$descriptor() { + return dot_product_i8.DESC; + } + + /** + * Downcall method handle for: + * {@snippet lang=c : + * float dot_product_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static MethodHandle dot_product_i8$handle() { + return dot_product_i8.HANDLE; + } + + /** + * Address for: + * {@snippet lang=c : + * float dot_product_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static MemorySegment dot_product_i8$address() { + return dot_product_i8.ADDR; + } + + /** + * {@snippet lang=c : + * float dot_product_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static float dot_product_i8(MemorySegment a, long aoffset, MemorySegment b, long boffset, long length) { + var mh$ = dot_product_i8.HANDLE; + try { + if (TRACE_DOWNCALLS) { + traceDowncall("dot_product_i8", a, aoffset, b, boffset, length); + } + return (float)mh$.invokeExact(a, aoffset, b, boffset, length); + } catch (Throwable ex$) { + throw new AssertionError("should not reach here", ex$); + } + } + + private static class euclidean_i8 { + public static final FunctionDescriptor DESC = FunctionDescriptor.of( + NativeSimdOps.C_FLOAT, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG, + NativeSimdOps.C_LONG + ); + + public static final MemorySegment ADDR = NativeSimdOps.findOrThrow("euclidean_i8"); + + public static final MethodHandle HANDLE = Linker.nativeLinker().downcallHandle(ADDR, DESC, Linker.Option.critical(true)); + } + + /** + * Function descriptor for: + * {@snippet lang=c : + * float euclidean_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static FunctionDescriptor euclidean_i8$descriptor() { + return euclidean_i8.DESC; + } + + /** + * Downcall method handle for: + * {@snippet lang=c : + * float euclidean_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static MethodHandle euclidean_i8$handle() { + return euclidean_i8.HANDLE; + } + + /** + * Address for: + * {@snippet lang=c : + * float euclidean_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static MemorySegment euclidean_i8$address() { + return euclidean_i8.ADDR; + } + + /** + * {@snippet lang=c : + * float euclidean_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static float euclidean_i8(MemorySegment a, long aoffset, MemorySegment b, long boffset, long length) { + var mh$ = euclidean_i8.HANDLE; + try { + if (TRACE_DOWNCALLS) { + traceDowncall("euclidean_i8", a, aoffset, b, boffset, length); + } + return (float)mh$.invokeExact(a, aoffset, b, boffset, length); + } catch (Throwable ex$) { + throw new AssertionError("should not reach here", ex$); + } + } + + private static class cosine_i8 { + public static final FunctionDescriptor DESC = FunctionDescriptor.of( + NativeSimdOps.C_FLOAT, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG, + NativeSimdOps.C_LONG + ); + + public static final MemorySegment ADDR = NativeSimdOps.findOrThrow("cosine_i8"); + + public static final MethodHandle HANDLE = Linker.nativeLinker().downcallHandle(ADDR, DESC, Linker.Option.critical(true)); + } + + /** + * Function descriptor for: + * {@snippet lang=c : + * float cosine_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static FunctionDescriptor cosine_i8$descriptor() { + return cosine_i8.DESC; + } + + /** + * Downcall method handle for: + * {@snippet lang=c : + * float cosine_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static MethodHandle cosine_i8$handle() { + return cosine_i8.HANDLE; + } + + /** + * Address for: + * {@snippet lang=c : + * float cosine_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static MemorySegment cosine_i8$address() { + return cosine_i8.ADDR; + } + + /** + * {@snippet lang=c : + * float cosine_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static float cosine_i8(MemorySegment a, long aoffset, MemorySegment b, long boffset, long length) { + var mh$ = cosine_i8.HANDLE; + try { + if (TRACE_DOWNCALLS) { + traceDowncall("cosine_i8", a, aoffset, b, boffset, length); + } + return (float)mh$.invokeExact(a, aoffset, b, boffset, length); + } catch (Throwable ex$) { + throw new AssertionError("should not reach here", ex$); + } + } + private static class jvector_simd_get_active_isa { public static final FunctionDescriptor DESC = FunctionDescriptor.of( NativeSimdOps.C_POINTER ); diff --git a/jvector-native/src/main/native/benchmarks/bench_similarity_i8.cpp b/jvector-native/src/main/native/benchmarks/bench_similarity_i8.cpp new file mode 100644 index 000000000..fd9af793a --- /dev/null +++ b/jvector-native/src/main/native/benchmarks/bench_similarity_i8.cpp @@ -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. + */ + +// Google Benchmark micro-benchmarks for the int8 vector similarity kernels: +// dot_product_i8, euclidean_i8, cosine_i8 +// +// Parameterised over the realistic embedding dimensions used in production: +// 128, 256, 512, 1024, 1536, 3072 +// +// Build (requires google-benchmark installed or available via pkg-config): +// meson setup build && ninja -C build bench_simd_kernels +// +// Run: +// ./build/bench_simd_kernels [--benchmark_filter=] + +#include +#include +#include + +#include "jvector_simd.h" + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +// Deterministic, non-zero int8 vector: values cycle through a signed range to +// avoid degenerate all-zero inputs while staying within [-128, 127]. +static std::vector make_i8_vec(size_t n, int8_t seed) +{ + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + int val = seed + static_cast(i % 127); + if (i % 3 == 0) val = -val; + // clamp to [-127, 127] to keep vectors non-degenerate for cosine + if (val > 127) val = 127; + if (val < -127) val = -127; + v[i] = static_cast(val); + } + return v; +} + +// Benchmark sizes matching production embedding dimensions. +static const std::vector kBenchSizes = {128, 256, 512, 1024, 1536, 3072}; + +// --------------------------------------------------------------------------- +// dot_product_i8 +// --------------------------------------------------------------------------- + +static void BM_dot_product_i8(benchmark::State& state) +{ + const size_t n = static_cast(state.range(0)); + auto a = make_i8_vec(n, 7); + auto b = make_i8_vec(n, 13); + + for (auto _ : state) { + float result = dot_product_i8(a.data(), 0, b.data(), 0, n); + benchmark::DoNotOptimize(result); + } + + state.SetItemsProcessed(state.iterations() * static_cast(n)); + state.SetBytesProcessed(state.iterations() * static_cast(n) * 2 * sizeof(int8_t)); +} +BENCHMARK(BM_dot_product_i8)->ArgsProduct({kBenchSizes}); + +// --------------------------------------------------------------------------- +// euclidean_i8 +// --------------------------------------------------------------------------- + +static void BM_euclidean_i8(benchmark::State& state) +{ + const size_t n = static_cast(state.range(0)); + auto a = make_i8_vec(n, 7); + auto b = make_i8_vec(n, 13); + + for (auto _ : state) { + float result = euclidean_i8(a.data(), 0, b.data(), 0, n); + benchmark::DoNotOptimize(result); + } + + state.SetItemsProcessed(state.iterations() * static_cast(n)); + state.SetBytesProcessed(state.iterations() * static_cast(n) * 2 * sizeof(int8_t)); +} +BENCHMARK(BM_euclidean_i8)->ArgsProduct({kBenchSizes}); + +// --------------------------------------------------------------------------- +// cosine_i8 +// --------------------------------------------------------------------------- + +static void BM_cosine_i8(benchmark::State& state) +{ + const size_t n = static_cast(state.range(0)); + auto a = make_i8_vec(n, 7); + auto b = make_i8_vec(n, 13); + + for (auto _ : state) { + float result = cosine_i8(a.data(), 0, b.data(), 0, n); + benchmark::DoNotOptimize(result); + } + + state.SetItemsProcessed(state.iterations() * static_cast(n)); + state.SetBytesProcessed(state.iterations() * static_cast(n) * 2 * sizeof(int8_t)); +} +BENCHMARK(BM_cosine_i8)->ArgsProduct({kBenchSizes}); + diff --git a/jvector-native/src/main/native/meson.build b/jvector-native/src/main/native/meson.build index ec85d86e4..921160fb0 100644 --- a/jvector-native/src/main/native/meson.build +++ b/jvector-native/src/main/native/meson.build @@ -198,6 +198,7 @@ if gtest_dep.found() sources : [ 'tests/test_helpers.cpp', 'tests/test_similarity.cpp', + 'tests/test_similarity_i8.cpp', 'tests/test_elementwise.cpp', cpu_features_test_src, ], @@ -221,7 +222,10 @@ gbench_dep = dependency('benchmark', required: false) if gbench_dep.found() executable( 'bench_simd_kernels', - sources : 'benchmarks/bench_similarity_f32.cpp', + sources : [ + 'benchmarks/bench_similarity_f32.cpp', + 'benchmarks/bench_similarity_i8.cpp', + ], dependencies: [vectorutil_dep, gbench_dep], cpp_args : ['-O3'], ) diff --git a/jvector-native/src/main/native/src/jvector_avx3_dl_kernels.cpp b/jvector-native/src/main/native/src/jvector_avx3_dl_kernels.cpp index ae8ab73bb..11e7ed33d 100644 --- a/jvector-native/src/main/native/src/jvector_avx3_dl_kernels.cpp +++ b/jvector-native/src/main/native/src/jvector_avx3_dl_kernels.cpp @@ -25,13 +25,358 @@ // VNNI, VBMI, VBMI2, IFMA, BITALG, VPOPCNTDQ, GFNI, VAES, VPCLMULQDQ // // Compiled with -march=icelake-server. -// Highway will select HWY_AVX3_DL as the static target. +// +// This file uses raw Intel AVX-512 intrinsics directly — NO Google Highway — +// so we get exactly the instructions we intend with zero abstraction overhead. + +#include +#include +#include +#include // AVX-512 + VNNI intrinsics #include "jvector_simd.h" -#include "hwy/highway.h" -#include "assert_hwy_targets.h" -namespace hn = hwy::HWY_NAMESPACE; +// ============================================================================= +// Register naming convention +// zmm = 512-bit (16 × int32, 32 × int16, 64 × int8) +// ymm = 256-bit (32 × int8, 16 × int16) +// xmm = 128-bit (16 × int8, 8 × int16) +// +// VNNI instructions used +// ───────────────────────────────────────────────────────────────────────────── +// VPDPBUSD zmm_acc, zmm_a, zmm_b +// For each group of 4 adjacent lanes (i×4 .. i×4+3): +// acc[i] += (u8)a[i×4+0] * (i8)b[i×4+0] +// + (u8)a[i×4+1] * (i8)b[i×4+1] +// + (u8)a[i×4+2] * (i8)b[i×4+2] +// + (u8)a[i×4+3] * (i8)b[i×4+3] +// → 16 i32 accumulations, 64 int8 products per zmm register per cycle. +// Latency: 3 cycles. Throughput: 1/cycle (two ports on Ice Lake). +// +// VPDPWSSD zmm_acc, zmm_a, zmm_b +// For each group of 2 adjacent i16 lanes (i×2, i×2+1): +// acc[i] += (i16)a[i×2+0] * (i16)b[i×2+0] +// + (i16)a[i×2+1] * (i16)b[i×2+1] +// → 16 i32 accumulations, 32 int16 products per zmm per cycle. +// Latency: 3 cycles. Throughput: 1/cycle. +// +// Signed i8 × signed i8 using VPDPBUSD +// ───────────────────────────────────────────────────────────────────────────── +// VPDPBUSD requires operand A to be unsigned. For signed inputs we apply the +// standard bias trick: +// (a + 128) is always non-negative, so we use it as the unsigned operand. +// (a+128) * b = a*b + 128*b → a*b = VPDPBUSD(a+128, b) - 128 * sum(b) +// +// The bias (128*sum(b)) is constant per zmm load of b, computed as: +// _mm512_dpwssd_epi32(zero, b, set1_epi16(128)) [reuse VPDPWSSD] +// and subtracted once per iteration from the accumulator. +// +// This adds one VPDPWSSD + one VPADDD per iteration, which is negligible +// compared to the main VPDPBUSD throughput. +// +// Unrolling strategy +// ───────────────────────────────────────────────────────────────────────────── +// With 3-cycle VPDPBUSD latency and 1/cycle throughput (ports 0+5), we need +// at least 4 independent accumulator chains to keep the ports saturated: +// issued cycle 0: port 0 ← acc0 +// issued cycle 1: port 5 ← acc1 +// issued cycle 2: port 0 ← acc2 +// issued cycle 3: port 5 ← acc3 (acc0 writeback done, cycle 3) +// 4× unrolling fully hides the 3-cycle latency. +// ============================================================================= namespace AVX3_DL { +// --------------------------------------------------------------------------- +// Horizontal reduce: sum all 16 int32 lanes of a zmm register. +// _mm512_reduce_add_epi32 emits the optimal fold-down sequence; the compiler +// schedules it across surrounding instructions better than manual shuffles. +// --------------------------------------------------------------------------- +static inline int32_t hsum_epi32(__m512i v) +{ + return _mm512_reduce_add_epi32(v); +} + +// --------------------------------------------------------------------------- +// dot_product_i8 — VPDPBUSD with bias correction for signed i8 × signed i8 +// --------------------------------------------------------------------------- +// +// Algorithm +// acc = VPDPBUSD(acc, a_u8, b_i8) where a_u8 = a + 128 +// bias = VPDPWSSD(bias, b_i8, 128) accumulates 128 * sum(b) +// result = hsum(acc) - hsum(bias) +// +// 4× unrolled (256 bytes/iteration) to saturate both ICX VNNI ports and +// fully hide the 3-cycle VPDPBUSD latency. +float dot_product_i8(const int8_t * __restrict__ a, size_t aoffset, + const int8_t * __restrict__ b, size_t boffset, + size_t length) +{ + a += aoffset; + b += boffset; + + const __m512i bias128 = _mm512_set1_epi16(128); + + __m512i acc0 = _mm512_setzero_si512(), acc1 = _mm512_setzero_si512(); + __m512i acc2 = _mm512_setzero_si512(), acc3 = _mm512_setzero_si512(); + __m512i bias0 = _mm512_setzero_si512(), bias1 = _mm512_setzero_si512(); + __m512i bias2 = _mm512_setzero_si512(), bias3 = _mm512_setzero_si512(); + + size_t i = 0; + for (; i + 256 <= length; i += 256) { + __m512i va0 = _mm512_loadu_si512(a + i + 0); + __m512i va1 = _mm512_loadu_si512(a + i + 64); + __m512i va2 = _mm512_loadu_si512(a + i + 128); + __m512i va3 = _mm512_loadu_si512(a + i + 192); + __m512i vb0 = _mm512_loadu_si512(b + i + 0); + __m512i vb1 = _mm512_loadu_si512(b + i + 64); + __m512i vb2 = _mm512_loadu_si512(b + i + 128); + __m512i vb3 = _mm512_loadu_si512(b + i + 192); + + // Flip sign bit: maps signed [-128,127] → unsigned [0,255]. + const __m512i flip = _mm512_set1_epi8(-128); + __m512i au0 = _mm512_add_epi8(va0, flip); + __m512i au1 = _mm512_add_epi8(va1, flip); + __m512i au2 = _mm512_add_epi8(va2, flip); + __m512i au3 = _mm512_add_epi8(va3, flip); + + // VPDPBUSD: acc[i] += (u8)au[4i+k] * (i8)vb[4i+k], k=0..3 + acc0 = _mm512_dpbusd_epi32(acc0, au0, vb0); + acc1 = _mm512_dpbusd_epi32(acc1, au1, vb1); + acc2 = _mm512_dpbusd_epi32(acc2, au2, vb2); + acc3 = _mm512_dpbusd_epi32(acc3, au3, vb3); + + // Bias: promote vb to i16 then compute 128 * sum(vb) using VPDPWSSD. + // Each 64-byte zmm of int8 is split into two 512-bit i16 vectors. + __m512i vb0_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 0))); + __m512i vb0_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 32))); + __m512i vb1_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 64))); + __m512i vb1_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 96))); + __m512i vb2_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 128))); + __m512i vb2_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 160))); + __m512i vb3_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 192))); + __m512i vb3_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 224))); + + bias0 = _mm512_dpwssd_epi32(bias0, vb0_lo, bias128); + bias0 = _mm512_dpwssd_epi32(bias0, vb0_hi, bias128); + bias1 = _mm512_dpwssd_epi32(bias1, vb1_lo, bias128); + bias1 = _mm512_dpwssd_epi32(bias1, vb1_hi, bias128); + bias2 = _mm512_dpwssd_epi32(bias2, vb2_lo, bias128); + bias2 = _mm512_dpwssd_epi32(bias2, vb2_hi, bias128); + bias3 = _mm512_dpwssd_epi32(bias3, vb3_lo, bias128); + bias3 = _mm512_dpwssd_epi32(bias3, vb3_hi, bias128); + } + __m512i acc = _mm512_add_epi32(_mm512_add_epi32(acc0, acc1), + _mm512_add_epi32(acc2, acc3)); + __m512i bias = _mm512_add_epi32(_mm512_add_epi32(bias0, bias1), + _mm512_add_epi32(bias2, bias3)); + + // Single-zmm tail (residual 64-byte blocks). + for (; i + 64 <= length; i += 64) { + __m512i va = _mm512_loadu_si512(a + i); + __m512i vb = _mm512_loadu_si512(b + i); + __m512i au = _mm512_add_epi8(va, _mm512_set1_epi8(-128)); + acc = _mm512_dpbusd_epi32(acc, au, vb); + + // Promote vb to i16 for bias calculation + __m512i vb_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i))); + __m512i vb_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 32))); + bias = _mm512_dpwssd_epi32(bias, vb_lo, bias128); + bias = _mm512_dpwssd_epi32(bias, vb_hi, bias128); + } + + int32_t result = hsum_epi32(acc) - hsum_epi32(bias); + + // Scalar tail. + for (; i < length; i++) + result += (int32_t)a[i] * (int32_t)b[i]; + + return (float)result; +} + +// --------------------------------------------------------------------------- +// euclidean_i8 — VPMOVSXBW sign-extend + VPDPWSSD squared differences +// --------------------------------------------------------------------------- +// +// Each 64-byte zmm block is processed as two 32-byte halves: +// _mm512_cvtepi8_epi16(__m256i) = VPMOVSXBW: sign-extends 32×i8 → 32×i16 +// diff = da - db (i16 subtraction, no overflow since range is [-255,255]) +// acc = VPDPWSSD(acc, diff, diff) +// +// 4× unrolled (256 bytes/iteration). +float euclidean_i8(const int8_t * __restrict__ a, size_t aoffset, + const int8_t * __restrict__ b, size_t boffset, + size_t length) +{ + a += aoffset; + b += boffset; + + __m512i acc0 = _mm512_setzero_si512(), acc1 = _mm512_setzero_si512(); + __m512i acc2 = _mm512_setzero_si512(), acc3 = _mm512_setzero_si512(); + + size_t i = 0; + for (; i + 256 <= length; i += 256) { +#define EUCL_BLOCK(off, acc_var) \ + { \ + const int8_t *ap = a + i + (off), *bp = b + i + (off); \ + __m512i da_lo = _mm512_cvtepi8_epi16( \ + _mm256_loadu_si256(reinterpret_cast(ap))); \ + __m512i db_lo = _mm512_cvtepi8_epi16( \ + _mm256_loadu_si256(reinterpret_cast(bp))); \ + __m512i da_hi = _mm512_cvtepi8_epi16( \ + _mm256_loadu_si256(reinterpret_cast(ap + 32))); \ + __m512i db_hi = _mm512_cvtepi8_epi16( \ + _mm256_loadu_si256(reinterpret_cast(bp + 32))); \ + __m512i diff_lo = _mm512_sub_epi16(da_lo, db_lo); \ + __m512i diff_hi = _mm512_sub_epi16(da_hi, db_hi); \ + acc_var = _mm512_dpwssd_epi32(acc_var, diff_lo, diff_lo); \ + acc_var = _mm512_dpwssd_epi32(acc_var, diff_hi, diff_hi); \ + } + EUCL_BLOCK( 0, acc0) + EUCL_BLOCK( 64, acc1) + EUCL_BLOCK(128, acc2) + EUCL_BLOCK(192, acc3) +#undef EUCL_BLOCK + } + __m512i acc = _mm512_add_epi32(_mm512_add_epi32(acc0, acc1), + _mm512_add_epi32(acc2, acc3)); + + // Single 64-byte tail blocks. + for (; i + 64 <= length; i += 64) { + __m512i da_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i))); + __m512i db_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i))); + __m512i da_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i + 32))); + __m512i db_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 32))); + acc = _mm512_dpwssd_epi32(acc, _mm512_sub_epi16(da_lo, db_lo), _mm512_sub_epi16(da_lo, db_lo)); + acc = _mm512_dpwssd_epi32(acc, _mm512_sub_epi16(da_hi, db_hi), _mm512_sub_epi16(da_hi, db_hi)); + } + + // 32-byte tail (one ymm → one 512-bit i16 vector). + if (i + 32 <= length) { + __m512i da = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i))); + __m512i db = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i))); + __m512i diff = _mm512_sub_epi16(da, db); + acc = _mm512_dpwssd_epi32(acc, diff, diff); + i += 32; + } + + int32_t result = hsum_epi32(acc); + + // Scalar tail. + for (; i < length; i++) { + int32_t d = (int32_t)a[i] - (int32_t)b[i]; + result += d * d; + } + return (float)result; +} + +// --------------------------------------------------------------------------- +// cosine_i8 — three parallel VPDPBUSD chains with bias correction +// --------------------------------------------------------------------------- +// +// Computes dot(a,b), ||a||², ||b||² in a single pass using VPDPBUSD. +// Bias trick: a_u = a+128 (unsigned), then subtract 128*sum(b) and 128*sum(a). +// +// dot(a,b) = hsum(VPDPBUSD(acc_dot, a_u, b)) - 128*sum(b) +// ||a||² = hsum(VPDPBUSD(acc_normA, a_u, a)) - 128*sum(a) +// ||b||² = hsum(VPDPBUSD(acc_normB, b_u, b)) - 128*sum(b) +// +// normB reuses the same biasAB accumulator as dot (both need 128*sum(b)). +// 2× unrolled (128 bytes/iteration) with 6 VPDPBUSD + 4 VPDPWSSD per iter. +float cosine_i8(const int8_t * __restrict__ a, size_t aoffset, + const int8_t * __restrict__ b, size_t boffset, + size_t length) +{ + a += aoffset; + b += boffset; + + const __m512i bias128 = _mm512_set1_epi16(128); + const __m512i flip = _mm512_set1_epi8(-128); + + __m512i dot0 = _mm512_setzero_si512(), dot1 = _mm512_setzero_si512(); + __m512i normA0 = _mm512_setzero_si512(), normA1 = _mm512_setzero_si512(); + __m512i normB0 = _mm512_setzero_si512(), normB1 = _mm512_setzero_si512(); + __m512i biasAB0 = _mm512_setzero_si512(), biasAB1 = _mm512_setzero_si512(); + __m512i biasA0 = _mm512_setzero_si512(), biasA1 = _mm512_setzero_si512(); + + size_t i = 0; + for (; i + 128 <= length; i += 128) { + __m512i va0 = _mm512_loadu_si512(a + i); + __m512i vb0 = _mm512_loadu_si512(b + i); + __m512i va1 = _mm512_loadu_si512(a + i + 64); + __m512i vb1 = _mm512_loadu_si512(b + i + 64); + __m512i au0 = _mm512_add_epi8(va0, flip); + __m512i bu0 = _mm512_add_epi8(vb0, flip); + __m512i au1 = _mm512_add_epi8(va1, flip); + __m512i bu1 = _mm512_add_epi8(vb1, flip); + + dot0 = _mm512_dpbusd_epi32(dot0, au0, vb0); + dot1 = _mm512_dpbusd_epi32(dot1, au1, vb1); + normA0 = _mm512_dpbusd_epi32(normA0, au0, va0); + normA1 = _mm512_dpbusd_epi32(normA1, au1, va1); + normB0 = _mm512_dpbusd_epi32(normB0, bu0, vb0); + normB1 = _mm512_dpbusd_epi32(normB1, bu1, vb1); + + // Promote to i16 for bias calculations + __m512i va0_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i))); + __m512i va0_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i + 32))); + __m512i vb0_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i))); + __m512i vb0_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 32))); + __m512i va1_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i + 64))); + __m512i va1_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i + 96))); + __m512i vb1_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 64))); + __m512i vb1_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 96))); + + biasAB0 = _mm512_dpwssd_epi32(biasAB0, vb0_lo, bias128); + biasAB0 = _mm512_dpwssd_epi32(biasAB0, vb0_hi, bias128); + biasAB1 = _mm512_dpwssd_epi32(biasAB1, vb1_lo, bias128); + biasAB1 = _mm512_dpwssd_epi32(biasAB1, vb1_hi, bias128); + biasA0 = _mm512_dpwssd_epi32(biasA0, va0_lo, bias128); + biasA0 = _mm512_dpwssd_epi32(biasA0, va0_hi, bias128); + biasA1 = _mm512_dpwssd_epi32(biasA1, va1_lo, bias128); + biasA1 = _mm512_dpwssd_epi32(biasA1, va1_hi, bias128); + } + __m512i dot = _mm512_add_epi32(dot0, dot1); + __m512i normA = _mm512_add_epi32(normA0, normA1); + __m512i normB = _mm512_add_epi32(normB0, normB1); + __m512i biasAB = _mm512_add_epi32(biasAB0, biasAB1); + __m512i biasA = _mm512_add_epi32(biasA0, biasA1); + + // Single-zmm tail. + for (; i + 64 <= length; i += 64) { + __m512i va = _mm512_loadu_si512(a + i); + __m512i vb = _mm512_loadu_si512(b + i); + __m512i au = _mm512_add_epi8(va, flip); + __m512i bu = _mm512_add_epi8(vb, flip); + dot = _mm512_dpbusd_epi32(dot, au, vb); + normA = _mm512_dpbusd_epi32(normA, au, va); + normB = _mm512_dpbusd_epi32(normB, bu, vb); + + // Promote to i16 for bias calculations + __m512i va_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i))); + __m512i va_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i + 32))); + __m512i vb_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i))); + __m512i vb_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 32))); + + biasAB = _mm512_dpwssd_epi32(biasAB, vb_lo, bias128); + biasAB = _mm512_dpwssd_epi32(biasAB, vb_hi, bias128); + biasA = _mm512_dpwssd_epi32(biasA, va_lo, bias128); + biasA = _mm512_dpwssd_epi32(biasA, va_hi, bias128); + } + + // Apply bias corrections before scalar tail. + int64_t dotResult = (int64_t)hsum_epi32(dot) - (int64_t)hsum_epi32(biasAB); + int64_t normAResult = (int64_t)hsum_epi32(normA) - (int64_t)hsum_epi32(biasA); + int64_t normBResult = (int64_t)hsum_epi32(normB) - (int64_t)hsum_epi32(biasAB); + + // Scalar tail. + for (; i < length; i++) { + int32_t ai = a[i], bi = b[i]; + dotResult += (int64_t)ai * bi; + normAResult += (int64_t)ai * ai; + normBResult += (int64_t)bi * bi; + } + + return (float)(dotResult / sqrt((double)normAResult * (double)normBResult)); +} + } // namespace AVX3_DL diff --git a/jvector-native/src/main/native/src/jvector_simd.cpp b/jvector-native/src/main/native/src/jvector_simd.cpp index ec69d5cfa..82b3ceae2 100644 --- a/jvector-native/src/main/native/src/jvector_simd.cpp +++ b/jvector-native/src/main/native/src/jvector_simd.cpp @@ -107,11 +107,14 @@ static const KernelVTable AVX3_vtable = { }; #undef KERNEL_ENTRY -// AVX3_DL (Ice Lake) inherits all slots from AVX3 unchanged for now. -// To override a slot: t.kernel_name = AVX3_DL::kernel_name; -// The implementation must exist in jvector_avx3_dl_kernels.cpp. +// AVX3_DL (Ice Lake) inherits all slots from AVX3, then overrides the three +// int8 similarity kernels with VNNI-accelerated versions from +// jvector_avx3_dl_kernels.cpp. static const KernelVTable AVX3_DL_vtable = []() { KernelVTable t = AVX3_vtable; + t.dot_product_i8 = AVX3_DL::dot_product_i8; + t.euclidean_i8 = AVX3_DL::euclidean_i8; + t.cosine_i8 = AVX3_DL::cosine_i8; return t; }(); diff --git a/jvector-native/src/main/native/src/jvector_simd_kernel_list.h b/jvector-native/src/main/native/src/jvector_simd_kernel_list.h index 03ec7f809..d76141b45 100644 --- a/jvector-native/src/main/native/src/jvector_simd_kernel_list.h +++ b/jvector-native/src/main/native/src/jvector_simd_kernel_list.h @@ -68,7 +68,11 @@ KERNEL_ENTRY(float, nvq_square_l2_distance_8bit, (const float *vector, const unsigned char *quantized, size_t length, float alpha, float x0, float minValue, float maxValue), (vector, quantized, length, alpha, x0, minValue, maxValue)) \ KERNEL_ENTRY(float, nvq_dot_product_8bit, (const float *vector, const unsigned char *quantized, size_t length, float alpha, float x0, float minValue, float maxValue), (vector, quantized, length, alpha, x0, minValue, maxValue)) \ KERNEL_ENTRY(int64_t, nvq_cosine_8bit_packed, (const float *vector, const unsigned char *quantized, size_t length, float alpha, float x0, float minValue, float maxValue, const float *centroid), (vector, quantized, length, alpha, x0, minValue, maxValue, centroid)) \ - KERNEL_ENTRY(void, nvq_shuffle_query_in_place_8bit, (float *vector, size_t length), (vector, length)) + KERNEL_ENTRY(void, nvq_shuffle_query_in_place_8bit, (float *vector, size_t length), (vector, length)) \ + /* Int8 byte-vector similarity (VNNI-accelerated on AVX3_DL+) */ \ + KERNEL_ENTRY(float, dot_product_i8, (const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length), (a, aoffset, b, boffset, length)) \ + KERNEL_ENTRY(float, euclidean_i8, (const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length), (a, aoffset, b, boffset, length)) \ + KERNEL_ENTRY(float, cosine_i8, (const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length), (a, aoffset, b, boffset, length)) /* ── ADD NEW KERNEL_ENTRY LINES ABOVE THIS LINE ── */ // clang-format on diff --git a/jvector-native/src/main/native/src/jvector_simd_kernels.cpp b/jvector-native/src/main/native/src/jvector_simd_kernels.cpp index 69c7fc062..aecab40ff 100644 --- a/jvector-native/src/main/native/src/jvector_simd_kernels.cpp +++ b/jvector-native/src/main/native/src/jvector_simd_kernels.cpp @@ -1903,4 +1903,188 @@ HWY_FLATTEN int64_t nvq_cosine_8bit_packed(const float *HWY_RESTRICT vector, return ((int64_t)bmag_bits << 32) | (int64_t)(uint32_t)sum_bits; } +// ============================================================================= +// Int8 byte-vector similarity kernels +// ============================================================================= +// +// These kernels operate on signed int8 (int8_t) vectors — e.g. the output of +// scalar quantization. The generic path here (compiled for SSE4.2, AVX2, AVX3) +// widens i8→i16 using ReorderWidenMulAccumulate, then accumulates into i32. +// +// On the AVX3_DL (Ice Lake+) tier these implementations are overridden in +// jvector_avx3_dl_kernels.cpp with raw AVX-512 VNNI intrinsics — processing +// 64 bytes per VPDPBUSD clock in a single instruction. +// ============================================================================= + +// Horizontal dot product of two signed int8 vectors. +HWY_FLATTEN float dot_product_i8(const int8_t *HWY_RESTRICT a, size_t aoffset, + const int8_t *HWY_RESTRICT b, size_t boffset, + size_t length) +{ + a += aoffset; + b += boffset; + const hn::ScalableTag d8; + const hn::RepartitionToWideX2 d32; // int32, lanes = len(d8)/4 + const hn::Repartition d16; // int16 + const size_t lanes8 = hn::Lanes(d8); + + // Four independent accumulators hide the multi-cycle MADD latency. + auto acc0 = hn::Zero(d32), acc1 = hn::Zero(d32); + auto acc2 = hn::Zero(d32), acc3 = hn::Zero(d32); + auto acc0h = hn::Zero(d32), acc1h = hn::Zero(d32); + auto acc2h = hn::Zero(d32), acc3h = hn::Zero(d32); + size_t i = 0; + for (; i + 4 * lanes8 <= length; i += 4 * lanes8) { + auto va0 = hn::LoadU(d8, a + i); + auto vb0 = hn::LoadU(d8, b + i); + auto va1 = hn::LoadU(d8, a + i + lanes8); + auto vb1 = hn::LoadU(d8, b + i + lanes8); + auto va2 = hn::LoadU(d8, a + i + 2*lanes8); + auto vb2 = hn::LoadU(d8, b + i + 2*lanes8); + auto va3 = hn::LoadU(d8, a + i + 3*lanes8); + auto vb3 = hn::LoadU(d8, b + i + 3*lanes8); + + // Promote to i16 and accumulate using ReorderWidenMulAccumulate (2 i16s -> 1 i32) + acc0 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteLowerTo(d16, va0), hn::PromoteLowerTo(d16, vb0), acc0, acc0h); + acc0 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteUpperTo(d16, va0), hn::PromoteUpperTo(d16, vb0), acc0, acc0h); + acc1 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteLowerTo(d16, va1), hn::PromoteLowerTo(d16, vb1), acc1, acc1h); + acc1 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteUpperTo(d16, va1), hn::PromoteUpperTo(d16, vb1), acc1, acc1h); + acc2 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteLowerTo(d16, va2), hn::PromoteLowerTo(d16, vb2), acc2, acc2h); + acc2 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteUpperTo(d16, va2), hn::PromoteUpperTo(d16, vb2), acc2, acc2h); + acc3 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteLowerTo(d16, va3), hn::PromoteLowerTo(d16, vb3), acc3, acc3h); + acc3 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteUpperTo(d16, va3), hn::PromoteUpperTo(d16, vb3), acc3, acc3h); + } + auto acc = hn::Add(hn::Add(acc0, acc1), hn::Add(acc2, acc3)); + auto acch = hn::Add(hn::Add(acc0h, acc1h), hn::Add(acc2h, acc3h)); + acc = hn::Add(acc, acch); + + for (; i + lanes8 <= length; i += lanes8) { + auto va = hn::LoadU(d8, a + i); + auto vb = hn::LoadU(d8, b + i); + auto dummy_hi = hn::Zero(d32); + acc = hn::ReorderWidenMulAccumulate(d32, hn::PromoteLowerTo(d16, va), hn::PromoteLowerTo(d16, vb), acc, dummy_hi); + acc = hn::Add(acc, dummy_hi); + dummy_hi = hn::Zero(d32); + acc = hn::ReorderWidenMulAccumulate(d32, hn::PromoteUpperTo(d16, va), hn::PromoteUpperTo(d16, vb), acc, dummy_hi); + acc = hn::Add(acc, dummy_hi); + } + int32_t result = hn::ReduceSum(d32, acc); + for (; i < length; i++) result += (int32_t)a[i] * (int32_t)b[i]; + return (float)result; +} + +// Sum of squared differences of two signed int8 vectors. +// Promote i8→i16, subtract in i16, then ReorderWidenMulAccumulate into i32. +HWY_FLATTEN float euclidean_i8(const int8_t *HWY_RESTRICT a, size_t aoffset, + const int8_t *HWY_RESTRICT b, size_t boffset, + size_t length) +{ + a += aoffset; + b += boffset; + const hn::ScalableTag d8; + const hn::RepartitionToWideX2 d32; // int32, lanes = len(d8)/4 + const size_t lanes8 = hn::Lanes(d8); + + auto acc0 = hn::Zero(d32), acc0h = hn::Zero(d32); + auto acc1 = hn::Zero(d32), acc1h = hn::Zero(d32); + auto acc2 = hn::Zero(d32), acc2h = hn::Zero(d32); + auto acc3 = hn::Zero(d32), acc3h = hn::Zero(d32); + size_t i = 0; + for (; i + 4 * lanes8 <= length; i += 4 * lanes8) { +#define DO_EUCL_BLOCK(off, lo_var, hi_var) \ + { \ + const hn::RepartitionToWide _d16; \ + auto _va8 = hn::LoadU(d8, a + i + (off)); \ + auto _vb8 = hn::LoadU(d8, b + i + (off)); \ + auto _diff_lo = hn::Sub(hn::PromoteLowerTo(_d16, _va8), \ + hn::PromoteLowerTo(_d16, _vb8)); \ + auto _diff_hi = hn::Sub(hn::PromoteUpperTo(_d16, _va8), \ + hn::PromoteUpperTo(_d16, _vb8)); \ + lo_var = hn::ReorderWidenMulAccumulate(d32, _diff_lo, _diff_lo, lo_var, hi_var); \ + lo_var = hn::ReorderWidenMulAccumulate(d32, _diff_hi, _diff_hi, lo_var, hi_var); \ + } + DO_EUCL_BLOCK(0, acc0, acc0h) + DO_EUCL_BLOCK(lanes8, acc1, acc1h) + DO_EUCL_BLOCK(2*lanes8, acc2, acc2h) + DO_EUCL_BLOCK(3*lanes8, acc3, acc3h) +#undef DO_EUCL_BLOCK + } + auto acc = hn::Add(hn::Add(acc0, acc1), hn::Add(acc2, acc3)); + auto acch = hn::Add(hn::Add(acc0h, acc1h), hn::Add(acc2h, acc3h)); + acc = hn::Add(acc, acch); + for (; i + lanes8 <= length; i += lanes8) { + const hn::RepartitionToWide d16; + auto va8 = hn::LoadU(d8, a + i); + auto vb8 = hn::LoadU(d8, b + i); + auto diff_lo = hn::Sub(hn::PromoteLowerTo(d16, va8), hn::PromoteLowerTo(d16, vb8)); + auto diff_hi = hn::Sub(hn::PromoteUpperTo(d16, va8), hn::PromoteUpperTo(d16, vb8)); + auto dummy_hi = hn::Zero(d32); + acc = hn::ReorderWidenMulAccumulate(d32, diff_lo, diff_lo, acc, dummy_hi); + acc = hn::Add(acc, dummy_hi); + dummy_hi = hn::Zero(d32); + acc = hn::ReorderWidenMulAccumulate(d32, diff_hi, diff_hi, acc, dummy_hi); + acc = hn::Add(acc, dummy_hi); + } + int32_t result = hn::ReduceSum(d32, acc); + for (; i < length; i++) { + int32_t d = (int32_t)a[i] - (int32_t)b[i]; + result += d * d; + } + return (float)result; +} + +// Cosine similarity of two signed int8 vectors. +// Computes dot(a,b), dot(a,a), dot(b,b) in parallel over a single pass. +HWY_FLATTEN float cosine_i8(const int8_t *HWY_RESTRICT a, size_t aoffset, + const int8_t *HWY_RESTRICT b, size_t boffset, + size_t length) +{ + a += aoffset; + b += boffset; + const hn::ScalableTag d8; + const hn::RepartitionToWideX2 d32; + const hn::Repartition d16; + const size_t lanes8 = hn::Lanes(d8); + + auto dot = hn::Zero(d32); + auto doth = hn::Zero(d32); + auto normA = hn::Zero(d32); + auto normAh = hn::Zero(d32); + auto normB = hn::Zero(d32); + auto normBh = hn::Zero(d32); + + size_t i = 0; + for (; i + lanes8 <= length; i += lanes8) { + auto va = hn::LoadU(d8, a + i); + auto vb = hn::LoadU(d8, b + i); + auto va_lo = hn::PromoteLowerTo(d16, va); + auto va_hi = hn::PromoteUpperTo(d16, va); + auto vb_lo = hn::PromoteLowerTo(d16, vb); + auto vb_hi = hn::PromoteUpperTo(d16, vb); + + dot = hn::ReorderWidenMulAccumulate(d32, va_lo, vb_lo, dot, doth); + dot = hn::ReorderWidenMulAccumulate(d32, va_hi, vb_hi, dot, doth); + normA = hn::ReorderWidenMulAccumulate(d32, va_lo, va_lo, normA, normAh); + normA = hn::ReorderWidenMulAccumulate(d32, va_hi, va_hi, normA, normAh); + normB = hn::ReorderWidenMulAccumulate(d32, vb_lo, vb_lo, normB, normBh); + normB = hn::ReorderWidenMulAccumulate(d32, vb_hi, vb_hi, normB, normBh); + } + + dot = hn::Add(dot, doth); + normA = hn::Add(normA, normAh); + normB = hn::Add(normB, normBh); + + int64_t dotResult = (int64_t)hn::ReduceSum(d32, dot); + int64_t normAResult = (int64_t)hn::ReduceSum(d32, normA); + int64_t normBResult = (int64_t)hn::ReduceSum(d32, normB); + + for (; i < length; i++) { + int32_t ai = a[i], bi = b[i]; + dotResult += ai * bi; + normAResult += ai * ai; + normBResult += bi * bi; + } + return (float)(dotResult / sqrt((double)normAResult * (double)normBResult)); +} + } // namespace JV_ISA diff --git a/jvector-native/src/main/native/tests/test_helpers.cpp b/jvector-native/src/main/native/tests/test_helpers.cpp index 654e86734..7d6ba81ee 100644 --- a/jvector-native/src/main/native/tests/test_helpers.cpp +++ b/jvector-native/src/main/native/tests/test_helpers.cpp @@ -76,6 +76,10 @@ const std::vector kKernelTestParams = { {100, "large_mixed_tail"}, {128, "large_power_of_2"}, {255, "large_odd_tail_15"}, + // ---- i8 VNNI 256-byte unroll boundaries (dot_product_i8/euclidean_i8) - + {256, "i8_vnni_4x_exact"}, + {263, "i8_vnni_4x_tail_7"}, + {135, "i8_vnni_2zmm_tail_7"}, }; std::unordered_set parse_cpuinfo_line(const std::string& key) @@ -108,3 +112,21 @@ std::vector make_vec(size_t n, float seed) return v; } +// Produces n int8_t values with a mix of signs and magnitudes. +// The pattern ensures no element is zero (important for cosine tests). +std::vector make_vec_i8(size_t n, int8_t seed) +{ + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + // Scale seed by a small per-element factor to get variety, + // then clamp to [-100, 100] to keep products well within int16 range. + int val = static_cast(seed) + static_cast(i % 13) - 6; + if (i % 3 == 0) val = -val; // mix of signs + if (val == 0) val = 1; // never zero + if (val > 100) val = 100; + if (val < -100) val = -100; + v[i] = static_cast(val); + } + return v; +} + diff --git a/jvector-native/src/main/native/tests/test_helpers.h b/jvector-native/src/main/native/tests/test_helpers.h index 8127cd7b3..96f102c15 100644 --- a/jvector-native/src/main/native/tests/test_helpers.h +++ b/jvector-native/src/main/native/tests/test_helpers.h @@ -51,7 +51,11 @@ std::unordered_set parse_cpuinfo_line(const std::string& key); // so that no element is exactly zero (important for cosine tests). // --------------------------------------------------------------------------- -std::vector make_vec(size_t n, float seed); +std::vector make_vec(size_t n, float seed); + +// make_vec_i8(n, seed) produces n int8_t values with a mix of signs +// suitable for testing the i8 similarity kernels. +std::vector make_vec_i8(size_t n, int8_t seed); // --------------------------------------------------------------------------- // Shared test parameter — vector length + human-readable path description. diff --git a/jvector-native/src/main/native/tests/test_similarity_i8.cpp b/jvector-native/src/main/native/tests/test_similarity_i8.cpp new file mode 100644 index 000000000..51443a4bd --- /dev/null +++ b/jvector-native/src/main/native/tests/test_similarity_i8.cpp @@ -0,0 +1,229 @@ +/* + * 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. + */ + +// Tests for int8 vector similarity kernels: dot_product_i8, euclidean_i8, cosine_i8. +// +// The kernels operate on signed int8_t vectors and return a float result: +// dot_product_i8 — (float) sum(a[i] * b[i]) +// euclidean_i8 — (float) sum((a[i] - b[i])^2) (squared L2 distance) +// cosine_i8 — (float) dot(a,b) / sqrt(||a||^2 * ||b||^2) +// +// On AVX3_DL (Ice Lake+) these are overridden with VNNI (VPDPBUSD/VPDPWSSD) +// implementations; on all other tiers the generic Highway path is used. +// +// All tests are parametrised over kKernelTestParams (defined in test_helpers.cpp), +// which covers every ISA-tier loop-boundary for both f32 and i8 kernels, including +// the VNNI-specific 64/128/256-byte unroll boundaries added for the i8 suite. + +#include "test_helpers.h" + +// --------------------------------------------------------------------------- +// Reference scalar implementations +// --------------------------------------------------------------------------- + +static float ref_dot_i8(const std::vector& a, const std::vector& b) +{ + int64_t s = 0; + for (size_t i = 0; i < a.size(); ++i) + s += static_cast(a[i]) * static_cast(b[i]); + return static_cast(s); +} + +static float ref_euclidean_i8(const std::vector& a, const std::vector& b) +{ + int64_t s = 0; + for (size_t i = 0; i < a.size(); ++i) { + int32_t d = static_cast(a[i]) - static_cast(b[i]); + s += d * d; + } + return static_cast(s); +} + +static float ref_cosine_i8(const std::vector& a, const std::vector& b) +{ + int64_t dot = 0, normA = 0, normB = 0; + for (size_t i = 0; i < a.size(); ++i) { + int32_t ai = a[i], bi = b[i]; + dot += static_cast(ai) * bi; + normA += static_cast(ai) * ai; + normB += static_cast(bi) * bi; + } + return static_cast(dot / std::sqrt(static_cast(normA) + * static_cast(normB))); +} + +// --------------------------------------------------------------------------- +// Parametrised test fixture +// --------------------------------------------------------------------------- + +class SimilarityI8Test : public ::testing::TestWithParam {}; + +// --------------------------------------------------------------------------- +// dot_product_i8 — SIMD result must match the scalar reference +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, DotProduct) +{ + const size_t n = GetParam().length; + auto a = make_vec_i8(n, 7); + auto b = make_vec_i8(n, 11); + + const float want = ref_dot_i8(a, b); + const float got = dot_product_i8(a.data(), 0, b.data(), 0, n); + + // Integer accumulation with a single int64→float cast — result is exact. + EXPECT_EQ(got, want); +} + +// --------------------------------------------------------------------------- +// dot_product_i8 with non-zero offsets — exercises the aoffset/boffset path +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, DotProductWithOffset) +{ + const size_t n = GetParam().length; + const size_t prefix = 5; // arbitrary prefix that must be ignored + + std::vector a_pad(prefix + n, 0); + std::vector b_pad(prefix + n, 0); + auto a = make_vec_i8(n, 7); + auto b = make_vec_i8(n, 11); + std::copy(a.begin(), a.end(), a_pad.begin() + prefix); + std::copy(b.begin(), b.end(), b_pad.begin() + prefix); + + const float want = ref_dot_i8(a, b); + const float got = dot_product_i8(a_pad.data(), prefix, b_pad.data(), prefix, n); + + EXPECT_EQ(got, want); +} + +// --------------------------------------------------------------------------- +// dot_product_i8 — zero vector gives exactly 0.0 +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, DotProductZeroVector) +{ + const size_t n = GetParam().length; + auto a = make_vec_i8(n, 7); + std::vector z(n, 0); + + EXPECT_EQ(dot_product_i8(a.data(), 0, z.data(), 0, n), 0.0f); +} + +// --------------------------------------------------------------------------- +// euclidean_i8 — SIMD result must match the scalar reference +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, Euclidean) +{ + const size_t n = GetParam().length; + auto a = make_vec_i8(n, 7); + auto b = make_vec_i8(n, 11); + + const float want = ref_euclidean_i8(a, b); + const float got = euclidean_i8(a.data(), 0, b.data(), 0, n); + + // Integer accumulation with a single int64→float cast — result is exact. + EXPECT_EQ(got, want); +} + +// --------------------------------------------------------------------------- +// euclidean_i8 — identical vectors must give exactly 0 +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, EuclideanSameVector) +{ + const size_t n = GetParam().length; + auto a = make_vec_i8(n, 9); + + const float got = euclidean_i8(a.data(), 0, a.data(), 0, n); + + EXPECT_EQ(got, 0.0f); +} + +// --------------------------------------------------------------------------- +// cosine_i8 — SIMD result must match the scalar reference +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, Cosine) +{ + const size_t n = GetParam().length; + auto a = make_vec_i8(n, 7); + auto b = make_vec_i8(n, 11); + + const float want = ref_cosine_i8(a, b); + const float got = cosine_i8(a.data(), 0, b.data(), 0, n); + + EXPECT_NEAR(got, want, 1e-5f); +} + +// --------------------------------------------------------------------------- +// cosine_i8 — parallel vectors (b = k*a, k > 0) should give similarity ≈ 1.0 +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, CosineParallelVectors) +{ + const size_t n = GetParam().length; + // Use small magnitudes so that 2*val stays within int8 range. + auto a = make_vec_i8(n, 3); + std::vector b(n); + for (size_t i = 0; i < n; ++i) + b[i] = static_cast(std::max(-127, std::min(127, 2 * static_cast(a[i])))); + + const float got = cosine_i8(a.data(), 0, b.data(), 0, n); + + EXPECT_NEAR(got, 1.0f, 1e-5f); +} + +// --------------------------------------------------------------------------- +// cosine_i8 — orthogonal vectors should give similarity ≈ 0.0 +// +// Same analytic construction as the f32 test: for even n, +// a = [+1, +1, +1, ...] +// b = [+1, -1, +1, -1, ...] → dot(a,b) = 0. +// Odd n: the odd last element is zeroed out on b (unchanged on a) so the +// dot product remains zero without affecting the norms materially. +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, CosineOrthogonalVectors) +{ + const size_t n = GetParam().length; + if (n < 2) GTEST_SKIP() << "need at least 2 elements for orthogonality"; + + const size_t even_n = n - (n % 2); + + std::vector a(n, 0), b(n, 0); + for (size_t i = 0; i < even_n; ++i) { + a[i] = 1; + b[i] = (i % 2 == 0) ? 1 : -1; + } + + const float got = cosine_i8(a.data(), 0, b.data(), 0, n); + + EXPECT_NEAR(got, 0.0f, 1e-5f); +} + +// --------------------------------------------------------------------------- +// Instantiation — named using the description field +// --------------------------------------------------------------------------- + +INSTANTIATE_TEST_SUITE_P( + AllSizes, + SimilarityI8Test, + ::testing::ValuesIn(kKernelTestParams), + [](const ::testing::TestParamInfo& info) { + return info.param.description; + }); diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestDatasetPrefetching.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestDatasetPrefetching.java index 5379cb0ef..af6f85f4f 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestDatasetPrefetching.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/example/TestDatasetPrefetching.java @@ -49,18 +49,18 @@ private MultiConfig config(String build, String search, boolean fused, String re @Test public void yamlDefaultsToResidentAndOptInIsParsed() throws Exception { var file = Files.createTempFile("prefetch-flag", ".yml"); try { - Files.writeString(file, "yamlSchemaVersion: 1\nonDiskIndexVersion: 6\n"); + Files.writeString(file, "yamlSchemaVersion: 1\nonDiskIndexVersion: " + OnDiskGraphIndex.CURRENT_VERSION + "\n"); assertFalse(MultiConfig.getConfig(file.toString()).useDatasetPrefetching); - Files.writeString(file, "yamlSchemaVersion: 1\nonDiskIndexVersion: 6\nuseDatasetPrefetching: false\n"); + Files.writeString(file, "yamlSchemaVersion: 1\nonDiskIndexVersion: " + OnDiskGraphIndex.CURRENT_VERSION + "\nuseDatasetPrefetching: false\n"); assertFalse(MultiConfig.getConfig(file.toString()).useDatasetPrefetching); - Files.writeString(file, "yamlSchemaVersion: 1\nonDiskIndexVersion: 6\nuseDatasetPrefetching: true\n"); + Files.writeString(file, "yamlSchemaVersion: 1\nonDiskIndexVersion: " + OnDiskGraphIndex.CURRENT_VERSION + "\nuseDatasetPrefetching: true\n"); assertTrue(MultiConfig.getConfig(file.toString()).useDatasetPrefetching); } finally { Files.delete(file); } } @Test public void datasetPrefetchingAndCompactionOptionsAreIndependent() throws Exception { var file = Files.createTempFile("prefetch-compaction-flags", ".yml"); try { - String common = "yamlSchemaVersion: 1\nonDiskIndexVersion: 6\nuseDatasetPrefetching: true\n"; + String common = "yamlSchemaVersion: 1\nonDiskIndexVersion: " + OnDiskGraphIndex.CURRENT_VERSION + "\nuseDatasetPrefetching: true\n"; Files.writeString(file, common); var config = MultiConfig.getConfig(file.toString()); assertTrue(config.useDatasetPrefetching); @@ -78,7 +78,7 @@ private MultiConfig config(String build, String search, boolean fused, String re @Test public void versionedYamlRejectsUnknownSettings() throws Exception { var file = Files.createTempFile("prefetch-unknown-setting", ".yml"); try { - Files.writeString(file, "yamlSchemaVersion: 1\nonDiskIndexVersion: 6\nuseDatasetPrefetching: false\nunknownSetting: true\n"); + Files.writeString(file, "yamlSchemaVersion: 1\nonDiskIndexVersion: " + OnDiskGraphIndex.CURRENT_VERSION + "\nuseDatasetPrefetching: false\nunknownSetting: true\n"); assertThrows(org.yaml.snakeyaml.error.YAMLException.class, () -> MultiConfig.getConfig(file.toString())); } finally { Files.delete(file); } diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestListRandomAccessByteVectorValues.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestListRandomAccessByteVectorValues.java new file mode 100644 index 000000000..81cab7fee --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestListRandomAccessByteVectorValues.java @@ -0,0 +1,83 @@ +/* + * 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.ByteSequence; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; +import org.junit.Test; + +import java.util.ArrayList; +import java.util.List; + +import static org.junit.Assert.*; + +public class TestListRandomAccessByteVectorValues { + + private static final VectorTypeSupport vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + + private ByteSequence seq(byte... values) { + return vts.createByteSequence(values); + } + + @Test + public void testSizeAndDimension() { + int dim = 4; + List> vectors = List.of(seq((byte) 1, (byte) 2, (byte) 3, (byte) 4), + seq((byte) 5, (byte) 6, (byte) 7, (byte) 8)); + var rabvv = new ListRandomAccessByteVectorValues(vectors, dim); + assertEquals(2, rabvv.size()); + assertEquals(dim, rabvv.dimension()); + } + + @Test + public void testGetVectorReturnsCorrectEntry() { + var v0 = seq((byte) 10, (byte) 20); + var v1 = seq((byte) -1, (byte) -2); + var v2 = seq((byte) 127, (byte) -128); + var rabvv = new ListRandomAccessByteVectorValues(List.of(v0, v1, v2), 2); + + assertSame(v0, rabvv.getVector(0)); + assertSame(v1, rabvv.getVector(1)); + assertSame(v2, rabvv.getVector(2)); + } + + @Test + public void testIsValueShared() { + var rabvv = new ListRandomAccessByteVectorValues(List.of(seq((byte) 0)), 1); + assertFalse(rabvv.isValueShared()); + } + + @Test + public void testCopyReturnsSelf() { + var rabvv = new ListRandomAccessByteVectorValues(List.of(seq((byte) 1, (byte) 2)), 2); + assertSame(rabvv, rabvv.copy()); + } + + @Test + public void testMutableBackingListIsReflected() { + // ListRandomAccessByteVectorValues documents that additions to the backing list are visible + List> backing = new ArrayList<>(); + backing.add(seq((byte) 1)); + var rabvv = new ListRandomAccessByteVectorValues(backing, 1); + assertEquals(1, rabvv.size()); + + backing.add(seq((byte) 2)); + assertEquals(2, rabvv.size()); + assertSame(backing.get(1), rabvv.getVector(1)); + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestVectorGraph.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestVectorGraph.java index 52bdc872a..4f80eb24b 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestVectorGraph.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestVectorGraph.java @@ -33,10 +33,13 @@ import io.github.jbellis.jvector.util.Bits; import io.github.jbellis.jvector.util.BoundedLongHeap; import io.github.jbellis.jvector.util.FixedBitSet; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.ByteSequence; import io.github.jbellis.jvector.vector.types.VectorFloat; import io.github.jbellis.jvector.vector.types.VectorTypeSupport; +import io.github.jbellis.jvector.graph.VectorValues; import org.junit.Before; import org.junit.Test; @@ -433,7 +436,7 @@ public void testGraphIndexBuilderInvalid() { public void testGraphIndexBuilderInvalid(boolean addHierarchy) { assertThrows(NullPointerException.class, - () -> new GraphIndexBuilder(null, null, 0, 0, 1.0f, 1.0f, addHierarchy)); + () -> new GraphIndexBuilder((RandomAccessVectorValues) null, (VectorSimilarityFunction) null, 0, 0, 1.0f, 1.0f, addHierarchy)); // M must be > 0 assertThrows(IllegalArgumentException.class, () -> { @@ -800,4 +803,163 @@ protected static Bits createRandomAcceptOrds(int startIndex, int length) { } return bits; } + + // ----------------------------------------------------------------------- + // Byte-vector (int8) graph construction tests + // ----------------------------------------------------------------------- + + /** Build a small int8 graph using every similarity function and verify it is navigable. */ + @Test + public void testByteVectorBuildAllSimilarityFunctions() { + for (var bvsf : ByteVectorSimilarityFunction.values()) { + testByteVectorBuild(bvsf, false); + testByteVectorBuild(bvsf, true); + } + } + + private void testByteVectorBuild(ByteVectorSimilarityFunction bvsf, boolean addHierarchy) { + int n = 50, dim = 8; + var rabvv = randomByteVectorValues(n, dim); + var builder = GraphIndexBuilder.builder(rabvv, bvsf, 8) + .withBeamWidth(20) + .withNeighborOverflow(1.2f) + .withAlpha(1.2f) + .build(); + var graph = builder.build(rabvv); + validateIndex(graph); + assertNotNull("entry node must be set", graph.getView().entryNode()); + assertEquals(n, graph.size(0)); + } + + /** addGraphNode(int, ByteSequence) inserts nodes one by one and yields the right count. */ + @Test + public void testByteVectorAddGraphNodeOneByOne() { + int n = 30, dim = 4; + var rabvv = randomByteVectorValues(n, dim); + var builder = GraphIndexBuilder.builder(rabvv, ByteVectorSimilarityFunction.EUCLIDEAN, 4) + .withBeamWidth(10) + .withNeighborOverflow(1.2f) + .withAlpha(1.2f) + .build(); + for (int i = 0; i < n; i++) { + builder.addGraphNode(i, rabvv.getVector(i)); + } + builder.cleanup(); + var graph = builder.getGraph(); + validateIndex(graph); + assertEquals(n, graph.size(0)); + // every node's neighbor count is within the declared degree + var view = graph.getView(); + for (int i = 0; i < n; i++) { + assertTrue(view.getNeighborsIterator(0, i).size() <= graph.getDegree(0)); + } + } + + /** + * addGraphNode(int, ByteSequence) must throw when the builder was constructed with a + * float-vector score provider. + */ + @Test + public void testByteAddGraphNodeThrowsOnFloatBuilder() { + var floatRavv = circularVectorValues(10); + var builder = new GraphIndexBuilder(floatRavv, VectorSimilarityFunction.EUCLIDEAN, 4, 10, 1.2f, 1.2f, false); + var bs = vectorTypeSupport.createByteSequence(4); + assertThrows(UnsupportedOperationException.class, () -> builder.addGraphNode(0, bs)); + } + + /** + * build(VectorValues) dispatches correctly when given a RandomAccessByteVectorValues. + */ + @Test + public void testByteVectorBuildViaGenericBuildMethod() { + int n = 40, dim = 6; + var rabvv = randomByteVectorValues(n, dim); + var builder = GraphIndexBuilder.builder(rabvv, ByteVectorSimilarityFunction.DOT_PRODUCT, 6) + .withBeamWidth(20) + .withNeighborOverflow(1.2f) + .withAlpha(1.2f) + .build(); + // call the generic VectorValues overload explicitly + var graph = builder.build((VectorValues) rabvv); + validateIndex(graph); + assertEquals(n, graph.size(0)); + } + + /** refineFinalGraph=false and refineFinalGraph=true both produce valid graphs. */ + @Test + public void testByteVectorRefineFinalGraphVariants() { + int n = 40, dim = 4; + for (boolean refine : new boolean[]{false, true}) { + var rabvv = randomByteVectorValues(n, dim); + // Use the 7-arg constructor (refineFinalGraph defaults to true); test both via the + // BSP constructor which exposes the refineFinalGraph knob + var bsp = io.github.jbellis.jvector.graph.similarity.BuildScoreProvider.byteVectorScoreProvider( + rabvv, ByteVectorSimilarityFunction.EUCLIDEAN); + var builder = new GraphIndexBuilder(bsp, dim, 4, 20, 1.2f, 1.2f, false, refine); + var graph = builder.build(rabvv); + validateIndex(graph); + assertEquals(n, graph.size(0)); + var view = graph.getView(); + int ub = graph.getIdUpperBound(); + for (int i = 0; i < n; i++) { + for (var it = view.getNeighborsIterator(0, i); it.hasNext(); ) { + int nb = it.nextInt(); + assertTrue("neighbor " + nb + " out of bounds", nb >= 0 && nb < ub); + } + } + } + } + + /** Smoke-test: a built byte-vector graph achieves reasonable top-5 recall. */ + @Test + public void testByteVectorSearchRecall() { + int n = 200, dim = 16; + int topK = 5; + var rabvv = randomByteVectorValues(n, dim); + var bvsf = ByteVectorSimilarityFunction.EUCLIDEAN; + var builder = GraphIndexBuilder.builder(rabvv, bvsf, 16) + .withBeamWidth(50) + .withNeighborOverflow(1.2f) + .withAlpha(1.2f) + .build(); + var graph = builder.build(rabvv); + + int totalMatches = 0; + int trials = 100; + for (int t = 0; t < trials; t++) { + byte[] rawQ = new byte[dim]; + getRandom().nextBytes(rawQ); + ByteSequence query = vectorTypeSupport.createByteSequence(rawQ); + + // brute-force top-K + NodeQueue expected = new NodeQueue(new BoundedLongHeap(topK), NodeQueue.Order.MIN_HEAP); + for (int i = 0; i < n; i++) { + expected.push(i, bvsf.compare(query, rabvv.getVector(i))); + } + + // graph search with efSearch = 100 (larger than topK to explore more candidates) + var ssp = new DefaultSearchScoreProvider( + (io.github.jbellis.jvector.graph.similarity.ScoreFunction.ExactScoreFunction) + node -> bvsf.compare(query, rabvv.getVector(node))); + var result = new GraphSearcher(graph).search(ssp, topK, 100, 0f, 0f, Bits.ALL); + var actualNodeIds = Arrays.stream(result.getNodes(), 0, topK) + .mapToInt(ns -> ns.node).toArray(); + + totalMatches += computeOverlap(actualNodeIds, expected.nodesCopy()); + } + // with efSearch=100 over a 200-node graph, we can visit most nodes; expect >=90% recall + double overlap = totalMatches / (double) (trials * topK); + assertTrue("byte-vector recall " + overlap + " is too low", overlap > 0.9); + } + + /** Creates a list-backed RandomAccessByteVectorValues with random int8 vectors. */ + private ListRandomAccessByteVectorValues randomByteVectorValues(int n, int dim) { + var list = new java.util.ArrayList>(n); + for (int i = 0; i < n; i++) { + byte[] raw = new byte[dim]; + getRandom().nextBytes(raw); + list.add(vectorTypeSupport.createByteSequence(raw)); + } + return new ListRandomAccessByteVectorValues(list, dim); + } } diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestFeatureHeaderVersion.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestFeatureHeaderVersion.java new file mode 100644 index 000000000..8f9be0f12 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestFeatureHeaderVersion.java @@ -0,0 +1,111 @@ +/* + * 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.disk; + +import com.carrotsearch.randomizedtesting.RandomizedTest; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.disk.ByteBufferIndexWriter; +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.FusedPQ; +import io.github.jbellis.jvector.graph.disk.feature.NVQ; +import io.github.jbellis.jvector.graph.disk.feature.SeparatedNVQ; +import io.github.jbellis.jvector.quantization.NVQuantization; +import io.github.jbellis.jvector.quantization.ProductQuantization; +import org.junit.Test; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.EnumMap; +import java.util.List; + +import static org.junit.Assert.assertArrayEquals; + +/** + * A feature header embedded in a graph serializes its compressor at the graph's version, not at + * {@link OnDiskGraphIndex#CURRENT_VERSION}, so a graph written at an older version is laid out as + * that version wrote it. + */ +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +public class TestFeatureHeaderVersion extends RandomizedTest { + private static final int DIMENSION = 32, MAX_DEGREE = 16; + + private interface Body { + void write(IndexWriter out, int version) throws IOException; + } + + @Test + public void featureHeadersAreWrittenAtTheGraphVersion() throws IOException { + var ravv = new ListRandomAccessVectorValues(TestUtil.createRandomVectors(512, DIMENSION), DIMENSION); + var pq = ProductQuantization.compute(ravv, 8, 256, false); + var nvq = NVQuantization.compute(ravv, 2); + + // Each feature's header is its compressor at the graph's version, plus any fields of its own. + var cases = List.of( + new Case(FeatureId.FUSED_PQ, new FusedPQ(MAX_DEGREE, pq), pq::write), + new Case(FeatureId.NVQ_VECTORS, new NVQ(nvq), nvq::write), + new Case(FeatureId.SEPARATED_NVQ, new SeparatedNVQ(nvq, 1234L), (out, v) -> { + nvq.write(out, v); + out.writeLong(1234L); + })); + + for (int version = 2; version <= OnDiskGraphIndex.CURRENT_VERSION; version++) { + var format = GraphIndexFormatFactory.forVersion(version); + for (var c : cases) { + if (!format.getSupportedFeatures().contains(c.id)) continue; + var features = new EnumMap(FeatureId.class); + features.put(c.id, c.feature); + byte[] written = bytes(out -> format.writeHeaderFeatures(out, features)); + int v = version; + byte[] expected = bytes(out -> c.body.write(out, v)); + // The format's own framing (feature ids, counts) precedes the header. + byte[] header = Arrays.copyOfRange(written, written.length - expected.length, written.length); + assertArrayEquals(c.id + " at version " + version, expected, header); + } + } + } + + private static final class Case { + final FeatureId id; + final Feature feature; + final Body body; + + Case(FeatureId id, Feature feature, Body body) { + this.id = id; + this.feature = feature; + this.body = body; + } + } + + private interface Writes { + void to(IndexWriter out) throws IOException; + } + + private static byte[] bytes(Writes writes) throws IOException { + try (var out = ByteBufferIndexWriter.create(1024 * 1024, false)) { + writes.to(out); + ByteBuffer data = out.getWrittenData(); + byte[] result = new byte[data.remaining()]; + data.get(result); + return result; + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestGraphIndexFormatFactory.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestGraphIndexFormatFactory.java index 27c0a82e8..67d15e0f1 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestGraphIndexFormatFactory.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestGraphIndexFormatFactory.java @@ -89,9 +89,9 @@ public void testFormatCharacteristicsPerVersion() { assertFalse(v2.usesFooter()); assertEquals(EnumSet.of(FeatureId.INLINE_VECTORS), v2.getSupportedFeatures()); - // versions 3-5 support every feature except FUSED_PQ; multi-layer arrives at v4, - // footer-based metadata arrives at v5. - var nonFused = EnumSet.complementOf(EnumSet.of(FeatureId.FUSED_PQ)); + // versions 3-5 support every feature except FUSED_PQ (v6+) and INLINE_BYTE_VECTORS (v7+); + // multi-layer arrives at v4, footer-based metadata arrives at v5. + var nonFused = EnumSet.complementOf(EnumSet.of(FeatureId.FUSED_PQ, FeatureId.INLINE_BYTE_VECTORS)); for (int version : new int[]{3, 4, 5}) { var format = GraphIndexFormatFactory.forVersion(version); assertEquals("version " + version + " supported features", nonFused, format.getSupportedFeatures()); @@ -104,7 +104,15 @@ public void testFormatCharacteristicsPerVersion() { assertTrue(v6.supportsMultiLayer()); assertTrue(v6.usesFooter()); assertTrue(v6.supportsFeature(FeatureId.FUSED_PQ)); - assertEquals(EnumSet.allOf(FeatureId.class), v6.getSupportedFeatures()); + assertFalse(v6.supportsFeature(FeatureId.INLINE_BYTE_VECTORS)); + assertEquals(EnumSet.complementOf(EnumSet.of(FeatureId.INLINE_BYTE_VECTORS)), v6.getSupportedFeatures()); + + var v7 = GraphIndexFormatFactory.forVersion(7); + assertTrue(v7.supportsMultiLayer()); + assertTrue(v7.usesFooter()); + assertTrue(v7.supportsFeature(FeatureId.FUSED_PQ)); + assertTrue(v7.supportsFeature(FeatureId.INLINE_BYTE_VECTORS)); + assertEquals(EnumSet.allOf(FeatureId.class), v7.getSupportedFeatures()); } @Test diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndex.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndex.java index c76796bb4..d82c4e0b4 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndex.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndex.java @@ -23,6 +23,7 @@ import io.github.jbellis.jvector.graph.GraphIndexBuilder; import io.github.jbellis.jvector.graph.GraphSearcher; import io.github.jbellis.jvector.graph.ImmutableGraphIndex; +import io.github.jbellis.jvector.graph.ListRandomAccessByteVectorValues; import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; import io.github.jbellis.jvector.graph.NodesIterator; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; @@ -30,6 +31,7 @@ import io.github.jbellis.jvector.graph.disk.feature.Feature; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; import io.github.jbellis.jvector.graph.disk.feature.FusedPQ; +import io.github.jbellis.jvector.graph.disk.feature.InlineByteVectors; import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; import io.github.jbellis.jvector.graph.disk.feature.NVQ; import io.github.jbellis.jvector.graph.disk.feature.SeparatedNVQ; @@ -38,7 +40,13 @@ import io.github.jbellis.jvector.quantization.PQVectors; import io.github.jbellis.jvector.quantization.ProductQuantization; import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider; +import io.github.jbellis.jvector.graph.similarity.ScoreFunction; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.ByteSequence; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; import org.junit.After; import org.junit.Before; import org.junit.Test; @@ -340,7 +348,7 @@ public void testV0Read() throws IOException { var onDiskGraph = OnDiskGraphIndex.load(readerSupplier); var onDiskView = onDiskGraph.getView()) { - assertEquals(88, onDiskGraph.ramBytesUsed()); // Current size of graph that hasn't been read from + assertEquals(96, onDiskGraph.ramBytesUsed()); // Current size of graph that hasn't been read from assertEquals(32, onDiskGraph.getDegree(0)); assertEquals(2, onDiskGraph.version); assertEquals(100_000, onDiskGraph.size(0)); @@ -572,4 +580,282 @@ public void testIncrementalWrites() throws IOException { throw new RuntimeException(e); } } + + // ----------------------------------------------------------------------- + // InlineByteVectors (int8) tests + // ----------------------------------------------------------------------- + + private static final VectorTypeSupport vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + + /** Builds a list of n random int8 vectors of the given dimension. */ + private List> randomByteVectors(int n, int dim) { + var list = new ArrayList>(n); + for (int i = 0; i < n; i++) { + byte[] raw = new byte[dim]; + getRandom().nextBytes(raw); + list.add(vts.createByteSequence(raw)); + } + return list; + } + + /** + * Full round-trip: build with InlineByteVectors, write, reload, verify + * getByteVector() returns the original bytes for every node. + */ + @Test + public void testInlineByteVectorsRoundTrip() throws IOException { + int n = 50, dim = 8; + var vectors = randomByteVectors(n, dim); + var rabvv = new ListRandomAccessByteVectorValues(vectors, dim); + + // Build a graph + var builder = GraphIndexBuilder.builder(rabvv, ByteVectorSimilarityFunction.EUCLIDEAN, 8) + .withBeamWidth(20) + .withNeighborOverflow(1.2f) + .withAlpha(1.2f) + .build(); + var graph = builder.build(rabvv); + + // Write with InlineByteVectors feature + var outputPath = testDirectory.resolve("byte_vectors_roundtrip"); + try (var writer = new OnDiskGraphIndexWriter.Builder(graph, outputPath) + .with(new InlineByteVectors(dim)) + .build()) + { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_BYTE_VECTORS, + nodeId -> new InlineByteVectors.State(rabvv.getVector(nodeId)))); + } + + // Reload and verify every vector + try (var rs = new SimpleMappedReader.Supplier(outputPath.toAbsolutePath()); + var onDisk = OnDiskGraphIndex.load(rs); + var view = onDisk.getView()) + { + assertTrue("INLINE_BYTE_VECTORS feature missing", + onDisk.getFeatureSet().contains(FeatureId.INLINE_BYTE_VECTORS)); + assertEquals(dim, onDisk.dimension); + + for (int i = 0; i < n; i++) { + var expected = rabvv.getVector(i); + var actual = view.getByteVector(i); + assertEquals("byte vector length mismatch at node " + i, expected.length(), actual.length()); + for (int d = 0; d < dim; d++) { + assertEquals("byte mismatch at node " + i + " dim " + d, + expected.get(d), actual.get(d)); + } + } + } + } + + /** + * featureSize() of InlineByteVectors must equal the dimension (1 byte per component), + * confirming 4× compression vs float32 InlineVectors. + */ + @Test + public void testInlineByteVectorsFeatureSize() { + for (int dim : new int[]{1, 8, 128, 256}) { + var ibv = new InlineByteVectors(dim); + assertEquals("featureSize should equal dimension", dim, ibv.featureSize()); + // float32 InlineVectors uses 4 * dim bytes + assertEquals("byte feature 4x smaller than float32", + new InlineVectors(dim).featureSize(), 4 * ibv.featureSize()); + } + } + + @Test + public void testInlineByteVectorsRequiresVersion7() { + var graph = new TestUtil.RandomlyConnectedGraphIndex(32, 4, getRandom()); + var e = assertThrows(IllegalArgumentException.class, () -> new OnDiskGraphIndexWriter.Builder(graph, + testDirectory.resolve("version_6_inline_byte_vectors")) + .withVersion(6) + .with(new InlineByteVectors(8)) + .build()); + assertTrue(e.getMessage(), e.getMessage().startsWith("Feature INLINE_BYTE_VECTORS is not supported by version 6")); + } + + /** + * getByteVector throws UnsupportedOperationException when the graph has no INLINE_BYTE_VECTORS feature. + */ + @Test + public void testGetByteVectorThrowsWithoutFeature() throws IOException { + // Build and write a float-vector graph (no INLINE_BYTE_VECTORS) + var graph = new TestUtil.RandomlyConnectedGraphIndex(10, 4, getRandom()); + var ravv = new TestVectorGraph.CircularFloatVectorValues(10); + var outputPath = testDirectory.resolve("float_graph_no_bytes"); + TestUtil.writeGraph(graph, ravv, outputPath); + + try (var rs = new SimpleMappedReader.Supplier(outputPath.toAbsolutePath()); + var onDisk = OnDiskGraphIndex.load(rs); + var view = onDisk.getView()) + { + assertFalse(onDisk.getFeatureSet().contains(FeatureId.INLINE_BYTE_VECTORS)); + assertThrows(UnsupportedOperationException.class, () -> view.getByteVector(0)); + } + } + + /** + * byteVectorRerankerFor throws UnsupportedOperationException when the graph has no INLINE_BYTE_VECTORS feature. + */ + @Test + public void testByteVectorRerankerThrowsWithoutFeature() throws IOException { + var graph = new TestUtil.RandomlyConnectedGraphIndex(10, 4, getRandom()); + var ravv = new TestVectorGraph.CircularFloatVectorValues(10); + var outputPath = testDirectory.resolve("float_graph_reranker"); + TestUtil.writeGraph(graph, ravv, outputPath); + + try (var rs = new SimpleMappedReader.Supplier(outputPath.toAbsolutePath()); + var onDisk = OnDiskGraphIndex.load(rs); + var view = onDisk.getView()) + { + var qb = vts.createByteSequence(2); + assertThrows(UnsupportedOperationException.class, + () -> view.byteVectorRerankerFor(qb, ByteVectorSimilarityFunction.EUCLIDEAN)); + } + } + + /** + * byteVectorRerankerFor returns scores consistent with a direct bvsf.compare() call. + */ + @Test + public void testByteVectorRerankerScoresMatchDirectCompare() throws IOException { + int n = 30, dim = 4; + var vectors = randomByteVectors(n, dim); + var rabvv = new ListRandomAccessByteVectorValues(vectors, dim); + var bvsf = ByteVectorSimilarityFunction.EUCLIDEAN; + + var builder = GraphIndexBuilder.builder(rabvv, bvsf, 8) + .withBeamWidth(20) + .withNeighborOverflow(1.2f) + .withAlpha(1.2f) + .build(); + var graph = builder.build(rabvv); + + var outputPath = testDirectory.resolve("byte_reranker_test"); + try (var writer = new OnDiskGraphIndexWriter.Builder(graph, outputPath) + .with(new InlineByteVectors(dim)) + .build()) + { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_BYTE_VECTORS, + nodeId -> new InlineByteVectors.State(rabvv.getVector(nodeId)))); + } + + try (var rs = new SimpleMappedReader.Supplier(outputPath.toAbsolutePath()); + var onDisk = OnDiskGraphIndex.load(rs); + var view = onDisk.getView()) + { + byte[] rawQ = new byte[dim]; + getRandom().nextBytes(rawQ); + var query = vts.createByteSequence(rawQ); + var reranker = view.byteVectorRerankerFor(query, bvsf); + + for (int i = 0; i < n; i++) { + float expected = bvsf.compare(query, rabvv.getVector(i)); + float actual = reranker.similarityTo(i); + assertEquals("reranker score mismatch at node " + i, expected, actual, 1e-5f); + } + } + } + + /** + * The reranker reuses a single scratch vector; scores must not depend on which nodes were scored before. + */ + @Test + public void testByteVectorRerankerReusesScratchWithoutStaleData() throws IOException { + int n = 30, dim = 16; + var vectors = randomByteVectors(n, dim); + var rabvv = new ListRandomAccessByteVectorValues(vectors, dim); + var bvsf = ByteVectorSimilarityFunction.EUCLIDEAN; + + var builder = GraphIndexBuilder.builder(rabvv, bvsf, 8) + .withBeamWidth(20) + .withNeighborOverflow(1.2f) + .withAlpha(1.2f) + .build(); + var graph = builder.build(rabvv); + + var outputPath = testDirectory.resolve("byte_reranker_scratch_test"); + try (var writer = new OnDiskGraphIndexWriter.Builder(graph, outputPath) + .with(new InlineByteVectors(dim)) + .build()) + { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_BYTE_VECTORS, + nodeId -> new InlineByteVectors.State(rabvv.getVector(nodeId)))); + } + + try (var rs = new SimpleMappedReader.Supplier(outputPath.toAbsolutePath()); + var onDisk = OnDiskGraphIndex.load(rs); + var view = onDisk.getView()) + { + byte[] rawQ = new byte[dim]; + getRandom().nextBytes(rawQ); + var query = vts.createByteSequence(rawQ); + var reranker = view.byteVectorRerankerFor(query, bvsf); + + // forward, backward, then repeated: each result must match a direct compare + for (int i = 0; i < n; i++) { + assertEquals(bvsf.compare(query, rabvv.getVector(i)), reranker.similarityTo(i), 1e-5f); + } + for (int i = n - 1; i >= 0; i--) { + assertEquals(bvsf.compare(query, rabvv.getVector(i)), reranker.similarityTo(i), 1e-5f); + } + assertEquals(bvsf.compare(query, rabvv.getVector(7)), reranker.similarityTo(7), 1e-5f); + assertEquals(bvsf.compare(query, rabvv.getVector(7)), reranker.similarityTo(7), 1e-5f); + + // getByteVectorInto fills a caller-supplied vector with the same bytes getByteVector returns + var into = vts.createByteSequence(dim); + view.getByteVectorInto(3, into); + assertTrue(view.getByteVector(3).equalTo(into)); + assertTrue(rabvv.getVector(3).equalTo(into)); + } + } + + /** + * All three byte-vector similarity functions produce results in [0, 1] for all nodes + * after a full int8 end-to-end write/search cycle. + */ + @Test + public void testByteVectorEndToEndAllSimilarityFunctions() throws IOException { + int n = 40, dim = 8; + + for (var bvsf : ByteVectorSimilarityFunction.values()) { + var vectors = randomByteVectors(n, dim); + var rabvv = new ListRandomAccessByteVectorValues(vectors, dim); + + var builder = GraphIndexBuilder.builder(rabvv, bvsf, 8) + .withBeamWidth(20) + .withNeighborOverflow(1.2f) + .withAlpha(1.2f) + .build(); + var graph = builder.build(rabvv); + + var outputPath = testDirectory.resolve("byte_e2e_" + bvsf.name()); + try (var writer = new OnDiskGraphIndexWriter.Builder(graph, outputPath) + .with(new InlineByteVectors(dim)) + .build()) + { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_BYTE_VECTORS, + nodeId -> new InlineByteVectors.State(rabvv.getVector(nodeId)))); + } + + try (var rs = new SimpleMappedReader.Supplier(outputPath.toAbsolutePath()); + var onDisk = OnDiskGraphIndex.load(rs); + var view = onDisk.getView()) + { + byte[] rawQ = new byte[dim]; + getRandom().nextBytes(rawQ); + var query = vts.createByteSequence(rawQ); + var reranker = view.byteVectorRerankerFor(query, bvsf); + var ssp = new DefaultSearchScoreProvider(reranker); + var result = new GraphSearcher(onDisk).search(ssp, 5, Bits.ALL); + + for (var ns : result.getNodes()) { + float score = ns.score; + // DOT_PRODUCT is positive but unbounded above; the others are in [0,1] + float max = bvsf == ByteVectorSimilarityFunction.DOT_PRODUCT ? Float.MAX_VALUE : 1.0f; + assertTrue(bvsf + " score out of range for node " + ns.node + ": " + score, + score > 0f && score <= max); + } + } + } + } } diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProviderTest.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProviderTest.java index 4942b8efb..fcd5f0fee 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProviderTest.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProviderTest.java @@ -16,9 +16,13 @@ package io.github.jbellis.jvector.graph.similarity; +import io.github.jbellis.jvector.graph.ListRandomAccessByteVectorValues; +import io.github.jbellis.jvector.graph.RandomAccessByteVectorValues; import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.ByteSequence; import io.github.jbellis.jvector.vector.types.VectorFloat; import io.github.jbellis.jvector.vector.types.VectorTypeSupport; import org.junit.Test; @@ -26,7 +30,8 @@ import java.util.ArrayList; import java.util.List; -import static org.junit.Assert.assertEquals; +import static org.junit.Assert.*; +import static org.junit.Assert.assertThrows; public class BuildScoreProviderTest { private static final VectorTypeSupport vts = VectorizationProvider.getInstance().getVectorTypeSupport(); @@ -69,4 +74,190 @@ public void testOrdinalMapping() { var dsp0 = bsp.diversityProviderFor(0); assertEquals(vsf.compare(vectors.get(2), vectors.get(0)), dsp0.exactScoreFunction().similarityTo(1), 1e-6f); } + + // ----------------------------------------------------------------------- + // byteVectorScoreProvider tests + // ----------------------------------------------------------------------- + + private ByteSequence bseq(byte... values) { + return vts.createByteSequence(values); + } + + /** Helper: build a small RABVV with three 2-D signed-byte vectors. */ + private ListRandomAccessByteVectorValues byteRavv() { + return new ListRandomAccessByteVectorValues( + List.of(bseq((byte) 10, (byte) 0), // node 0 + bseq((byte) 0, (byte) 10), // node 1 + bseq((byte) -10, (byte) 0)), // node 2 + 2); + } + + @Test + public void testByteVectorIsExact() { + var bsp = BuildScoreProvider.byteVectorScoreProvider(byteRavv(), ByteVectorSimilarityFunction.EUCLIDEAN); + assertTrue(bsp.isExact()); + } + + @Test + public void testByteVectorSearchProviderForByteSequence() { + var bvsf = ByteVectorSimilarityFunction.EUCLIDEAN; + var rabvv = byteRavv(); + var bsp = BuildScoreProvider.byteVectorScoreProvider(rabvv, bvsf); + + // searchProviderFor(ByteSequence) scores all nodes against the given query + ByteSequence query = bseq((byte) 10, (byte) 0); // identical to node 0 + var ssp = bsp.searchProviderFor(query); + + // node 0 should be self-similar (score = 1.0 for EUCLIDEAN with zero distance) + assertEquals(1.0f, ssp.exactScoreFunction().similarityTo(0), 1e-5f); + // node 2 = [-10, 0], distance vs [10,0] = 400, not 1.0 + assertTrue(ssp.exactScoreFunction().similarityTo(2) < 1.0f); + } + + @Test + public void testByteVectorSearchProviderForNode() { + var bvsf = ByteVectorSimilarityFunction.DOT_PRODUCT; + var rabvv = byteRavv(); + var bsp = BuildScoreProvider.byteVectorScoreProvider(rabvv, bvsf); + + // searchProviderFor(int) should delegate to searchProviderFor(ByteSequence) + var sspByNode = bsp.searchProviderFor(0); + ByteSequence v0 = rabvv.getVector(0); + var sspBySeq = bsp.searchProviderFor(v0); + assertEquals(sspByNode.exactScoreFunction().similarityTo(1), + sspBySeq.exactScoreFunction().similarityTo(1), 1e-6f); + } + + @Test + public void testByteVectorDiversityProviderMatchesSearch() { + var bvsf = ByteVectorSimilarityFunction.COSINE; + var bsp = BuildScoreProvider.byteVectorScoreProvider(byteRavv(), bvsf); + + // diversityProviderFor delegates to searchProviderFor(int) + var search = bsp.searchProviderFor(1); + var diversity = bsp.diversityProviderFor(1); + assertEquals(search.exactScoreFunction().similarityTo(0), + diversity.exactScoreFunction().similarityTo(0), 1e-6f); + } + + @Test + public void testByteVectorDiversityScoreFunction() { + var bvsf = ByteVectorSimilarityFunction.EUCLIDEAN; + var rabvv = byteRavv(); + var bsp = BuildScoreProvider.byteVectorScoreProvider(rabvv, bvsf); + + // diversityScoreFunctionFor(n1).similarityTo(n2) == bvsf.compare(v_n1, v_n2) + var dsf = bsp.diversityScoreFunctionFor(0); + assertEquals(bvsf.compare(rabvv.getVector(0), rabvv.getVector(2)), + dsf.similarityTo(2), 1e-6f); + } + + @Test + public void testByteVectorApproximateCentroid() { + // centroid of [10,0], [0,10], [-10,0] should be [0, 10/3] + var bsp = BuildScoreProvider.byteVectorScoreProvider(byteRavv(), ByteVectorSimilarityFunction.EUCLIDEAN); + var centroid = bsp.approximateCentroid(); + assertEquals(2, centroid.length()); + assertEquals(0.0f, centroid.get(0), 1e-5f); + assertEquals(10.0f / 3.0f, centroid.get(1), 1e-5f); + } + + @Test + public void testByteVectorApproximateCentroidSkipsNullVectors() { + // Mirrors the float provider: null entries are skipped but still count towards the divisor (size = 3), + // so the centroid of [0,10], null, [0,20] is [0, 30/3] = [0,10]. + var withNull = new ArrayList>(); + withNull.add(bseq((byte) 0, (byte) 10)); + withNull.add(null); + withNull.add(bseq((byte) 0, (byte) 20)); + var bsp = BuildScoreProvider.byteVectorScoreProvider( + new ListRandomAccessByteVectorValues(withNull, 2), ByteVectorSimilarityFunction.EUCLIDEAN); + + var centroid = bsp.approximateCentroid(); + assertEquals(0.0f, centroid.get(0), 1e-5f); + assertEquals(10.0f, centroid.get(1), 1e-5f); + } + + @Test + public void testByteVectorThrowsForFloatQuery() { + var bsp = BuildScoreProvider.byteVectorScoreProvider(byteRavv(), ByteVectorSimilarityFunction.EUCLIDEAN); + VectorFloat floatQuery = vts.createFloatVector(new float[]{1.0f, 0.0f}); + assertThrows(UnsupportedOperationException.class, () -> bsp.searchProviderFor(floatQuery)); + } + + /** A byte RAVV whose getVector always returns the same per-instance buffer, like a disk-backed source. */ + private static class SharedBufferByteVectors implements RandomAccessByteVectorValues { + private final List> data; + private final int dimension; + private final ByteSequence buffer; + + SharedBufferByteVectors(List> data, int dimension) { + this.data = data; + this.dimension = dimension; + this.buffer = vts.createByteSequence(dimension); + } + + @Override + public int size() { + return data.size(); + } + + @Override + public int dimension() { + return dimension; + } + + @Override + public ByteSequence getVector(int nodeId) { + buffer.copyFrom(data.get(nodeId), 0, 0, dimension); + return buffer; + } + + @Override + public boolean isValueShared() { + return true; + } + + @Override + public RandomAccessByteVectorValues copy() { + return new SharedBufferByteVectors(data, dimension); + } + } + + private SharedBufferByteVectors sharedByteRavv() { + return new SharedBufferByteVectors( + List.of(bseq((byte) 10, (byte) 0), // node 0 + bseq((byte) 0, (byte) 10), // node 1 + bseq((byte) -10, (byte) 0)), // node 2 + 2); + } + + @Test + public void testByteVectorSearchProviderSurvivesSharedBufferReuse() { + var bvsf = ByteVectorSimilarityFunction.EUCLIDEAN; + var bsp = BuildScoreProvider.byteVectorScoreProvider(sharedByteRavv(), bvsf); + var reference = byteRavv(); + + // Creating a second scorer on the same thread must not change what the first one compares against. + var scorer0 = bsp.searchProviderFor(0).exactScoreFunction(); + var scorer1 = bsp.searchProviderFor(1).exactScoreFunction(); + + assertEquals(bvsf.compare(reference.getVector(1), reference.getVector(2)), scorer1.similarityTo(2), 1e-6f); + assertEquals(bvsf.compare(reference.getVector(0), reference.getVector(2)), scorer0.similarityTo(2), 1e-6f); + assertEquals(bvsf.compare(reference.getVector(1), reference.getVector(2)), scorer1.similarityTo(2), 1e-6f); + } + + @Test + public void testByteVectorDiversityScoreFunctionSurvivesSharedBufferReuse() { + var bvsf = ByteVectorSimilarityFunction.EUCLIDEAN; + var bsp = BuildScoreProvider.byteVectorScoreProvider(sharedByteRavv(), bvsf); + var reference = byteRavv(); + + var sf0 = bsp.diversityScoreFunctionFor(0); + var sf1 = bsp.diversityScoreFunctionFor(1); + + assertEquals(bvsf.compare(reference.getVector(1), reference.getVector(2)), sf1.similarityTo(2), 1e-6f); + assertEquals(bvsf.compare(reference.getVector(0), reference.getVector(2)), sf0.similarityTo(2), 1e-6f); + assertEquals(bvsf.compare(reference.getVector(1), reference.getVector(2)), sf1.similarityTo(2), 1e-6f); + } } \ No newline at end of file diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/TestByteVectorSimilarityFunction.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/TestByteVectorSimilarityFunction.java new file mode 100644 index 000000000..95986382b --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/TestByteVectorSimilarityFunction.java @@ -0,0 +1,323 @@ +/* + * 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 com.carrotsearch.randomizedtesting.RandomizedTest; +import io.github.jbellis.jvector.vector.types.ByteSequence; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; +import org.junit.Test; + +import static org.junit.Assert.*; + +public class TestByteVectorSimilarityFunction extends RandomizedTest { + + private static final VectorTypeSupport vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + + /** + * Fills {@code raw} with random bytes, guaranteeing it is not the zero vector (cosine is undefined for it). + * RandomizedTest's generator is biased toward runs of 0, so all-zero draws do happen. + */ + private void nextNonZeroBytes(byte[] raw) { + getRandom().nextBytes(raw); + boolean allZero = true; + for (byte b : raw) { + if (b != 0) { allZero = false; break; } + } + if (allZero) raw[0] = 1; + } + + private ByteSequence seq(byte... values) { + return vts.createByteSequence(values); + } + + // ----------------------------------------------------------------------- + // EUCLIDEAN + // ----------------------------------------------------------------------- + + @Test + public void testEuclideanKnownValue() { + // v1=[1,0], v2=[0,1] squared-L2 = 1+1 = 2 + // expected = 1 / (1 + 2) = 1/3 + var v1 = seq((byte) 1, (byte) 0); + var v2 = seq((byte) 0, (byte) 1); + assertEquals(1.0f / 3.0f, ByteVectorSimilarityFunction.EUCLIDEAN.compare(v1, v2), 1e-6f); + } + + @Test + public void testEuclideanDistinguishesSmallDistancesInHighDimensions() { + // At 1024 dimensions, squared distances 0, 1, 2 and 3 must not collapse to the same float score. + int dim = 1024; + var base = seq(new byte[dim]); + float previous = Float.POSITIVE_INFINITY; + for (int sq = 0; sq <= 3; sq++) { + byte[] other = new byte[dim]; + for (int i = 0; i < sq; i++) { + other[i] = 1; // each unit component adds 1 to the squared distance + } + float score = ByteVectorSimilarityFunction.EUCLIDEAN.compare(base, seq(other)); + assertEquals(1.0f / (1.0f + sq), score, 1e-6f); + assertTrue("score for squared distance " + sq + " should be strictly lower than for " + (sq - 1), + score < previous); + previous = score; + } + } + + @Test + public void testEuclideanIdenticalVectors() { + var v = seq((byte) 42, (byte) -7, (byte) 100); + // squaredL2 = 0, so score = 1/(1+0) = 1.0 + assertEquals(1.0f, ByteVectorSimilarityFunction.EUCLIDEAN.compare(v, v), 1e-6f); + } + + @Test + public void testEuclideanResultInRange() { + for (int trial = 0; trial < 50; trial++) { + byte[] raw1 = new byte[32]; + byte[] raw2 = new byte[32]; + getRandom().nextBytes(raw1); + getRandom().nextBytes(raw2); + float score = ByteVectorSimilarityFunction.EUCLIDEAN.compare(seq(raw1), seq(raw2)); + assertTrue("EUCLIDEAN score out of (0,1]: " + score, score > 0f && score <= 1.0f); + } + } + + @Test + public void testEuclideanSymmetry() { + byte[] raw1 = new byte[16]; + byte[] raw2 = new byte[16]; + getRandom().nextBytes(raw1); + getRandom().nextBytes(raw2); + assertEquals( + ByteVectorSimilarityFunction.EUCLIDEAN.compare(seq(raw1), seq(raw2)), + ByteVectorSimilarityFunction.EUCLIDEAN.compare(seq(raw2), seq(raw1)), + 1e-6f); + } + + // ----------------------------------------------------------------------- + // DOT_PRODUCT + // ----------------------------------------------------------------------- + + @Test + public void testDotProductKnownValue() { + // v1=[1,0], v2=[0,1] dot = 0 + // expected = dot + 1 = 1 + var v1 = seq((byte) 1, (byte) 0); + var v2 = seq((byte) 0, (byte) 1); + assertEquals(1.0f, ByteVectorSimilarityFunction.DOT_PRODUCT.compare(v1, v2), 1e-6f); + // negative dot = -2: 1 / (1 + 2) + assertEquals(1f / 3f, ByteVectorSimilarityFunction.DOT_PRODUCT.compare(seq((byte) 1, (byte) 1), seq((byte) -1, (byte) -1)), 1e-6f); + } + + @Test + public void testDotProductResultInRange() { + for (int trial = 0; trial < 50; trial++) { + byte[] raw1 = new byte[32]; + byte[] raw2 = new byte[32]; + getRandom().nextBytes(raw1); + getRandom().nextBytes(raw2); + float score = ByteVectorSimilarityFunction.DOT_PRODUCT.compare(seq(raw1), seq(raw2)); + assertTrue("DOT_PRODUCT score not positive: " + score, score > 0f); + } + } + + @Test + public void testDotProductKeepsDistinctScoresAtHighDimension() { + // At 1,024 dimensions, adjacent integer dot products must not collapse to the same float score. + int dim = 1024; + byte[] q = new byte[dim]; + java.util.Arrays.fill(q, (byte) 1); + float prev = Float.NEGATIVE_INFINITY; + for (int dot = -3; dot <= 3; dot++) { + byte[] v = new byte[dim]; + v[0] = (byte) dot; + float score = ByteVectorSimilarityFunction.DOT_PRODUCT.compare(seq(q), seq(v)); + assertTrue("scores must strictly increase with dot product: " + prev + " -> " + score, score > prev); + prev = score; + } + } + + @Test + public void testDotProductSymmetry() { + byte[] raw1 = new byte[16]; + byte[] raw2 = new byte[16]; + getRandom().nextBytes(raw1); + getRandom().nextBytes(raw2); + assertEquals( + ByteVectorSimilarityFunction.DOT_PRODUCT.compare(seq(raw1), seq(raw2)), + ByteVectorSimilarityFunction.DOT_PRODUCT.compare(seq(raw2), seq(raw1)), + 1e-6f); + } + + // ----------------------------------------------------------------------- + // COSINE + // ----------------------------------------------------------------------- + + @Test + public void testCosineParallelVectors() { + // v1 == v2 → cosine = 1.0, score = (1+1)/2 = 1.0 + var v = seq((byte) 3, (byte) 4); + assertEquals(1.0f, ByteVectorSimilarityFunction.COSINE.compare(v, v), 1e-5f); + } + + @Test + public void testCosineOrthogonalVectors() { + // [1,0] · [0,1] = 0, cosine = 0, score = (1+0)/2 = 0.5 + var v1 = seq((byte) 1, (byte) 0); + var v2 = seq((byte) 0, (byte) 1); + assertEquals(0.5f, ByteVectorSimilarityFunction.COSINE.compare(v1, v2), 1e-5f); + } + + @Test + public void testCosineResultInRange() { + for (int trial = 0; trial < 50; trial++) { + byte[] raw1 = new byte[32]; + byte[] raw2 = new byte[32]; + nextNonZeroBytes(raw1); + nextNonZeroBytes(raw2); + float score = ByteVectorSimilarityFunction.COSINE.compare(seq(raw1), seq(raw2)); + assertTrue("COSINE score out of [0,1]: " + score, score >= 0f && score <= 1.0f); + } + } + + @Test + public void testCosineSymmetry() { + byte[] raw1 = new byte[16]; + byte[] raw2 = new byte[16]; + nextNonZeroBytes(raw1); + nextNonZeroBytes(raw2); + assertEquals( + ByteVectorSimilarityFunction.COSINE.compare(seq(raw1), seq(raw2)), + ByteVectorSimilarityFunction.COSINE.compare(seq(raw2), seq(raw1)), + 1e-5f); + } + + @Test + public void testCosineDimensionMismatchThrows() { + var v1 = seq((byte) 1, (byte) 2, (byte) 3); + var v2 = seq((byte) 1, (byte) 2); + assertThrows(IllegalArgumentException.class, () -> VectorUtil.cosine(v1, v2)); + } + + @Test + public void testCosineZeroVectorIsRejectedByAssertion() { + // Mirrors the float overload: zero vectors are unsupported and caught by an assertion, not defined. + boolean assertionsEnabled = false; + assert assertionsEnabled = true; + org.junit.Assume.assumeTrue("requires -ea", assertionsEnabled); + + var zero = seq(new byte[8]); + var nonZero = seq((byte) 1, (byte) 2, (byte) 3, (byte) 4, (byte) 5, (byte) 6, (byte) 7, (byte) 8); + assertThrows(AssertionError.class, () -> VectorUtil.cosine(zero, zero)); + assertThrows(AssertionError.class, () -> VectorUtil.cosine(zero, nonZero)); + assertThrows(AssertionError.class, () -> VectorUtil.cosine(nonZero, zero)); + } + + // ----------------------------------------------------------------------- + // Boundary values + // ----------------------------------------------------------------------- + + @Test + public void testAllMaxValues() { + // All -128 vectors — EUCLIDEAN identity = 1, DOT_PRODUCT = n*128^2 + 1, COSINE = 1 + byte[] rawMin = new byte[8]; + java.util.Arrays.fill(rawMin, (byte) -128); + var vMin = seq(rawMin); + assertEquals(1.0f, ByteVectorSimilarityFunction.EUCLIDEAN.compare(vMin, vMin), 1e-5f); + assertEquals(8 * 128f * 128f + 1f, ByteVectorSimilarityFunction.DOT_PRODUCT.compare(vMin, vMin), 1e-3f); + assertEquals(1.0f, ByteVectorSimilarityFunction.COSINE.compare(vMin, vMin), 1e-5f); + + // All 127 vectors — EUCLIDEAN identity = 1, COSINE = 1 + byte[] rawMax = new byte[8]; + java.util.Arrays.fill(rawMax, (byte) 127); + var vMax = seq(rawMax); + assertEquals(1.0f, ByteVectorSimilarityFunction.EUCLIDEAN.compare(vMax, vMax), 1e-5f); + float expectedDot = 8 * 127f * 127f + 1f; + assertEquals(expectedDot, ByteVectorSimilarityFunction.DOT_PRODUCT.compare(vMax, vMax), 1e-3f); + assertEquals(1.0f, ByteVectorSimilarityFunction.COSINE.compare(vMax, vMax), 1e-5f); + } + + // ----------------------------------------------------------------------- + // Default implementation accumulation + // ----------------------------------------------------------------------- + + @Test + public void testDefaultImplementationAccumulatesExactlyInHighDimensions() { + // Float accumulation stops being exact above 2^24; these sums exceed it, so a float accumulator + // would round differently from the integer-accumulating Panama and native implementations. + var impl = new DefaultVectorUtilSupport(); + int dim = 4096; + byte[][][] cases = new byte[3][2][dim]; + java.util.Arrays.fill(cases[0][0], (byte) 127); + java.util.Arrays.fill(cases[0][1], (byte) 127); + java.util.Arrays.fill(cases[1][0], (byte) 127); + java.util.Arrays.fill(cases[1][1], (byte) -128); + getRandom().nextBytes(cases[2][0]); + getRandom().nextBytes(cases[2][1]); + + for (byte[][] c : cases) { + long dot = 0, normA = 0, normB = 0, sq = 0; + for (int i = 0; i < dim; i++) { + int ai = c[0][i], bi = c[1][i]; + dot += ai * bi; + normA += ai * ai; + normB += bi * bi; + sq += (ai - bi) * (ai - bi); + } + var a = seq(c[0]); + var b = seq(c[1]); + assertEquals((float) dot, impl.dotProduct(a, b), 0f); + assertEquals((float) sq, impl.squareDistance(a, b), 0f); + assertEquals((float) (dot / Math.sqrt((double) normA * normB)), impl.cosine(a, b), 0f); + } + } + + // ----------------------------------------------------------------------- + // Sliced inputs + // ----------------------------------------------------------------------- + + @Test + public void testSlicedInputsMatchScalarReference() { + // Lengths cover: shorter than any SIMD step, exact multiples, and lengths with a scalar tail. + for (int length : new int[] {5, 8, 16, 37, 64, 100}) { + for (int[] offsets : new int[][] {{0, 0}, {3, 0}, {0, 5}, {7, 11}}) { + byte[] raw1 = new byte[offsets[0] + length + 9]; + byte[] raw2 = new byte[offsets[1] + length + 9]; + getRandom().nextBytes(raw1); + getRandom().nextBytes(raw2); + var a = seq(raw1).slice(offsets[0], length); + var b = seq(raw2).slice(offsets[1], length); + + long dot = 0, normA = 0, normB = 0, sq = 0; + for (int i = 0; i < length; i++) { + int ai = raw1[offsets[0] + i], bi = raw2[offsets[1] + i]; + dot += ai * bi; + normA += ai * ai; + normB += bi * bi; + sq += (ai - bi) * (ai - bi); + } + String msg = "length=" + length + " offsets=" + offsets[0] + "," + offsets[1]; + assertEquals(msg, (float) dot, VectorUtil.dotProduct(a, b), 0f); + assertEquals(msg, (float) sq, VectorUtil.squareL2Distance(a, b), 0f); + // RandomizedTest's nextBytes is biased toward 0/-1 runs, so a short slice can be all zeros; + // cosine is undefined for zero vectors (VectorUtil.cosine asserts on them), so skip it then. + if (normA != 0 && normB != 0) { + assertEquals(msg, (float) (dot / Math.sqrt((double) normA * normB)), VectorUtil.cosine(a, b), 1e-5f); + } + } + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/TestVectorizationProvider.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/TestVectorizationProvider.java index b8d0b5c3e..a21d2efcf 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/TestVectorizationProvider.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/TestVectorizationProvider.java @@ -19,6 +19,7 @@ import com.carrotsearch.randomizedtesting.RandomizedTest; import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.vector.types.ByteSequence; import io.github.jbellis.jvector.vector.types.FloatArray; import io.github.jbellis.jvector.vector.types.VectorFloat; import io.github.jbellis.jvector.vector.types.VectorTypeSupport; @@ -60,6 +61,121 @@ public void testSimilarityMetricsFloat() { Assert.assertEquals(a.getVectorUtilSupport().squareDistance(v1a, v2a), b.getVectorUtilSupport().squareDistance(v1b, v2b), 0.0001f); } + @Test + public void testSimilarityMetricsByte() { + Assume.assumeTrue(hasSimd); + + VectorizationProvider a = new DefaultVectorizationProvider(); + VectorizationProvider b = VectorizationProvider.getInstance(); + + // Use a prime-length vector that is not a multiple of 8 or 16 + int dim = 107; + byte[] rawA = new byte[dim]; + byte[] rawB = new byte[dim]; + getRandom().nextBytes(rawA); + getRandom().nextBytes(rawB); + + ByteSequence bsA_scalar = a.getVectorTypeSupport().createByteSequence(rawA); + ByteSequence bsB_scalar = a.getVectorTypeSupport().createByteSequence(rawB); + ByteSequence bsA_simd = b.getVectorTypeSupport().createByteSequence(rawA); + ByteSequence bsB_simd = b.getVectorTypeSupport().createByteSequence(rawB); + + Assert.assertEquals( + a.getVectorUtilSupport().dotProduct(bsA_scalar, bsB_scalar), + b.getVectorUtilSupport().dotProduct(bsA_simd, bsB_simd), + 0.0001f); + Assert.assertEquals( + a.getVectorUtilSupport().squareDistance(bsA_scalar, bsB_scalar), + b.getVectorUtilSupport().squareDistance(bsA_simd, bsB_simd), + 0.0001f); + Assert.assertEquals( + a.getVectorUtilSupport().cosine(bsA_scalar, bsB_scalar), + b.getVectorUtilSupport().cosine(bsA_simd, bsB_simd), + 0.0001f); + } + + /** + * Verifies that the SIMD byte-vector kernels agree with the scalar baseline for + * several dimensions that stress boundary conditions in the vectorised loops: + *

    + *
  • dim=1 — the absolute minimum case
  • + *
  • dim=8 — exact multiple of the smallest SIMD lane width
  • + *
  • dim=128 — exact multiple of larger lane widths
  • + *
  • dim=256 — exact multiple of AVX2/AVX512 lane widths
  • + *
  • dim=255 — one less than 256, exposes tail-loop handling
  • + *
+ */ + @Test + public void testSimilarityMetricsByteEdgeDimensions() { + Assume.assumeTrue(hasSimd); + + VectorizationProvider scalar = new DefaultVectorizationProvider(); + VectorizationProvider simd = VectorizationProvider.getInstance(); + + for (int dim : new int[]{1, 8, 128, 255, 256}) { + byte[] rawA = new byte[dim]; + byte[] rawB = new byte[dim]; + getRandom().nextBytes(rawA); + getRandom().nextBytes(rawB); + + ByteSequence sA = scalar.getVectorTypeSupport().createByteSequence(rawA); + ByteSequence sB = scalar.getVectorTypeSupport().createByteSequence(rawB); + ByteSequence vA = simd.getVectorTypeSupport().createByteSequence(rawA); + ByteSequence vB = simd.getVectorTypeSupport().createByteSequence(rawB); + + Assert.assertEquals("dotProduct dim=" + dim, + scalar.getVectorUtilSupport().dotProduct(sA, sB), + simd.getVectorUtilSupport().dotProduct(vA, vB), + 0.0001f); + Assert.assertEquals("squareDistance dim=" + dim, + scalar.getVectorUtilSupport().squareDistance(sA, sB), + simd.getVectorUtilSupport().squareDistance(vA, vB), + 0.0001f); + Assert.assertEquals("cosine dim=" + dim, + scalar.getVectorUtilSupport().cosine(sA, sB), + simd.getVectorUtilSupport().cosine(vA, vB), + 0.0001f); + } + } + + /** + * SIMD kernels must treat bytes as signed. This test uses vectors whose + * correct result depends on negative byte values to catch sign-extension bugs. + */ + @Test + public void testSimilarityMetricsByteSignHandling() { + Assume.assumeTrue(hasSimd); + + VectorizationProvider scalar = new DefaultVectorizationProvider(); + VectorizationProvider simd = VectorizationProvider.getInstance(); + + // mix of extreme signed values: max positive 127, min negative -128 + byte[] rawA = new byte[16]; + byte[] rawB = new byte[16]; + for (int i = 0; i < 16; i++) { + rawA[i] = (i % 2 == 0) ? (byte) 127 : (byte) -128; + rawB[i] = (i % 2 == 0) ? (byte) -128 : (byte) 127; + } + + ByteSequence sA = scalar.getVectorTypeSupport().createByteSequence(rawA); + ByteSequence sB = scalar.getVectorTypeSupport().createByteSequence(rawB); + ByteSequence vA = simd.getVectorTypeSupport().createByteSequence(rawA); + ByteSequence vB = simd.getVectorTypeSupport().createByteSequence(rawB); + + Assert.assertEquals("signed dotProduct", + scalar.getVectorUtilSupport().dotProduct(sA, sB), + simd.getVectorUtilSupport().dotProduct(vA, vB), + 0.0001f); + Assert.assertEquals("signed squareDistance", + scalar.getVectorUtilSupport().squareDistance(sA, sB), + simd.getVectorUtilSupport().squareDistance(vA, vB), + 0.0001f); + Assert.assertEquals("signed cosine", + scalar.getVectorUtilSupport().cosine(sA, sB), + simd.getVectorUtilSupport().cosine(vA, vB), + 0.0001f); + } + @Test public void testAssembleAndSum() { Assume.assumeTrue(hasSimd); diff --git a/jvector-tests/src/test/resources/provider-selection/abi-library.c b/jvector-tests/src/test/resources/provider-selection/abi-library.c index d5a8a2558..28f899f48 100644 --- a/jvector-tests/src/test/resources/provider-selection/abi-library.c +++ b/jvector-tests/src/test/resources/provider-selection/abi-library.c @@ -72,3 +72,6 @@ void closest_centroid_f32(void) {} void nvq_growth_rate(void) {} void scale_to_f32(void) {} void subtract_add_in_place_f32(void) {} +void dot_product_i8(void) {} +void euclidean_i8(void) {} +void cosine_i8(void) {} diff --git a/jvector-tests/src/test/resources/provider-selection/missing-kernel.c b/jvector-tests/src/test/resources/provider-selection/missing-kernel.c index 9d26ac6f6..b07618886 100644 --- a/jvector-tests/src/test/resources/provider-selection/missing-kernel.c +++ b/jvector-tests/src/test/resources/provider-selection/missing-kernel.c @@ -61,3 +61,6 @@ void closest_centroid_f32(void) {} void nvq_growth_rate(void) {} void scale_to_f32(void) {} void subtract_add_in_place_f32(void) {} +void dot_product_i8(void) {} +void euclidean_i8(void) {} +void cosine_i8(void) {} 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 fc6dbc4b2..87accacc2 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 @@ -56,8 +56,11 @@ protected void intoVectorFloat(FloatVector vector, VectorFloat v, int offset) vector.intoArray(((ArrayVectorFloat) v).get(), offset); } + /// Loads from the backing array of an array-backed sequence (plain or sliced). + /// `offset` is an index into the backing array, so callers working with a slice + /// must include `vector.offset()`. protected ByteVector fromByteSequence(VectorSpecies SPEC, ByteSequence vector, int offset) { - return ByteVector.fromArray(SPEC, ((ArrayByteSequence) vector).get(), offset); + return ByteVector.fromArray(SPEC, (byte[]) vector.get(), offset); } protected void intoByteSequence(ByteVector vector, ByteSequence v, int offset) { @@ -1028,14 +1031,250 @@ float assembleAndSumPQ_512( return res; } + // ----------------------------------------------------------------------- + // ByteSequence similarity metrics – Panama SIMD implementations + // + // Strategy: widen signed bytes to int32 via B2I (no AND-mask needed for + // signed arithmetic), accumulate products in IntVector lanes, then reduce. + // The byte-vector species is 1/4 the width of the int species: + // 512-bit int (16 lanes) <- SPECIES_128 bytes + // 256-bit int (8 lanes) <- SPECIES_64 bytes + // 128-bit preferred <- scalar (ByteVector.SPECIES_32 does not exist; + // 128-bit SIMD shows no benefit for this workload) + // ----------------------------------------------------------------------- + + /** + * Vectorized dot product of two signed int8 byte vectors. + */ + @Override + public float dotProduct(ByteSequence a, ByteSequence b) { + return switch (PREFERRED_BIT_SIZE) { + case 512 -> dotProductBytes512(a, b); + case 256 -> dotProductBytes256(a, b); + default -> dotProductBytes128(a, b); + }; + } + + float dotProductBytes512(ByteSequence a, ByteSequence b) { + final int length = a.length(); + final int step = ByteVector.SPECIES_128.length(); // 16 + final int limit = ByteVector.SPECIES_128.loopBound(length); + IntVector acc = IntVector.zero(IntVector.SPECIES_512); + + for (int i = 0; i < limit; i += step) { + IntVector va = fromByteSequence(ByteVector.SPECIES_128, a, a.offset() + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_512, 0) + .reinterpretAsInts(); + IntVector vb = fromByteSequence(ByteVector.SPECIES_128, b, b.offset() + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_512, 0) + .reinterpretAsInts(); + acc = acc.add(va.mul(vb)); + } + + int result = acc.reduceLanes(VectorOperators.ADD); + for (int i = limit; i < length; i++) { + result += a.get(i) * b.get(i); + } + return result; + } + + float dotProductBytes256(ByteSequence a, ByteSequence b) { + final int length = a.length(); + final int step = ByteVector.SPECIES_64.length(); // 8 + final int limit = ByteVector.SPECIES_64.loopBound(length); + IntVector acc = IntVector.zero(IntVector.SPECIES_256); + + for (int i = 0; i < limit; i += step) { + IntVector va = fromByteSequence(ByteVector.SPECIES_64, a, a.offset() + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_256, 0) + .reinterpretAsInts(); + IntVector vb = fromByteSequence(ByteVector.SPECIES_64, b, b.offset() + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_256, 0) + .reinterpretAsInts(); + acc = acc.add(va.mul(vb)); + } + + int result = acc.reduceLanes(VectorOperators.ADD); + for (int i = limit; i < length; i++) { + result += a.get(i) * b.get(i); + } + return result; + } + + float dotProductBytes128(ByteSequence a, ByteSequence b) { + // ByteVector.SPECIES_32 does not exist; scalar is fastest at 128-bit width + final int length = a.length(); + int result = 0; + for (int i = 0; i < length; i++) { + result += a.get(i) * b.get(i); + } + return result; + } + /** - * Vectorized calculation of Hamming distance for two arrays of long integers. - * Both arrays should have the same length. - * - * @param a The first array - * @param b The second array - * @return The Hamming distance + * Vectorized sum of squared differences between two signed int8 byte vectors. */ + @Override + public float squareDistance(ByteSequence a, ByteSequence b) { + return switch (PREFERRED_BIT_SIZE) { + case 512 -> squareDistanceBytes512(a, b); + case 256 -> squareDistanceBytes256(a, b); + default -> squareDistanceBytes128(a, b); + }; + } + + float squareDistanceBytes512(ByteSequence a, ByteSequence b) { + final int length = a.length(); + final int step = ByteVector.SPECIES_128.length(); + final int limit = ByteVector.SPECIES_128.loopBound(length); + IntVector acc = IntVector.zero(IntVector.SPECIES_512); + + for (int i = 0; i < limit; i += step) { + IntVector va = fromByteSequence(ByteVector.SPECIES_128, a, a.offset() + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_512, 0) + .reinterpretAsInts(); + IntVector vb = fromByteSequence(ByteVector.SPECIES_128, b, b.offset() + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_512, 0) + .reinterpretAsInts(); + IntVector diff = va.sub(vb); + acc = acc.add(diff.mul(diff)); + } + + int result = acc.reduceLanes(VectorOperators.ADD); + for (int i = limit; i < length; i++) { + int diff = a.get(i) - b.get(i); + result += diff * diff; + } + return result; + } + + float squareDistanceBytes256(ByteSequence a, ByteSequence b) { + final int length = a.length(); + final int step = ByteVector.SPECIES_64.length(); + final int limit = ByteVector.SPECIES_64.loopBound(length); + IntVector acc = IntVector.zero(IntVector.SPECIES_256); + + for (int i = 0; i < limit; i += step) { + IntVector va = fromByteSequence(ByteVector.SPECIES_64, a, a.offset() + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_256, 0) + .reinterpretAsInts(); + IntVector vb = fromByteSequence(ByteVector.SPECIES_64, b, b.offset() + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_256, 0) + .reinterpretAsInts(); + IntVector diff = va.sub(vb); + acc = acc.add(diff.mul(diff)); + } + + int result = acc.reduceLanes(VectorOperators.ADD); + for (int i = limit; i < length; i++) { + int diff = a.get(i) - b.get(i); + result += diff * diff; + } + return result; + } + + float squareDistanceBytes128(ByteSequence a, ByteSequence b) { + // ByteVector.SPECIES_32 does not exist; scalar is fastest at 128-bit width + final int length = a.length(); + int result = 0; + for (int i = 0; i < length; i++) { + int diff = a.get(i) - b.get(i); + result += diff * diff; + } + return result; + } + + /** + * Vectorized cosine similarity between two signed int8 byte vectors. + */ + @Override + public float cosine(ByteSequence a, ByteSequence b) { + return switch (PREFERRED_BIT_SIZE) { + case 512 -> cosineBytes512(a, b); + case 256 -> cosineBytes256(a, b); + default -> cosineBytes128(a, b); + }; + } + + float cosineBytes512(ByteSequence a, ByteSequence b) { + final int length = a.length(); + final int step = ByteVector.SPECIES_128.length(); + final int limit = ByteVector.SPECIES_128.loopBound(length); + IntVector dot = IntVector.zero(IntVector.SPECIES_512); + IntVector normA = IntVector.zero(IntVector.SPECIES_512); + IntVector normB = IntVector.zero(IntVector.SPECIES_512); + + for (int i = 0; i < limit; i += step) { + IntVector va = fromByteSequence(ByteVector.SPECIES_128, a, a.offset() + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_512, 0) + .reinterpretAsInts(); + IntVector vb = fromByteSequence(ByteVector.SPECIES_128, b, b.offset() + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_512, 0) + .reinterpretAsInts(); + dot = dot.add(va.mul(vb)); + normA = normA.add(va.mul(va)); + normB = normB.add(vb.mul(vb)); + } + + long dotResult = dot.reduceLanes(VectorOperators.ADD); + long normAResult = normA.reduceLanes(VectorOperators.ADD); + long normBResult = normB.reduceLanes(VectorOperators.ADD); + + for (int i = limit; i < length; i++) { + int ai = a.get(i), bi = b.get(i); + dotResult += ai * bi; + normAResult += ai * ai; + normBResult += bi * bi; + } + return (float) (dotResult / Math.sqrt((double) normAResult * normBResult)); + } + + float cosineBytes256(ByteSequence a, ByteSequence b) { + final int length = a.length(); + final int step = ByteVector.SPECIES_64.length(); + final int limit = ByteVector.SPECIES_64.loopBound(length); + IntVector dot = IntVector.zero(IntVector.SPECIES_256); + IntVector normA = IntVector.zero(IntVector.SPECIES_256); + IntVector normB = IntVector.zero(IntVector.SPECIES_256); + + for (int i = 0; i < limit; i += step) { + IntVector va = fromByteSequence(ByteVector.SPECIES_64, a, a.offset() + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_256, 0) + .reinterpretAsInts(); + IntVector vb = fromByteSequence(ByteVector.SPECIES_64, b, b.offset() + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_256, 0) + .reinterpretAsInts(); + dot = dot.add(va.mul(vb)); + normA = normA.add(va.mul(va)); + normB = normB.add(vb.mul(vb)); + } + + long dotResult = dot.reduceLanes(VectorOperators.ADD); + long normAResult = normA.reduceLanes(VectorOperators.ADD); + long normBResult = normB.reduceLanes(VectorOperators.ADD); + + for (int i = limit; i < length; i++) { + int ai = a.get(i), bi = b.get(i); + dotResult += ai * bi; + normAResult += ai * ai; + normBResult += bi * bi; + } + return (float) (dotResult / Math.sqrt((double) normAResult * normBResult)); + } + + float cosineBytes128(ByteSequence a, ByteSequence b) { + // ByteVector.SPECIES_32 does not exist; scalar is fastest at 128-bit width + final int length = a.length(); + long dotResult = 0, normAResult = 0, normBResult = 0; + for (int i = 0; i < length; i++) { + int ai = a.get(i), bi = b.get(i); + dotResult += ai * bi; + normAResult += ai * ai; + normBResult += bi * bi; + } + return (float) (dotResult / Math.sqrt((double) normAResult * normBResult)); + } + @Override public int hammingDistance(long[] a, long[] b) { var sum = LongVector.zero(LongVector.SPECIES_PREFERRED);