Skip to content

Commit a58a9dc

Browse files
author
Ryan Bogan
authored
Add quantization state reader and writer (#1997)
Add quantization state reader and writer Signed-off-by: Ryan Bogan <rbogan@amazon.com>
1 parent 435e417 commit a58a9dc

18 files changed

Lines changed: 1341 additions & 37 deletions

‎CHANGELOG.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,4 +21,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
2121
### Infrastructure
2222
### Documentation
2323
### Maintenance
24-
### Refactoring
24+
### Refactoring

‎release-notes/opensearch-knn.release-notes-2.17.0.0.md‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,4 +30,5 @@ Compatible with OpenSearch 2.17.0
3030
* Restructure mappers to better handle null cases and avoid branching in parsing [#1939](https://github.com/opensearch-project/k-NN/pull/1939)
3131
* Added Quantization Framework and implemented 1Bit and multibit quantizer[#1889](https://github.com/opensearch-project/k-NN/issues/1889)
3232
* Encapsulate dimension, vector data type validation/processing inside Library [#1957](https://github.com/opensearch-project/k-NN/pull/1957)
33-
* Add quantization state cache [#1960](https://github.com/opensearch-project/k-NN/pull/1960)
33+
* Add quantization state cache [#1960](https://github.com/opensearch-project/k-NN/pull/1960)
34+
* Add quantization state reader and writer [#1997](https://github.com/opensearch-project/k-NN/pull/1997)

‎src/main/java/org/opensearch/knn/common/KNNConstants.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@ public class KNNConstants {
7474
public static final String MINIMAL_MODE_AND_COMPRESSION_FEATURE = "mode_and_compression_feature";
7575

7676
public static final String RADIAL_SEARCH_KEY = "radial_search";
77+
public static final String QUANTIZATION_STATE_FILE_SUFFIX = "osknnqstate";
7778

7879
// Lucene specific constants
7980
public static final String LUCENE_NAME = "lucene";

‎src/main/java/org/opensearch/knn/index/KNNSettings.java‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
import org.opensearch.knn.index.memory.NativeMemoryCacheManager;
2525
import org.opensearch.knn.index.memory.NativeMemoryCacheManagerDto;
2626
import org.opensearch.knn.index.util.IndexHyperParametersUtil;
27-
import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateCache;
27+
import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateCacheManager;
2828
import org.opensearch.monitor.jvm.JvmInfo;
2929
import org.opensearch.monitor.os.OsProbe;
3030

@@ -60,6 +60,7 @@ public class KNNSettings {
6060
private static final OsProbe osProbe = OsProbe.getInstance();
6161

6262
private static final int INDEX_THREAD_QTY_MAX = 32;
63+
private static final QuantizationStateCacheManager quantizationStateCacheManager = QuantizationStateCacheManager.getInstance();
6364

6465
/**
6566
* Settings name
@@ -379,11 +380,11 @@ private void setSettingsUpdateConsumers() {
379380
NativeMemoryCacheManager.getInstance().rebuildCache(builder.build());
380381
}, Stream.concat(dynamicCacheSettings.values().stream(), FEATURE_FLAGS.values().stream()).collect(Collectors.toUnmodifiableList()));
381382
clusterService.getClusterSettings().addSettingsUpdateConsumer(QUANTIZATION_STATE_CACHE_SIZE_LIMIT_SETTING, it -> {
382-
QuantizationStateCache.getInstance().setMaxCacheSizeInKB(it.getKb());
383-
QuantizationStateCache.getInstance().rebuildCache();
383+
quantizationStateCacheManager.setMaxCacheSizeInKB(it.getKb());
384+
quantizationStateCacheManager.rebuildCache();
384385
});
385386
clusterService.getClusterSettings().addSettingsUpdateConsumer(QUANTIZATION_STATE_CACHE_EXPIRY_TIME_MINUTES_SETTING, it -> {
386-
QuantizationStateCache.getInstance().rebuildCache();
387+
quantizationStateCacheManager.rebuildCache();
387388
});
388389
}
389390

‎src/main/java/org/opensearch/knn/index/codec/KNN990Codec/KNN990QuantizationStateReader.java‎

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.knn.index.codec.KNN990Codec;
7+
8+
import com.google.common.annotations.VisibleForTesting;
9+
import lombok.extern.log4j.Log4j2;
10+
import org.apache.lucene.codecs.CodecUtil;
11+
import org.apache.lucene.index.IndexFileNames;
12+
import org.apache.lucene.index.SegmentReadState;
13+
import org.apache.lucene.store.IOContext;
14+
import org.apache.lucene.store.IndexInput;
15+
import org.opensearch.knn.common.KNNConstants;
16+
import org.opensearch.knn.quantization.enums.ScalarQuantizationType;
17+
import org.opensearch.knn.quantization.models.quantizationParams.ScalarQuantizationParams;
18+
import org.opensearch.knn.quantization.models.quantizationState.MultiBitScalarQuantizationState;
19+
import org.opensearch.knn.quantization.models.quantizationState.OneBitScalarQuantizationState;
20+
import org.opensearch.knn.quantization.models.quantizationState.QuantizationState;
21+
import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateReadConfig;
22+
23+
import java.io.IOException;
24+
import java.util.Collections;
25+
import java.util.HashMap;
26+
import java.util.Map;
27+
28+
/**
29+
* Reads quantization states
30+
*/
31+
@Log4j2
32+
public final class KNN990QuantizationStateReader {
33+
34+
/**
35+
* Read quantization states and return list of fieldNames and bytes
36+
* File format:
37+
* Header
38+
* QS1 state bytes
39+
* QS2 state bytes
40+
* Number of quantization states
41+
* QS1 field number
42+
* QS1 state bytes length
43+
* QS1 position of state bytes
44+
* QS2 field number
45+
* QS2 state bytes length
46+
* QS2 position of state bytes
47+
* Position of index section (where QS1 field name is located)
48+
* -1 (marker)
49+
* Footer
50+
*
51+
* @param state the read state to read from
52+
*/
53+
public static Map<String, byte[]> read(SegmentReadState state) throws IOException {
54+
String quantizationStateFileName = getQuantizationStateFileName(state);
55+
Map<String, byte[]> readQuantizationStateInfos = null;
56+
57+
try (IndexInput input = state.directory.openInput(quantizationStateFileName, IOContext.READ)) {
58+
CodecUtil.retrieveChecksum(input);
59+
60+
int numFields = getNumFields(input);
61+
62+
readQuantizationStateInfos = new HashMap<>();
63+
64+
// Read each field's metadata from the index section and then read bytes
65+
for (int i = 0; i < numFields; i++) {
66+
int fieldNumber = input.readInt();
67+
int length = input.readInt();
68+
long position = input.readVLong();
69+
byte[] stateBytes = readStateBytes(input, position, length);
70+
String fieldName = state.fieldInfos.fieldInfo(fieldNumber).getName();
71+
readQuantizationStateInfos.put(fieldName, stateBytes);
72+
}
73+
} catch (Exception e) {
74+
log.warn(String.format("Unable to read the quantization state file for segment %s", state.segmentInfo.name), e);
75+
return Collections.emptyMap();
76+
}
77+
return readQuantizationStateInfos;
78+
}
79+
80+
/**
81+
* Reads an individual quantization state for a given field
82+
* @param readConfig a config class that contains necessary information for reading the state
83+
* @return quantization state
84+
*/
85+
public static QuantizationState read(QuantizationStateReadConfig readConfig) throws IOException {
86+
SegmentReadState segmentReadState = readConfig.getSegmentReadState();
87+
String field = readConfig.getField();
88+
String quantizationStateFileName = getQuantizationStateFileName(segmentReadState);
89+
int fieldNumber = segmentReadState.fieldInfos.fieldInfo(field).getFieldNumber();
90+
91+
try (IndexInput input = segmentReadState.directory.openInput(quantizationStateFileName, IOContext.READ)) {
92+
CodecUtil.retrieveChecksum(input);
93+
int numFields = getNumFields(input);
94+
95+
long position = -1;
96+
int length = 0;
97+
98+
// Read each field's metadata from the index section, break when correct field is found
99+
for (int i = 0; i < numFields; i++) {
100+
int tempFieldNumber = input.readInt();
101+
int tempLength = input.readInt();
102+
long tempPosition = input.readVLong();
103+
if (tempFieldNumber == fieldNumber) {
104+
position = tempPosition;
105+
length = tempLength;
106+
break;
107+
}
108+
}
109+
110+
if (position == -1 || length == 0) {
111+
throw new IllegalArgumentException(String.format("Field %s not found", field));
112+
}
113+
114+
byte[] stateBytes = readStateBytes(input, position, length);
115+
116+
// Deserialize the byte array to a quantization state object
117+
ScalarQuantizationType scalarQuantizationType = ((ScalarQuantizationParams) readConfig.getQuantizationParams()).getSqType();
118+
switch (scalarQuantizationType) {
119+
case ONE_BIT:
120+
return OneBitScalarQuantizationState.fromByteArray(stateBytes);
121+
case TWO_BIT:
122+
case FOUR_BIT:
123+
return MultiBitScalarQuantizationState.fromByteArray(stateBytes);
124+
default:
125+
throw new IllegalArgumentException(String.format("Unexpected scalar quantization type: %s", scalarQuantizationType));
126+
}
127+
} catch (Exception e) {
128+
log.warn(String.format("Unable to read the quantization state file for segment %s", segmentReadState.segmentInfo.name), e);
129+
return null;
130+
}
131+
}
132+
133+
@VisibleForTesting
134+
static int getNumFields(IndexInput input) throws IOException {
135+
long footerStart = input.length() - CodecUtil.footerLength();
136+
long markerAndIndexPosition = footerStart - Integer.BYTES - Long.BYTES;
137+
input.seek(markerAndIndexPosition);
138+
long indexStartPosition = input.readLong();
139+
input.seek(indexStartPosition);
140+
return input.readInt();
141+
}
142+
143+
@VisibleForTesting
144+
static byte[] readStateBytes(IndexInput input, long position, int length) throws IOException {
145+
input.seek(position);
146+
byte[] stateBytes = new byte[length];
147+
input.readBytes(stateBytes, 0, length);
148+
return stateBytes;
149+
}
150+
151+
@VisibleForTesting
152+
static String getQuantizationStateFileName(SegmentReadState state) {
153+
return IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, KNNConstants.QUANTIZATION_STATE_FILE_SUFFIX);
154+
}
155+
}

‎src/main/java/org/opensearch/knn/index/codec/KNN990Codec/KNN990QuantizationStateWriter.java‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.knn.index.codec.KNN990Codec;
7+
8+
import lombok.AllArgsConstructor;
9+
import lombok.Setter;
10+
import org.apache.lucene.codecs.CodecUtil;
11+
import org.apache.lucene.index.IndexFileNames;
12+
import org.apache.lucene.index.SegmentWriteState;
13+
import org.apache.lucene.store.IndexOutput;
14+
import org.opensearch.knn.common.KNNConstants;
15+
import org.opensearch.knn.quantization.models.quantizationState.QuantizationState;
16+
17+
import java.io.IOException;
18+
import java.util.ArrayList;
19+
import java.util.List;
20+
21+
/**
22+
* Writes quantization states to off heap memory
23+
*/
24+
public final class KNN990QuantizationStateWriter {
25+
26+
private final IndexOutput output;
27+
private List<FieldQuantizationState> fieldQuantizationStates = new ArrayList<>();
28+
static final String NATIVE_ENGINES_990_KNN_VECTORS_FORMAT_QS_DATA = "NativeEngines990KnnVectorsFormatQSData";
29+
30+
/**
31+
* Constructor
32+
* Overall file format for writer:
33+
* Header
34+
* QS1 state bytes
35+
* QS2 state bytes
36+
* Number of quantization states
37+
* QS1 field number
38+
* QS1 state bytes length
39+
* QS1 position of state bytes
40+
* QS2 field number
41+
* QS2 state bytes length
42+
* QS2 position of state bytes
43+
* Position of index section (where QS1 field name is located)
44+
* -1 (marker)
45+
* Footer
46+
* @param segmentWriteState segment write state containing segment information
47+
* @throws IOException exception could be thrown while creating the output
48+
*/
49+
public KNN990QuantizationStateWriter(SegmentWriteState segmentWriteState) throws IOException {
50+
String quantizationStateFileName = IndexFileNames.segmentFileName(
51+
segmentWriteState.segmentInfo.name,
52+
segmentWriteState.segmentSuffix,
53+
KNNConstants.QUANTIZATION_STATE_FILE_SUFFIX
54+
);
55+
56+
output = segmentWriteState.directory.createOutput(quantizationStateFileName, segmentWriteState.context);
57+
}
58+
59+
/**
60+
* Writes an index header
61+
* @param segmentWriteState state containing segment information
62+
* @throws IOException exception could be thrown while writing header
63+
*/
64+
public void writeHeader(SegmentWriteState segmentWriteState) throws IOException {
65+
CodecUtil.writeIndexHeader(
66+
output,
67+
NATIVE_ENGINES_990_KNN_VECTORS_FORMAT_QS_DATA,
68+
0,
69+
segmentWriteState.segmentInfo.getId(),
70+
segmentWriteState.segmentSuffix
71+
);
72+
}
73+
74+
/**
75+
* Writes a quantization state as bytes
76+
*
77+
* @param fieldNumber field number
78+
* @param quantizationState quantization state
79+
* @throws IOException could be thrown while writing
80+
*/
81+
public void writeState(int fieldNumber, QuantizationState quantizationState) throws IOException {
82+
byte[] stateBytes = quantizationState.toByteArray();
83+
long position = output.getFilePointer();
84+
output.writeBytes(stateBytes, stateBytes.length);
85+
fieldQuantizationStates.add(new FieldQuantizationState(fieldNumber, stateBytes, position));
86+
}
87+
88+
/**
89+
* Writes index footer and other index information for parsing later
90+
* @throws IOException could be thrown while writing
91+
*/
92+
public void writeFooter() throws IOException {
93+
long indexStartPosition = output.getFilePointer();
94+
output.writeInt(fieldQuantizationStates.size());
95+
for (FieldQuantizationState fieldQuantizationState : fieldQuantizationStates) {
96+
output.writeInt(fieldQuantizationState.fieldNumber);
97+
output.writeInt(fieldQuantizationState.stateBytes.length);
98+
output.writeVLong(fieldQuantizationState.position);
99+
}
100+
output.writeLong(indexStartPosition);
101+
output.writeInt(-1);
102+
CodecUtil.writeFooter(output);
103+
}
104+
105+
@AllArgsConstructor
106+
private static class FieldQuantizationState {
107+
final int fieldNumber;
108+
final byte[] stateBytes;
109+
@Setter
110+
Long position;
111+
}
112+
113+
public void closeOutput() throws IOException {
114+
output.close();
115+
}
116+
}

0 commit comments

Comments
 (0)