Adding main segment implementation for API and indexing. #2653

oaganesh · 2025-04-08T21:32:59Z

Description

Adding implementation for profiling and statistical analysis of KNN vector segments in OpenSearch. Provides functionality to collect, process, and store statistical information about vector dimensions across different shards and segments.

Example input:

Creating the index

curl -X PUT "http://localhost:9200/my-index" -H 'Content-Type: application/json' -d '     
{
  "settings": {                          
    "index": {
      "knn": true
    }
  },
  "mappings": {
    "properties": {
      "my_vector_field": {
        "type": "knn_vector",
        "dimension": 4
      }
    }
  }
}'

Adding vectors to the index

curl -X POST "http://localhost:9200/my-index/_doc" -H 'Content-Type: application/json' -d'
{
  "my_vector_field": [1.0, 2.0, 3.0, 4.0]
}'

curl -X POST "http://localhost:9200/my-index/_doc" -H 'Content-Type: application/json' -d'
{
  "my_vector_field": [1.5, 2.0, 3.0, 4.0]
}'

Flushing the index

curl -X POST "http://localhost:9200/my-index/_flush"

Observed example output:

curl -X GET "localhost:9200/_plugins/_knn/sampling/my-index/stats?pretty"
{
  "index_summary" : {
    "doc_count" : 2,
    "size_in_bytes" : 5763,
    "timestamp" : "2025-04-07T21:38:31.814948Z"
  },
  "vector_stats" : {
    "sample_size" : 2,
    "summary_stats" : {
      "dimensions" : [
        {
          "dimension" : 0,
          "count" : 2,
          "min" : 1.0,
          "max" : 1.5,
          "sum" : 2.5,
          "mean" : 1.25,
          "standardDeviation" : 0.3536,
          "variance" : 0.125
        },
        {
          "dimension" : 1,
          "count" : 2,
          "min" : 2.0,
          "max" : 2.0,
          "sum" : 4.0,
          "mean" : 2.0,
          "standardDeviation" : 0.0,
          "variance" : 0.0
        },
        {
          "dimension" : 2,
          "count" : 2,
          "min" : 3.0,
          "max" : 3.0,
          "sum" : 6.0,
          "mean" : 3.0,
          "standardDeviation" : 0.0,
          "variance" : 0.0
        },
        {
          "dimension" : 3,
          "count" : 2,
          "min" : 4.0,
          "max" : 4.0,
          "sum" : 8.0,
          "mean" : 4.0,
          "standardDeviation" : 0.0,
          "variance" : 0.0
        }
      ]
    }
  }
}

Related Issues

Implements #2687

Check List

[ ✔️ ] New functionality includes testing.
[ ✔️ ] New functionality has been documented.
[ ✔️] API changes companion pull request created.
[ ✔️] Commits are signed per the DCO using --signoff.
[ ✔️] Public documentation issue/PR created.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check ✔️.

…earch-project#2624) Signed-off-by: Arun Ganesh <[email protected]>

markwu-sde · 2025-04-09T07:00:50Z

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

+                log.info("Starting vector profiling for field: {}", fieldInfo.getName());
+                SegmentProfilerState.profileVectors(knnVectorValuesSupplier, segmentWriteState, fieldInfo.getName());
+                log.info("Completed vector profiling for field: {}", fieldInfo.getName());
+            } catch (Exception e) {


Should we catch the exception in our SegmentProfileState class instead? We want to keep the business logic here relatively untouched.

markwu-sde · 2025-04-09T07:00:53Z

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


    private KNNStats knnStats;
    private ClusterService clusterService;
+    // private IndicesService indicesService;


uncomment or remove this line. Why do we need IndicesService here? It's a pretty heavy object to be using inside KnnPlugin

markwu-sde · 2025-04-09T07:00:54Z

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

    private KNNStats knnStats;
    private ClusterService clusterService;
+    // private IndicesService indicesService;
+    private Environment environment;


Do we need Environment here? What information are we trying to get here?

It's used to access the data directory paths and node-specific items. It's needed to pass to the RestKNNSamplingStatsHandler to locate and read statistics files stored on disk. The environment provides access to the node's data paths where vector statistics files are stored under the indices directory structure.

We don't need this - we can directly follow what's being used in the QuantizationState

markwu-sde · 2025-04-09T07:00:56Z

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

        Supplier<RepositoriesService> repositoriesServiceSupplier
    ) {
        this.clusterService = clusterService;
+        // this.indicesService = client.getInstanceFromNode(IndicesService.class);


Let's not leave commented lines in the PR.

markwu-sde · 2025-04-09T07:00:59Z

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

        this.clusterService = clusterService;
+        // this.indicesService = client.getInstanceFromNode(IndicesService.class);
        this.repositoriesServiceSupplier = repositoriesServiceSupplier;
+        this.environment = environment;


Why do we need the environment here?

It's used to access the data directory paths and node-specific items. It's needed to pass to the RestKNNSamplingStatsHandler to locate and read statistics files stored on disk. The environment provides access to the node's data paths where vector statistics files are stored under the indices directory structure.

We don't need this - we can directly follow what's being used in the QuantizationState

markwu-sde · 2025-04-09T07:01:02Z

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

+            clusterService,
+            indexNameExpressionResolver,
+            this.environment
+            // indicesService


Remove comment.

markwu-sde · 2025-04-09T07:01:04Z

src/main/java/org/opensearch/knn/plugin/rest/RestKNNSamplingStatsHandler.java

+ * Rest handler for sampling stats endpoint
+ */
+public class RestKNNSamplingStatsHandler extends BaseRestHandler {
+    // private final IndicesService indicesService;


Remove comments.

markwu-sde · 2025-04-09T07:01:06Z

src/main/java/org/opensearch/knn/plugin/rest/RestKNNSamplingStatsHandler.java

+     * @param indexNameExpressionResolver Resolver for index names
+     * @param environment OpenSearch environment configuration
+     */
+    public RestKNNSamplingStatsHandler(


NIT: use lombok for @AllArgsConstructor

markwu-sde · 2025-04-09T07:01:10Z

src/main/java/org/opensearch/knn/plugin/rest/RestKNNSamplingStatsHandler.java

+        Environment environment
+        // IndicesService indicesService
+    ) {
+        // this.indicesService = indicesService;


Remove comments.

markwu-sde · 2025-04-09T07:01:12Z

src/main/java/org/opensearch/knn/plugin/rest/RestKNNSamplingStatsHandler.java

+     */
+    @Override
+    protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
+        String indexName = request.param("index");


Let's not hardcode parameter strings here.

An OS index can consist of multiple vector fields so getting just the index could return multiple fields. Do we need to query by field?

markwu-sde · 2025-04-09T07:01:17Z