Skip to content

Commit

Permalink
Merge branch 'main' into FinalMain
Browse files Browse the repository at this point in the history
Signed-off-by: Vikasht34 <[email protected]>
  • Loading branch information
Vikasht34 authored Sep 9, 2024
2 parents a390c1d + 8f6b177 commit f90a221
Show file tree
Hide file tree
Showing 14 changed files with 237 additions and 93 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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[]
Expand Down Expand Up @@ -90,7 +97,7 @@ public void close() {
*/
public void reset() {
vectorAddress = 0;
vectorsToTransfer.clear();
vectorsToTransfer = null;
}

protected abstract void deallocate();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -782,6 +782,8 @@ public ParametrizedFieldMapper.Builder getMergeBuilder() {
.vectorDataType(vectorDataType)
.versionCreated(indexCreatedVersion)
.dimension(fieldType().getKnnMappingConfig().getDimension())
.compressionLevel(fieldType().getKnnMappingConfig().getCompressionLevel())
.mode(fieldType().getKnnMappingConfig().getMode())
.build();
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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
Expand All @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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));
}
}
}
Loading

0 comments on commit f90a221

Please sign in to comment.