Bounds the transferLimit in OffheapVectorTransfer (opensearch-project…

…#2070) The array list buffer was unnecessarily allocated a large memory irrespective of the vectors to transfer. This change considers total vectors Signed-off-by: Tejas Shah <[email protected]>
heemin32 · Sep 9, 2024 · 0492fb3 · 0492fb3
1 parent cbc6343
commit 0492fb3
Show file tree

Hide file tree

Showing 11 changed files with 126 additions and 90 deletions.
diff --git a/src/main/java/org/opensearch/knn/index/codec/nativeindex/DefaultIndexBuildStrategy.java b/src/main/java/org/opensearch/knn/index/codec/nativeindex/DefaultIndexBuildStrategy.java
@@ -8,7 +8,6 @@
 import lombok.AccessLevel;
 import lombok.NoArgsConstructor;
 import org.opensearch.knn.common.KNNConstants;
-import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams;
 import org.opensearch.knn.index.codec.transfer.OffHeapVectorTransfer;
 import org.opensearch.knn.index.vectorvalues.KNNVectorValues;
@@ -56,8 +55,13 @@ public void buildAndWriteIndex(final BuildIndexParams indexInfo) throws IOExcept
         iterateVectorValuesOnce(knnVectorValues);
         IndexBuildSetup indexBuildSetup = QuantizationIndexUtils.prepareIndexBuild(knnVectorValues, indexInfo);
 
-        int transferLimit = (int) Math.max(1, KNNSettings.getVectorStreamingMemoryLimit().getBytes() / indexBuildSetup.getBytesPerVector());
-        try (final OffHeapVectorTransfer vectorTransfer = getVectorTransfer(indexInfo.getVectorDataType(), transferLimit)) {
+        try (
+            final OffHeapVectorTransfer vectorTransfer = getVectorTransfer(
+                indexInfo.getVectorDataType(),
+                indexBuildSetup.getBytesPerVector(),
+                indexInfo.getTotalLiveDocs()
+            )
+        ) {
             final List<Integer> transferredDocIds = new ArrayList<>(indexInfo.getTotalLiveDocs());
 
             while (knnVectorValues.docId() != NO_MORE_DOCS) {

diff --git a/...java/org/opensearch/knn/index/codec/nativeindex/MemOptimizedNativeIndexBuildStrategy.java b/...java/org/opensearch/knn/index/codec/nativeindex/MemOptimizedNativeIndexBuildStrategy.java
@@ -7,7 +7,6 @@
 
 import lombok.AccessLevel;
 import lombok.NoArgsConstructor;
-import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams;
 import org.opensearch.knn.index.codec.transfer.OffHeapVectorTransfer;
 import org.opensearch.knn.index.engine.KNNEngine;
@@ -70,10 +69,15 @@ public void buildAndWriteIndex(final BuildIndexParams indexInfo) throws IOExcept
             )
         );
 
-        int transferLimit = (int) Math.max(1, KNNSettings.getVectorStreamingMemoryLimit().getBytes() / indexBuildSetup.getBytesPerVector());
-        try (final OffHeapVectorTransfer vectorTransfer = getVectorTransfer(indexInfo.getVectorDataType(), transferLimit)) {
+        try (
+            final OffHeapVectorTransfer vectorTransfer = getVectorTransfer(
+                indexInfo.getVectorDataType(),
+                indexBuildSetup.getBytesPerVector(),
+                indexInfo.getTotalLiveDocs()
+            )
+        ) {
 
-            final List<Integer> transferredDocIds = new ArrayList<>(transferLimit);
+            final List<Integer> transferredDocIds = new ArrayList<>(vectorTransfer.getTransferLimit());
 
             while (knnVectorValues.docId() != NO_MORE_DOCS) {
                 Object vector = QuantizationIndexUtils.processAndReturnVector(knnVectorValues, indexBuildSetup);

diff --git a/src/main/java/org/opensearch/knn/index/codec/transfer/OffHeapBinaryVectorTransfer.java b/src/main/java/org/opensearch/knn/index/codec/transfer/OffHeapBinaryVectorTransfer.java
@@ -17,8 +17,8 @@
  */
 public final class OffHeapBinaryVectorTransfer extends OffHeapVectorTransfer<byte[]> {
 
-    public OffHeapBinaryVectorTransfer(int transferLimit) {
-        super(transferLimit);
+    public OffHeapBinaryVectorTransfer(int bytesPerVector, int totalVectorsToTransfer) {
+        super(bytesPerVector, totalVectorsToTransfer);
     }
 
     @Override

diff --git a/src/main/java/org/opensearch/knn/index/codec/transfer/OffHeapByteVectorTransfer.java b/src/main/java/org/opensearch/knn/index/codec/transfer/OffHeapByteVectorTransfer.java
@@ -17,8 +17,8 @@
  */
 public final class OffHeapByteVectorTransfer extends OffHeapVectorTransfer<byte[]> {
 
-    public OffHeapByteVectorTransfer(int transferLimit) {
-        super(transferLimit);
+    public OffHeapByteVectorTransfer(int bytesPerVector, int totalVectorsToTransfer) {
+        super(bytesPerVector, totalVectorsToTransfer);
     }
 
     @Override

diff --git a/src/main/java/org/opensearch/knn/index/codec/transfer/OffHeapFloatVectorTransfer.java b/src/main/java/org/opensearch/knn/index/codec/transfer/OffHeapFloatVectorTransfer.java
@@ -15,8 +15,8 @@
  */
 public final class OffHeapFloatVectorTransfer extends OffHeapVectorTransfer<float[]> {
 
-    public OffHeapFloatVectorTransfer(int transferLimit) {
-        super(transferLimit);
+    public OffHeapFloatVectorTransfer(int bytesPerVector, int totalVectorsToTransfer) {
+        super(bytesPerVector, totalVectorsToTransfer);
     }
 
     @Override

diff --git a/src/main/java/org/opensearch/knn/index/codec/transfer/OffHeapVectorTransfer.java b/src/main/java/org/opensearch/knn/index/codec/transfer/OffHeapVectorTransfer.java
@@ -6,6 +6,7 @@
 package org.opensearch.knn.index.codec.transfer;
 
 import lombok.Getter;
+import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.vectorvalues.KNNVectorValues;
 
 import java.io.Closeable;
@@ -27,16 +28,22 @@ public abstract class OffHeapVectorTransfer<T> implements Closeable {
 
     @Getter
     private long vectorAddress;
+    @Getter
     protected final int transferLimit;
 
-    private final List<T> vectorsToTransfer;
+    private List<T> vectorsToTransfer;
 
-    public OffHeapVectorTransfer(final int transferLimit) {
-        this.transferLimit = transferLimit;
-        this.vectorsToTransfer = new ArrayList<>(transferLimit);
+    public OffHeapVectorTransfer(int bytesPerVector, int totalVectorsToTransfer) {
+        this.transferLimit = computeTransferLimit(bytesPerVector, totalVectorsToTransfer);
+        this.vectorsToTransfer = new ArrayList<>(this.transferLimit);
         this.vectorAddress = 0;
     }
 
+    private int computeTransferLimit(int bytesPerVector, int totalVectorsToTransfer) {
+        int limit = (int) Math.max(1, KNNSettings.getVectorStreamingMemoryLimit().getBytes() / bytesPerVector);
+        return Math.min(limit, totalVectorsToTransfer);
+    }
+
     /**
      * Transfer vectors to off-heap
      * @param vector float[] or byte[]
@@ -90,7 +97,7 @@ public void close() {
      */
     public void reset() {
         vectorAddress = 0;
-        vectorsToTransfer.clear();
+        vectorsToTransfer = null;
     }
 
     protected abstract void deallocate();

diff --git a/src/main/java/org/opensearch/knn/index/codec/transfer/OffHeapVectorTransferFactory.java b/src/main/java/org/opensearch/knn/index/codec/transfer/OffHeapVectorTransferFactory.java
@@ -18,18 +18,23 @@ public final class OffHeapVectorTransferFactory {
     /**
      * Gets the right vector transfer object based on vector data type
      * @param vectorDataType {@link VectorDataType}
-     * @param transferLimit max number of vectors that can be transferred to off heap in one transfer
+     * @param bytesPerVector Bytes used per vector
+     * @param totalVectorsToTransfer total number of vectors that will be transferred off heap
      * @return Correct implementation of {@link OffHeapVectorTransfer}
      * @param <T> float[] or byte[]
      */
-    public static <T> OffHeapVectorTransfer<T> getVectorTransfer(final VectorDataType vectorDataType, final int transferLimit) {
+    public static <T> OffHeapVectorTransfer<T> getVectorTransfer(
+        final VectorDataType vectorDataType,
+        int bytesPerVector,
+        int totalVectorsToTransfer
+    ) {
         switch (vectorDataType) {
             case FLOAT:
-                return (OffHeapVectorTransfer<T>) new OffHeapFloatVectorTransfer(transferLimit);
+                return (OffHeapVectorTransfer<T>) new OffHeapFloatVectorTransfer(bytesPerVector, totalVectorsToTransfer);
             case BINARY:
-                return (OffHeapVectorTransfer<T>) new OffHeapBinaryVectorTransfer(transferLimit);
+                return (OffHeapVectorTransfer<T>) new OffHeapBinaryVectorTransfer(bytesPerVector, totalVectorsToTransfer);
             case BYTE:
-                return (OffHeapVectorTransfer<T>) new OffHeapByteVectorTransfer(transferLimit);
+                return (OffHeapVectorTransfer<T>) new OffHeapByteVectorTransfer(bytesPerVector, totalVectorsToTransfer);
             default:
                 throw new IllegalArgumentException("Unsupported vector data type: " + vectorDataType);
         }

diff --git a/src/test/java/org/opensearch/knn/index/codec/nativeindex/DefaultIndexBuildStrategyTests.java b/src/test/java/org/opensearch/knn/index/codec/nativeindex/DefaultIndexBuildStrategyTests.java
@@ -57,14 +57,11 @@ public void testBuildAndWrite() {
         final KNNVectorValues<byte[]> knnVectorValues = KNNVectorValuesFactory.getVectorValues(VectorDataType.FLOAT, randomVectorValues);
 
         try (
-            MockedStatic<KNNSettings> mockedKNNSettings = mockStatic(KNNSettings.class);
             MockedStatic<JNIService> mockedJNIService = mockStatic(JNIService.class);
             MockedStatic<OffHeapVectorTransferFactory> mockedOffHeapVectorTransferFactory = mockStatic(OffHeapVectorTransferFactory.class)
         ) {
-
-            mockedKNNSettings.when(KNNSettings::getVectorStreamingMemoryLimit).thenReturn(new ByteSizeValue(16));
             OffHeapVectorTransfer offHeapVectorTransfer = mock(OffHeapVectorTransfer.class);
-            mockedOffHeapVectorTransferFactory.when(() -> OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.FLOAT, 2))
+            mockedOffHeapVectorTransferFactory.when(() -> OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.FLOAT, 8, 3))
                 .thenReturn(offHeapVectorTransfer);
 
             when(offHeapVectorTransfer.getVectorAddress()).thenReturn(200L);
@@ -131,7 +128,7 @@ public void testBuildAndWrite_withQuantization() {
             mockedJNIService.when(() -> JNIService.initIndex(3, 2, Map.of("index", "param"), KNNEngine.FAISS)).thenReturn(100L);
 
             OffHeapVectorTransfer offHeapVectorTransfer = mock(OffHeapVectorTransfer.class);
-            mockedOffHeapVectorTransferFactory.when(() -> OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.FLOAT, 2))
+            mockedOffHeapVectorTransferFactory.when(() -> OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.FLOAT, 8, 3))
                 .thenReturn(offHeapVectorTransfer);
 
             QuantizationService quantizationService = mock(QuantizationService.class);
@@ -237,14 +234,12 @@ public void testBuildAndWriteWithModel() {
             docs
         );
         try (
-            MockedStatic<KNNSettings> mockedKNNSettings = mockStatic(KNNSettings.class);
             MockedStatic<JNIService> mockedJNIService = mockStatic(JNIService.class);
             MockedStatic<OffHeapVectorTransferFactory> mockedOffHeapVectorTransferFactory = mockStatic(OffHeapVectorTransferFactory.class)
         ) {
 
-            mockedKNNSettings.when(KNNSettings::getVectorStreamingMemoryLimit).thenReturn(new ByteSizeValue(16));
             OffHeapVectorTransfer offHeapVectorTransfer = mock(OffHeapVectorTransfer.class);
-            mockedOffHeapVectorTransferFactory.when(() -> OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.FLOAT, 2))
+            mockedOffHeapVectorTransferFactory.when(() -> OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.FLOAT, 8, 3))
                 .thenReturn(offHeapVectorTransfer);
 
             when(offHeapVectorTransfer.getVectorAddress()).thenReturn(200L);

diff --git a/...org/opensearch/knn/index/codec/nativeindex/MemOptimizedNativeIndexBuildStrategyTests.java b/...org/opensearch/knn/index/codec/nativeindex/MemOptimizedNativeIndexBuildStrategyTests.java
@@ -9,8 +9,6 @@
 import org.mockito.ArgumentCaptor;
 import org.mockito.MockedStatic;
 import org.mockito.Mockito;
-import org.opensearch.core.common.unit.ByteSizeValue;
-import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.VectorDataType;
 import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams;
 import org.opensearch.knn.index.codec.transfer.OffHeapVectorTransfer;
@@ -49,21 +47,20 @@ public void testBuildAndWrite() {
         final KNNVectorValues<byte[]> knnVectorValues = KNNVectorValuesFactory.getVectorValues(VectorDataType.FLOAT, randomVectorValues);
 
         try (
-            MockedStatic<KNNSettings> mockedKNNSettings = Mockito.mockStatic(KNNSettings.class);
             MockedStatic<JNIService> mockedJNIService = Mockito.mockStatic(JNIService.class);
             MockedStatic<OffHeapVectorTransferFactory> mockedOffHeapVectorTransferFactory = Mockito.mockStatic(
                 OffHeapVectorTransferFactory.class
             );
         ) {
 
             // Limits transfer to 2 vectors
-            mockedKNNSettings.when(KNNSettings::getVectorStreamingMemoryLimit).thenReturn(new ByteSizeValue(16));
             mockedJNIService.when(() -> JNIService.initIndex(3, 2, Map.of("index", "param"), KNNEngine.FAISS)).thenReturn(100L);
 
             OffHeapVectorTransfer offHeapVectorTransfer = mock(OffHeapVectorTransfer.class);
-            mockedOffHeapVectorTransferFactory.when(() -> OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.FLOAT, 2))
+            mockedOffHeapVectorTransferFactory.when(() -> OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.FLOAT, 8, 3))
                 .thenReturn(offHeapVectorTransfer);
 
+            when(offHeapVectorTransfer.getTransferLimit()).thenReturn(2);
             when(offHeapVectorTransfer.transfer(vectorTransferCapture.capture(), eq(false))).thenReturn(false)
                 .thenReturn(true)
                 .thenReturn(false);
@@ -145,7 +142,6 @@ public void testBuildAndWrite_withQuantization() {
         final KNNVectorValues<byte[]> knnVectorValues = KNNVectorValuesFactory.getVectorValues(VectorDataType.FLOAT, randomVectorValues);
 
         try (
-            MockedStatic<KNNSettings> mockedKNNSettings = Mockito.mockStatic(KNNSettings.class);
             MockedStatic<JNIService> mockedJNIService = Mockito.mockStatic(JNIService.class);
             MockedStatic<OffHeapVectorTransferFactory> mockedOffHeapVectorTransferFactory = Mockito.mockStatic(
                 OffHeapVectorTransferFactory.class
@@ -154,11 +150,11 @@ public void testBuildAndWrite_withQuantization() {
         ) {
 
             // Limits transfer to 2 vectors
-            mockedKNNSettings.when(KNNSettings::getVectorStreamingMemoryLimit).thenReturn(new ByteSizeValue(16));
             mockedJNIService.when(() -> JNIService.initIndex(3, 2, Map.of("index", "param"), KNNEngine.FAISS)).thenReturn(100L);
 
             OffHeapVectorTransfer offHeapVectorTransfer = mock(OffHeapVectorTransfer.class);
-            mockedOffHeapVectorTransferFactory.when(() -> OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.FLOAT, 2))
+            when(offHeapVectorTransfer.getTransferLimit()).thenReturn(2);
+            mockedOffHeapVectorTransferFactory.when(() -> OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.FLOAT, 8, 3))
                 .thenReturn(offHeapVectorTransfer);
 
             QuantizationService quantizationService = mock(QuantizationService.class);

diff --git a/src/test/java/org/opensearch/knn/index/codec/transfer/OffHeapVectorTransferFactoryTests.java b/src/test/java/org/opensearch/knn/index/codec/transfer/OffHeapVectorTransferFactoryTests.java
@@ -5,22 +5,30 @@
 
 package org.opensearch.knn.index.codec.transfer;
 
+import org.mockito.MockedStatic;
+import org.opensearch.core.common.unit.ByteSizeValue;
+import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.VectorDataType;
 import org.opensearch.test.OpenSearchTestCase;
 
+import static org.mockito.Mockito.mockStatic;
+
 public class OffHeapVectorTransferFactoryTests extends OpenSearchTestCase {
 
     public void testOffHeapVectorTransferFactory() {
-        var floatVectorTransfer = OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.FLOAT, 10);
-        assertEquals(OffHeapFloatVectorTransfer.class, floatVectorTransfer.getClass());
-        assertNotSame(floatVectorTransfer, OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.FLOAT, 10));
+        try (MockedStatic<KNNSettings> mockedKNNSettings = mockStatic(KNNSettings.class)) {
+            mockedKNNSettings.when(KNNSettings::getVectorStreamingMemoryLimit).thenReturn(new ByteSizeValue(16));
+            var floatVectorTransfer = OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.FLOAT, 10, 10);
+            assertEquals(OffHeapFloatVectorTransfer.class, floatVectorTransfer.getClass());
+            assertNotSame(floatVectorTransfer, OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.FLOAT, 10, 10));
 
-        var byteVectorTransfer = OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.BYTE, 10);
-        assertEquals(OffHeapByteVectorTransfer.class, byteVectorTransfer.getClass());
-        assertNotSame(byteVectorTransfer, OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.BYTE, 10));
+            var byteVectorTransfer = OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.BYTE, 10, 10);
+            assertEquals(OffHeapByteVectorTransfer.class, byteVectorTransfer.getClass());
+            assertNotSame(byteVectorTransfer, OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.BYTE, 10, 10));
 
-        var binaryVectorTransfer = OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.BINARY, 10);
-        assertEquals(OffHeapBinaryVectorTransfer.class, binaryVectorTransfer.getClass());
-        assertNotSame(binaryVectorTransfer, OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.BINARY, 10));
+            var binaryVectorTransfer = OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.BINARY, 10, 10);
+            assertEquals(OffHeapBinaryVectorTransfer.class, binaryVectorTransfer.getClass());
+            assertNotSame(binaryVectorTransfer, OffHeapVectorTransferFactory.getVectorTransfer(VectorDataType.BINARY, 10, 10));
+        }
     }
 }