Add quantization state reader and writer

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

@@ -71,6 +71,7 @@ public class KNNConstants {
     public static final VectorDataType DEFAULT_VECTOR_DATA_TYPE_FIELD = VectorDataType.FLOAT;
 
     public static final String RADIAL_SEARCH_KEY = "radial_search";
+    public static final String QUANTIZATION_STATE_FILE_SUFFIX = "qs";
 
     // Lucene specific constants
     public static final String LUCENE_NAME = "lucene";

src/main/java/org/opensearch/knn/index/codec/KNN990Codec/KNNQuantizationStateReader.java

@@ -0,0 +1,143 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.codec.KNN990Codec;
+
+import org.apache.lucene.codecs.CodecUtil;
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.IndexFileNames;
+import org.apache.lucene.index.SegmentReadState;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexInput;
+import org.opensearch.knn.common.KNNConstants;
+import org.opensearch.knn.quantization.models.quantizationState.QuantizationState;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Reads quantization states
+ */
+public class KNNQuantizationStateReader {
+
+    /**
+     * Read quantization states and return list of fieldNames and bytes
+     *
+     * File format:
+     * Header
+     * QS1 state bytes
+     * QS2 state bytes
+     * Number of quantization states
+     * QS1 field name
+     * QS1 state bytes length
+     * QS1 position of state bytes
+     * QS2 field name
+     * QS2 state bytes length
+     * QS2 position of state bytes
+     * Position of index section (where QS1 field name is located)
+     * -1 (marker)
+     * Footer
+     *
+     * @param state the read state to read from
+     */
+    public Map<String, byte[]> read(SegmentReadState state) {
+        String quantizationStateFileName = IndexFileNames.segmentFileName(
+            state.segmentInfo.name,
+            state.segmentSuffix,
+            KNNConstants.QUANTIZATION_STATE_FILE_SUFFIX
+        );
+        Map<String, byte[]> readQuantizationStateInfos = new HashMap<>();
+
+        try (IndexInput input = state.directory.openInput(quantizationStateFileName, IOContext.READ)) {
+
+            int numFields = getNumFields(input);
+
+            List<String> fieldNames = new ArrayList<>();
+            List<Long> positions = new ArrayList<>();
+            List<Integer> lengths = new ArrayList<>();
+
+            // Read each field's metadata from the index section
+            for (int i = 0; i < numFields; i++) {
+                fieldNames.add(input.readString());
+                int length = input.readInt();
+                lengths.add(length);
+                long position = input.readVLong();
+                positions.add(position);
+            }
+            // Read each field's bytes
+            for (int i = 0; i < numFields; i++) {
+                input.seek(positions.get(i));
+                byte[] stateBytes = new byte[lengths.get(i)];
+                input.readBytes(stateBytes, 0, lengths.get(i));
+                readQuantizationStateInfos.put(fieldNames.get(i), stateBytes);
+            }
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+        return readQuantizationStateInfos;
+    }
+
+    /**
+     * Reads an individual quantization state for a given field
+     * @param directory directory to open input
+     * @param segmentName segment name
+     * @param segmentSuffix segment suffix
+     * @param fieldInfo field information
+     * @return quantization state
+     */
+    public QuantizationState read(Directory directory, String segmentName, String segmentSuffix, FieldInfo fieldInfo) throws IOException {
+        String quantizationStateFileName = IndexFileNames.segmentFileName(
+            segmentName,
+            segmentSuffix,
+            KNNConstants.QUANTIZATION_STATE_FILE_SUFFIX
+        );
+        String fieldName = fieldInfo.getName();
+
+        IndexInput input = directory.openInput(quantizationStateFileName, IOContext.READ);
+        CodecUtil.retrieveChecksum(input);
+        int numFields = getNumFields(input);
+
+        long position = -1;
+        int length = 0;
+
+        // Read each field's metadata from the index section
+        for (int i = 0; i < numFields; i++) {
+            String tempFieldName = input.readString();
+            int tempLength = input.readInt();
+            long tempPosition = input.readVLong();
+            if (tempFieldName.equals(fieldName)) {
+                position = tempPosition;
+                length = tempLength;
+                break;
+            }
+        }
+
+        if (position == -1 || length == 0) {
+            throw new IllegalArgumentException(String.format("Field %s not found", fieldName));
+        }
+
+        input.seek(position);
+        byte[] stateBytes = new byte[length];
+        input.readBytes(stateBytes, 0, length);
+        input.close();
+        // Deserialize the byte array to a quantization state object
+        // TODO: Get params from field info and deserialize
+        return null;
+    }
+
+    private int getNumFields(IndexInput input) throws IOException {
+        long footerStart = input.length() - CodecUtil.footerLength();
+        long markerAndIndexPosition = footerStart - Integer.BYTES - Long.BYTES;
+        input.seek(markerAndIndexPosition);
+        long indexStartPosition = input.readLong();
+        input.readInt();
+        input.seek(indexStartPosition);
+        return input.readInt();
+    }
+}

src/main/java/org/opensearch/knn/index/codec/KNN990Codec/KNNQuantizationStateWriter.java

@@ -0,0 +1,90 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.codec.KNN990Codec;
+
+import lombok.AllArgsConstructor;
+import org.apache.lucene.codecs.CodecUtil;
+import org.apache.lucene.index.IndexFileNames;
+import org.apache.lucene.index.SegmentWriteState;
+import org.apache.lucene.store.IndexOutput;
+import org.opensearch.knn.common.KNNConstants;
+import org.opensearch.knn.quantization.models.quantizationState.QuantizationState;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Writes quantization states to off heap memory
+ */
+public class KNNQuantizationStateWriter {
+
+    private final IndexOutput output;
+    private List<FieldQuantizationState> fieldQuantizationStates = new ArrayList<>();
+
+    /**
+     * Constructor
+     * @param segmentWriteState segment write state containing segment information
+     * @throws IOException exception could be thrown while creating the output
+     */
+    public KNNQuantizationStateWriter(SegmentWriteState segmentWriteState) throws IOException {
+        String quantizationStateFileName = IndexFileNames.segmentFileName(
+            segmentWriteState.segmentInfo.name,
+            segmentWriteState.segmentSuffix,
+            KNNConstants.QUANTIZATION_STATE_FILE_SUFFIX
+        );
+
+        output = segmentWriteState.directory.createOutput(quantizationStateFileName, segmentWriteState.context);
+    }
+
+    /**
+     * Writes an index header
+     * @param segmentWriteState state containing segment information
+     * @throws IOException exception could be thrown while writing header
+     */
+    public void writeHeader(SegmentWriteState segmentWriteState) throws IOException {
+        CodecUtil.writeIndexHeader(output, "QuantizationCodec", 0, segmentWriteState.segmentInfo.getId(), segmentWriteState.segmentSuffix);
+    }
+
+    /**
+     * Writes a quantization state as bytes
+     * @param fieldName field name
+     * @param quantizationState quantization state
+     * @throws IOException could be thrown while writing
+     */
+    public void writeState(String fieldName, QuantizationState quantizationState) throws IOException {
+        byte[] stateBytes = quantizationState.toByteArray();
+        long position = output.getFilePointer();
+        output.writeBytes(stateBytes, stateBytes.length);
+        fieldQuantizationStates.add(new FieldQuantizationState(fieldName, stateBytes, position));
+    }
+
+    /**
+     * Writes index footer and other index information for parsing later
+     * @throws IOException could be thrown while writing
+     */
+    public void writeFooter() throws IOException {
+        long indexStartPosition = output.getFilePointer();
+        output.writeInt(fieldQuantizationStates.size());
+        for (FieldQuantizationState fieldQuantizationState : fieldQuantizationStates) {
+            output.writeString(fieldQuantizationState.fieldName);
+            output.writeInt(fieldQuantizationState.stateBytes.length);
+            output.writeVLong(fieldQuantizationState.position);
+        }
+        output.writeLong(indexStartPosition);
+        output.writeInt(-1);
+        CodecUtil.writeFooter(output);
+        output.close();
+        fieldQuantizationStates = new ArrayList<>();
+    }
+
+    @AllArgsConstructor
+    private static class FieldQuantizationState {
+        final String fieldName;
+        final byte[] stateBytes;
+        final Long position;
+    }
+}

src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsWriter.java

@@ -11,7 +11,6 @@
 
 package org.opensearch.knn.index.codec.KNN990Codec;
 
-import lombok.RequiredArgsConstructor;
 import lombok.extern.log4j.Log4j2;
 import org.apache.lucene.codecs.KnnFieldVectorsWriter;
 import org.apache.lucene.codecs.KnnVectorsWriter;
@@ -39,14 +38,21 @@
  * A KNNVectorsWriter class for writing the vector data strcutures and flat vectors for Native Engines.
  */
 @Log4j2
-@RequiredArgsConstructor
 public class NativeEngines990KnnVectorsWriter extends KnnVectorsWriter {
     private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(NativeEngines990KnnVectorsWriter.class);
     private final SegmentWriteState segmentWriteState;
     private final FlatVectorsWriter flatVectorsWriter;
+    private final KNNQuantizationStateWriter quantizationStateWriter;
     private final List<NativeEngineFieldVectorsWriter<?>> fields = new ArrayList<>();
     private boolean finished;
 
+    public NativeEngines990KnnVectorsWriter(SegmentWriteState segmentWriteState, FlatVectorsWriter flatVectorsWriter) throws IOException {
+        this.segmentWriteState = segmentWriteState;
+        this.flatVectorsWriter = flatVectorsWriter;
+        this.quantizationStateWriter = new KNNQuantizationStateWriter(segmentWriteState);
+
+    }
+
     /**
      * Add new field for indexing.
      * In Lucene, we use single file for all the vector fields so here we need to see how we are going to make things
@@ -70,6 +76,9 @@ public KnnFieldVectorsWriter<?> addField(final FieldInfo fieldInfo) throws IOExc
     public void flush(int maxDoc, final Sorter.DocMap sortMap) throws IOException {
         // simply write data in the flat file
         flatVectorsWriter.flush(maxDoc, sortMap);
+
+        quantizationStateWriter.writeHeader(segmentWriteState);
+
         for (final NativeEngineFieldVectorsWriter<?> field : fields) {
             final VectorDataType vectorDataType = extractVectorDataType(field.getFieldInfo());
             final KNNVectorValues<?> knnVectorValues = KNNVectorValuesFactory.getVectorValues(
@@ -78,8 +87,12 @@ public void flush(int maxDoc, final Sorter.DocMap sortMap) throws IOException {
                 field.getVectors()
             );
 
+            // TODO: Extract quantization state here, uncomment below line once implemented
+            // quantizationStateWriter.writeState(field.getFieldInfo().getName(), quantizationState);
+
             NativeIndexWriter.getWriter(field.getFieldInfo(), segmentWriteState).flushIndex(knnVectorValues);
         }
+        quantizationStateWriter.writeFooter();
     }
 
     @Override