Adding serializer and api implementation for segment profiler state

CHANGELOG.md

@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ## [Unreleased 2.x](https://github.com/opensearch-project/k-NN/compare/2.19...2.x)
 ### Features
 * [Vector Profiler] Adding basic generic vector profiler implementation and tests. [#2624](https://github.com/opensearch-project/k-NN/pull/2624)
+* [Vector Profiler] Adding serializer and api implementation for segment profiler state. [#2687](https://github.com/opensearch-project/k-NN/pull/2687)
 
 ### Enhancements
 ### Bug Fixes

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

@@ -84,6 +84,7 @@ public class KNNConstants {
     public static final String RADIAL_SEARCH_KEY = "radial_search";
     public static final String MODEL_VERSION = "model_version";
     public static final String QUANTIZATION_STATE_FILE_SUFFIX = "osknnqstate";
+    public static final String SEGMENT_PROFILE_STATE_FILE_SUFFIX = "segpfstate";
 
     // Lucene specific constants
     public static final String LUCENE_NAME = "lucene";

src/main/java/org/opensearch/knn/index/KNNIndexShard.java

@@ -28,6 +28,9 @@
 import org.opensearch.knn.index.memory.NativeMemoryEntryContext;
 import org.opensearch.knn.index.memory.NativeMemoryLoadStrategy;
 import org.opensearch.knn.index.engine.KNNEngine;
+import org.opensearch.knn.index.query.SegmentProfilerUtil;
+import org.opensearch.knn.plugin.transport.KNNIndexShardProfileResult;
+import org.opensearch.knn.profiler.SegmentProfilerState;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -82,6 +85,42 @@ public String getIndexName() {
         return indexShard.shardId().getIndexName();
     }
 
+    public KNNIndexShardProfileResult profile(final String field) {
+        try (Engine.Searcher searcher = indexShard.acquireSearcher("knn-profile")) {
+            log.info("[KNN] Profiling field [{}] in index [{}]", field, getIndexName());
+
+            List<SegmentProfilerState> segmentLevelProfilerStates = new ArrayList<>();
+
+            // For each leaf, collect the profile
+            for (LeafReaderContext leaf : searcher.getIndexReader().leaves()) {
+                try {
+                    SegmentProfilerState state = SegmentProfilerUtil.getSegmentProfileState(leaf.reader(), field);
+                    segmentLevelProfilerStates.add(state);
+                    log.debug("[KNN] Successfully profiled segment for field [{}]", field);
+                } catch (IOException e) {
+                    log.error("[KNN] Failed to profile segment for field [{}]: {}", field, e.getMessage(), e);
+                    throw new RuntimeException("Failed to profile segment: " + e.getMessage(), e);
+                }
+            }
+
+            if (segmentLevelProfilerStates.isEmpty()) {
+                log.warn("[KNN] No segments found with field [{}] in index [{}]", field, getIndexName());
+            } else {
+                log.info(
+                    "[KNN] Successfully profiled [{}] segments for field [{}] in index [{}]",
+                    segmentLevelProfilerStates.size(),
+                    field,
+                    getIndexName()
+                );
+            }
+
+            return new KNNIndexShardProfileResult(segmentLevelProfilerStates, indexShard.shardId().toString());
+        } catch (Exception e) {
+            log.error("[KNN] Error profiling field [{}] in index [{}]: {}", field, getIndexName(), e.getMessage(), e);
+            throw new RuntimeException("Error during profiling: " + e.getMessage(), e);
+        }
+    }
+
     /**
      * Load all of the k-NN segments for this shard into the cache.
      *

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

@@ -12,6 +12,7 @@
 import org.apache.lucene.index.SegmentWriteState;
 import org.apache.lucene.store.IndexOutput;
 import org.opensearch.knn.common.KNNConstants;
+import org.opensearch.knn.profiler.SegmentProfilerState;
 import org.opensearch.knn.quantization.models.quantizationState.QuantizationState;
 
 import java.io.IOException;
@@ -46,14 +47,18 @@ public final class KNN990QuantizationStateWriter {
      * @param segmentWriteState segment write state containing segment information
      * @throws IOException exception could be thrown while creating the output
      */
-    public KNN990QuantizationStateWriter(SegmentWriteState segmentWriteState) throws IOException {
-        String quantizationStateFileName = IndexFileNames.segmentFileName(
+    public KNN990QuantizationStateWriter(SegmentWriteState segmentWriteState, String fileSuffix) throws IOException {
+        String stateFileName = IndexFileNames.segmentFileName(
             segmentWriteState.segmentInfo.name,
             segmentWriteState.segmentSuffix,
-            KNNConstants.QUANTIZATION_STATE_FILE_SUFFIX
+            fileSuffix
         );
 
-        output = segmentWriteState.directory.createOutput(quantizationStateFileName, segmentWriteState.context);
+        output = segmentWriteState.directory.createOutput(stateFileName, segmentWriteState.context);
+    }
+
+    public KNN990QuantizationStateWriter(SegmentWriteState segmentWriteState) throws IOException {
+        this(segmentWriteState, KNNConstants.QUANTIZATION_STATE_FILE_SUFFIX);
     }
 
     /**
@@ -85,6 +90,20 @@ public void writeState(int fieldNumber, QuantizationState quantizationState) thr
         fieldQuantizationStates.add(new FieldQuantizationState(fieldNumber, stateBytes, position));
     }
 
+    /**
+     * Writes a segment profile state as bytes
+     *
+     * @param fieldNumber field number
+     * @param segmentProfilerState segment profiler state
+     * @throws IOException could be thrown while writing
+     */
+    public void writeState(int fieldNumber, SegmentProfilerState segmentProfilerState) throws IOException {
+        byte[] stateBytes = segmentProfilerState.toByteArray();
+        long position = output.getFilePointer();
+        output.writeBytes(stateBytes, stateBytes.length);
+        fieldQuantizationStates.add(new FieldQuantizationState(fieldNumber, stateBytes, position));
+    }
+
     /**
      * Writes index footer and other index information for parsing later
      * @throws IOException could be thrown while writing

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

@@ -35,6 +35,10 @@
 import org.opensearch.knn.index.quantizationservice.QuantizationService;
 import org.opensearch.knn.memoryoptsearch.VectorSearcher;
 import org.opensearch.knn.memoryoptsearch.VectorSearcherFactory;
+import org.opensearch.knn.profiler.KNN990ProfileStateReader;
+import org.opensearch.knn.profiler.SegmentProfileKNNCollector;
+import org.opensearch.knn.profiler.SegmentProfileStateReadConfig;
+import org.opensearch.knn.profiler.SegmentProfilerState;
 import org.opensearch.knn.quantization.models.quantizationState.QuantizationState;
 import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateCacheManager;
 import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateReadConfig;
@@ -163,6 +167,14 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
             return;
         }
 
+        if (knnCollector instanceof SegmentProfileKNNCollector) {
+            SegmentProfilerState segmentProfileState = KNN990ProfileStateReader.read(
+                new SegmentProfileStateReadConfig(segmentReadState, field)
+            );
+            ((SegmentProfileKNNCollector) knnCollector).setSegmentProfilerState(segmentProfileState);
+            return;
+        }
+
         if (trySearchWithMemoryOptimizedSearch(field, target, knnCollector, acceptDocs, true)) {
             return;
         }

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

@@ -15,14 +15,12 @@
 import org.apache.lucene.codecs.KnnFieldVectorsWriter;
 import org.apache.lucene.codecs.KnnVectorsWriter;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
-import org.apache.lucene.index.FieldInfo;
-import org.apache.lucene.index.MergeState;
-import org.apache.lucene.index.SegmentWriteState;
-import org.apache.lucene.index.Sorter;
+import org.apache.lucene.index.*;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.util.IOUtils;
 import org.apache.lucene.util.RamUsageEstimator;
 import org.opensearch.common.StopWatch;
+import org.opensearch.knn.common.KNNConstants;
 import org.opensearch.knn.index.VectorDataType;
 import org.opensearch.knn.index.codec.nativeindex.NativeIndexBuildStrategyFactory;
 import org.opensearch.knn.index.codec.nativeindex.NativeIndexWriter;
@@ -52,6 +50,7 @@ public class NativeEngines990KnnVectorsWriter extends KnnVectorsWriter {
     private final SegmentWriteState segmentWriteState;
     private final FlatVectorsWriter flatVectorsWriter;
     private KNN990QuantizationStateWriter quantizationStateWriter;
+    private KNN990QuantizationStateWriter segmentStateWriter;
     private final List<NativeEngineFieldVectorsWriter<?>> fields = new ArrayList<>();
     private boolean finished;
     private final Integer approximateThreshold;
@@ -108,7 +107,7 @@ public void flush(int maxDoc, final Sorter.DocMap sortMap) throws IOException {
                 field.getVectors()
             );
             final QuantizationState quantizationState = train(field.getFieldInfo(), knnVectorValuesSupplier, totalLiveDocs);
-            SegmentProfilerState.profileVectors(knnVectorValuesSupplier);
+            profile(field.getFieldInfo(), knnVectorValuesSupplier, totalLiveDocs);
             // should skip graph building only for non quantization use case and if threshold is met
             if (quantizationState == null && shouldSkipBuildingVectorDataStructure(totalLiveDocs)) {
                 log.debug(
@@ -152,9 +151,13 @@ public void mergeOneField(final FieldInfo fieldInfo, final MergeState mergeState
         }
 
         final QuantizationState quantizationState = train(fieldInfo, knnVectorValuesSupplier, totalLiveDocs);
+
+        // Write the segment profile state to the directory
+        profile(fieldInfo, knnVectorValuesSupplier, totalLiveDocs);
+
         // should skip graph building only for non quantization use case and if threshold is met
         if (quantizationState == null && shouldSkipBuildingVectorDataStructure(totalLiveDocs)) {
-            log.debug(
+            log.info(
                 "Skip building vector data structure for field: {}, as liveDoc: {} is less than the threshold {} during merge",
                 fieldInfo.name,
                 totalLiveDocs,
@@ -190,6 +193,10 @@ public void finish() throws IOException {
         if (quantizationStateWriter != null) {
             quantizationStateWriter.writeFooter();
         }
+        if (segmentStateWriter != null) {
+            segmentStateWriter.writeFooter();
+        }
+
         flatVectorsWriter.finish();
     }
 
@@ -211,6 +218,9 @@ public void close() throws IOException {
         if (quantizationStateWriter != null) {
             quantizationStateWriter.closeOutput();
         }
+        if (segmentStateWriter != null) {
+            segmentStateWriter.closeOutput();
+        }
         IOUtils.close(flatVectorsWriter);
     }
 
@@ -243,6 +253,23 @@ private QuantizationState train(
         return quantizationState;
     }
 
+    private SegmentProfilerState profile(
+        final FieldInfo fieldInfo,
+        final Supplier<KNNVectorValues<?>> knnVectorValuesSupplier,
+        final int totalLiveDocs
+    ) throws IOException {
+
+        SegmentProfilerState segmentProfilerState = null;
+        if (totalLiveDocs > 0) {
+            initSegmentStateWriterIfNecessary();
+            String segmentId = segmentWriteState.segmentInfo.name;
+            SegmentProfilerState profileResultForSegment = SegmentProfilerState.profileVectors(knnVectorValuesSupplier, segmentId);
+            segmentStateWriter.writeState(fieldInfo.getFieldNumber(), profileResultForSegment);
+        }
+
+        return segmentProfilerState;
+    }
+
     /**
      * The {@link KNNVectorValues} will be exhausted after this function run. So make sure that you are not sending the
      * vectorsValues object which you plan to use later
@@ -260,11 +287,18 @@ private int getLiveDocs(KNNVectorValues<?> vectorValues) throws IOException {
 
     private void initQuantizationStateWriterIfNecessary() throws IOException {
         if (quantizationStateWriter == null) {
-            quantizationStateWriter = new KNN990QuantizationStateWriter(segmentWriteState);
+            quantizationStateWriter = new KNN990QuantizationStateWriter(segmentWriteState, KNNConstants.QUANTIZATION_STATE_FILE_SUFFIX);
             quantizationStateWriter.writeHeader(segmentWriteState);
         }
     }
 
+    private void initSegmentStateWriterIfNecessary() throws IOException {
+        if (segmentStateWriter == null) {
+            segmentStateWriter = new KNN990QuantizationStateWriter(segmentWriteState, KNNConstants.SEGMENT_PROFILE_STATE_FILE_SUFFIX);
+            segmentStateWriter.writeHeader(segmentWriteState);
+        }
+    }
+
     private boolean shouldSkipBuildingVectorDataStructure(final long docCount) {
         if (approximateThreshold < 0) {
             return true;

src/main/java/org/opensearch/knn/index/query/SegmentProfilerUtil.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.query;
+
+import lombok.experimental.UtilityClass;
+import org.apache.lucene.index.LeafReader;
+import org.opensearch.knn.profiler.SegmentProfileKNNCollector;
+import org.opensearch.knn.profiler.SegmentProfilerState;
+
+import java.io.IOException;
+import java.util.Locale;
+
+/**
+ * Utility class to get segment profiler state for a given field
+ */
+@UtilityClass
+public class SegmentProfilerUtil {
+
+    /**
+     * Gets the segment profile state for a given field
+     * @param leafReader The leaf reader to query
+     * @param fieldName The field name to profile
+     * @return The segment profiler state
+     * @throws IOException If there's an error reading the segment
+     */
+    public static SegmentProfilerState getSegmentProfileState(final LeafReader leafReader, String fieldName) throws IOException {
+        final SegmentProfileKNNCollector tempCollector = new SegmentProfileKNNCollector();
+        leafReader.searchNearestVectors(fieldName, new float[0], tempCollector, null);
+        if (tempCollector.getSegmentProfilerState() == null) {
+            throw new IllegalStateException(String.format(Locale.ROOT, "No segment state found for field %s", fieldName));
+        }
+        return tempCollector.getSegmentProfilerState();
+    }
+}

src/main/java/org/opensearch/knn/plugin/KNNPlugin.java

@@ -84,6 +84,9 @@
 import org.opensearch.knn.plugin.transport.UpdateModelGraveyardTransportAction;
 import org.opensearch.knn.plugin.transport.UpdateModelMetadataAction;
 import org.opensearch.knn.plugin.transport.UpdateModelMetadataTransportAction;
+import org.opensearch.knn.profiler.RestKNNProfileHandler;
+import org.opensearch.knn.plugin.transport.KNNProfileTransportAction;
+import org.opensearch.knn.plugin.transport.KNNProfileAction;
 import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateCache;
 import org.opensearch.knn.training.TrainingJobClusterStateListener;
 import org.opensearch.knn.training.TrainingJobRunner;
@@ -269,6 +272,7 @@ public List<RestHandler> getRestHandlers(
     ) {
 
         RestKNNStatsHandler restKNNStatsHandler = new RestKNNStatsHandler();
+        RestKNNProfileHandler restKNNProfileHandler = new RestKNNProfileHandler();
         RestKNNWarmupHandler restKNNWarmupHandler = new RestKNNWarmupHandler(
             settings,
             restController,
@@ -284,6 +288,7 @@ public List<RestHandler> getRestHandlers(
         return ImmutableList.of(
             restKNNStatsHandler,
             restKNNWarmupHandler,
+            restKNNProfileHandler,
             restGetModelHandler,
             restDeleteModelHandler,
             restTrainModelHandler,
@@ -300,6 +305,7 @@ public List<RestHandler> getRestHandlers(
         return Arrays.asList(
             new ActionHandler<>(KNNStatsAction.INSTANCE, KNNStatsTransportAction.class),
             new ActionHandler<>(KNNWarmupAction.INSTANCE, KNNWarmupTransportAction.class),
+            new ActionHandler<>(KNNProfileAction.INSTANCE, KNNProfileTransportAction.class),
             new ActionHandler<>(UpdateModelMetadataAction.INSTANCE, UpdateModelMetadataTransportAction.class),
             new ActionHandler<>(TrainingJobRouteDecisionInfoAction.INSTANCE, TrainingJobRouteDecisionInfoTransportAction.class),
             new ActionHandler<>(GetModelAction.INSTANCE, GetModelTransportAction.class),

src/main/java/org/opensearch/knn/plugin/transport/KNNIndexShardProfileResult.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.plugin.transport;
+
+import lombok.AllArgsConstructor;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.common.io.stream.Writeable;
+import org.opensearch.knn.profiler.SegmentProfilerState;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+@AllArgsConstructor
+public class KNNIndexShardProfileResult implements Writeable {
+    List<SegmentProfilerState> segmentProfilerStateList;
+    String shardId;
+
+    /**
+     * Constructor to deserialize from StreamInput
+     *
+     * @param in StreamInput to read from
+     * @throws IOException if there's an error reading from the stream
+     */
+    public KNNIndexShardProfileResult(StreamInput in) throws IOException {
+        this.shardId = in.readString();
+
+        // Read the number of SegmentProfilerStates
+        int size = in.readInt();
+        this.segmentProfilerStateList = new ArrayList<>(size);
+
+        // Read each SegmentProfilerState
+        for (int i = 0; i < size; i++) {
+            byte[] stateBytes = in.readByteArray();
+            this.segmentProfilerStateList.add(SegmentProfilerState.fromBytes(stateBytes));
+        }
+    }
+
+    @Override
+    public void writeTo(StreamOutput streamOutput) throws IOException {
+        streamOutput.writeString(shardId);
+
+        // Write the segment profiler state list size
+        streamOutput.writeInt(segmentProfilerStateList.size());
+
+        // Write each SegmentProfilerState as bytes
+        for (SegmentProfilerState state : segmentProfilerStateList) {
+            byte[] stateBytes = state.toByteArray();
+            streamOutput.writeByteArray(stateBytes);
+        }
+    }
+}