diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ArrayVectorFloat.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ArrayVectorFloat.java index 32dce1f35..46c7a59dc 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ArrayVectorFloat.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ArrayVectorFloat.java @@ -18,6 +18,7 @@ import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.util.RamUsageEstimator; +import io.github.jbellis.jvector.vector.types.FloatArray; import io.github.jbellis.jvector.vector.types.VectorFloat; import java.io.IOException; @@ -26,7 +27,7 @@ /** * VectorFloat implementation backed by an on-heap float array. */ -final public class ArrayVectorFloat implements VectorFloat +final public class ArrayVectorFloat implements VectorFloat, FloatArray { private final float[] data; @@ -123,5 +124,10 @@ public int hashCode() { return this.getHashCode(); } + + @Override + public float[] array() { + return get(); + } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/types/FloatArray.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/types/FloatArray.java new file mode 100644 index 000000000..8577d6c1f --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/types/FloatArray.java @@ -0,0 +1,27 @@ +/* + * 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.types; + +/** + * Return the {@code float[]} representation of the {@link VectorFloat} if its type supports + * such conversions. + * + * @apiNote this is an experimental API and may change in the future releases. + */ +public interface FloatArray { + float[] array(); +} diff --git a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/MemorySegmentVectorFloat.java b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/MemorySegmentVectorFloat.java index 3e3f67e1f..4227629aa 100644 --- a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/MemorySegmentVectorFloat.java +++ b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/MemorySegmentVectorFloat.java @@ -18,6 +18,7 @@ import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.util.RamUsageEstimator; +import io.github.jbellis.jvector.vector.types.FloatArray; import io.github.jbellis.jvector.vector.types.VectorFloat; import java.io.IOException; @@ -27,7 +28,7 @@ /** * VectorFloat implementation backed by an on-heap MemorySegment. */ -final public class MemorySegmentVectorFloat implements VectorFloat +final public class MemorySegmentVectorFloat implements VectorFloat, FloatArray { private final MemorySegment segment; @@ -141,4 +142,9 @@ public boolean equals(Object o) public int hashCode() { return this.getHashCode(); } + + @Override + public float[] array() { + return (float[])segment.heapBase().get(); + } } diff --git a/jvector-native/src/test/java/io/github/jbellis/jvector/vector/MemorySegmentVectorProviderTest.java b/jvector-native/src/test/java/io/github/jbellis/jvector/vector/MemorySegmentVectorProviderTest.java index 29b415a86..5afa4b2bf 100644 --- a/jvector-native/src/test/java/io/github/jbellis/jvector/vector/MemorySegmentVectorProviderTest.java +++ b/jvector-native/src/test/java/io/github/jbellis/jvector/vector/MemorySegmentVectorProviderTest.java @@ -18,8 +18,15 @@ import io.github.jbellis.jvector.disk.IndexWriter; 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 static org.hamcrest.CoreMatchers.instanceOf; import java.io.IOException; +import java.util.Random; + +import org.junit.Assert; import org.junit.jupiter.api.Test; public class MemorySegmentVectorProviderTest { @@ -75,6 +82,36 @@ void testWriteByteSequenceZeroLength() throws IOException { org.junit.jupiter.api.Assertions.assertArrayEquals(new byte[0], dummyWriter.toByteArray()); } + @Test + void testFloatVectorAsArray() { + final MemorySegmentVectorProvider provider = new MemorySegmentVectorProvider(); + + final Random random = new Random(); + final VectorFloat vf = randomVector(provider, random, 1021); + + Assert.assertThat(vf, instanceOf(FloatArray.class)); + Assert.assertArrayEquals(((FloatArray) vf).array(), getVector(vf), 0.0001f); + } + + private static VectorFloat randomVector(MemorySegmentVectorProvider provider, Random random, int dim) { + var vec = provider.createFloatVector(dim); + for (int i = 0; i < dim; i++) { + vec.set(i, random.nextFloat()); + if (random.nextBoolean()) { + vec.set(i, -vec.get(i)); + } + } + return vec; + } + + private static float[] getVector(VectorFloat vf) { + final float[] arr = new float[vf.length()]; + for (int i = 0; i < vf.length(); ++i) { + arr[i] = vf.get(i); + } + return arr; + } + /** * A lightweight mock to capture IndexWriter output without boilerplate. */ 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 4a3b69e93..81a99aafc 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 @@ -17,7 +17,9 @@ package io.github.jbellis.jvector.vector; import com.carrotsearch.randomizedtesting.RandomizedTest; + import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.vector.types.FloatArray; import io.github.jbellis.jvector.vector.types.VectorFloat; import io.github.jbellis.jvector.vector.types.VectorTypeSupport; import org.junit.Assert; @@ -25,6 +27,7 @@ import org.junit.Test; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.hamcrest.CoreMatchers.instanceOf; public class TestVectorizationProvider extends RandomizedTest { @@ -47,6 +50,11 @@ public void testSimilarityMetricsFloat() { v2a.set(i, v2b.get(i)); } + Assert.assertThat(v1a, instanceOf(FloatArray.class)); + Assert.assertThat(v2a, instanceOf(FloatArray.class)); + Assert.assertArrayEquals(((FloatArray) v1a).array(), getVector(v1a), 0.0001f); + Assert.assertArrayEquals(((FloatArray) v2a).array(), getVector(v2a), 0.0001f); + Assert.assertEquals(a.getVectorUtilSupport().dotProduct(v1a,v2a), b.getVectorUtilSupport().dotProduct(v1b, v2b), 0.0001f); Assert.assertEquals(a.getVectorUtilSupport().cosine(v1a,v2a), b.getVectorUtilSupport().cosine(v1b, v2b), 0.0001f); Assert.assertEquals(a.getVectorUtilSupport().squareDistance(v1a, v2a), b.getVectorUtilSupport().squareDistance(v1b, v2b), 0.0001f); @@ -71,13 +79,18 @@ public void testAssembleAndSum() { offsets[c] = (byte) (c * skipSize); } + Assert.assertThat(v2, instanceOf(FloatArray.class)); + Assert.assertThat(v3, instanceOf(FloatArray.class)); + Assert.assertArrayEquals(((FloatArray) v2).array(), getVector(v2), 0.0001f); + Assert.assertArrayEquals(((FloatArray) v3).array(), getVector(v3), 0.0001f); + Assert.assertEquals(a.getVectorUtilSupport().sum(v3), b.getVectorUtilSupport().sum(v3), 0.0001); Assert.assertEquals(a.getVectorUtilSupport().sum(v3), a.getVectorUtilSupport().assembleAndSum(v2, 0, vectorTypeSupport.createByteSequence(offsets)), 0.0001); Assert.assertEquals(b.getVectorUtilSupport().sum(v3), b.getVectorUtilSupport().assembleAndSum(v2, 0, vectorTypeSupport.createByteSequence(offsets)), 0.0001); } } - public static String REQUIRE_SPECIFIC_VECTORIZATION_PROVIDER="Test_RequireSpecificVectorizationProvider"; + public static String REQUIRE_SPECIFIC_VECTORIZATION_PROVIDER="Test_RequireSpecificVectorizationProvider"; /** * To run with native-access vector support, use @@ -117,4 +130,11 @@ public void testVectorSupportTypeIsExpected() { } } + private static float[] getVector(VectorFloat vf) { + final float[] arr = new float[vf.length()]; + for (int i = 0; i < vf.length(); ++i) { + arr[i] = vf.get(i); + } + return arr; + } } diff --git a/pom.xml b/pom.xml index a80eedf31..a502864b9 100644 --- a/pom.xml +++ b/pom.xml @@ -109,7 +109,7 @@ org.apache.maven.plugins maven-javadoc-plugin - 3.5.0 + 3.12.0 attach-javadocs @@ -123,6 +123,14 @@ --add-modules=jdk.incubator.vector 22 + + + + apiNote + a + API Note: + +