add stats for text embedding processor with different settings

Signed-off-by: will-hwang <[email protected]>

Author: will-hwang

CHANGELOG.md

@@ -15,7 +15,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Performance Improvement] Add custom bulk scorer for hybrid query (2-3x faster) ([#1289](https://github.com/opensearch-project/neural-search/pull/1289))
 - [Stats] Add stats for text chunking processor algorithms ([#1308](https://github.com/opensearch-project/neural-search/pull/1308))
 - Support custom weights in RRF normalization processor ([#1322](https://github.com/opensearch-project/neural-search/pull/1322))
-
+- [Stats] Add stats for text embedding processor with different settings ([#1332](https://github.com/opensearch-project/neural-search/pull/1332))
 ### Bug Fixes
 - Fix score value as null for single shard when sorting is not done on score field ([#1277](https://github.com/opensearch-project/neural-search/pull/1277))
 - Return bad request for stats API calls with invalid stat names instead of ignoring them ([#1291](https://github.com/opensearch-project/neural-search/pull/1291))

src/main/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessor.java

@@ -72,6 +72,7 @@ public void doExecute(
             generateAndSetInference(ingestDocument, processMap, inferenceList, handler);
             return;
         }
+        EventStatsManager.increment(EventStatName.TEXT_EMBEDDING_PROCESSOR_SKIP_EXISTING_EXECUTIONS);
         // if skipExisting flag is turned on, eligible inference texts will be compared and filtered after embeddings are copied
         Object index = ingestDocument.getSourceAndMetadata().get(INDEX_FIELD);
         Object id = ingestDocument.getSourceAndMetadata().get(ID_FIELD);
@@ -109,6 +110,7 @@ public void doBatchExecute(List<String> inferenceList, Consumer<List<?>> handler
     @Override
     public void subBatchExecute(List<IngestDocumentWrapper> ingestDocumentWrappers, Consumer<List<IngestDocumentWrapper>> handler) {
         try {
+            EventStatsManager.increment(EventStatName.TEXT_EMBEDDING_PROCESSOR_EXECUTIONS);
             if (CollectionUtils.isEmpty(ingestDocumentWrappers)) {
                 handler.accept(ingestDocumentWrappers);
                 return;
@@ -126,6 +128,7 @@ public void subBatchExecute(List<IngestDocumentWrapper> ingestDocumentWrappers,
             }
             // skipExisting flag is turned on, eligible inference texts in dataForInferences will be compared and filtered after embeddings
             // are copied
+            EventStatsManager.increment(EventStatName.TEXT_EMBEDDING_PROCESSOR_SKIP_EXISTING_EXECUTIONS);
             openSearchClient.execute(
                 MultiGetAction.INSTANCE,
                 buildMultiGetRequest(dataForInferences),

src/main/java/org/opensearch/neuralsearch/stats/events/EventStatName.java

@@ -19,6 +19,11 @@
 @Getter
 public enum EventStatName implements StatName {
     TEXT_EMBEDDING_PROCESSOR_EXECUTIONS("text_embedding_executions", "processors.ingest", EventStatType.TIMESTAMPED_EVENT_COUNTER),
+    TEXT_EMBEDDING_PROCESSOR_SKIP_EXISTING_EXECUTIONS(
+        "text_embedding_skip_existing_executions",
+        "processors.ingest",
+        EventStatType.TIMESTAMPED_EVENT_COUNTER
+    ),
     TEXT_CHUNKING_PROCESSOR_EXECUTIONS("text_chunking_executions", "processors.ingest", EventStatType.TIMESTAMPED_EVENT_COUNTER),
     TEXT_CHUNKING_FIXED_LENGTH_EXECUTIONS(
         "text_chunking_fixed_length_executions",

src/main/java/org/opensearch/neuralsearch/stats/info/InfoStatName.java

@@ -21,6 +21,7 @@ public enum InfoStatName implements StatName {
     // Cluster info
     CLUSTER_VERSION("cluster_version", "", InfoStatType.INFO_STRING),
     TEXT_EMBEDDING_PROCESSORS("text_embedding_processors_in_pipelines", "processors.ingest", InfoStatType.INFO_COUNTER),
+    TEXT_EMBEDDING_SKIP_EXISTING_PROCESSORS("text_embedding_skip_existing_processors", "processors.ingest", InfoStatType.INFO_COUNTER),
     TEXT_CHUNKING_PROCESSORS("text_chunking_processors", "processors.ingest", InfoStatType.INFO_COUNTER),
     TEXT_CHUNKING_DELIMITER_PROCESSORS("text_chunking_delimiter_processors", "processors.ingest", InfoStatType.INFO_COUNTER),
     TEXT_CHUNKING_FIXED_LENGTH_PROCESSORS("text_chunking_fixed_length_processors", "processors.ingest", InfoStatType.INFO_COUNTER);

src/main/java/org/opensearch/neuralsearch/stats/info/InfoStatsManager.java

@@ -17,6 +17,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.stream.Collectors;
 
 /**
@@ -130,7 +131,7 @@ private void addIngestProcessorStats(Map<InfoStatName, CountableInfoStatSnapshot
                     Map<String, Object> processorConfig = asMap(entry.getValue());
                     switch (processorType) {
                         case TextEmbeddingProcessor.TYPE:
-                            increment(stats, InfoStatName.TEXT_EMBEDDING_PROCESSORS);
+                            countTextEmbeddingProcessorStats(stats, processorConfig);
                             break;
                         case TextChunkingProcessor.TYPE:
                             countTextChunkingProcessorStats(stats, processorConfig);
@@ -141,6 +142,19 @@ private void addIngestProcessorStats(Map<InfoStatName, CountableInfoStatSnapshot
         }
     }
 
+    /**
+     * Counts text embedding processor stats based on processor config
+     * @param stats map containing the stat to increment
+     * @param processorConfig map of the processor config, parsed to add stats
+     */
+    private void countTextEmbeddingProcessorStats(Map<InfoStatName, CountableInfoStatSnapshot> stats, Map<String, Object> processorConfig) {
+        increment(stats, InfoStatName.TEXT_EMBEDDING_PROCESSORS);
+        Object skipExisting = processorConfig.get(TextEmbeddingProcessor.SKIP_EXISTING);
+        if (Objects.nonNull(skipExisting) && skipExisting.equals(Boolean.TRUE)) {
+            increment(stats, InfoStatName.TEXT_EMBEDDING_SKIP_EXISTING_PROCESSORS);
+        }
+    }
+
     /**
      * Counts text chunking processor stats based on processor config
      * @param stats map containing the stat to increment

src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java

@@ -9,6 +9,7 @@
 import java.net.URL;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -31,6 +32,8 @@
 
 import com.google.common.collect.ImmutableList;
 import org.opensearch.neuralsearch.query.NeuralQueryBuilder;
+import org.opensearch.neuralsearch.stats.events.EventStatName;
+import org.opensearch.neuralsearch.stats.info.InfoStatName;
 
 public class TextEmbeddingProcessorIT extends BaseNeuralSearchIT {
 
@@ -478,4 +481,53 @@ public void testTextEmbeddingProcessorWithReindexOperation() throws Exception {
         reindex(fromIndexName, toIndexName);
         assertEquals(1, getDocCount(toIndexName));
     }
+
+    public void testTextEmbeddingProcessor_processorStats_successful() throws Exception {
+        updateClusterSettings("plugins.neural_search.stats_enabled", true);
+        String modelId = uploadTextEmbeddingModel();
+        loadModel(modelId);
+        createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_EMBEDDING_WITH_SKIP_EXISTING);
+        createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME);
+        ingestDocument(INDEX_NAME, INGEST_DOC1, "1");
+        updateDocument(INDEX_NAME, UPDATE_DOC1, "1");
+        assertEquals(1, getDocCount(INDEX_NAME));
+        assertEquals(2, getDocById(INDEX_NAME, "1").get("_version"));
+        // Get stats
+        String responseBody = executeNeuralStatRequest(new ArrayList<>(), new ArrayList<>());
+        Map<String, Object> stats = parseInfoStatsResponse(responseBody);
+        Map<String, Object> allNodesStats = parseAggregatedNodeStatsResponse(responseBody);
+
+        // Parse json to get stats
+        assertEquals(2, getNestedValue(allNodesStats, EventStatName.TEXT_EMBEDDING_PROCESSOR_EXECUTIONS));
+        assertEquals(2, getNestedValue(allNodesStats, EventStatName.TEXT_EMBEDDING_PROCESSOR_SKIP_EXISTING_EXECUTIONS));
+
+        assertEquals(1, getNestedValue(stats, InfoStatName.TEXT_EMBEDDING_PROCESSORS));
+        assertEquals(1, getNestedValue(stats, InfoStatName.TEXT_EMBEDDING_SKIP_EXISTING_PROCESSORS));
+        // Reset stats
+        updateClusterSettings("plugins.neural_search.stats_enabled", false);
+    }
+
+    public void testTextEmbeddingProcessor_batch__processorStats_successful() throws Exception {
+        updateClusterSettings("plugins.neural_search.stats_enabled", true);
+        String modelId = uploadTextEmbeddingModel();
+        loadModel(modelId);
+        createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_EMBEDDING_WITH_SKIP_EXISTING);
+        createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME);
+        ingestBatchDocumentWithBulk("batch_", 2, Collections.emptySet(), Collections.emptySet(), List.of(INGEST_DOC1, INGEST_DOC2));
+        assertEquals(2, getDocCount(INDEX_NAME));
+
+        // Get stats
+        String responseBody = executeNeuralStatRequest(new ArrayList<>(), new ArrayList<>());
+        Map<String, Object> stats = parseInfoStatsResponse(responseBody);
+        Map<String, Object> allNodesStats = parseAggregatedNodeStatsResponse(responseBody);
+
+        // Parse json to get stats
+        assertEquals(2, getNestedValue(allNodesStats, EventStatName.TEXT_EMBEDDING_PROCESSOR_EXECUTIONS));
+        assertEquals(2, getNestedValue(allNodesStats, EventStatName.TEXT_EMBEDDING_PROCESSOR_SKIP_EXISTING_EXECUTIONS));
+
+        assertEquals(1, getNestedValue(stats, InfoStatName.TEXT_EMBEDDING_PROCESSORS));
+        assertEquals(1, getNestedValue(stats, InfoStatName.TEXT_EMBEDDING_SKIP_EXISTING_PROCESSORS));
+        updateClusterSettings("plugins.neural_search.stats_enabled", false);
+    }
+
 }